// Ray tracing acceleration structure management.
// Wraps wgpu's experimental BLAS/TLAS API for scene-level ray queries.
// Requires Features::EXPERIMENTAL_RAY_QUERY.

use crate::dreamlet_catalog::GpuDreamlet;

/// Ray tracing context — manages acceleration structures for the scene.
#[derive(Default)]
pub struct RayTracingContext {
    pub enabled: bool,
    pub blas_list: Vec<BlasEntry>,
    pub tlas: Option<wgpu::Tlas>,
    pub max_instances: u32,
}

/// A bottom-level acceleration structure with its mesh metadata.
pub struct BlasEntry {
    pub blas: wgpu::Blas,
    pub label: String,
    pub vertex_count: u32,
    pub index_count: u32,
    /// Cached size descriptor for BLAS rebuilds.
    pub size_desc: wgpu::BlasTriangleGeometrySizeDescriptor,
}

impl RayTracingContext {
    /// Initialize ray tracing if the device supports it.
    pub fn init_if_supported(device: &wgpu::Device, features: wgpu::Features) -> Self {
        let enabled = features.contains(wgpu::Features::EXPERIMENTAL_RAY_QUERY);
        if !enabled {
            return Self::default();
        }

        let max_instances = 1024;
        let tlas = device.create_tlas(&wgpu::CreateTlasDescriptor {
            label: Some("tlas"),
            max_instances,
            flags: wgpu::AccelerationStructureFlags::PREFER_FAST_TRACE,
            update_mode: wgpu::AccelerationStructureUpdateMode::Build,
        });

        Self {
            enabled,
            blas_list: Vec::with_capacity(64),
            tlas: Some(tlas),
            max_instances,
        }
    }

    /// Re-initialize with a larger TLAS capacity. Call when instance count exceeds max_instances.
    pub fn resize_tlas(&mut self, device: &wgpu::Device, new_max: u32) {
        if !self.enabled {
            return;
        }
        self.max_instances = new_max;
        self.tlas = Some(device.create_tlas(&wgpu::CreateTlasDescriptor {
            label: Some("tlas"),
            max_instances: new_max,
            flags: wgpu::AccelerationStructureFlags::PREFER_FAST_TRACE,
            update_mode: wgpu::AccelerationStructureUpdateMode::Build,
        }));
    }

    /// Create a BLAS for triangle geometry.
    pub fn create_triangle_blas(
        &mut self,
        device: &wgpu::Device,
        label: &str,
        vertex_count: u32,
        index_count: Option<u32>,
        vertex_format: wgpu::VertexFormat,
    ) -> Option<usize> {
        if !self.enabled {
            return None;
        }

        if self.blas_list.len() >= self.max_instances as usize {
            return None;
        }

        let size_desc = wgpu::BlasTriangleGeometrySizeDescriptor {
            vertex_format,
            vertex_count,
            index_format: index_count.map(|_| wgpu::IndexFormat::Uint32),
            index_count,
            flags: wgpu::AccelerationStructureGeometryFlags::OPAQUE,
        };

        let blas = device.create_blas(
            &wgpu::CreateBlasDescriptor {
                label: Some(label),
                flags: wgpu::AccelerationStructureFlags::PREFER_FAST_TRACE,
                update_mode: wgpu::AccelerationStructureUpdateMode::Build,
            },
            wgpu::BlasGeometrySizeDescriptors::Triangles {
                descriptors: vec![size_desc.clone()],
            },
        );

        let idx = self.blas_list.len();
        self.blas_list.push(BlasEntry {
            blas,
            label: label.to_string(),
            vertex_count,
            index_count: index_count.unwrap_or(0),
            size_desc,
        });
        Some(idx)
    }

    /// Find an existing BLAS by label. Returns the index if found.
    pub fn find_blas(&self, label: &str) -> Option<usize> {
        self.blas_list.iter().position(|b| b.label == label)
    }

    /// Register a BLAS from existing GPU buffers (vertex + optional index).
    /// If a BLAS with the same label already exists, returns its existing index (deduplication).
    /// Builds the BLAS immediately using the provided encoder.
    /// Returns the BLAS index for use with `set_instance`.
    pub fn register_primitive_blas(
        &mut self,
        device: &wgpu::Device,
        encoder: &mut wgpu::CommandEncoder,
        label: &str,
        vertex_buffer: &wgpu::Buffer,
        vertex_count: u32,
        vertex_stride: u64,
        index_buffer: Option<&wgpu::Buffer>,
        index_count: Option<u32>,
    ) -> Option<usize> {
        if !self.enabled {
            return None;
        }

        // Deduplication: return existing BLAS if one with this label already exists.
        // Prevents duplicate VRAM allocation when the same primitive is registered multiple times.
        if let Some(existing) = self.find_blas(label) {
            return Some(existing);
        }

        let size_desc = wgpu::BlasTriangleGeometrySizeDescriptor {
            vertex_format: wgpu::VertexFormat::Float32x3,
            vertex_count,
            index_format: index_count.map(|_| wgpu::IndexFormat::Uint32),
            index_count,
            flags: wgpu::AccelerationStructureGeometryFlags::OPAQUE,
        };

        let blas = device.create_blas(
            &wgpu::CreateBlasDescriptor {
                label: Some(label),
                flags: wgpu::AccelerationStructureFlags::PREFER_FAST_TRACE,
                update_mode: wgpu::AccelerationStructureUpdateMode::Build,
            },
            wgpu::BlasGeometrySizeDescriptors::Triangles {
                descriptors: vec![size_desc.clone()],
            },
        );

        // Build immediately
        let geometry = wgpu::BlasTriangleGeometry {
            size: &size_desc,
            vertex_buffer,
            first_vertex: 0,
            vertex_stride,
            index_buffer,
            first_index: index_count.map(|_| 0),
            transform_buffer: None,
            transform_buffer_offset: None,
        };

        let build_entry = wgpu::BlasBuildEntry {
            blas: &blas,
            geometry: wgpu::BlasGeometries::TriangleGeometries(vec![geometry]),
        };

        if let Some(tlas) = &self.tlas {
            encoder.build_acceleration_structures(std::iter::once(&build_entry), std::iter::once(tlas));
        }

        let idx = self.blas_list.len();
        self.blas_list.push(BlasEntry {
            blas,
            label: label.to_string(),
            vertex_count,
            index_count: index_count.unwrap_or(0),
            size_desc,
        });
        Some(idx)
    }

    /// Populate TLAS instances from a slice of GpuDreamlets.
    /// Each dreamlet's position, rotation (quaternion), and scale are converted
    /// to a 3x4 row-major affine matrix for the TLAS instance.
    ///
    /// When `blas_index` is provided, all dreamlets use that single BLAS.
    /// When `None`, each dreamlet's `mesh_id / 3` selects its BLAS (per-primitive mapping).
    /// Returns the number of instances set.
    pub fn sync_tlas_from_dreamlets(&mut self, dreamlets: &[GpuDreamlet], blas_index: usize) -> u32 {
        self.sync_tlas_from_dreamlets_mapped(dreamlets, Some(blas_index))
    }

    /// Populate TLAS instances with per-dreamlet BLAS selection.
    /// Each dreamlet's `mesh_id / 3` maps to a BLAS index (matching PrimitiveKind ordering).
    /// Falls back to BLAS 0 if the computed index exceeds the BLAS list.
    pub fn sync_tlas_from_dreamlets_mapped(&mut self, dreamlets: &[GpuDreamlet], default_blas: Option<usize>) -> u32 {
        if !self.enabled || self.blas_list.is_empty() {
            return 0;
        }

        let count = dreamlets.len().min(self.max_instances as usize);
        let blas_count = self.blas_list.len();

        if let Some(tlas) = &mut self.tlas {
            for (i, d) in dreamlets.iter().take(count).enumerate() {
                let transform = quat_pos_scale_to_affine(d.rotation, d.position, d.scale);
                // Per-primitive BLAS: mesh_id / 3 = kind_index (matches PrimitiveKind::ALL ordering)
                let bi = default_blas.unwrap_or_else(|| {
                    let kind_index = (d.mesh_id / 3) as usize;
                    kind_index.min(blas_count - 1)
                });
                let bi = bi.min(blas_count - 1);
                if let Some(slot) = tlas.get_mut_single(i) {
                    *slot = Some(wgpu::TlasInstance::new(
                        &self.blas_list[bi].blas,
                        transform,
                        i as u32,
                        0xFF,
                    ));
                }
            }
            // Clear remaining slots
            for i in count..self.max_instances as usize {
                if let Some(slot) = tlas.get_mut_single(i) {
                    *slot = None;
                }
            }
        }

        count as u32
    }

    /// Create a BLAS that supports per-frame vertex updates (for animated/skinned content).
    /// Uses `AccelerationStructureUpdateMode::Update` for efficient refitting.
    pub fn create_updatable_blas(
        &mut self,
        device: &wgpu::Device,
        label: &str,
        vertex_count: u32,
        index_count: Option<u32>,
        vertex_format: wgpu::VertexFormat,
    ) -> Option<usize> {
        if !self.enabled {
            return None;
        }

        let size_desc = wgpu::BlasTriangleGeometrySizeDescriptor {
            vertex_format,
            vertex_count,
            index_format: index_count.map(|_| wgpu::IndexFormat::Uint32),
            index_count,
            flags: wgpu::AccelerationStructureGeometryFlags::OPAQUE,
        };

        let blas = device.create_blas(
            &wgpu::CreateBlasDescriptor {
                label: Some(label),
                flags: wgpu::AccelerationStructureFlags::PREFER_FAST_BUILD,
                update_mode: wgpu::AccelerationStructureUpdateMode::PreferUpdate,
            },
            wgpu::BlasGeometrySizeDescriptors::Triangles {
                descriptors: vec![size_desc.clone()],
            },
        );

        let idx = self.blas_list.len();
        self.blas_list.push(BlasEntry {
            blas,
            label: label.to_string(),
            vertex_count,
            index_count: index_count.unwrap_or(0),
            size_desc,
        });
        Some(idx)
    }

    /// Update an existing BLAS with new vertex data (for animated content).
    /// The BLAS must have been created with `create_updatable_blas`.
    pub fn update_blas(
        &self,
        encoder: &mut wgpu::CommandEncoder,
        blas_index: usize,
        vertex_buffer: &wgpu::Buffer,
        vertex_stride: u64,
        index_buffer: Option<&wgpu::Buffer>,
    ) {
        if !self.enabled || blas_index >= self.blas_list.len() {
            return;
        }
        let entry = &self.blas_list[blas_index];
        let geometry = wgpu::BlasTriangleGeometry {
            size: &entry.size_desc,
            vertex_buffer,
            first_vertex: 0,
            vertex_stride,
            index_buffer,
            first_index: entry.size_desc.index_count.map(|_| 0),
            transform_buffer: None,
            transform_buffer_offset: None,
        };
        let build_entry = wgpu::BlasBuildEntry {
            blas: &entry.blas,
            geometry: wgpu::BlasGeometries::TriangleGeometries(vec![geometry]),
        };
        if let Some(tlas) = &self.tlas {
            encoder.build_acceleration_structures(std::iter::once(&build_entry), std::iter::once(tlas));
        }
    }

    /// Rebuild TLAS after instance updates. Call after `sync_tlas_from_dreamlets`.
    pub fn rebuild_tlas(&self, encoder: &mut wgpu::CommandEncoder) {
        if !self.enabled {
            return;
        }
        if let Some(tlas) = &self.tlas {
            encoder.build_acceleration_structures(std::iter::empty(), std::iter::once(tlas));
        }
    }

    /// Remove a BLAS entry by index. Returns `true` if the entry was removed.
    pub fn remove_blas(&mut self, index: usize) -> bool {
        if index < self.blas_list.len() {
            self.blas_list.remove(index);
            true
        } else {
            false
        }
    }

    /// Build acceleration structures. Call after geometry buffers are uploaded.
    pub fn build_structures(&self, encoder: &mut wgpu::CommandEncoder, build_entries: &[wgpu::BlasBuildEntry<'_>]) {
        if !self.enabled {
            return;
        }
        if let Some(tlas) = &self.tlas {
            encoder.build_acceleration_structures(build_entries.iter(), std::iter::once(tlas));
        }
    }

    /// Set a TLAS instance. Transform is a row-major 3x4 affine matrix.
    pub fn set_instance(&mut self, index: usize, blas_index: usize, transform: [f32; 12], custom_data: u32, mask: u8) {
        if let Some(tlas) = &mut self.tlas {
            if let Some(slot) = tlas.get_mut_single(index) {
                if blas_index < self.blas_list.len() {
                    *slot = Some(wgpu::TlasInstance::new(
                        &self.blas_list[blas_index].blas,
                        transform,
                        custom_data,
                        mask,
                    ));
                }
            }
        }
    }

    /// Clear all TLAS instances.
    pub fn clear_instances(&mut self) {
        if let Some(tlas) = &mut self.tlas {
            if let Some(slice) = tlas.get_mut_slice(0..self.max_instances as usize) {
                for slot in slice.iter_mut() {
                    *slot = None;
                }
            }
        }
    }

    /// Get TLAS as a binding resource for use in bind groups (ray query in shaders).
    pub fn tlas_binding(&self) -> Option<wgpu::BindingResource<'_>> {
        self.tlas.as_ref().map(|t| t.as_binding())
    }

    /// Get a reference to the TLAS (for RT passes that need it).
    pub fn tlas(&self) -> Option<&wgpu::Tlas> {
        self.tlas.as_ref()
    }

    /// DL-09: Register LOD geometry in TLAS for far-field traces.
    /// Simplified geometry (low-poly) for distant objects improves RT perf.
    pub fn register_lod_blas(
        &mut self,
        device: &wgpu::Device,
        encoder: &mut wgpu::CommandEncoder,
        label: &str,
        vertex_buf: &wgpu::Buffer,
        vertex_count: u32,
        vertex_stride: u64,
        index_buf: Option<&wgpu::Buffer>,
        index_count: Option<u32>,
    ) {
        // Delegate to existing register_primitive_blas with "lod_" prefix
        self.register_primitive_blas(
            device,
            encoder,
            &format!("lod_{}", label),
            vertex_buf,
            vertex_count,
            vertex_stride,
            index_buf,
            index_count,
        );
    }

    /// Estimated VRAM usage for acceleration structures.
    /// Approximation: ~256 bytes per BLAS primitive + 64 bytes per TLAS instance + scratch.
    pub fn vram_estimate(&self) -> u64 {
        let blas_vram: u64 = self
            .blas_list
            .iter()
            .map(|b| {
                let prims = if b.index_count > 0 {
                    b.index_count / 3
                } else {
                    b.vertex_count / 3
                };
                prims as u64 * 256
            })
            .sum();
        let tlas_vram = self.max_instances as u64 * 64;
        blas_vram + tlas_vram
    }
}

/// Convert quaternion (x, y, z, w) + position + uniform scale to 3x4 row-major affine.
pub fn quat_pos_scale_to_affine(quat: [f32; 4], pos: [f32; 3], scale: f32) -> [f32; 12] {
    let [qx, qy, qz, qw] = quat;

    // Rotation matrix from quaternion (scaled by uniform scale)
    let x2 = qx + qx;
    let y2 = qy + qy;
    let z2 = qz + qz;
    let xx = qx * x2;
    let xy = qx * y2;
    let xz = qx * z2;
    let yy = qy * y2;
    let yz = qy * z2;
    let zz = qz * z2;
    let wx = qw * x2;
    let wy = qw * y2;
    let wz = qw * z2;

    [
        (1.0 - yy - zz) * scale,
        (xy - wz) * scale,
        (xz + wy) * scale,
        pos[0],
        (xy + wz) * scale,
        (1.0 - xx - zz) * scale,
        (yz - wx) * scale,
        pos[1],
        (xz - wy) * scale,
        (yz + wx) * scale,
        (1.0 - xx - yy) * scale,
        pos[2],
    ]
}

/// Identity transform for TLAS instances (3x4 row-major).
pub const IDENTITY_TRANSFORM: [f32; 12] = [
    1.0, 0.0, 0.0, 0.0, // row 0
    0.0, 1.0, 0.0, 0.0, // row 1
    0.0, 0.0, 1.0, 0.0, // row 2
];

/// Build a 3x4 translation transform for a TLAS instance.
pub fn translation_transform(x: f32, y: f32, z: f32) -> [f32; 12] {
    [
        1.0, 0.0, 0.0, x, // row 0
        0.0, 1.0, 0.0, y, // row 1
        0.0, 0.0, 1.0, z, // row 2
    ]
}

/// Build a 3x4 scale + translation transform for a TLAS instance.
pub fn scale_translation_transform(sx: f32, sy: f32, sz: f32, tx: f32, ty: f32, tz: f32) -> [f32; 12] {
    [sx, 0.0, 0.0, tx, 0.0, sy, 0.0, ty, 0.0, 0.0, sz, tz]
}

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

    #[test]
    fn default_disabled() {
        let ctx = RayTracingContext::default();
        assert!(!ctx.enabled);
        assert!(ctx.tlas.is_none());
    }

    #[test]
    fn identity_transform_values() {
        assert_eq!(IDENTITY_TRANSFORM[0], 1.0);
        assert_eq!(IDENTITY_TRANSFORM[5], 1.0);
        assert_eq!(IDENTITY_TRANSFORM[10], 1.0);
    }

    #[test]
    fn translation_transform_values() {
        let t = translation_transform(5.0, 10.0, 15.0);
        assert_eq!(t[3], 5.0);
        assert_eq!(t[7], 10.0);
        assert_eq!(t[11], 15.0);
    }

    #[test]
    fn scale_translation_transform_values() {
        let t = scale_translation_transform(2.0, 3.0, 4.0, 1.0, 2.0, 3.0);
        assert_eq!(t[0], 2.0);
        assert_eq!(t[5], 3.0);
        assert_eq!(t[10], 4.0);
        assert_eq!(t[3], 1.0);
        assert_eq!(t[7], 2.0);
        assert_eq!(t[11], 3.0);
    }

    #[test]
    fn quat_identity_produces_identity_rotation() {
        let t = quat_pos_scale_to_affine([0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0], 1.0);
        // Should be identity rotation with zero translation
        assert!((t[0] - 1.0).abs() < 1e-6);
        assert!((t[5] - 1.0).abs() < 1e-6);
        assert!((t[10] - 1.0).abs() < 1e-6);
        assert!((t[3]).abs() < 1e-6);
        assert!((t[7]).abs() < 1e-6);
        assert!((t[11]).abs() < 1e-6);
    }

    #[test]
    fn quat_with_scale_and_translation() {
        let t = quat_pos_scale_to_affine([0.0, 0.0, 0.0, 1.0], [1.0, 2.0, 3.0], 2.0);
        assert!((t[0] - 2.0).abs() < 1e-6); // scale applied
        assert!((t[5] - 2.0).abs() < 1e-6);
        assert!((t[10] - 2.0).abs() < 1e-6);
        assert!((t[3] - 1.0).abs() < 1e-6); // translation
        assert!((t[7] - 2.0).abs() < 1e-6);
        assert!((t[11] - 3.0).abs() < 1e-6);
    }

    #[test]
    fn vram_estimate_empty() {
        let ctx = RayTracingContext {
            enabled: true,
            blas_list: vec![],
            tlas: None,
            max_instances: 1024,
        };
        assert_eq!(ctx.vram_estimate(), 1024 * 64);
    }
}
