// Surface Cache — orthographic material capture atlas for RT GI hit lighting.
//
// Captures per-mesh material data into a 4096x4096 Rgba8Unorm atlas.
// Each card is a 32x32 tile representing one mesh's material from an
// orthographic view. Amortized capture: 1-2 meshes per frame to avoid
// frame spikes. Dreamlet content bypasses capture (128B GpuDreamlet
// already has material data indexed by instance_custom_data).
//
// The atlas is bound as a texture in the RT GI hit lighting path (DL-04)
// for non-dreamlet meshes that require full material evaluation.

/// Per-card metadata (96 bytes, #[repr(C)], Pod).
///
/// Describes one 32x32 tile in the atlas: which mesh it captures, where
/// in the atlas it lives, and the orthographic view-projection used to
/// rasterize the mesh's material into the tile.
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct CardMetadata {
    /// Orthographic capture view-projection matrix.
    pub capture_view_proj: [f32; 16],
    /// UV offset in atlas (normalized 0..1).
    pub atlas_uv_offset: [f32; 2],
    /// UV scale in atlas (normalized).
    pub atlas_uv_scale: [f32; 2],
    /// Mesh ID this card represents.
    pub mesh_id: u32,
    /// Frame when this card was last updated.
    pub last_update_frame: u32,
    pub _pad: [f32; 2],
}

/// Maximum number of cards in the atlas (128 columns x 128 rows of 32x32 tiles).
const MAX_CARDS: u32 = 128 * 128; // 16384 cards at 32x32 in a 4096x4096 atlas

/// Default card resolution in pixels (each tile is 32x32).
const CARD_RESOLUTION: u32 = 32;

/// Atlas texture size in pixels.
const ATLAS_SIZE: u32 = 4096;

/// Surface cache atlas manager.
///
/// Owns the atlas texture, per-card metadata buffer, and the compute
/// pipeline that captures mesh materials into atlas tiles. Amortized
/// capture dispatches 1-2 cards per frame via a round-robin index.
pub struct SurfaceCacheAtlas {
    pub enabled: bool,
    atlas_texture: Option<wgpu::Texture>,
    atlas_view: Option<wgpu::TextureView>,
    card_buffer: Option<wgpu::Buffer>,
    capture_pipeline: Option<wgpu::ComputePipeline>,
    bgl: Option<wgpu::BindGroupLayout>,
    max_cards: u32,
    card_resolution: u32,
    atlas_size: u32,
    next_capture_index: u32,
    frame_index: u32,
    initialized: bool,
}

impl Default for SurfaceCacheAtlas {
    fn default() -> Self {
        Self {
            enabled: false,
            atlas_texture: None,
            atlas_view: None,
            card_buffer: None,
            capture_pipeline: None,
            bgl: None,
            max_cards: MAX_CARDS,
            card_resolution: CARD_RESOLUTION,
            atlas_size: ATLAS_SIZE,
            next_capture_index: 0,
            frame_index: 0,
            initialized: false,
        }
    }
}

impl SurfaceCacheAtlas {
    /// Create a new surface cache atlas (disabled by default).
    pub fn new() -> Self {
        Self::default()
    }

    /// Initialize GPU resources: atlas texture, card metadata buffer.
    /// Does not create the capture pipeline (call `create_pipeline` separately).
    pub fn init(&mut self, device: &wgpu::Device) {
        if self.initialized {
            return;
        }

        // Atlas texture: 4096x4096 Rgba8Unorm (storage + sampled).
        let texture = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("surface_cache_atlas"),
            size: wgpu::Extent3d {
                width: self.atlas_size,
                height: self.atlas_size,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8Unorm,
            usage: wgpu::TextureUsages::STORAGE_BINDING
                | wgpu::TextureUsages::TEXTURE_BINDING
                | wgpu::TextureUsages::COPY_DST,
            view_formats: &[],
        });
        self.atlas_view = Some(texture.create_view(&wgpu::TextureViewDescriptor::default()));
        self.atlas_texture = Some(texture);

        // Card metadata buffer: max_cards * 96 bytes.
        self.card_buffer = Some(device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("surface_cache_cards"),
            size: self.max_cards as u64 * std::mem::size_of::<CardMetadata>() as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        }));

        self.initialized = true;
        log::info!(
            "Surface cache initialized ({}x{} atlas, {} max cards, {}x{} per card)",
            self.atlas_size,
            self.atlas_size,
            self.max_cards,
            self.card_resolution,
            self.card_resolution
        );
    }

    /// Create the capture compute pipeline from surface_cache_capture.wgsl.
    pub fn create_pipeline(&mut self, device: &wgpu::Device) {
        // Bind group layout: uniform params + material storage + atlas storage texture.
        let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("surface_cache_bgl"),
            entries: &[
                // binding 0: CaptureParams uniform
                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: materials storage (read)
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
                // binding 2: atlas storage texture (write)
                wgpu::BindGroupLayoutEntry {
                    binding: 2,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::StorageTexture {
                        access: wgpu::StorageTextureAccess::WriteOnly,
                        format: wgpu::TextureFormat::Rgba8Unorm,
                        view_dimension: wgpu::TextureViewDimension::D2,
                    },
                    count: None,
                },
            ],
        });

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

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

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

        self.bgl = Some(bgl);

        log::info!("Surface cache capture pipeline created");
    }

    /// Amortized capture: dispatch `count` card captures this frame.
    ///
    /// Each capture writes a 32x32 tile into the atlas at the next
    /// round-robin slot. The caller provides the command encoder;
    /// this method records compute dispatches into it.
    ///
    /// Typical usage: call with `count = 1` or `count = 2` per frame
    /// to spread capture cost across many frames.
    pub fn capture_cards(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
        count: u32,
    ) {
        if !self.enabled || !self.initialized {
            return;
        }
        let Some(ref pipeline) = self.capture_pipeline else {
            return;
        };
        let Some(ref bgl) = self.bgl else { return };
        let Some(ref atlas_view) = self.atlas_view else {
            return;
        };
        let Some(ref card_buffer) = self.card_buffer else {
            return;
        };

        let cards_per_row = self.atlas_size / self.card_resolution;

        for i in 0..count {
            let card_index = (self.next_capture_index + i) % self.max_cards;
            let tile_col = card_index % cards_per_row;
            let tile_row = card_index / cards_per_row;
            let atlas_offset_x = tile_col * self.card_resolution;
            let atlas_offset_y = tile_row * self.card_resolution;

            // Upload per-card capture params as a small uniform buffer.
            // CaptureParams layout: mat4x4 (64B) + vec2<u32> (8B) + u32 (4B) + u32 (4B) = 80 bytes.
            #[repr(C)]
            #[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
            struct CaptureParams {
                view_proj: [f32; 16],
                atlas_offset: [u32; 2],
                tile_size: u32,
                mesh_id: u32,
            }

            // Identity orthographic capture (placeholder — real capture uses
            // per-mesh bounding box to compute the orthographic projection).
            #[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,
            ];

            let params = CaptureParams {
                view_proj: identity,
                atlas_offset: [atlas_offset_x, atlas_offset_y],
                tile_size: self.card_resolution,
                mesh_id: card_index,
            };

            // Create a temporary uniform buffer for this dispatch.
            let param_buf = device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("surface_cache_params"),
                size: std::mem::size_of::<CaptureParams>() as u64,
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            queue.write_buffer(&param_buf, 0, bytemuck::bytes_of(&params));

            let bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("surface_cache_bg"),
                layout: bgl,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: param_buf.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: card_buffer.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 2,
                        resource: wgpu::BindingResource::TextureView(atlas_view),
                    },
                ],
            });

            // Dispatch: ceil(tile_size / 8) workgroups in x and y.
            let wg = self.card_resolution.div_ceil(8);
            {
                let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                    label: Some("surface_cache_capture"),
                    timestamp_writes: None,
                });
                pass.set_pipeline(pipeline);
                pass.set_bind_group(0, &bg, &[]);
                pass.dispatch_workgroups(wg, wg, 1);
            }

            // Update card metadata in the buffer.
            let uv_scale_x = self.card_resolution as f32 / self.atlas_size as f32;
            let uv_scale_y = self.card_resolution as f32 / self.atlas_size as f32;
            let uv_offset_x = atlas_offset_x as f32 / self.atlas_size as f32;
            let uv_offset_y = atlas_offset_y as f32 / self.atlas_size as f32;

            let metadata = CardMetadata {
                capture_view_proj: identity,
                atlas_uv_offset: [uv_offset_x, uv_offset_y],
                atlas_uv_scale: [uv_scale_x, uv_scale_y],
                mesh_id: card_index,
                last_update_frame: self.frame_index,
                _pad: [0.0; 2],
            };

            let offset = card_index as u64 * std::mem::size_of::<CardMetadata>() as u64;
            queue.write_buffer(card_buffer, offset, bytemuck::bytes_of(&metadata));
        }

        self.next_capture_index = (self.next_capture_index + count) % self.max_cards;
        self.frame_index = self.frame_index.wrapping_add(1);
    }

    /// Get the atlas texture view for external binding (RT GI hit lighting).
    pub fn atlas_view(&self) -> Option<&wgpu::TextureView> {
        self.atlas_view.as_ref()
    }

    /// Get the card metadata buffer for external binding.
    pub fn card_buffer(&self) -> Option<&wgpu::Buffer> {
        self.card_buffer.as_ref()
    }

    /// Estimate total VRAM usage in bytes.
    pub fn vram_estimate(&self) -> u64 {
        // Atlas: 4096 x 4096 x 4 bytes (Rgba8Unorm)
        let atlas_bytes = self.atlas_size as u64 * self.atlas_size as u64 * 4;
        // Card metadata: max_cards x 96 bytes
        let card_bytes = self.max_cards as u64 * std::mem::size_of::<CardMetadata>() as u64;
        atlas_bytes + card_bytes
    }

    /// Drop all GPU resources.
    pub fn destroy(&mut self) {
        if let Some(t) = self.atlas_texture.take() {
            t.destroy();
        }
        if let Some(b) = self.card_buffer.take() {
            b.destroy();
        }
        self.atlas_view = None;
        self.capture_pipeline = None;
        self.bgl = None;
        self.next_capture_index = 0;
        self.initialized = false;
    }
}

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

    #[test]
    fn card_metadata_size() {
        assert_eq!(std::mem::size_of::<CardMetadata>(), 96);
    }

    #[test]
    fn card_metadata_alignment() {
        assert_eq!(std::mem::align_of::<CardMetadata>(), 4);
    }

    #[test]
    fn default_state() {
        let cache = SurfaceCacheAtlas::default();
        assert!(!cache.enabled);
        assert!(!cache.initialized);
        assert!(cache.atlas_view().is_none());
        assert!(cache.card_buffer().is_none());
        assert_eq!(cache.max_cards, MAX_CARDS);
        assert_eq!(cache.card_resolution, CARD_RESOLUTION);
        assert_eq!(cache.atlas_size, ATLAS_SIZE);
        assert_eq!(cache.next_capture_index, 0);
        assert_eq!(cache.frame_index, 0);
    }

    #[test]
    fn new_returns_default() {
        let a = SurfaceCacheAtlas::new();
        let b = SurfaceCacheAtlas::default();
        assert_eq!(a.enabled, b.enabled);
        assert_eq!(a.max_cards, b.max_cards);
        assert_eq!(a.atlas_size, b.atlas_size);
    }

    #[test]
    fn vram_estimate_correct() {
        let cache = SurfaceCacheAtlas::default();
        let estimate = cache.vram_estimate();

        // Atlas: 4096 * 4096 * 4 = 67,108,864 bytes (64 MB)
        let atlas_bytes = ATLAS_SIZE as u64 * ATLAS_SIZE as u64 * 4;
        assert_eq!(atlas_bytes, 67_108_864);

        // Cards: 16384 * 96 = 1,572,864 bytes (~1.5 MB)
        let card_bytes = MAX_CARDS as u64 * 96;
        assert_eq!(card_bytes, 1_572_864);

        assert_eq!(estimate, atlas_bytes + card_bytes);

        // Total should be under 70 MB
        let budget = 70 * 1024 * 1024;
        assert!(
            estimate < budget,
            "VRAM estimate {} bytes ({:.1} MB) exceeds 70 MB budget",
            estimate,
            estimate as f64 / (1024.0 * 1024.0)
        );
    }

    #[test]
    fn atlas_tile_coverage() {
        // Verify that max_cards tiles exactly cover the atlas
        let cards_per_row = ATLAS_SIZE / CARD_RESOLUTION;
        let cards_per_col = ATLAS_SIZE / CARD_RESOLUTION;
        assert_eq!(cards_per_row, 128);
        assert_eq!(cards_per_col, 128);
        assert_eq!(MAX_CARDS, cards_per_row * cards_per_col);
    }

    #[test]
    fn card_metadata_pod_zeroed() {
        use bytemuck::Zeroable;
        // Verify Pod + Zeroable derive works (all-zeros is valid)
        let zeroed = CardMetadata::zeroed();
        assert_eq!(zeroed.mesh_id, 0);
        assert_eq!(zeroed.last_update_frame, 0);
        assert_eq!(zeroed.atlas_uv_offset, [0.0, 0.0]);
        assert_eq!(zeroed.atlas_uv_scale, [0.0, 0.0]);
    }
}
