// SSAO — Screen-Space Ambient Occlusion compute pass manager.
// Dispatches half-resolution AO computation + bilateral blur.
// Output: R8Unorm texture sampled in PBR ambient term.

/// SSAO uniform parameters (matches ssao.wgsl SsaoParams).
#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
pub struct SsaoParams {
    pub inv_proj: [f32; 16],
    pub proj: [f32; 16],
    pub screen_size: [f32; 2],
    pub radius: f32,
    pub bias: f32,
}

/// SSAO pass: compute AO from depth buffer, blur, output R8Unorm texture.
pub struct SsaoPass {
    ssao_pipeline: Option<wgpu::ComputePipeline>,
    blur_pipeline: Option<wgpu::ComputePipeline>,
    ssao_bgl: Option<wgpu::BindGroupLayout>,
    blur_bgl: Option<wgpu::BindGroupLayout>,
    params_buf: Option<wgpu::Buffer>,
    ao_texture: Option<wgpu::Texture>,
    ao_view: Option<wgpu::TextureView>,
    blur_texture: Option<wgpu::Texture>,
    blur_view: Option<wgpu::TextureView>,
    sampler: Option<wgpu::Sampler>,
    width: u32,
    height: u32,
    pub enabled: bool,
    initialized: bool,
    /// Whether bind groups need to be recreated (set on resize/init).
    pub bg_dirty: bool,
}

impl Default for SsaoPass {
    fn default() -> Self {
        Self {
            ssao_pipeline: None,
            blur_pipeline: None,
            ssao_bgl: None,
            blur_bgl: None,
            params_buf: None,
            ao_texture: None,
            ao_view: None,
            blur_texture: None,
            blur_view: None,
            sampler: None,
            width: 0,
            height: 0,
            enabled: true,
            initialized: false,
            bg_dirty: true,
        }
    }
}

impl SsaoPass {
    /// Initialize pipelines and allocate textures for the given resolution.
    pub fn init(&mut self, device: &wgpu::Device, width: u32, height: u32) {
        if self.initialized && self.width == width && self.height == height {
            return;
        }

        let half_w = width.max(2) / 2;
        let half_h = height.max(2) / 2;

        // SSAO bind group layout: params (uniform) + depth (texture_depth_2d) + output (storage)
        // No sampler needed — shader uses textureLoad for depth reads.
        let ssao_bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("ssao_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::Depth,
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 2,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::StorageTexture {
                        access: wgpu::StorageTextureAccess::WriteOnly,
                        format: wgpu::TextureFormat::Rgba16Float,
                        view_dimension: wgpu::TextureViewDimension::D2,
                    },
                    count: None,
                },
            ],
        });

        // Blur bind group layout: same structure but reads AO texture
        let blur_bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("ssao_blur_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::D2,
                        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: wgpu::TextureFormat::Rgba16Float,
                        view_dimension: wgpu::TextureViewDimension::D2,
                    },
                    count: None,
                },
            ],
        });

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

        let ssao_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("ssao_layout"),
            bind_group_layouts: &[Some(&ssao_bgl)],
            immediate_size: 0,
        });
        let blur_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("ssao_blur_layout"),
            bind_group_layouts: &[Some(&blur_bgl)],
            immediate_size: 0,
        });

        self.ssao_pipeline = Some(device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("ssao_compute"),
            layout: Some(&ssao_layout),
            module: &shader,
            entry_point: Some("main"),
            compilation_options: Default::default(),
            cache: None,
        }));
        self.blur_pipeline = Some(device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("ssao_blur"),
            layout: Some(&blur_layout),
            module: &shader,
            entry_point: Some("blur"),
            compilation_options: Default::default(),
            cache: None,
        }));

        // AO textures at half resolution
        let ao_tex = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("ssao_ao"),
            size: wgpu::Extent3d {
                width: half_w,
                height: half_h,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba16Float,
            usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::TEXTURE_BINDING,
            view_formats: &[],
        });
        let blur_tex = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("ssao_blur"),
            size: wgpu::Extent3d {
                width: half_w,
                height: half_h,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba16Float,
            usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::TEXTURE_BINDING,
            view_formats: &[],
        });

        self.ao_view = Some(ao_tex.create_view(&wgpu::TextureViewDescriptor::default()));
        self.blur_view = Some(blur_tex.create_view(&wgpu::TextureViewDescriptor::default()));
        self.ao_texture = Some(ao_tex);
        self.blur_texture = Some(blur_tex);

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

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

        self.ssao_bgl = Some(ssao_bgl);
        self.blur_bgl = Some(blur_bgl);
        self.width = width;
        self.height = height;
        self.initialized = true;
        self.bg_dirty = true;

        log::info!("SSAO initialized ({}x{} → {}x{} AO)", width, height, half_w, half_h);
    }

    /// Dispatch SSAO compute + blur. Returns the blurred AO texture view for PBR binding.
    pub fn dispatch(
        &self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
        depth_view: &wgpu::TextureView,
        camera: &crate::camera::Camera,
    ) -> Option<&wgpu::TextureView> {
        if !self.enabled || !self.initialized {
            return None;
        }
        let Some(ref ssao_pipe) = self.ssao_pipeline else {
            return None;
        };
        let Some(ref blur_pipe) = self.blur_pipeline else {
            return None;
        };
        let Some(ref ssao_bgl) = self.ssao_bgl else { return None };
        let Some(ref blur_bgl) = self.blur_bgl else { return None };
        let Some(ref params_buf) = self.params_buf else {
            return None;
        };
        let Some(ref ao_view) = self.ao_view else { return None };
        let Some(ref blur_view) = self.blur_view else {
            return None;
        };
        let Some(ref sampler) = self.sampler else { return None };

        let half_w = self.width.max(2) / 2;
        let half_h = self.height.max(2) / 2;

        // Upload params
        let params = SsaoParams {
            inv_proj: camera.projection_matrix.inverse().to_cols_array(),
            proj: camera.projection_matrix.to_cols_array(),
            screen_size: [self.width as f32, self.height as f32],
            radius: 0.5,
            bias: 0.025,
        };
        queue.write_buffer(params_buf, 0, bytemuck::bytes_of(&params));

        // SSAO pass: depth → AO (no sampler — uses textureLoad in shader)
        let ssao_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("ssao_bg"),
            layout: ssao_bgl,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: params_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(depth_view),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::TextureView(ao_view),
                },
            ],
        });
        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("ssao_compute"),
                timestamp_writes: None,
            });
            pass.set_pipeline(ssao_pipe);
            pass.set_bind_group(0, &ssao_bg, &[]);
            pass.dispatch_workgroups(half_w.div_ceil(8), half_h.div_ceil(8), 1);
        }

        // Blur pass: AO → blurred AO
        let blur_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("ssao_blur_bg"),
            layout: blur_bgl,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: params_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(ao_view),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::Sampler(sampler),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: wgpu::BindingResource::TextureView(blur_view),
                },
            ],
        });
        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("ssao_blur_pass"),
                timestamp_writes: None,
            });
            pass.set_pipeline(blur_pipe);
            pass.set_bind_group(0, &blur_bg, &[]);
            pass.dispatch_workgroups(half_w.div_ceil(8), half_h.div_ceil(8), 1);
        }

        Some(blur_view)
    }

    /// Get the blurred AO texture view for external binding.
    pub fn ao_view(&self) -> Option<&wgpu::TextureView> {
        self.blur_view.as_ref()
    }

    pub fn destroy(&mut self) {
        if let Some(t) = self.ao_texture.take() {
            t.destroy();
        }
        if let Some(t) = self.blur_texture.take() {
            t.destroy();
        }
        if let Some(b) = self.params_buf.take() {
            b.destroy();
        }
        self.ao_view = None;
        self.blur_view = None;
        self.initialized = false;
    }
}

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

    #[test]
    fn ssao_params_size() {
        assert_eq!(std::mem::size_of::<SsaoParams>(), 144);
    }

    #[test]
    fn ssao_default_enabled() {
        let pass = SsaoPass::default();
        assert!(pass.enabled);
        assert!(!pass.initialized);
    }
}
