use std::collections::HashMap;

use ash::vk::{self};

use crate::{RenderCtx, decode_rif, game::Camera};

use super::{MemType, loading_world::WorldComponent};

pub struct Skybox {
    pub buf: vk::Buffer,
    pub image: vk::Image,
    pub imageview: vk::ImageView,
    pub sampler: vk::Sampler,
    pipeline: vk::Pipeline,
    pipe_layout: vk::PipelineLayout,
}

impl WorldComponent for Skybox {
    fn descriptors(&self) -> HashMap<vk::DescriptorType, u32> {
        let mut map = HashMap::new();
        map.insert(vk::DescriptorType::COMBINED_IMAGE_SAMPLER, 1);
        map
    }

    fn desc_layout(&self, ctx: &RenderCtx) -> vk::DescriptorSetLayout {
        let image_binding = vk::DescriptorSetLayoutBinding::default()
            .binding(0)
            .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
            .descriptor_count(1)
            .stage_flags(vk::ShaderStageFlags::FRAGMENT);
        let layouts = [image_binding];

        let layout_info = vk::DescriptorSetLayoutCreateInfo::default().bindings(&layouts);
        let layout = unsafe {
            ctx.dev
                .create_descriptor_set_layout(&layout_info, None)
                .expect("Failed to create descriptor set layout")
        };
        layout
    }

    fn write_desc_set(
        &self,
        ctx: &RenderCtx,
        desc_set: vk::DescriptorSet,
        _: (vk::ImageView, vk::Sampler),
    ) {
        let img_info = vk::DescriptorImageInfo::default()
            .image_layout(vk::ImageLayout::GENERAL)
            .image_view(self.imageview)
            .sampler(self.sampler);
        let img_infos = &[img_info];
        let img_desc = vk::WriteDescriptorSet::default()
            .dst_set(desc_set)
            .dst_binding(0)
            .dst_array_element(0)
            .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
            .image_info(img_infos);

        unsafe {
            ctx.dev.update_descriptor_sets(&[img_desc], &[]);
        }
    }

    fn pipeline(&self) -> (vk::Pipeline, vk::PipelineLayout) {
        (self.pipeline, self.pipe_layout)
    }

    fn push_constants(
        &self,
        camera: &Camera,
        res: [u32; 2],
        base_color: [f32; 4],
    ) -> Vec<(vk::ShaderStageFlags, u32, Vec<u8>)> {
        let mut constants = Vec::new();
        let cam_data: Vec<u8> = (Vec::<f32>::from([
            camera.origin.x,
            camera.origin.y,
            camera.origin.z,
            0.0,
            camera.rotation.x,
            camera.rotation.y,
            camera.rotation.z,
        ]))
        .iter()
        .map(|f| f.to_ne_bytes())
        .collect::<Vec<_>>()
        .into_flattened();
        constants.push((vk::ShaderStageFlags::VERTEX, 0, cam_data));

        let screen_res: Vec<u8> = res
            .iter()
            .map(|f| f.to_ne_bytes())
            .collect::<Vec<_>>()
            .into_flattened();
        constants.push((vk::ShaderStageFlags::VERTEX, 32, screen_res));

        constants
    }
}

impl Skybox {}

impl RenderCtx {
    pub(super) fn create_skybox(&self, pass: vk::RenderPass) -> Skybox {
        let skybox_mem = self.mem_alloc(MemType::HostVisibile, 3 * 2048 * 2048 * 6);
        let skybox_buf = self.alloc_buf(
            skybox_mem,
            3 * 2048 * 2048 * 6,
            vk::BufferUsageFlags::TRANSFER_SRC,
        );
        let skybox_ptr = self.buf_to_ptr(skybox_mem, skybox_buf);
        let img = std::fs::read("resources/skybox.rif").unwrap();
        let data = decode_rif(&img);
        unsafe {
            std::ptr::copy(data.as_ptr(), skybox_ptr as *mut u8, data.len());
        }

        let image = make_skybox_image(self);
        let imageview = make_skybox_image_view(self, image);
        let sampler = make_skybox_sampler(self);

        let pipe_layout = setup_pipe_layout(self);
        let pipeline = setup_pipeline(self, pass, pipe_layout);

        Skybox {
            buf: skybox_buf,
            image,
            imageview,
            sampler,
            pipeline,
            pipe_layout,
        }
    }
}

fn make_skybox_image(ctx: &RenderCtx) -> vk::Image {
    let qf_idxs = [ctx.queue_idx];
    let create_info = vk::ImageCreateInfo::default()
        .flags(vk::ImageCreateFlags::CUBE_COMPATIBLE)
        .image_type(vk::ImageType::TYPE_2D)
        .format(vk::Format::R8G8B8_UNORM)
        .extent(vk::Extent3D::default().width(2048).height(2048).depth(1))
        .mip_levels(1)
        .array_layers(6)
        .samples(vk::SampleCountFlags::TYPE_1)
        .tiling(vk::ImageTiling::OPTIMAL)
        .usage(vk::ImageUsageFlags::TRANSFER_DST | vk::ImageUsageFlags::SAMPLED)
        .sharing_mode(vk::SharingMode::EXCLUSIVE)
        .initial_layout(vk::ImageLayout::UNDEFINED)
        .queue_family_indices(&qf_idxs);
    let image =
        unsafe { ctx.dev.create_image(&create_info, None) }.expect("Failed to create image");

    let skybox_mem_size: vk::MemoryRequirements =
        unsafe { ctx.dev.get_image_memory_requirements(image) };
    let skybox_dev_mem = ctx.mem_alloc(MemType::HostInvisible, skybox_mem_size.size);
    unsafe { ctx.dev.bind_image_memory(image, skybox_dev_mem, 0) }
        .expect("Failed to bind image memory");

    image
}

fn make_skybox_image_view(ctx: &RenderCtx, image: vk::Image) -> vk::ImageView {
    let create_info = vk::ImageViewCreateInfo::default()
        .image(image)
        .view_type(vk::ImageViewType::CUBE)
        .format(vk::Format::R8G8B8_UNORM)
        .components(vk::ComponentMapping::default())
        .subresource_range(
            vk::ImageSubresourceRange::default()
                .aspect_mask(vk::ImageAspectFlags::COLOR)
                .base_mip_level(0)
                .level_count(1)
                .base_array_layer(0)
                .layer_count(6),
        );
    unsafe { ctx.dev.create_image_view(&create_info, None) }.expect("Failed to create image view")
}

fn make_skybox_sampler(ctx: &RenderCtx) -> vk::Sampler {
    let create_info = vk::SamplerCreateInfo::default()
        .mag_filter(vk::Filter::LINEAR)
        .min_filter(vk::Filter::LINEAR)
        .mipmap_mode(vk::SamplerMipmapMode::LINEAR)
        .address_mode_u(vk::SamplerAddressMode::CLAMP_TO_EDGE)
        .address_mode_v(vk::SamplerAddressMode::CLAMP_TO_EDGE)
        .address_mode_w(vk::SamplerAddressMode::CLAMP_TO_EDGE)
        .max_anisotropy(1.0)
        .max_lod(1.0)
        .border_color(vk::BorderColor::FLOAT_OPAQUE_WHITE);
    unsafe { ctx.dev.create_sampler(&create_info, None) }.expect("Failed to create sampler")
}

fn load_skybox_shaders(ctx: &RenderCtx) -> (vk::ShaderModule, vk::ShaderModule) {
    let skybox_vert_shader = ctx.shader_mod_from_file(
        "shaders/skybox_vert.spv",
        vk::ShaderModuleCreateFlags::empty(),
    );
    let skybox_frag_shader = ctx.shader_mod_from_file(
        "shaders/skybox_frag.spv",
        vk::ShaderModuleCreateFlags::empty(),
    );

    (skybox_vert_shader, skybox_frag_shader)
}

fn setup_desc_layout(ctx: &RenderCtx) -> vk::DescriptorSetLayout {
    let image_binding = vk::DescriptorSetLayoutBinding::default()
        .binding(0)
        .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
        .descriptor_count(1)
        .stage_flags(vk::ShaderStageFlags::FRAGMENT);
    let layouts = [image_binding];

    let layout_info = vk::DescriptorSetLayoutCreateInfo::default().bindings(&layouts);
    let layout = unsafe { ctx.dev.create_descriptor_set_layout(&layout_info, None) }
        .expect("Failed to create descriptor set layout");
    layout
}

fn setup_pipe_layout(ctx: &RenderCtx) -> vk::PipelineLayout {
    let layout = setup_desc_layout(ctx);
    let push_range: [vk::PushConstantRange; 1] = [
        vk::PushConstantRange::default()
            .stage_flags(vk::ShaderStageFlags::VERTEX)
            .offset(0)
            .size(48), // vec4 camera_orig, camera_rot; uvec2 screen_res
    ];

    let layouts = &[layout];

    let pipe_layout = vk::PipelineLayoutCreateInfo::default()
        .set_layouts(layouts.as_ref())
        .push_constant_ranges(&push_range);
    let pipe_layout = unsafe {
        ctx.dev
            .create_pipeline_layout(&pipe_layout, None)
            .expect("Failed to create pipeline layout")
    };
    pipe_layout
}

fn setup_pipeline(
    ctx: &RenderCtx,
    pass: vk::RenderPass,
    layout: vk::PipelineLayout,
) -> vk::Pipeline {
    let (vert_shader, frag_shader) = load_skybox_shaders(ctx);
    let shader_stages = ctx.setup_simple_shader_stage(vert_shader, frag_shader);
    let vert_input = vk::PipelineVertexInputStateCreateInfo::default();
    let input_assembly = vk::PipelineInputAssemblyStateCreateInfo::default()
        .topology(vk::PrimitiveTopology::TRIANGLE_LIST);
    let rasterization = vk::PipelineRasterizationStateCreateInfo::default()
        .polygon_mode(vk::PolygonMode::FILL)
        .cull_mode(vk::CullModeFlags::NONE)
        .front_face(vk::FrontFace::CLOCKWISE)
        .line_width(1.0);
    let multisample = vk::PipelineMultisampleStateCreateInfo::default()
        .rasterization_samples(vk::SampleCountFlags::TYPE_1);
    let depth_stencil = vk::PipelineDepthStencilStateCreateInfo::default()
        .depth_test_enable(true)
        .depth_write_enable(true)
        .depth_compare_op(vk::CompareOp::LESS)
        .depth_bounds_test_enable(true)
        .stencil_test_enable(false)
        .min_depth_bounds(0.0)
        .max_depth_bounds(1.0);

    let blend = vk::PipelineColorBlendAttachmentState::default()
        .src_color_blend_factor(vk::BlendFactor::SRC_COLOR)
        .dst_color_blend_factor(vk::BlendFactor::SRC_COLOR)
        .color_blend_op(vk::BlendOp::ADD)
        .src_alpha_blend_factor(vk::BlendFactor::SRC_COLOR)
        .dst_alpha_blend_factor(vk::BlendFactor::SRC_COLOR)
        .alpha_blend_op(vk::BlendOp::ADD)
        .color_write_mask(
            vk::ColorComponentFlags::R
                | vk::ColorComponentFlags::G
                | vk::ColorComponentFlags::B
                | vk::ColorComponentFlags::A,
        );

    let blend_attachments = [blend];

    let color_blend = vk::PipelineColorBlendStateCreateInfo::default()
        .attachments(blend_attachments.as_ref())
        .logic_op_enable(false)
        .logic_op(vk::LogicOp::COPY)
        .blend_constants([1.0, 1.0, 1.0, 1.0]);

    let dynamic = vk::PipelineDynamicStateCreateInfo::default().dynamic_states(&[
        vk::DynamicState::VIEWPORT_WITH_COUNT,
        vk::DynamicState::SCISSOR_WITH_COUNT,
    ]);

    let pipe_info = vk::GraphicsPipelineCreateInfo::default()
        .stages(&shader_stages)
        .vertex_input_state(&vert_input)
        .input_assembly_state(&input_assembly)
        .rasterization_state(&rasterization)
        .multisample_state(&multisample)
        .depth_stencil_state(&depth_stencil)
        .color_blend_state(&color_blend)
        .layout(layout)
        .dynamic_state(&dynamic)
        .render_pass(pass)
        .subpass(0);

    let pipe = unsafe {
        ctx.dev
            .create_graphics_pipelines(vk::PipelineCache::null(), &[pipe_info], None)
            .expect("Failed to create graphics pipeline")
            .remove(0)
    };
    println!("Created pipeline");
    pipe
}
