use crate::ecs::mesh::components::Vertex;
use crate::render::wgpu::rendergraph::{PassExecutionContext, PassNode};
use std::collections::HashMap;

#[repr(C)]
#[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct ShadowUniforms {
    light_view_projection: [[f32; 4]; 4],
    _padding: [[f32; 4]; 4],
    _padding2: [[f32; 4]; 4],
    _padding3: [[f32; 4]; 4],
    _padding4: [f32; 4],
}

pub struct ShadowDepthPass {
    pipeline: wgpu::RenderPipeline,

    uniform_buffer: wgpu::Buffer,
    uniform_bind_group: wgpu::BindGroup,

    transform_buffer: wgpu::Buffer,
    transform_bind_group: wgpu::BindGroup,
    transform_bind_group_layout: wgpu::BindGroupLayout,

    transform_buffer_size: usize,

    vertex_buffer: wgpu::Buffer,
    index_buffer: wgpu::Buffer,
    vertex_buffer_size: u64,
    index_buffer_size: u64,

    meshes: HashMap<String, (u32, u32, u32, u32)>,

    current_vertex_offset: u32,
    current_index_offset: u32,

    shadow_caster_entities: Vec<crate::ecs::world::Entity>,
}

impl ShadowDepthPass {
    pub fn new(device: &wgpu::Device, queue: &wgpu::Queue) -> Self {
        let shader =
            device.create_shader_module(wgpu::include_wgsl!("../shaders/shadow_depth.wgsl"));

        let uniform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Shadow Uniform Buffer"),
            size: std::mem::size_of::<ShadowUniforms>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let uniform_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("Shadow Uniform Bind Group Layout"),
                entries: &[wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::VERTEX,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                }],
            });

        let uniform_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Shadow Uniform Bind Group"),
            layout: &uniform_bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: uniform_buffer.as_entire_binding(),
            }],
        });

        let initial_transform_count = 1000;
        let transform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Shadow Transform Buffer"),
            size: (std::mem::size_of::<[[f32; 4]; 4]>() * initial_transform_count) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let transform_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("Shadow Transform Bind Group Layout"),
                entries: &[wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::VERTEX,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                }],
            });

        let transform_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Shadow Transform Bind Group"),
            layout: &transform_bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: transform_buffer.as_entire_binding(),
            }],
        });

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("Shadow Pipeline Layout"),
            bind_group_layouts: &[&uniform_bind_group_layout, &transform_bind_group_layout],
            push_constant_ranges: &[],
        });

        let vertex_attributes = wgpu::vertex_attr_array![
            0 => Float32x3,
            1 => Float32x3,
            2 => Float32x2,
        ];

        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("Shadow Depth Pipeline"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: Some("vertex_main"),
                buffers: &[wgpu::VertexBufferLayout {
                    array_stride: std::mem::size_of::<Vertex>() as wgpu::BufferAddress,
                    step_mode: wgpu::VertexStepMode::Vertex,
                    attributes: &vertex_attributes,
                }],
                compilation_options: Default::default(),
            },
            fragment: None,
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: Some(wgpu::Face::Back),
                polygon_mode: wgpu::PolygonMode::Fill,
                unclipped_depth: false,
                conservative: false,
            },
            depth_stencil: Some(wgpu::DepthStencilState {
                format: wgpu::TextureFormat::Depth32Float,
                depth_write_enabled: true,
                depth_compare: wgpu::CompareFunction::LessEqual,
                stencil: wgpu::StencilState::default(),
                bias: wgpu::DepthBiasState {
                    constant: 2,
                    slope_scale: 2.0,
                    clamp: 0.0,
                },
            }),
            multisample: wgpu::MultisampleState::default(),
            multiview: None,
            cache: None,
        });

        let initial_vertex_count = 10000;
        let initial_index_count = 30000;

        let vertex_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Shadow Vertex Buffer"),
            size: (std::mem::size_of::<Vertex>() * initial_vertex_count) as u64,
            usage: wgpu::BufferUsages::VERTEX
                | wgpu::BufferUsages::COPY_DST
                | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });

        let index_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Shadow Index Buffer"),
            size: (std::mem::size_of::<u32>() * initial_index_count) as u64,
            usage: wgpu::BufferUsages::INDEX
                | wgpu::BufferUsages::COPY_DST
                | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });

        let mut pass = Self {
            pipeline,
            uniform_buffer,
            uniform_bind_group,
            transform_buffer,
            transform_bind_group,
            transform_bind_group_layout,
            transform_buffer_size: initial_transform_count,
            vertex_buffer,
            index_buffer,
            vertex_buffer_size: (std::mem::size_of::<Vertex>() * initial_vertex_count) as u64,
            index_buffer_size: (std::mem::size_of::<u32>() * initial_index_count) as u64,
            meshes: HashMap::new(),
            current_vertex_offset: 0,
            current_index_offset: 0,
            shadow_caster_entities: Vec::new(),
        };

        use crate::ecs::mesh::components::{
            create_cone_mesh, create_cube_mesh, create_cylinder_mesh, create_plane_mesh,
            create_sphere_mesh, create_torus_mesh,
        };
        let cube_mesh = create_cube_mesh();
        pass.add_mesh(
            device,
            queue,
            "Cube".to_string(),
            &cube_mesh.vertices,
            &cube_mesh.indices,
        );
        let sphere_mesh = create_sphere_mesh(1.0, 16);
        pass.add_mesh(
            device,
            queue,
            "Sphere".to_string(),
            &sphere_mesh.vertices,
            &sphere_mesh.indices,
        );
        let torus_mesh = create_torus_mesh(1.0, 0.3, 16, 16);
        pass.add_mesh(
            device,
            queue,
            "Torus".to_string(),
            &torus_mesh.vertices,
            &torus_mesh.indices,
        );
        let plane_mesh = create_plane_mesh(2.0);
        pass.add_mesh(
            device,
            queue,
            "Plane".to_string(),
            &plane_mesh.vertices,
            &plane_mesh.indices,
        );
        let cylinder_mesh = create_cylinder_mesh(0.5, 1.0, 16);
        pass.add_mesh(
            device,
            queue,
            "Cylinder".to_string(),
            &cylinder_mesh.vertices,
            &cylinder_mesh.indices,
        );
        let cone_mesh = create_cone_mesh(0.5, 1.0, 16);
        pass.add_mesh(
            device,
            queue,
            "Cone".to_string(),
            &cone_mesh.vertices,
            &cone_mesh.indices,
        );

        pass
    }

    fn add_mesh(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        name: String,
        vertices: &[Vertex],
        indices: &[u32],
    ) {
        let vertex_size = std::mem::size_of_val(vertices) as u64;
        let index_size = std::mem::size_of_val(indices) as u64;

        let required_vertex_size = (self.current_vertex_offset as u64
            * std::mem::size_of::<Vertex>() as u64)
            + vertex_size;
        let required_index_size =
            (self.current_index_offset as u64 * std::mem::size_of::<u32>() as u64) + index_size;

        if required_vertex_size > self.vertex_buffer_size {
            let new_size = (required_vertex_size * 2).max(self.vertex_buffer_size * 2);
            let new_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("Shadow Vertex Buffer (Resized)"),
                size: new_size,
                usage: wgpu::BufferUsages::VERTEX
                    | wgpu::BufferUsages::COPY_DST
                    | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            });

            let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("Shadow Vertex Buffer Copy"),
            });
            encoder.copy_buffer_to_buffer(
                &self.vertex_buffer,
                0,
                &new_buffer,
                0,
                self.vertex_buffer_size,
            );
            queue.submit(Some(encoder.finish()));

            self.vertex_buffer = new_buffer;
            self.vertex_buffer_size = new_size;
        }

        if required_index_size > self.index_buffer_size {
            let new_size = (required_index_size * 2).max(self.index_buffer_size * 2);
            let new_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("Shadow Index Buffer (Resized)"),
                size: new_size,
                usage: wgpu::BufferUsages::INDEX
                    | wgpu::BufferUsages::COPY_DST
                    | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            });

            let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("Shadow Index Buffer Copy"),
            });
            encoder.copy_buffer_to_buffer(
                &self.index_buffer,
                0,
                &new_buffer,
                0,
                self.index_buffer_size,
            );
            queue.submit(Some(encoder.finish()));

            self.index_buffer = new_buffer;
            self.index_buffer_size = new_size;
        }

        queue.write_buffer(
            &self.vertex_buffer,
            self.current_vertex_offset as u64 * std::mem::size_of::<Vertex>() as u64,
            bytemuck::cast_slice(vertices),
        );

        queue.write_buffer(
            &self.index_buffer,
            self.current_index_offset as u64 * std::mem::size_of::<u32>() as u64,
            bytemuck::cast_slice(indices),
        );

        self.meshes.insert(
            name,
            (
                self.current_vertex_offset,
                vertices.len() as u32,
                self.current_index_offset,
                indices.len() as u32,
            ),
        );

        self.current_vertex_offset += vertices.len() as u32;
        self.current_index_offset += indices.len() as u32;
    }

    fn sync_meshes(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        mesh_cache: &crate::ecs::prefab::resources::MeshCache,
    ) {
        for (name, mesh) in mesh_cache.iter() {
            if !self.meshes.contains_key(name.as_str()) {
                let vertex_size = (mesh.vertices.len() * std::mem::size_of::<Vertex>()) as u64;
                let index_size = (mesh.indices.len() * std::mem::size_of::<u32>()) as u64;

                let required_vertex_size = (self.current_vertex_offset as u64
                    * std::mem::size_of::<Vertex>() as u64)
                    + vertex_size;
                let required_index_size = (self.current_index_offset as u64
                    * std::mem::size_of::<u32>() as u64)
                    + index_size;

                if required_vertex_size > self.vertex_buffer_size {
                    let new_size = (required_vertex_size * 2).max(self.vertex_buffer_size * 2);
                    let new_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                        label: Some("Shadow Vertex Buffer (Resized)"),
                        size: new_size,
                        usage: wgpu::BufferUsages::VERTEX
                            | wgpu::BufferUsages::COPY_DST
                            | wgpu::BufferUsages::COPY_SRC,
                        mapped_at_creation: false,
                    });

                    let mut encoder =
                        device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
                            label: Some("Shadow Vertex Buffer Copy"),
                        });
                    encoder.copy_buffer_to_buffer(
                        &self.vertex_buffer,
                        0,
                        &new_buffer,
                        0,
                        self.vertex_buffer_size,
                    );
                    queue.submit(Some(encoder.finish()));

                    self.vertex_buffer = new_buffer;
                    self.vertex_buffer_size = new_size;
                }

                if required_index_size > self.index_buffer_size {
                    let new_size = (required_index_size * 2).max(self.index_buffer_size * 2);
                    let new_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                        label: Some("Shadow Index Buffer (Resized)"),
                        size: new_size,
                        usage: wgpu::BufferUsages::INDEX
                            | wgpu::BufferUsages::COPY_DST
                            | wgpu::BufferUsages::COPY_SRC,
                        mapped_at_creation: false,
                    });

                    let mut encoder =
                        device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
                            label: Some("Shadow Index Buffer Copy"),
                        });
                    encoder.copy_buffer_to_buffer(
                        &self.index_buffer,
                        0,
                        &new_buffer,
                        0,
                        self.index_buffer_size,
                    );
                    queue.submit(Some(encoder.finish()));

                    self.index_buffer = new_buffer;
                    self.index_buffer_size = new_size;
                }

                queue.write_buffer(
                    &self.vertex_buffer,
                    self.current_vertex_offset as u64 * std::mem::size_of::<Vertex>() as u64,
                    bytemuck::cast_slice(&mesh.vertices),
                );

                queue.write_buffer(
                    &self.index_buffer,
                    self.current_index_offset as u64 * std::mem::size_of::<u32>() as u64,
                    bytemuck::cast_slice(&mesh.indices),
                );

                self.meshes.insert(
                    name.clone(),
                    (
                        self.current_vertex_offset,
                        mesh.vertices.len() as u32,
                        self.current_index_offset,
                        mesh.indices.len() as u32,
                    ),
                );

                self.current_vertex_offset += mesh.vertices.len() as u32;
                self.current_index_offset += mesh.indices.len() as u32;
            }
        }
    }
}

impl PassNode<crate::ecs::world::World> for ShadowDepthPass {
    fn name(&self) -> &str {
        "shadow_depth_pass"
    }

    fn reads(&self) -> Vec<&str> {
        vec![]
    }

    fn writes(&self) -> Vec<&str> {
        vec!["shadow_depth"]
    }

    fn reads_writes(&self) -> Vec<&str> {
        vec![]
    }

    fn prepare(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        world: &crate::ecs::world::World,
    ) {
        self.sync_meshes(device, queue, &world.resources.mesh_cache);

        let directional_light = world
            .query_entities(crate::ecs::world::LIGHT | crate::ecs::world::GLOBAL_TRANSFORM)
            .find_map(|entity| {
                if let (Some(light), Some(transform)) =
                    (world.get_light(entity), world.get_global_transform(entity))
                {
                    if matches!(
                        light.light_type,
                        crate::ecs::light::components::LightType::Directional
                    ) && light.cast_shadows
                    {
                        Some((light.clone(), *transform))
                    } else {
                        None
                    }
                } else {
                    None
                }
            });

        if let Some((_light, light_transform)) = directional_light {
            let light_direction = light_transform.forward_vector();
            let scene_center = nalgebra_glm::vec3(0.0, 0.0, 0.0);
            let light_distance = 100.0;
            let light_position = scene_center - light_direction * light_distance;

            let up = if light_direction.y.abs() > 0.99 {
                nalgebra_glm::vec3(1.0, 0.0, 0.0)
            } else {
                nalgebra_glm::vec3(0.0, 1.0, 0.0)
            };

            let light_view = nalgebra_glm::look_at(&light_position, &scene_center, &up);

            let light_projection = nalgebra_glm::ortho_zo(-50.0, 50.0, -50.0, 50.0, 0.1, 200.0);

            let light_view_projection = light_projection * light_view;

            let uniforms = ShadowUniforms {
                light_view_projection: light_view_projection.into(),
                _padding: [[0.0; 4]; 4],
                _padding2: [[0.0; 4]; 4],
                _padding3: [[0.0; 4]; 4],
                _padding4: [0.0; 4],
            };

            queue.write_buffer(&self.uniform_buffer, 0, bytemuck::cast_slice(&[uniforms]));
        }

        self.shadow_caster_entities.clear();
        let mut transforms = Vec::new();

        for entity in world.query_entities(
            crate::ecs::world::RENDER_MESH
                | crate::ecs::world::GLOBAL_TRANSFORM
                | crate::ecs::world::CASTS_SHADOW,
        ) {
            if let (Some(render_mesh), Some(transform)) = (
                world.get_render_mesh(entity),
                world.get_global_transform(entity),
            ) && self.meshes.contains_key(&render_mesh.name)
            {
                let model_matrix: [[f32; 4]; 4] = transform.0.into();
                transforms.push(model_matrix);
                self.shadow_caster_entities.push(entity);
            }
        }

        if !transforms.is_empty() {
            if transforms.len() > self.transform_buffer_size {
                let new_size = (transforms.len() as f32 * 2.0).ceil() as usize;
                let new_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("Shadow Transform Buffer (Resized)"),
                    size: (std::mem::size_of::<[[f32; 4]; 4]>() * new_size) as u64,
                    usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                    mapped_at_creation: false,
                });

                self.transform_buffer = new_buffer;
                self.transform_buffer_size = new_size;

                self.transform_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("Shadow Transform Bind Group"),
                    layout: &self.transform_bind_group_layout,
                    entries: &[wgpu::BindGroupEntry {
                        binding: 0,
                        resource: self.transform_buffer.as_entire_binding(),
                    }],
                });
            }

            queue.write_buffer(&self.transform_buffer, 0, bytemuck::cast_slice(&transforms));
        }
    }

    fn execute<'r, 'e>(
        &mut self,
        context: PassExecutionContext<'r, 'e, crate::ecs::world::World>,
    ) -> crate::render::wgpu::rendergraph::Result<
        Vec<crate::render::wgpu::rendergraph::SubGraphRunCommand<'r>>,
    > {
        if self.shadow_caster_entities.is_empty() {
            return Ok(vec![]);
        }

        let has_shadow_casting_light = context
            .configs
            .query_entities(crate::ecs::world::LIGHT)
            .any(|entity| {
                if let Some(light) = context.configs.get_light(entity) {
                    matches!(
                        light.light_type,
                        crate::ecs::light::components::LightType::Directional
                    ) && light.cast_shadows
                } else {
                    false
                }
            });

        if !has_shadow_casting_light {
            return Ok(vec![]);
        }

        let (shadow_depth_view, shadow_depth_load, _shadow_depth_store) =
            context.get_depth_attachment("shadow_depth")?;

        let mut shadow_pass = context
            .encoder
            .begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("Shadow Depth Pass"),
                color_attachments: &[],
                depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
                    view: shadow_depth_view,
                    depth_ops: Some(wgpu::Operations {
                        load: shadow_depth_load,
                        store: wgpu::StoreOp::Store,
                    }),
                    stencil_ops: None,
                }),
                timestamp_writes: None,
                occlusion_query_set: None,
            });

        shadow_pass.set_pipeline(&self.pipeline);
        shadow_pass.set_bind_group(0, &self.uniform_bind_group, &[]);
        shadow_pass.set_bind_group(1, &self.transform_bind_group, &[]);
        shadow_pass.set_vertex_buffer(0, self.vertex_buffer.slice(..));
        shadow_pass.set_index_buffer(self.index_buffer.slice(..), wgpu::IndexFormat::Uint32);

        for (instance_index, &entity) in self.shadow_caster_entities.iter().enumerate() {
            if let Some(render_mesh) = context.configs.get_render_mesh(entity)
                && let Some(&(vertex_offset, _vertex_count, index_offset, index_count)) =
                    self.meshes.get(&render_mesh.name)
            {
                shadow_pass.draw_indexed(
                    index_offset..(index_offset + index_count),
                    vertex_offset as i32,
                    instance_index as u32..(instance_index as u32 + 1),
                );
            }
        }

        drop(shadow_pass);

        Ok(vec![])
    }
}
