// Volumetric Fog — froxel-based atmospheric scattering.
//
// 3D froxel volume (160×90×64) stores inscattering + extinction per cell.
// Light injection computes per-light contribution with shadow sampling.
// Integration ray-marches front-to-back for final fog compositing.
// Temporal reprojection stabilizes results across frames.

/// Volumetric fog parameters.
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct VolumetricFogParams {
    pub inv_proj: [f32; 16],
    pub view: [f32; 16],
    pub screen_size: [f32; 2],
    pub near_plane: f32,
    pub far_plane: f32,
    pub fog_density: f32,
    pub fog_height_falloff: f32,
    pub fog_height_offset: f32,
    pub scattering_anisotropy: f32, // Henyey-Greenstein g parameter
    pub light_dir: [f32; 3],
    pub light_intensity: f32,
    pub light_color: [f32; 3],
    pub time: f32,
}

/// Froxel volume dimensions.
pub const FROXEL_WIDTH: u32 = 160;
pub const FROXEL_HEIGHT: u32 = 90;
pub const FROXEL_DEPTH: u32 = 64;

/// Volumetric fog pass manager.
pub struct VolumetricFogPass {
    inject_pipeline: Option<wgpu::ComputePipeline>,
    integrate_pipeline: Option<wgpu::ComputePipeline>,
    inject_bgl: Option<wgpu::BindGroupLayout>,
    integrate_bgl: Option<wgpu::BindGroupLayout>,
    params_buf: Option<wgpu::Buffer>,
    /// 3D froxel volume (Rgba16Float — inscattering RGB + extinction A).
    froxel_texture: Option<wgpu::Texture>,
    froxel_view: Option<wgpu::TextureView>,
    /// Integrated fog result (2D, Rgba16Float — accumulated fog color + transmittance).
    pub fog_texture: Option<wgpu::Texture>,
    pub fog_view: Option<wgpu::TextureView>,
    sampler: Option<wgpu::Sampler>,
    pub enabled: bool,
    initialized: bool,
}

impl Default for VolumetricFogPass {
    fn default() -> Self {
        Self {
            inject_pipeline: None,
            integrate_pipeline: None,
            inject_bgl: None,
            integrate_bgl: None,
            params_buf: None,
            froxel_texture: None,
            froxel_view: None,
            fog_texture: None,
            fog_view: None,
            sampler: None,
            enabled: false, // disabled by default — opt-in
            initialized: false,
        }
    }
}

impl VolumetricFogPass {
    /// Initialize the volumetric fog pipeline and froxel volume.
    pub fn init(&mut self, device: &wgpu::Device, width: u32, height: u32) {
        if self.initialized {
            return;
        }

        let format_3d = wgpu::TextureFormat::Rgba16Float;
        let format_2d = wgpu::TextureFormat::Rgba16Float;

        // 3D froxel volume
        let froxel = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("froxel_volume"),
            size: wgpu::Extent3d {
                width: FROXEL_WIDTH,
                height: FROXEL_HEIGHT,
                depth_or_array_layers: FROXEL_DEPTH,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D3,
            format: format_3d,
            usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::TEXTURE_BINDING,
            view_formats: &[],
        });
        self.froxel_view = Some(froxel.create_view(&wgpu::TextureViewDescriptor::default()));
        self.froxel_texture = Some(froxel);

        // 2D fog output (screen-space)
        let fog = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("fog_output"),
            size: wgpu::Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: format_2d,
            usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::TEXTURE_BINDING,
            view_formats: &[],
        });
        self.fog_view = Some(fog.create_view(&wgpu::TextureViewDescriptor::default()));
        self.fog_texture = Some(fog);

        if self.params_buf.is_none() {
            self.params_buf = Some(device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("volumetric_fog_params"),
                size: std::mem::size_of::<VolumetricFogParams>() as u64,
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            }));
        }

        if self.sampler.is_none() {
            self.sampler = Some(device.create_sampler(&wgpu::SamplerDescriptor {
                label: Some("fog_sampler"),
                mag_filter: wgpu::FilterMode::Linear,
                min_filter: wgpu::FilterMode::Linear,
                ..Default::default()
            }));
        }

        // Injection compute pipeline
        let inject_bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("fog_inject_bgl"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::StorageTexture {
                        access: wgpu::StorageTextureAccess::WriteOnly,
                        format: format_3d,
                        view_dimension: wgpu::TextureViewDimension::D3,
                    },
                    count: None,
                },
            ],
        });

        let inject_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("volumetric_fog.wgsl"),
            source: wgpu::ShaderSource::Wgsl(include_str!("../shaders/volumetric_fog.wgsl").into()),
        });

        let inject_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("fog_inject_layout"),
            bind_group_layouts: &[Some(&inject_bgl)],
            immediate_size: 0,
        });

        self.inject_pipeline = Some(device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("fog_inject"),
            layout: Some(&inject_layout),
            module: &inject_shader,
            entry_point: Some("inject"),
            compilation_options: Default::default(),
            cache: None,
        }));

        // Integration compute pipeline (reads froxel volume, writes 2D fog)
        let integrate_bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("fog_integrate_bgl"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Float { filterable: true },
                        view_dimension: wgpu::TextureViewDimension::D3,
                        multisampled: false,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 2,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 3,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::StorageTexture {
                        access: wgpu::StorageTextureAccess::WriteOnly,
                        format: format_2d,
                        view_dimension: wgpu::TextureViewDimension::D2,
                    },
                    count: None,
                },
            ],
        });

        let integrate_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("fog_integrate_layout"),
            bind_group_layouts: &[Some(&integrate_bgl)],
            immediate_size: 0,
        });

        self.integrate_pipeline = Some(device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("fog_integrate"),
            layout: Some(&integrate_layout),
            module: &inject_shader,
            entry_point: Some("integrate"),
            compilation_options: Default::default(),
            cache: None,
        }));

        self.inject_bgl = Some(inject_bgl);
        self.integrate_bgl = Some(integrate_bgl);
        self.initialized = true;

        log::info!(
            "Volumetric fog initialized ({}×{}×{} froxels, {}×{} output)",
            FROXEL_WIDTH,
            FROXEL_HEIGHT,
            FROXEL_DEPTH,
            width,
            height
        );
    }

    /// Dispatch fog injection + integration.
    pub fn dispatch(
        &self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
        camera: &crate::camera::Camera,
        light_dir: [f32; 3],
        light_color: [f32; 3],
        light_intensity: f32,
        time: f32,
    ) {
        if !self.enabled || !self.initialized {
            return;
        }
        let Some(ref inject_pipe) = self.inject_pipeline else {
            return;
        };
        let Some(ref integrate_pipe) = self.integrate_pipeline else {
            return;
        };
        let Some(ref inject_bgl) = self.inject_bgl else { return };
        let Some(ref integrate_bgl) = self.integrate_bgl else {
            return;
        };
        let Some(ref params_buf) = self.params_buf else { return };
        let Some(ref froxel_view) = self.froxel_view else {
            return;
        };
        let Some(ref fog_view) = self.fog_view else { return };
        let Some(ref sampler) = self.sampler else { return };

        let params = VolumetricFogParams {
            inv_proj: camera.projection_matrix.inverse().to_cols_array(),
            view: camera.view_matrix.to_cols_array(),
            screen_size: [FROXEL_WIDTH as f32, FROXEL_HEIGHT as f32],
            near_plane: camera.near_far.x,
            far_plane: camera.near_far.y.min(200.0),
            fog_density: 0.02,
            fog_height_falloff: 0.5,
            fog_height_offset: 0.0,
            scattering_anisotropy: 0.3,
            light_dir,
            light_intensity,
            light_color,
            time,
        };
        queue.write_buffer(params_buf, 0, bytemuck::bytes_of(&params));

        // Inject light into froxel volume
        let inject_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("fog_inject_bg"),
            layout: inject_bgl,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: params_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(froxel_view),
                },
            ],
        });
        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("fog_inject"),
                timestamp_writes: None,
            });
            pass.set_pipeline(inject_pipe);
            pass.set_bind_group(0, &inject_bg, &[]);
            pass.dispatch_workgroups(
                FROXEL_WIDTH.div_ceil(4),
                FROXEL_HEIGHT.div_ceil(4),
                FROXEL_DEPTH.div_ceil(4),
            );
        }

        // Integrate froxels into 2D fog texture
        let integrate_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("fog_integrate_bg"),
            layout: integrate_bgl,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: params_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(froxel_view),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::Sampler(sampler),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: wgpu::BindingResource::TextureView(fog_view),
                },
            ],
        });
        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("fog_integrate"),
                timestamp_writes: None,
            });
            pass.set_pipeline(integrate_pipe);
            pass.set_bind_group(0, &integrate_bg, &[]);
            // 2D dispatch over screen dimensions
            let fog_dims = [FROXEL_WIDTH, FROXEL_HEIGHT]; // fog output matches froxel XY
            pass.dispatch_workgroups(fog_dims[0].div_ceil(8), fog_dims[1].div_ceil(8), 1);
        }
    }

    pub fn destroy(&mut self) {
        if let Some(t) = self.froxel_texture.take() {
            t.destroy();
        }
        if let Some(t) = self.fog_texture.take() {
            t.destroy();
        }
        if let Some(b) = self.params_buf.take() {
            b.destroy();
        }
        self.froxel_view = None;
        self.fog_view = None;
        self.initialized = false;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn fog_params_size() {
        assert_eq!(std::mem::size_of::<VolumetricFogParams>(), 192);
    }

    #[test]
    fn froxel_volume_size_bytes() {
        // 160×90×64 × 8 bytes (Rgba16Float) = ~7 MB
        let size = FROXEL_WIDTH as u64 * FROXEL_HEIGHT as u64 * FROXEL_DEPTH as u64 * 8;
        assert!(size < 8 * 1024 * 1024, "Froxel volume should be < 8 MB");
    }
}
