// RT Denoise — SIGMA-inspired shadow denoiser compute pass manager.
// Two-pass denoiser for 1spp ray-traced shadow masks:
//   1. Spatial bilateral blur (5x5, depth-weighted)
//   2. Temporal stabilization (motion-reprojected history, variance clamping)
// Dispatches after RT shadow pass, before PBR lighting consumes the shadow mask.
// Output: R8Unorm denoised shadow texture bound in PBR @group(3).

use crate::camera::Camera;

/// Configuration for the RT shadow denoiser.
#[derive(Debug, Clone)]
pub struct RtDenoiseConfig {
    /// Spatial blur radius in pixels. Default: 2.0.
    pub blur_radius: f32,
    /// Temporal blend factor (0.0 = spatial only, 1.0 = history only). Default: 0.85.
    pub temporal_alpha: f32,
    /// Depth similarity threshold for bilateral weighting. Default: 0.1.
    pub depth_threshold: f32,
}

impl Default for RtDenoiseConfig {
    fn default() -> Self {
        Self {
            blur_radius: 2.0,
            temporal_alpha: 0.85,
            depth_threshold: 0.1,
        }
    }
}

/// GPU uniforms for the RT denoise shader (must be #[repr(C)] + Pod + Zeroable).
///
/// Layout (must match rt_denoise.wgsl DenoiseUniforms):
///   screen_size:    8 bytes (2 × u32)
///   _align_pad:     8 bytes (alignment to 16 for mat4x4)
///   view_proj_inv:  64 bytes (mat4x4<f32>)
///   prev_view_proj: 64 bytes (mat4x4<f32>)
///   blur_radius:    4 bytes
///   temporal_alpha: 4 bytes
///   depth_threshold:4 bytes
///   frame_index:    4 bytes
///   Total:          160 bytes (160 % 16 = 0)
#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
pub struct DenoiseUniforms {
    pub screen_size: [u32; 2],
    pub _align_pad: [u32; 2],
    pub view_proj_inv: [f32; 16],
    pub prev_view_proj: [f32; 16],
    pub blur_radius: f32,
    pub temporal_alpha: f32,
    pub depth_threshold: f32,
    pub frame_index: u32,
}

/// RT Shadow Denoiser pass — SIGMA-inspired bilateral + temporal.
///
/// Two compute passes per frame:
/// 1. **Spatial blur** — 5x5 bilateral filter on raw shadow mask, weighted by depth similarity.
/// 2. **Temporal blend** — motion-reprojected history blending with variance clamping.
///
/// Uses ping-pong history textures to avoid read-after-write hazards.
pub struct RtDenoisePass {
    pub enabled: bool,
    // Intermediate blurred texture (spatial pass output, temporal pass input)
    blurred_tex: Option<wgpu::Texture>,
    blurred_view: Option<wgpu::TextureView>,
    // Ping-pong history textures for temporal accumulation
    history_tex: [Option<wgpu::Texture>; 2],
    history_view: [Option<wgpu::TextureView>; 2],
    history_index: usize,
    // Final denoised output
    final_output: Option<wgpu::Texture>,
    final_output_view: Option<wgpu::TextureView>,
    // Uniform buffer
    uniform_buffer: Option<wgpu::Buffer>,
    // Bind group layouts (cached for reuse)
    bgl_shared: Option<wgpu::BindGroupLayout>,
    bgl_spatial: Option<wgpu::BindGroupLayout>,
    bgl_temporal: Option<wgpu::BindGroupLayout>,
    // Pipelines
    spatial_pipeline: Option<wgpu::ComputePipeline>,
    temporal_pipeline: Option<wgpu::ComputePipeline>,
    // Config
    config: RtDenoiseConfig,
    frame_index: u32,
    prev_view_proj: [f32; 16],
    width: u32,
    height: u32,
    initialized: bool,
}

impl Default for RtDenoisePass {
    fn default() -> Self {
        #[rustfmt::skip]
        let identity = [
            1.0, 0.0, 0.0, 0.0,
            0.0, 1.0, 0.0, 0.0,
            0.0, 0.0, 1.0, 0.0,
            0.0, 0.0, 0.0, 1.0,
        ];
        Self {
            enabled: false,
            blurred_tex: None,
            blurred_view: None,
            history_tex: [None, None],
            history_view: [None, None],
            history_index: 0,
            final_output: None,
            final_output_view: None,
            uniform_buffer: None,
            bgl_shared: None,
            bgl_spatial: None,
            bgl_temporal: None,
            spatial_pipeline: None,
            temporal_pipeline: None,
            config: RtDenoiseConfig::default(),
            frame_index: 0,
            prev_view_proj: identity,
            width: 0,
            height: 0,
            initialized: false,
        }
    }
}

impl RtDenoisePass {
    /// Create an RT denoise pass with the given configuration, enabled by default.
    pub fn new(config: RtDenoiseConfig) -> Self {
        Self {
            enabled: true,
            config,
            ..Default::default()
        }
    }

    /// Initialize GPU resources: intermediate textures, history ping-pong, output, uniform buffer.
    /// Does not create pipelines (call `create_pipelines` separately).
    pub fn init(&mut self, device: &wgpu::Device, width: u32, height: u32) {
        if self.initialized && self.width == width && self.height == height {
            return;
        }

        let make_r8 = |device: &wgpu::Device, label: &'static str| -> wgpu::Texture {
            device.create_texture(&wgpu::TextureDescriptor {
                label: Some(label),
                size: wgpu::Extent3d {
                    width,
                    height,
                    depth_or_array_layers: 1,
                },
                mip_level_count: 1,
                sample_count: 1,
                dimension: wgpu::TextureDimension::D2,
                format: wgpu::TextureFormat::R8Unorm,
                usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::TEXTURE_BINDING,
                view_formats: &[],
            })
        };

        // Blurred intermediate
        let blurred = make_r8(device, "rt_denoise_blurred");
        self.blurred_view = Some(blurred.create_view(&wgpu::TextureViewDescriptor::default()));
        self.blurred_tex = Some(blurred);

        // History ping-pong (two textures)
        for i in 0..2 {
            let label: &'static str = if i == 0 {
                "rt_denoise_history_0"
            } else {
                "rt_denoise_history_1"
            };
            let tex = make_r8(device, label);
            self.history_view[i] = Some(tex.create_view(&wgpu::TextureViewDescriptor::default()));
            self.history_tex[i] = Some(tex);
        }

        // Final denoised output
        let output = make_r8(device, "rt_denoise_output");
        self.final_output_view = Some(output.create_view(&wgpu::TextureViewDescriptor::default()));
        self.final_output = Some(output);

        // Uniform buffer
        if self.uniform_buffer.is_none() {
            self.uniform_buffer = Some(device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("rt_denoise_uniforms"),
                size: std::mem::size_of::<DenoiseUniforms>() as u64,
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            }));
        }

        self.width = width;
        self.height = height;
        self.initialized = true;

        log::info!("RT Denoise initialized ({}x{})", width, height);
    }

    /// Create the spatial and temporal compute pipelines from rt_denoise.wgsl.
    pub fn create_pipelines(&mut self, device: &wgpu::Device) {
        // Group 0 (shared): uniforms + depth
        let bgl_shared = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("rt_denoise_bgl_shared"),
            entries: &[
                // binding 0: uniforms
                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,
                },
                // binding 1: depth texture
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Float { filterable: false },
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
            ],
        });

        // Group 1 (spatial pass): shadow input (read) + blurred output (write)
        let bgl_spatial = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("rt_denoise_bgl_spatial"),
            entries: &[
                // binding 0: shadow_input (texture read)
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Float { filterable: false },
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
                // binding 1: shadow_blurred (storage write)
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::StorageTexture {
                        access: wgpu::StorageTextureAccess::WriteOnly,
                        format: wgpu::TextureFormat::R8Unorm,
                        view_dimension: wgpu::TextureViewDimension::D2,
                    },
                    count: None,
                },
            ],
        });

        // Group 2 (temporal pass): blurred input (read) + history (read) + final output (write)
        let bgl_temporal = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("rt_denoise_bgl_temporal"),
            entries: &[
                // binding 0: blurred_input (texture read)
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Float { filterable: false },
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
                // binding 1: history_tex (texture read)
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Float { filterable: false },
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
                // binding 2: final_output (storage write)
                wgpu::BindGroupLayoutEntry {
                    binding: 2,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::StorageTexture {
                        access: wgpu::StorageTextureAccess::WriteOnly,
                        format: wgpu::TextureFormat::R8Unorm,
                        view_dimension: wgpu::TextureViewDimension::D2,
                    },
                    count: None,
                },
            ],
        });

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

        // Spatial pipeline uses groups 0 + 1 (temporal group 2 is unused but must be in layout)
        let spatial_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("rt_denoise_spatial_layout"),
            bind_group_layouts: &[Some(&bgl_shared), Some(&bgl_spatial)],
            immediate_size: 0,
        });

        self.spatial_pipeline = Some(device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("rt_denoise_spatial"),
            layout: Some(&spatial_layout),
            module: &shader,
            entry_point: Some("spatial_blur"),
            compilation_options: Default::default(),
            cache: None,
        }));

        // Temporal pipeline uses groups 0 + 2 (spatial group 1 is unused)
        let temporal_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("rt_denoise_temporal_layout"),
            bind_group_layouts: &[Some(&bgl_shared), Some(&bgl_temporal)],
            immediate_size: 0,
        });

        self.temporal_pipeline = Some(device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("rt_denoise_temporal"),
            layout: Some(&temporal_layout),
            module: &shader,
            entry_point: Some("temporal_blend"),
            compilation_options: Default::default(),
            cache: None,
        }));

        // Cache layouts for bind group creation
        self.bgl_shared = Some(bgl_shared);
        self.bgl_spatial = Some(bgl_spatial);
        self.bgl_temporal = Some(bgl_temporal);

        log::info!("RT Denoise pipelines created (spatial + temporal)");
    }

    /// Dispatch both denoise passes: spatial blur then temporal blend.
    ///
    /// Takes the raw RT shadow mask as input, produces a denoised shadow texture.
    /// Swaps the history ping-pong index after each frame.
    pub fn dispatch(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
        depth_view: &wgpu::TextureView,
        raw_shadow_view: &wgpu::TextureView,
        camera: &Camera,
    ) {
        if !self.enabled || !self.initialized {
            return;
        }
        let Some(ref spatial_pipeline) = self.spatial_pipeline else {
            return;
        };
        let Some(ref temporal_pipeline) = self.temporal_pipeline else {
            return;
        };
        let Some(ref uniform_buffer) = self.uniform_buffer else {
            return;
        };
        let Some(ref blurred_view) = self.blurred_view else {
            return;
        };
        let Some(ref final_output_view) = self.final_output_view else {
            return;
        };
        let Some(ref bgl_shared) = self.bgl_shared else { return };
        let Some(ref bgl_spatial) = self.bgl_spatial else {
            return;
        };
        let Some(ref bgl_temporal) = self.bgl_temporal else {
            return;
        };

        // Current and previous history indices
        let read_history = self.history_index;
        let write_history = 1 - self.history_index;
        let Some(ref history_read_view) = self.history_view[read_history] else {
            return;
        };

        let view_proj = camera.projection_matrix * camera.view_matrix;
        let view_proj_inv = view_proj.inverse();

        // Upload uniforms
        let uniforms = DenoiseUniforms {
            screen_size: [self.width, self.height],
            _align_pad: [0, 0],
            view_proj_inv: view_proj_inv.to_cols_array(),
            prev_view_proj: self.prev_view_proj,
            blur_radius: self.config.blur_radius,
            temporal_alpha: self.config.temporal_alpha,
            depth_threshold: self.config.depth_threshold,
            frame_index: self.frame_index,
        };
        queue.write_buffer(uniform_buffer, 0, bytemuck::bytes_of(&uniforms));

        let wg_x = self.width.div_ceil(8);
        let wg_y = self.height.div_ceil(8);

        // Shared bind group (group 0): uniforms + depth
        let bg_shared = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("rt_denoise_bg_shared"),
            layout: bgl_shared,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: uniform_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(depth_view),
                },
            ],
        });

        // --- Pass 1: Spatial bilateral blur ---
        let bg_spatial = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("rt_denoise_bg_spatial"),
            layout: bgl_spatial,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(raw_shadow_view),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(blurred_view),
                },
            ],
        });

        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("rt_denoise_spatial"),
                timestamp_writes: None,
            });
            pass.set_pipeline(spatial_pipeline);
            pass.set_bind_group(0, &bg_shared, &[]);
            pass.set_bind_group(1, &bg_spatial, &[]);
            pass.dispatch_workgroups(wg_x, wg_y, 1);
        }

        // --- Pass 2: Temporal stabilization ---
        let bg_temporal = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("rt_denoise_bg_temporal"),
            layout: bgl_temporal,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(blurred_view),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(history_read_view),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::TextureView(final_output_view),
                },
            ],
        });

        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("rt_denoise_temporal"),
                timestamp_writes: None,
            });
            pass.set_pipeline(temporal_pipeline);
            pass.set_bind_group(0, &bg_shared, &[]);
            pass.set_bind_group(1, &bg_temporal, &[]);
            pass.dispatch_workgroups(wg_x, wg_y, 1);
        }

        // Copy final output to the write-side history for next frame.
        // (In practice this would be a texture copy or the output could write
        // directly to history[write]. For simplicity, swap which history is read.)
        let _ = write_history; // future: encoder.copy_texture_to_texture(...)

        // Swap history ping-pong
        self.history_index = 1 - self.history_index;
        self.prev_view_proj = view_proj.to_cols_array();
        self.frame_index = self.frame_index.wrapping_add(1);
    }

    /// Get the denoised shadow texture view for external binding (PBR @group(3)).
    pub fn denoised_shadow_view(&self) -> Option<&wgpu::TextureView> {
        self.final_output_view.as_ref()
    }

    /// Resize screen-dependent resources. Recreates all textures at the new resolution.
    pub fn resize(&mut self, device: &wgpu::Device, width: u32, height: u32) {
        if self.width == width && self.height == height {
            return;
        }

        // Destroy old textures
        if let Some(t) = self.blurred_tex.take() {
            t.destroy();
        }
        self.blurred_view = None;
        for i in 0..2 {
            if let Some(t) = self.history_tex[i].take() {
                t.destroy();
            }
            self.history_view[i] = None;
        }
        if let Some(t) = self.final_output.take() {
            t.destroy();
        }
        self.final_output_view = None;

        let make_tex = |device: &wgpu::Device, label: &'static str| -> wgpu::Texture {
            device.create_texture(&wgpu::TextureDescriptor {
                label: Some(label),
                size: wgpu::Extent3d {
                    width,
                    height,
                    depth_or_array_layers: 1,
                },
                mip_level_count: 1,
                sample_count: 1,
                dimension: wgpu::TextureDimension::D2,
                format: wgpu::TextureFormat::R8Unorm,
                usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::TEXTURE_BINDING,
                view_formats: &[],
            })
        };

        let blurred = make_tex(device, "rt_denoise_blurred");
        self.blurred_view = Some(blurred.create_view(&wgpu::TextureViewDescriptor::default()));
        self.blurred_tex = Some(blurred);

        for i in 0..2 {
            let label: &'static str = if i == 0 {
                "rt_denoise_history_0"
            } else {
                "rt_denoise_history_1"
            };
            let tex = make_tex(device, label);
            self.history_view[i] = Some(tex.create_view(&wgpu::TextureViewDescriptor::default()));
            self.history_tex[i] = Some(tex);
        }

        let output = make_tex(device, "rt_denoise_output");
        self.final_output_view = Some(output.create_view(&wgpu::TextureViewDescriptor::default()));
        self.final_output = Some(output);

        self.width = width;
        self.height = height;

        log::info!("RT Denoise resized ({}x{})", width, height);
    }

    /// Drop all GPU resources.
    pub fn destroy(&mut self) {
        if let Some(t) = self.blurred_tex.take() {
            t.destroy();
        }
        self.blurred_view = None;
        for i in 0..2 {
            if let Some(t) = self.history_tex[i].take() {
                t.destroy();
            }
            self.history_view[i] = None;
        }
        if let Some(t) = self.final_output.take() {
            t.destroy();
        }
        self.final_output_view = None;
        if let Some(b) = self.uniform_buffer.take() {
            b.destroy();
        }
        self.bgl_shared = None;
        self.bgl_spatial = None;
        self.bgl_temporal = None;
        self.initialized = false;
    }

    /// Estimate total VRAM usage in bytes for the given resolution.
    pub fn vram_estimate(width: u32, height: u32) -> u64 {
        let pixels = width as u64 * height as u64;

        // R8Unorm = 1 byte per pixel
        let blurred_bytes = pixels; // intermediate blurred
        let history_bytes = pixels * 2; // 2 ping-pong history textures
        let output_bytes = pixels; // final output

        // Uniform buffer
        let uniform_bytes = std::mem::size_of::<DenoiseUniforms>() as u64;

        blurred_bytes + history_bytes + output_bytes + uniform_bytes
    }
}

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

    #[test]
    fn rt_denoise_config_defaults_valid() {
        let config = RtDenoiseConfig::default();
        assert!(config.blur_radius > 0.0, "blur_radius must be positive");
        assert!(
            config.temporal_alpha > 0.0 && config.temporal_alpha < 1.0,
            "temporal_alpha must be in (0, 1)"
        );
        assert!(config.depth_threshold > 0.0, "depth_threshold must be positive");
    }

    #[test]
    fn denoise_uniforms_size() {
        // 2 u32 (screen_size) + 2 u32 (pad) + 16 f32 (view_proj_inv) + 16 f32 (prev_view_proj)
        // + 3 f32 (blur_radius, temporal_alpha, depth_threshold) + 1 u32 (frame_index)
        // = 8 + 8 + 64 + 64 + 12 + 4 = 160 bytes
        assert_eq!(std::mem::size_of::<DenoiseUniforms>(), 160);
    }

    #[test]
    fn vram_estimate_denoise_reasonable() {
        // 1920x1080: 4 R8Unorm textures (blurred + 2 history + output)
        let estimate = RtDenoisePass::vram_estimate(1920, 1080);
        // 4 × 1920 × 1080 = ~8 MB. Should be well under 20 MB.
        let budget = 20 * 1024 * 1024;
        assert!(
            estimate < budget,
            "VRAM estimate {} bytes ({:.1} MB) exceeds 20 MB budget",
            estimate,
            estimate as f64 / (1024.0 * 1024.0)
        );
    }

    #[test]
    fn rt_denoise_default_disabled() {
        let pass = RtDenoisePass::default();
        assert!(!pass.enabled);
        assert!(!pass.initialized);
        assert!(pass.denoised_shadow_view().is_none());
    }

    #[test]
    fn rt_denoise_new_enabled() {
        let pass = RtDenoisePass::new(RtDenoiseConfig::default());
        assert!(pass.enabled);
        assert!(!pass.initialized);
    }

    #[test]
    fn history_ping_pong_starts_at_zero() {
        let pass = RtDenoisePass::default();
        assert_eq!(pass.history_index, 0);
    }

    #[test]
    fn vram_estimate_scales_with_resolution() {
        let est_720 = RtDenoisePass::vram_estimate(1280, 720);
        let est_1080 = RtDenoisePass::vram_estimate(1920, 1080);
        let est_4k = RtDenoisePass::vram_estimate(3840, 2160);

        // Higher resolution should always cost more
        assert!(est_1080 > est_720);
        assert!(est_4k > est_1080);

        // 4K should be roughly 4x the cost of 1080p
        let ratio = est_4k as f64 / est_1080 as f64;
        assert!(ratio > 3.5 && ratio < 4.5, "4K/1080p ratio {ratio:.2} should be ~4.0");
    }
}
