// RT Shadows — Ray-Traced Shadow compute pass manager.
// Replaces CSM with per-pixel ray-traced shadow queries against the TLAS.
// Dispatches between main render pass and post-processing.
// Requires EXPERIMENTAL_RAY_QUERY feature (TLAS acceleration structure).
// Output: R8Unorm shadow mask texture sampled in PBR lighting pass.

use crate::camera::Camera;

/// Configuration for RT shadows.
#[derive(Debug, Clone)]
pub struct RtShadowConfig {
    /// Maximum ray distance for shadow queries. Default: 1000.0.
    pub max_ray_distance: f32,
    /// Bias to prevent shadow acne. Default: 0.001.
    pub shadow_bias: f32,
    /// Shadow rays per pixel (1 = hard shadows, >1 = soft). Default: 1.
    pub samples_per_pixel: u32,
}

impl Default for RtShadowConfig {
    fn default() -> Self {
        Self {
            max_ray_distance: 1000.0,
            shadow_bias: 0.001,
            samples_per_pixel: 1,
        }
    }
}

/// GPU uniforms for the RT shadow shader (must be #[repr(C)] + Pod + Zeroable).
#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
pub struct RtShadowUniforms {
    pub view_proj_inv: [f32; 16],
    pub light_direction: [f32; 3],
    pub max_ray_distance: f32,
    pub shadow_bias: f32,
    pub light_intensity: f32,
    pub screen_size: [u32; 2],
}

/// RT Shadow pass — replaces CSM with ray-traced shadow queries.
///
/// Single compute pass per frame: for each pixel, reconstructs world position
/// from depth, traces a shadow ray toward the light direction against the TLAS,
/// and writes a visibility value to the shadow mask.
pub struct RtShadowPass {
    pub enabled: bool,
    shadow_mask: Option<wgpu::Texture>,
    shadow_mask_view: Option<wgpu::TextureView>,
    uniform_buffer: Option<wgpu::Buffer>,
    bind_group: Option<wgpu::BindGroup>,
    cached_bgl: Option<wgpu::BindGroupLayout>,
    pipeline: Option<wgpu::ComputePipeline>,
    config: RtShadowConfig,
    width: u32,
    height: u32,
    initialized: bool,
    /// Whether bind groups need to be recreated (set on resize/init).
    pub bg_dirty: bool,
}

impl Default for RtShadowPass {
    fn default() -> Self {
        Self {
            enabled: false,
            shadow_mask: None,
            shadow_mask_view: None,
            uniform_buffer: None,
            bind_group: None,
            cached_bgl: None,
            pipeline: None,
            config: RtShadowConfig::default(),
            width: 0,
            height: 0,
            initialized: false,
            bg_dirty: true,
        }
    }
}

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

    /// Initialize GPU resources: shadow mask texture and uniform buffer.
    /// Does not create the pipeline (call `create_pipeline` separately after TLAS is ready).
    pub fn init(&mut self, device: &wgpu::Device, width: u32, height: u32) {
        if self.initialized && self.width == width && self.height == height {
            return;
        }

        // Shadow mask: R8Unorm at full resolution (0.0 = fully shadowed, 1.0 = fully lit)
        let mask_tex = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("rt_shadow_mask"),
            size: wgpu::Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8Unorm, // universally supported for storage
            usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::TEXTURE_BINDING,
            view_formats: &[],
        });
        self.shadow_mask_view = Some(mask_tex.create_view(&wgpu::TextureViewDescriptor::default()));
        self.shadow_mask = Some(mask_tex);

        // Uniform buffer
        if self.uniform_buffer.is_none() {
            self.uniform_buffer = Some(device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("rt_shadow_uniforms"),
                size: std::mem::size_of::<RtShadowUniforms>() 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 Shadows initialized ({}x{})", width, height);
    }

    /// Create an RT shadow pass with feature detection.
    /// If EXPERIMENTAL_RAY_QUERY is not supported, the pass is created but disabled.
    pub fn new_checked(device: &wgpu::Device, config: RtShadowConfig) -> Self {
        let supported = device.features().contains(wgpu::Features::EXPERIMENTAL_RAY_QUERY);
        if !supported {
            log::warn!("rt_shadows: EXPERIMENTAL_RAY_QUERY not supported, disabled");
        }
        Self {
            enabled: supported,
            config,
            ..Default::default()
        }
    }

    /// Create the compute pipeline from the rt_shadows.wgsl shader.
    /// Requires a valid TLAS for the acceleration structure bind group layout entry.
    /// Silently returns if EXPERIMENTAL_RAY_QUERY is not supported by the device.
    pub fn create_pipeline(&mut self, device: &wgpu::Device, _tlas: &wgpu::Tlas) {
        if !device.features().contains(wgpu::Features::EXPERIMENTAL_RAY_QUERY) {
            log::warn!("rt_shadows: skipping pipeline creation (no RT support)");
            self.enabled = false;
            return;
        }
        let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("rt_shadow_bgl"),
            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: acceleration structure (TLAS)
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::AccelerationStructure { vertex_return: false },
                    count: None,
                },
                // binding 2: depth texture
                wgpu::BindGroupLayoutEntry {
                    binding: 2,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Depth,
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
                // binding 3: shadow mask output
                wgpu::BindGroupLayoutEntry {
                    binding: 3,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::StorageTexture {
                        access: wgpu::StorageTextureAccess::WriteOnly,
                        format: wgpu::TextureFormat::Rgba8Unorm,
                        view_dimension: wgpu::TextureViewDimension::D2,
                    },
                    count: None,
                },
            ],
        });

        // SAFETY: RT shaders use `enable wgpu_ray_query;` — experimental WGSL extension.
        // GPU driver validates at pipeline creation. Same pattern as wgpu RT examples.
        let shader = unsafe {
            device.create_shader_module_trusted(
                wgpu::ShaderModuleDescriptor {
                    label: Some("rt_shadows.wgsl"),
                    source: wgpu::ShaderSource::Wgsl(include_str!("../shaders/rt_shadows.wgsl").into()),
                },
                wgpu::ShaderRuntimeChecks::unchecked(),
            )
        };

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

        self.pipeline = Some(device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("rt_shadow_compute"),
            layout: Some(&layout),
            module: &shader,
            entry_point: Some("rt_shadow_trace"),
            compilation_options: Default::default(),
            cache: None,
        }));

        // Cache layout for per-frame bind group recreation (avoids layout leak)
        self.cached_bgl = Some(bgl);

        log::info!("RT Shadow pipeline created");
    }

    /// Dispatch the shadow compute pass.
    /// Uploads uniforms, creates bind group, dispatches full-screen workgroups.
    pub fn dispatch(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
        tlas: &wgpu::Tlas,
        depth_view: &wgpu::TextureView,
        camera: &Camera,
        light_direction: [f32; 3],
        light_intensity: f32,
    ) {
        if !self.enabled || !self.initialized {
            return;
        }
        let Some(ref pipeline) = self.pipeline else { return };
        let Some(ref uniform_buffer) = self.uniform_buffer else {
            return;
        };
        let Some(ref shadow_mask_view) = self.shadow_mask_view else {
            return;
        };

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

        // Upload uniforms
        let uniforms = RtShadowUniforms {
            view_proj_inv: view_proj_inv.to_cols_array(),
            light_direction,
            max_ray_distance: self.config.max_ray_distance,
            shadow_bias: self.config.shadow_bias,
            light_intensity,
            screen_size: [self.width, self.height],
        };
        queue.write_buffer(uniform_buffer, 0, bytemuck::bytes_of(&uniforms));

        // Recreate bind group (depth view may change per frame). Layout is cached.
        let Some(ref bgl) = self.cached_bgl else { return };

        let bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("rt_shadow_bg"),
            layout: bgl,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: uniform_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: tlas.as_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::TextureView(depth_view),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: wgpu::BindingResource::TextureView(shadow_mask_view),
                },
            ],
        });

        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("rt_shadow_compute"),
                timestamp_writes: None,
            });
            pass.set_pipeline(pipeline);
            pass.set_bind_group(0, &bg, &[]);
            pass.dispatch_workgroups(self.width.div_ceil(8), self.height.div_ceil(8), 1);
        }
    }

    /// Get the shadow mask texture view for external binding (PBR lighting).
    pub fn shadow_mask_view(&self) -> Option<&wgpu::TextureView> {
        self.shadow_mask_view.as_ref()
    }

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

        // Destroy old texture
        if let Some(t) = self.shadow_mask.take() {
            t.destroy();
        }
        self.shadow_mask_view = None;
        self.bind_group = None;

        // Recreate shadow mask at new resolution
        let mask_tex = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("rt_shadow_mask"),
            size: wgpu::Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8Unorm, // universally supported for storage
            usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::TEXTURE_BINDING,
            view_formats: &[],
        });
        self.shadow_mask_view = Some(mask_tex.create_view(&wgpu::TextureViewDescriptor::default()));
        self.shadow_mask = Some(mask_tex);

        self.width = width;
        self.height = height;
        self.bg_dirty = true;

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

    /// Drop all GPU resources.
    pub fn destroy(&mut self) {
        if let Some(t) = self.shadow_mask.take() {
            t.destroy();
        }
        if let Some(b) = self.uniform_buffer.take() {
            b.destroy();
        }
        self.shadow_mask_view = None;
        self.bind_group = None;
        self.initialized = false;
    }

    /// Estimate total VRAM usage in bytes for the given resolution.
    pub fn vram_estimate(width: u32, height: u32) -> u64 {
        // Shadow mask: R8Unorm = 1 byte per pixel
        let mask_bytes = width as u64 * height as u64;

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

        mask_bytes + uniform_bytes
    }
}

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

    #[test]
    fn rt_shadow_config_defaults_valid() {
        let config = RtShadowConfig::default();
        assert!(config.max_ray_distance > 0.0);
        assert!(config.shadow_bias > 0.0);
        assert!(config.shadow_bias < 1.0);
        assert!(config.samples_per_pixel >= 1);
    }

    #[test]
    fn rt_shadow_uniforms_size() {
        // 16 floats (view_proj_inv) + 3 floats (light_dir) + 1 float (max_ray_dist)
        // + 1 float (bias) + 1 float (intensity) + 2 u32 (screen_size)
        // = 16*4 + 3*4 + 4 + 4 + 4 + 8 = 64 + 12 + 4 + 4 + 4 + 8 = 96 bytes
        assert_eq!(std::mem::size_of::<RtShadowUniforms>(), 96);
    }

    #[test]
    fn vram_estimate_shadow_reasonable() {
        // 1920x1080 shadow mask
        let estimate = RtShadowPass::vram_estimate(1920, 1080);
        // R8Unorm at 1080p = ~2 MB, should be well under 10 MB
        let budget = 10 * 1024 * 1024;
        assert!(
            estimate < budget,
            "VRAM estimate {} bytes ({:.1} MB) exceeds 10 MB budget",
            estimate,
            estimate as f64 / (1024.0 * 1024.0)
        );
    }

    #[test]
    fn rt_shadow_default_disabled() {
        let pass = RtShadowPass::default();
        assert!(!pass.enabled);
        assert!(!pass.initialized);
        assert!(pass.shadow_mask_view().is_none());
    }

    #[test]
    fn rt_shadow_new_enabled() {
        let pass = RtShadowPass::new(RtShadowConfig::default());
        assert!(pass.enabled);
        assert!(!pass.initialized);
    }
}
