// Hi-Z — Hierarchical depth buffer for GPU occlusion culling.
//
// After the depth pre-pass, downsample the depth buffer into a mip chain
// where each texel stores the MAX depth of its 2×2 parent region.
// The GPU cull shader samples the appropriate mip level to test if an
// object's bounding sphere is fully occluded by closer geometry.

/// Hi-Z mip chain manager.
pub struct HizPass {
    pipeline: Option<wgpu::ComputePipeline>,
    bgl: Option<wgpu::BindGroupLayout>,
    /// Mip chain texture (R32Float, full mip chain).
    hiz_texture: Option<wgpu::Texture>,
    /// Views for each mip level (for writing).
    mip_views: Vec<wgpu::TextureView>,
    /// Full texture view (for reading in cull shader).
    pub hiz_view: Option<wgpu::TextureView>,
    sampler: Option<wgpu::Sampler>,
    width: u32,
    height: u32,
    mip_count: u32,
    pub enabled: bool,
    initialized: bool,
}

impl Default for HizPass {
    fn default() -> Self {
        Self {
            pipeline: None,
            bgl: None,
            hiz_texture: None,
            mip_views: Vec::new(),
            hiz_view: None,
            sampler: None,
            width: 0,
            height: 0,
            mip_count: 0,
            enabled: true,
            initialized: false,
        }
    }
}

impl HizPass {
    /// Initialize the Hi-Z mip chain for the given resolution.
    pub fn init(&mut self, device: &wgpu::Device, width: u32, height: u32) {
        if self.initialized && self.width == width && self.height == height {
            return;
        }

        let mip_count = (width.max(height) as f32).log2().ceil() as u32;
        let format = wgpu::TextureFormat::R32Float;

        let texture = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("hiz_depth"),
            size: wgpu::Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
            mip_level_count: mip_count,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format,
            usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::TEXTURE_BINDING,
            view_formats: &[],
        });

        // Create per-mip views for writing
        let mut mip_views = Vec::with_capacity(mip_count as usize);
        for mip in 0..mip_count {
            mip_views.push(texture.create_view(&wgpu::TextureViewDescriptor {
                label: Some(&format!("hiz_mip_{mip}")),
                base_mip_level: mip,
                mip_level_count: Some(1),
                ..Default::default()
            }));
        }

        // Full view for reading (all mips)
        let hiz_view = texture.create_view(&wgpu::TextureViewDescriptor {
            label: Some("hiz_full"),
            ..Default::default()
        });

        if self.sampler.is_none() {
            self.sampler = Some(device.create_sampler(&wgpu::SamplerDescriptor {
                label: Some("hiz_sampler"),
                mag_filter: wgpu::FilterMode::Nearest,
                min_filter: wgpu::FilterMode::Nearest,
                mipmap_filter: wgpu::MipmapFilterMode::Nearest,
                ..Default::default()
            }));
        }

        // Build compute pipeline for mip downsample
        let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("hiz_bgl"),
            entries: &[
                // binding 0: source mip (texture)
                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: destination mip (storage write)
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::StorageTexture {
                        access: wgpu::StorageTextureAccess::WriteOnly,
                        format,
                        view_dimension: wgpu::TextureViewDimension::D2,
                    },
                    count: None,
                },
            ],
        });

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

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

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

        self.hiz_texture = Some(texture);
        self.mip_views = mip_views;
        self.hiz_view = Some(hiz_view);
        self.bgl = Some(bgl);
        self.width = width;
        self.height = height;
        self.mip_count = mip_count;
        self.initialized = true;

        log::info!("Hi-Z initialized ({}×{}, {} mip levels)", width, height, mip_count);
    }

    /// Build the Hi-Z mip chain from the depth buffer.
    /// Call after the depth pre-pass or shadow pass.
    pub fn build_mip_chain(
        &self,
        device: &wgpu::Device,
        encoder: &mut wgpu::CommandEncoder,
        depth_mip0_view: &wgpu::TextureView,
    ) {
        if !self.enabled || !self.initialized {
            return;
        }
        let Some(ref pipeline) = self.pipeline else { return };
        let Some(ref bgl) = self.bgl else { return };

        // First pass: depth buffer → mip 0
        // Subsequent passes: mip N → mip N+1
        let sources: Vec<&wgpu::TextureView> = std::iter::once(depth_mip0_view).chain(self.mip_views.iter()).collect();

        for mip in 0..self.mip_count.saturating_sub(1) as usize {
            let src = sources[mip];
            let dst = &self.mip_views[mip + 1];

            let bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some(&format!("hiz_mip_{}_bg", mip)),
                layout: bgl,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: wgpu::BindingResource::TextureView(src),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: wgpu::BindingResource::TextureView(dst),
                    },
                ],
            });

            let dst_w = (self.width >> (mip as u32 + 1)).max(1);
            let dst_h = (self.height >> (mip as u32 + 1)).max(1);

            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some(&format!("hiz_mip_{}", mip + 1)),
                timestamp_writes: None,
            });
            pass.set_pipeline(pipeline);
            pass.set_bind_group(0, &bg, &[]);
            pass.dispatch_workgroups(dst_w.div_ceil(8), dst_h.div_ceil(8), 1);
        }
    }

    pub fn destroy(&mut self) {
        if let Some(t) = self.hiz_texture.take() {
            t.destroy();
        }
        self.mip_views.clear();
        self.hiz_view = None;
        self.initialized = false;
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn mip_count_1080p() {
        let mips = (1920u32.max(1080) as f32).log2().ceil() as u32;
        assert_eq!(mips, 11); // 2^11 = 2048 > 1920
    }

    #[test]
    fn mip_count_4k() {
        let mips = (3840u32.max(2160) as f32).log2().ceil() as u32;
        assert_eq!(mips, 12); // 2^12 = 4096 > 3840
    }
}
