// Dreamwell GPU — IBL Precompute Manager
//
// Manages BRDF integration LUT, diffuse irradiance cubemap, and specular
// prefiltered environment cubemap. All textures are computed on the GPU
// via compute shaders at init time.

/// Expected BRDF LUT dimensions.
pub const BRDF_LUT_SIZE: u32 = 512;

/// Expected irradiance cubemap face dimensions.
pub const IRRADIANCE_SIZE: u32 = 64;

/// Expected specular prefilter mip chain base dimensions.
pub const SPECULAR_BASE_SIZE: u32 = 128;

/// Number of mip levels for specular prefiltered cubemap.
pub const SPECULAR_MIP_LEVELS: u32 = 6;

/// IBL precomputation resource manager.
///
/// Holds GPU textures for the split-sum PBR IBL pipeline:
/// - BRDF integration LUT (512x512, Rg16Float)
/// - Diffuse irradiance cubemap (64x64 per face, Rgba16Float)
/// - Specular prefiltered cubemap (128 base, 6 mip levels, Rgba16Float)
pub struct IblPrecompute {
    brdf_lut: Option<wgpu::Texture>,
    brdf_lut_view: Option<wgpu::TextureView>,
    irradiance_cube: Option<wgpu::Texture>,
    irradiance_view: Option<wgpu::TextureView>,
    specular_cube: Option<wgpu::Texture>,
    specular_view: Option<wgpu::TextureView>,
    brdf_pipeline: Option<wgpu::ComputePipeline>,
    brdf_bgl: Option<wgpu::BindGroupLayout>,
    initialized: bool,
}

impl IblPrecompute {
    /// Create an uninitialized IBL precompute manager.
    /// Call `compute_brdf_lut` to generate the BRDF integration LUT.
    pub fn new() -> Self {
        Self {
            brdf_lut: None,
            brdf_lut_view: None,
            irradiance_cube: None,
            irradiance_view: None,
            specular_cube: None,
            specular_view: None,
            brdf_pipeline: None,
            brdf_bgl: None,
            initialized: false,
        }
    }

    /// Compute the BRDF integration LUT on the GPU.
    ///
    /// Creates a 512x512 Rg16Float storage texture and dispatches the
    /// `brdf_lut.wgsl` compute shader to fill it. This is a one-time
    /// operation at init; the result is reused for all PBR materials.
    pub fn compute_brdf_lut(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) {
        // Create storage texture
        let texture = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("ibl_brdf_lut"),
            size: wgpu::Extent3d {
                width: BRDF_LUT_SIZE,
                height: BRDF_LUT_SIZE,
                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 view = texture.create_view(&wgpu::TextureViewDescriptor::default());

        // Create bind group layout
        let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("ibl_brdf_lut_bgl"),
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::COMPUTE,
                ty: wgpu::BindingType::StorageTexture {
                    access: wgpu::StorageTextureAccess::WriteOnly,
                    format: wgpu::TextureFormat::Rgba16Float,
                    view_dimension: wgpu::TextureViewDimension::D2,
                },
                count: None,
            }],
        });

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

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("ibl_brdf_lut_pipeline_layout"),
            bind_group_layouts: &[Some(&bgl)],
            immediate_size: 0,
        });

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

        // Create bind group
        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("ibl_brdf_lut_bind_group"),
            layout: &bgl,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: wgpu::BindingResource::TextureView(&view),
            }],
        });

        // Dispatch compute
        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("ibl_brdf_lut_encoder"),
        });

        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("ibl_brdf_lut_pass"),
                timestamp_writes: None,
            });
            pass.set_pipeline(&pipeline);
            pass.set_bind_group(0, &bind_group, &[]);
            pass.dispatch_workgroups(BRDF_LUT_SIZE / 8, BRDF_LUT_SIZE / 8, 1);
        }

        queue.submit(std::iter::once(encoder.finish()));

        self.brdf_lut = Some(texture);
        self.brdf_lut_view = Some(view);
        self.brdf_pipeline = Some(pipeline);
        self.brdf_bgl = Some(bgl);
        self.initialized = true;

        log::info!(
            "IBL: BRDF LUT computed ({}x{}, Rg16Float)",
            BRDF_LUT_SIZE,
            BRDF_LUT_SIZE,
        );
    }

    /// Compute the diffuse irradiance cubemap on the GPU.
    ///
    /// Creates a 64x64 Rgba16Float cubemap (6 array layers) and dispatches the
    /// `irradiance.wgsl` compute shader once per face. The environment source
    /// is sampled as an equirectangular 2D texture.
    pub fn compute_irradiance(&mut self, device: &wgpu::Device, queue: &wgpu::Queue, env_view: &wgpu::TextureView) {
        let texture = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("ibl_irradiance_cube"),
            size: wgpu::Extent3d {
                width: IRRADIANCE_SIZE,
                height: IRRADIANCE_SIZE,
                depth_or_array_layers: 6,
            },
            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 cube_view = texture.create_view(&wgpu::TextureViewDescriptor {
            dimension: Some(wgpu::TextureViewDimension::Cube),
            array_layer_count: Some(6),
            ..Default::default()
        });

        // Bind group layout: env source (texture_2d) + sampler + output face (storage texture)
        let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("ibl_irradiance_bgl"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    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: 1,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                    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::D2Array,
                    },
                    count: None,
                },
            ],
        });

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

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("ibl_irradiance_pipeline_layout"),
            bind_group_layouts: &[Some(&bgl)],
            immediate_size: 0,
        });

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

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

        // Create a 2D array view for storage output (all 6 layers)
        let output_view = texture.create_view(&wgpu::TextureViewDescriptor {
            dimension: Some(wgpu::TextureViewDimension::D2Array),
            array_layer_count: Some(6),
            ..Default::default()
        });

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("ibl_irradiance_bg"),
            layout: &bgl,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(env_view),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::Sampler(&sampler),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::TextureView(&output_view),
                },
            ],
        });

        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("ibl_irradiance_encoder"),
        });

        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("ibl_irradiance_pass"),
                timestamp_writes: None,
            });
            pass.set_pipeline(&pipeline);
            pass.set_bind_group(0, &bind_group, &[]);
            // Dispatch all 6 faces: shader uses global_invocation_id.z as face index
            let wg = IRRADIANCE_SIZE / 8;
            pass.dispatch_workgroups(wg, wg, 6);
        }

        queue.submit(std::iter::once(encoder.finish()));

        self.irradiance_cube = Some(texture);
        self.irradiance_view = Some(cube_view);

        log::info!(
            "IBL: irradiance cubemap computed ({}x{} per face, Rgba16Float)",
            IRRADIANCE_SIZE,
            IRRADIANCE_SIZE,
        );
    }

    /// Compute the specular prefiltered environment cubemap on the GPU.
    ///
    /// Creates a 128x128 base Rgba16Float cubemap with 6 mip levels and dispatches
    /// the `prefilter_env.wgsl` compute shader for each face×mip combination.
    /// Roughness increases per mip: `mip / (SPECULAR_MIP_LEVELS - 1)`.
    pub fn compute_specular_prefilter(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        env_view: &wgpu::TextureView,
    ) {
        let texture = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("ibl_specular_cube"),
            size: wgpu::Extent3d {
                width: SPECULAR_BASE_SIZE,
                height: SPECULAR_BASE_SIZE,
                depth_or_array_layers: 6,
            },
            mip_level_count: SPECULAR_MIP_LEVELS,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba16Float,
            usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::TEXTURE_BINDING,
            view_formats: &[],
        });

        let cube_view = texture.create_view(&wgpu::TextureViewDescriptor {
            dimension: Some(wgpu::TextureViewDimension::Cube),
            array_layer_count: Some(6),
            ..Default::default()
        });

        // Bind group layout: env source + sampler + output mip face (storage) + roughness uniform
        let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("ibl_prefilter_bgl"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    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: 1,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                    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::D2Array,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 3,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
            ],
        });

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

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("ibl_prefilter_pipeline_layout"),
            bind_group_layouts: &[Some(&bgl)],
            immediate_size: 0,
        });

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

        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            label: Some("ibl_prefilter_sampler"),
            mag_filter: wgpu::FilterMode::Linear,
            min_filter: wgpu::FilterMode::Linear,
            mipmap_filter: wgpu::MipmapFilterMode::Linear,
            ..Default::default()
        });

        // Roughness uniform buffer (4 bytes, updated per-mip dispatch)
        let roughness_buf = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("ibl_prefilter_roughness"),
            size: 16, // 4 bytes + padding to 16-byte alignment
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("ibl_prefilter_encoder"),
        });

        for mip in 0..SPECULAR_MIP_LEVELS {
            let mip_size = SPECULAR_BASE_SIZE >> mip;
            if mip_size == 0 {
                break;
            }

            let roughness = mip as f32 / (SPECULAR_MIP_LEVELS - 1) as f32;
            queue.write_buffer(&roughness_buf, 0, bytemuck::bytes_of(&[roughness, 0.0f32, 0.0, 0.0]));

            // Create per-mip output view (all 6 faces at this mip level)
            let mip_view = texture.create_view(&wgpu::TextureViewDescriptor {
                dimension: Some(wgpu::TextureViewDimension::D2Array),
                base_mip_level: mip,
                mip_level_count: Some(1),
                array_layer_count: Some(6),
                ..Default::default()
            });

            let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("ibl_prefilter_bg"),
                layout: &bgl,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: wgpu::BindingResource::TextureView(env_view),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: wgpu::BindingResource::Sampler(&sampler),
                    },
                    wgpu::BindGroupEntry {
                        binding: 2,
                        resource: wgpu::BindingResource::TextureView(&mip_view),
                    },
                    wgpu::BindGroupEntry {
                        binding: 3,
                        resource: roughness_buf.as_entire_binding(),
                    },
                ],
            });

            {
                let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                    label: Some("ibl_prefilter_pass"),
                    timestamp_writes: None,
                });
                pass.set_pipeline(&pipeline);
                pass.set_bind_group(0, &bind_group, &[]);
                let wg = (mip_size).max(8) / 8;
                pass.dispatch_workgroups(wg, wg, 6);
            }
        }

        queue.submit(std::iter::once(encoder.finish()));
        roughness_buf.destroy();

        self.specular_cube = Some(texture);
        self.specular_view = Some(cube_view);

        log::info!(
            "IBL: specular prefilter computed ({}x{} base, {} mip levels, Rgba16Float)",
            SPECULAR_BASE_SIZE,
            SPECULAR_BASE_SIZE,
            SPECULAR_MIP_LEVELS,
        );
    }

    /// Compute all IBL textures: BRDF LUT, irradiance cubemap, specular prefilter.
    ///
    /// Creates a 1x1 flat-white fallback environment texture when no HDRI is loaded.
    /// IBL with flat-white env produces physically correct uniform ambient light.
    /// No-op if already initialized.
    pub fn compute_all(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) {
        if self.initialized {
            return;
        }

        self.compute_brdf_lut(device, queue);

        // Create a 1x1 flat-white fallback environment texture.
        // When a real HDRI is loaded, call compute_irradiance/compute_specular_prefilter
        // with the environment texture view directly.
        let fallback_env = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("ibl_fallback_env"),
            size: wgpu::Extent3d {
                width: 1,
                height: 1,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba16Float,
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
            view_formats: &[],
        });
        // Write white pixel (1.0, 1.0, 1.0, 1.0) as IEEE 754 half-precision.
        // 0x3C00 = 1.0 in f16.
        let white_f16: [u16; 4] = [0x3C00, 0x3C00, 0x3C00, 0x3C00];
        queue.write_texture(
            wgpu::TexelCopyTextureInfo {
                texture: &fallback_env,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            bytemuck::cast_slice(&white_f16),
            wgpu::TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: Some(8), // 4 x f16 = 8 bytes
                rows_per_image: None,
            },
            wgpu::Extent3d {
                width: 1,
                height: 1,
                depth_or_array_layers: 1,
            },
        );
        let fallback_view = fallback_env.create_view(&wgpu::TextureViewDescriptor::default());

        self.compute_irradiance(device, queue, &fallback_view);
        self.compute_specular_prefilter(device, queue, &fallback_view);

        fallback_env.destroy();
        self.initialized = true;

        log::info!("IBL: all precompute textures initialized (fallback flat-white environment)");
    }

    /// Returns the BRDF LUT texture view, if computed.
    pub fn brdf_lut_view(&self) -> Option<&wgpu::TextureView> {
        self.brdf_lut_view.as_ref()
    }

    /// Returns the diffuse irradiance cubemap view, if computed.
    pub fn irradiance_view(&self) -> Option<&wgpu::TextureView> {
        self.irradiance_view.as_ref()
    }

    /// Returns the specular prefiltered cubemap view, if computed.
    pub fn specular_view(&self) -> Option<&wgpu::TextureView> {
        self.specular_view.as_ref()
    }

    /// Whether all IBL textures have been computed.
    pub fn is_initialized(&self) -> bool {
        self.initialized
    }

    /// Destroy all GPU resources and reset to uninitialized state.
    pub fn destroy(&mut self) {
        if let Some(tex) = self.brdf_lut.take() {
            tex.destroy();
        }
        if let Some(tex) = self.irradiance_cube.take() {
            tex.destroy();
        }
        if let Some(tex) = self.specular_cube.take() {
            tex.destroy();
        }
        self.brdf_lut_view = None;
        self.irradiance_view = None;
        self.specular_view = None;
        self.brdf_pipeline = None;
        self.brdf_bgl = None;
        self.initialized = false;
    }
}

impl Default for IblPrecompute {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn ibl_default_not_initialized() {
        let ibl = IblPrecompute::new();
        assert!(!ibl.is_initialized());
        assert!(ibl.brdf_lut_view().is_none());
        assert!(ibl.irradiance_view().is_none());
        assert!(ibl.specular_view().is_none());
    }

    #[test]
    fn brdf_lut_size() {
        // Compile-time constant assertions on expected dimensions.
        const _: () = assert!(BRDF_LUT_SIZE == 512);
        const _: () = assert!(BRDF_LUT_SIZE % 8 == 0); // Divisible by workgroup size
        const _: () = assert!(IRRADIANCE_SIZE == 64);
        const _: () = assert!(SPECULAR_BASE_SIZE == 128);
        const _: () = assert!(SPECULAR_MIP_LEVELS == 6);

        // Runtime check that dispatch would produce correct workgroup counts.
        assert_eq!(BRDF_LUT_SIZE / 8, 64);
    }
}
