use std::rc::Rc;

use exr::prelude::{ReadChannels, ReadLayers};
use glam::{Vec2, Vec3};
use wgpu::util::DeviceExt;

use crate::{
    loaders::mesh::{InstanceRaw, Mesh, Vertex},
    paths,
    render::{self, DrawCall, Pipeline},
    texture::HeightmapTexture,
};

pub struct SnowConfig
{
    pub depth_map_path: String,
    pub heightmap_path: String,
    pub terrain_size: Vec2,
    pub resolution: (u32, u32),
}

impl SnowConfig
{
    pub fn default() -> Self
    {
        Self {
            depth_map_path: paths::textures::snow_depth(),
            heightmap_path: paths::textures::terrain_heightmap(),
            terrain_size: Vec2::new(1000.0, 1000.0),
            resolution: (1000, 1000),
        }
    }
}

pub struct ClipmapConfig
{
    pub num_levels: u32,
    pub grid_size: u32,
    pub base_cell_size: f32,
}

impl Default for ClipmapConfig
{
    fn default() -> Self
    {
        Self {
            num_levels: 3,
            grid_size: 192,
            base_cell_size: 0.25,
        }
    }
}

struct ClipmapLevel
{
    mesh: Rc<Mesh>,
    instance_buffer: wgpu::Buffer,
}

pub struct SnowLayer
{
    pub depth_texture: wgpu::Texture,
    pub depth_texture_view: wgpu::TextureView,
    pub depth_bind_group: wgpu::BindGroup,
    pub width: u32,
    pub height: u32,
    pub deform_bind_group: wgpu::BindGroup,
    pub deform_pipeline: wgpu::ComputePipeline,
    pub deform_params_buffer: wgpu::Buffer,
    pub terrain_size: Vec2,

    clipmap_config: ClipmapConfig,
    levels: Vec<ClipmapLevel>,
    displacement_bind_group: wgpu::BindGroup,
    heightmap_texture: HeightmapTexture,
    depth_sampler: wgpu::Sampler,
}

fn create_instance_buffer() -> wgpu::Buffer
{
    render::with_device(|device| {
        device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Snow Clipmap Instance Buffer"),
            size: std::mem::size_of::<InstanceRaw>() as u64,
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        })
    })
}

impl SnowLayer
{
    pub fn load(config: &SnowConfig) -> anyhow::Result<Self>
    {
        println!("\n=== Loading Snow Layer ===");
        println!("Depth map path: {}", config.depth_map_path);
        println!("Heightmap path: {}", config.heightmap_path);
        println!("Terrain size: {:?}", config.terrain_size);

        let (depth_data, width, height) = Self::load_depth_map(&config.depth_map_path)?;

        println!("Using EXR dimensions: {}×{}", width, height);

        let (depth_texture, depth_texture_view, depth_bind_group) =
            Self::create_depth_texture(&depth_data, width, height);

        let (deform_pipeline, deform_bind_group, deform_params_buffer) =
            Self::create_deform_pipeline(&depth_texture_view);

        let heightmap_texture = render::with_device(|device| {
            render::with_queue(|queue| {
                HeightmapTexture::load(device, queue, &config.heightmap_path)
            })
        })?;

        let depth_sampler = render::with_device(|device| {
            device.create_sampler(&wgpu::SamplerDescriptor {
                label: Some("Snow Depth Sampler"),
                address_mode_u: wgpu::AddressMode::ClampToEdge,
                address_mode_v: wgpu::AddressMode::ClampToEdge,
                mag_filter: wgpu::FilterMode::Nearest,
                min_filter: wgpu::FilterMode::Nearest,
                ..Default::default()
            })
        });

        let displacement_bind_group = Self::create_displacement_bind_group(
            &heightmap_texture,
            &depth_texture_view,
            &depth_sampler,
        );

        let clipmap_config = ClipmapConfig::default();
        let levels = Vec::new();

        Ok(Self {
            depth_texture,
            depth_texture_view,
            depth_bind_group,
            width,
            height,
            deform_bind_group,
            deform_pipeline,
            deform_params_buffer,
            terrain_size: config.terrain_size,
            clipmap_config,
            levels,
            displacement_bind_group,
            heightmap_texture,
            depth_sampler,
        })
    }

    fn load_depth_map(path: &str) -> anyhow::Result<(Vec<f32>, u32, u32)>
    {
        println!("Loading snow depth map from: {}", path);

        let image = exr::prelude::read()
            .no_deep_data()
            .largest_resolution_level()
            .all_channels()
            .all_layers()
            .all_attributes()
            .from_file(path)?;

        let layer = &image.layer_data[0];
        let width = layer.size.width() as u32;
        let height = layer.size.height() as u32;

        println!("  Layer size: {}×{}", width, height);
        println!(
            "  Available channels: {:?}",
            layer
                .channel_data
                .list
                .iter()
                .map(|c| &c.name)
                .collect::<Vec<_>>()
        );

        let channel = layer
            .channel_data
            .list
            .iter()
            .find(|c| format!("{:?}", c.name).contains("\"R\""))
            .or_else(|| layer.channel_data.list.first())
            .ok_or_else(|| anyhow::anyhow!("No channels found in EXR"))?;

        println!("  Using channel: {:?}", channel.name);

        let depths: Vec<f32> = channel.sample_data.values_as_f32().collect();

        let min_value = depths.iter().cloned().fold(f32::INFINITY, f32::min);
        let max_value = depths.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let avg_value = depths.iter().sum::<f32>() / depths.len() as f32;
        let non_zero_count = depths.iter().filter(|&&v| v > 0.0001).count();

        println!("  Total values: {}", depths.len());
        println!(
            "  Min: {:.6}, Max: {:.6}, Avg: {:.6}",
            min_value, max_value, avg_value
        );
        println!(
            "  Non-zero values: {} ({:.1}%)",
            non_zero_count,
            (non_zero_count as f32 / depths.len() as f32) * 100.0
        );

        if max_value < 0.0001
        {
            println!(
                "  ⚠️  WARNING: All values are effectively zero! Snow depth map may be invalid."
            );
        }
        else
        {
            println!("  ✓ Snow depth data loaded successfully");
        }

        Ok((depths, width, height))
    }

    fn create_depth_texture(
        depth_data: &[f32],
        width: u32,
        height: u32,
    ) -> (wgpu::Texture, wgpu::TextureView, wgpu::BindGroup)
    {
        render::with_device(|device| {
            let size = wgpu::Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            };

            let texture = device.create_texture(&wgpu::TextureDescriptor {
                label: Some("Snow Depth Texture"),
                size,
                mip_level_count: 1,
                sample_count: 1,
                dimension: wgpu::TextureDimension::D2,
                format: wgpu::TextureFormat::R32Float,
                usage: wgpu::TextureUsages::TEXTURE_BINDING
                    | wgpu::TextureUsages::COPY_DST
                    | wgpu::TextureUsages::STORAGE_BINDING,
                view_formats: &[],
            });

            let data_bytes: &[u8] = bytemuck::cast_slice(depth_data);

            render::with_queue(|queue| {
                queue.write_texture(
                    wgpu::TexelCopyTextureInfo {
                        texture: &texture,
                        mip_level: 0,
                        origin: wgpu::Origin3d::ZERO,
                        aspect: wgpu::TextureAspect::All,
                    },
                    data_bytes,
                    wgpu::TexelCopyBufferLayout {
                        offset: 0,
                        bytes_per_row: Some(width * 4),
                        rows_per_image: Some(height),
                    },
                    size,
                );
            });

            let texture_view = texture.create_view(&wgpu::TextureViewDescriptor::default());

            let bind_group_layout =
                device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                    label: Some("Snow Depth Bind Group Layout"),
                    entries: &[wgpu::BindGroupLayoutEntry {
                        binding: 0,
                        visibility: wgpu::ShaderStages::FRAGMENT | wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::Texture {
                            sample_type: wgpu::TextureSampleType::Float { filterable: false },
                            view_dimension: wgpu::TextureViewDimension::D2,
                            multisampled: false,
                        },
                        count: None,
                    }],
                });

            let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("Snow Depth Bind Group"),
                layout: &bind_group_layout,
                entries: &[wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(&texture_view),
                }],
            });

            (texture, texture_view, bind_group)
        })
    }

    fn create_deform_pipeline(
        depth_texture_view: &wgpu::TextureView,
    ) -> (wgpu::ComputePipeline, wgpu::BindGroup, wgpu::Buffer)
    {
        render::with_device(|device| {
            let shader_source = std::fs::read_to_string(&paths::shaders::snow_deform())
                .expect("Failed to load snow deform shader");

            let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
                label: Some("Snow Deform Shader"),
                source: wgpu::ShaderSource::Wgsl(shader_source.into()),
            });

            let params_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("Snow Deform Params"),
                size: 32,
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });

            let bind_group_layout =
                device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                    label: Some("Snow Deform Bind Group Layout"),
                    entries: &[
                        wgpu::BindGroupLayoutEntry {
                            binding: 0,
                            visibility: wgpu::ShaderStages::COMPUTE,
                            ty: wgpu::BindingType::StorageTexture {
                                access: wgpu::StorageTextureAccess::ReadWrite,
                                format: wgpu::TextureFormat::R32Float,
                                view_dimension: wgpu::TextureViewDimension::D2,
                            },
                            count: None,
                        },
                        wgpu::BindGroupLayoutEntry {
                            binding: 1,
                            visibility: wgpu::ShaderStages::COMPUTE,
                            ty: wgpu::BindingType::Buffer {
                                ty: wgpu::BufferBindingType::Uniform,
                                has_dynamic_offset: false,
                                min_binding_size: None,
                            },
                            count: None,
                        },
                    ],
                });

            let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("Snow Deform Bind Group"),
                layout: &bind_group_layout,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: wgpu::BindingResource::TextureView(depth_texture_view),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: params_buffer.as_entire_binding(),
                    },
                ],
            });

            let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("Snow Deform Pipeline Layout"),
                bind_group_layouts: &[&bind_group_layout],
                immediate_size: 0,
            });

            let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("Snow Deform Pipeline"),
                layout: Some(&pipeline_layout),
                module: &shader,
                entry_point: Some("deform"),
                compilation_options: Default::default(),
                cache: None,
            });

            (pipeline, bind_group, params_buffer)
        })
    }

    fn create_displacement_bind_group(
        heightmap: &HeightmapTexture,
        depth_view: &wgpu::TextureView,
        depth_sampler: &wgpu::Sampler,
    ) -> wgpu::BindGroup
    {
        render::with_device(|device| {
            let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("Snow Displacement Bind Group Layout"),
                entries: &[
                    wgpu::BindGroupLayoutEntry {
                        binding: 0,
                        visibility: wgpu::ShaderStages::VERTEX,
                        ty: wgpu::BindingType::Texture {
                            sample_type: wgpu::TextureSampleType::Float { filterable: false },
                            view_dimension: wgpu::TextureViewDimension::D2,
                            multisampled: false,
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 1,
                        visibility: wgpu::ShaderStages::VERTEX,
                        ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::NonFiltering),
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 2,
                        visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
                        ty: wgpu::BindingType::Texture {
                            sample_type: wgpu::TextureSampleType::Float { filterable: false },
                            view_dimension: wgpu::TextureViewDimension::D2,
                            multisampled: false,
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 3,
                        visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
                        ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::NonFiltering),
                        count: None,
                    },
                ],
            });

            device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("Snow Displacement Bind Group"),
                layout: &layout,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: wgpu::BindingResource::TextureView(&heightmap.view),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: wgpu::BindingResource::Sampler(&heightmap.sampler),
                    },
                    wgpu::BindGroupEntry {
                        binding: 2,
                        resource: wgpu::BindingResource::TextureView(depth_view),
                    },
                    wgpu::BindGroupEntry {
                        binding: 3,
                        resource: wgpu::BindingResource::Sampler(depth_sampler),
                    },
                ],
            })
        })
    }

    fn generate_clipmap_grid(&self, level: u32) -> Mesh
    {
        let cell_size = self.clipmap_config.base_cell_size * (1 << level) as f32;
        let grid_size = self.clipmap_config.grid_size;
        let half_extent = (grid_size as f32 * cell_size) / 2.0;

        let mut vertices = Vec::new();
        let mut indices = Vec::new();

        for z in 0..=grid_size
        {
            for x in 0..=grid_size
            {
                vertices.push(Vertex {
                    position: [
                        -half_extent + x as f32 * cell_size,
                        0.0,
                        -half_extent + z as f32 * cell_size,
                    ],
                    normal: [0.0, 1.0, 0.0],
                    uv: [0.0, 0.0],
                });
            }
        }

        let row = grid_size + 1;
        for z in 0..grid_size
        {
            for x in 0..grid_size
            {
                let i0 = z * row + x;
                let i1 = i0 + 1;
                let i2 = i0 + row;
                let i3 = i2 + 1;
                indices.push(i0);
                indices.push(i2);
                indices.push(i1);
                indices.push(i1);
                indices.push(i2);
                indices.push(i3);
            }
        }

        let vertex_buffer = render::with_device(|device| {
            device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some(&format!("Snow Clipmap Level {} Vertex Buffer", level)),
                contents: bytemuck::cast_slice(&vertices),
                usage: wgpu::BufferUsages::VERTEX,
            })
        });

        let index_buffer = render::with_device(|device| {
            device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some(&format!("Snow Clipmap Level {} Index Buffer", level)),
                contents: bytemuck::cast_slice(&indices),
                usage: wgpu::BufferUsages::INDEX,
            })
        });

        let mut aabb_min = Vec3::splat(f32::MAX);
        let mut aabb_max = Vec3::splat(f32::MIN);
        for v in &vertices
        {
            let p = Vec3::from(v.position);
            aabb_min = aabb_min.min(p);
            aabb_max = aabb_max.max(p);
        }

        Mesh {
            vertex_buffer,
            index_buffer,
            num_indices: indices.len() as u32,
            aabb_min,
            aabb_max,
            cpu_positions: Vec::new(),
            cpu_indices: Vec::new(),
        }
    }

    pub fn update(&mut self, center: Vec3)
    {
        if self.levels.is_empty()
        {
            for level in 0..self.clipmap_config.num_levels
            {
                let mesh = self.generate_clipmap_grid(level);
                let instance_buffer = create_instance_buffer();

                self.levels.push(ClipmapLevel {
                    mesh: Rc::new(mesh),
                    instance_buffer,
                });
            }
        }

        render::with_queue(|queue| {
            for (level, clipmap_level) in self.levels.iter().enumerate()
            {
                let cell_size = self.clipmap_config.base_cell_size * (1u32 << level) as f32;
                let snapped_x = (center.x / cell_size).round() * cell_size;
                let snapped_z = (center.z / cell_size).round() * cell_size;
                let y_bias = -(level as f32) * 0.005;
                let model =
                    glam::Mat4::from_translation(glam::Vec3::new(snapped_x, y_bias, snapped_z));
                let instance_data = InstanceRaw {
                    model: model.to_cols_array_2d(),
                    dissolve_amount: 0.0,
                    _padding: [0.0; 3],
                };
                queue.write_buffer(
                    &clipmap_level.instance_buffer,
                    0,
                    bytemuck::cast_slice(&[instance_data]),
                );
            }
        });
    }

    pub fn get_draw_calls(&self) -> Vec<DrawCall>
    {
        self.levels
            .iter()
            .rev()
            .map(|level| DrawCall {
                vertex_buffer: level.mesh.vertex_buffer.clone(),
                index_buffer: level.mesh.index_buffer.clone(),
                num_indices: level.mesh.num_indices,
                model: glam::Mat4::IDENTITY,
                pipeline: Pipeline::SnowClipmap,
                instance_buffer: Some(level.instance_buffer.clone()),
                num_instances: 1,
                tile_scale: 2.0,
                enable_dissolve: false,
                enable_snow_light: true,
                displacement_bind_group: Some(self.displacement_bind_group.clone()),
                entity: None,
            })
            .collect()
    }

    pub fn deform_at_position(&self, position: Vec3, radius: f32, depth: f32)
    {
        render::with_queue(|queue| {
            let params_data = [
                position.x,
                position.z,
                radius,
                depth,
                self.terrain_size.x,
                self.terrain_size.y,
                0.0,
                0.0,
            ];
            let params_bytes: &[u8] = bytemuck::cast_slice(&params_data);

            queue.write_buffer(&self.deform_params_buffer, 0, params_bytes);
        });

        render::with_device(|device| {
            let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("Snow Deform Encoder"),
            });

            {
                let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                    label: Some("Snow Deform Pass"),
                    timestamp_writes: None,
                });

                compute_pass.set_pipeline(&self.deform_pipeline);
                compute_pass.set_bind_group(0, &self.deform_bind_group, &[]);

                let workgroup_size = 16;
                let dispatch_x = (self.width + workgroup_size - 1) / workgroup_size;
                let dispatch_y = (self.height + workgroup_size - 1) / workgroup_size;
                compute_pass.dispatch_workgroups(dispatch_x, dispatch_y, 1);
            }

            render::with_queue(|queue| {
                queue.submit(Some(encoder.finish()));
            });
        });
    }
}