// TAA — Temporal Anti-Aliasing pass manager.
// Jitters projection matrix per frame, resolves with history clamping.
// Ping-pong history buffers for stable accumulation.

/// TAA uniform parameters (matches taa.wgsl TaaParams).
#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
pub struct TaaParams {
    pub screen_size: [f32; 2],
    pub jitter: [f32; 2],
    pub prev_vp: [f32; 16],
    pub curr_inv_vp: [f32; 16],
}

/// Halton sequence for sub-pixel jitter (low-discrepancy, avoids patterns).
pub fn halton(index: u32, base: u32) -> f32 {
    let mut result = 0.0f32;
    let mut f = 1.0f32;
    let mut i = index;
    while i > 0 {
        f /= base as f32;
        result += f * (i % base) as f32;
        i /= base;
    }
    result
}

/// TAA pass: jittered render → temporal resolve → stable output.
pub struct TaaPass {
    pipeline: Option<wgpu::ComputePipeline>,
    bgl: Option<wgpu::BindGroupLayout>,
    params_buf: Option<wgpu::Buffer>,
    // Ping-pong history buffers
    history_a: Option<wgpu::Texture>,
    history_a_view: Option<wgpu::TextureView>,
    history_b: Option<wgpu::Texture>,
    history_b_view: Option<wgpu::TextureView>,
    /// Which history buffer was last written (false=A, true=B).
    write_b: bool,
    sampler: Option<wgpu::Sampler>,
    /// 1x1 fallback for motion vectors when not available (shader detects via dimensions).
    motion_fallback: Option<wgpu::Texture>,
    motion_fallback_view: Option<wgpu::TextureView>,
    width: u32,
    height: u32,
    frame_index: u32,
    prev_vp: [f32; 16],
    pub enabled: bool,
    initialized: bool,
    /// Cached resolve bind group (invalidated on resize).
    cached_resolve_bg: Option<wgpu::BindGroup>,
}

impl Default for TaaPass {
    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 {
            pipeline: None,
            bgl: None,
            params_buf: None,
            history_a: None,
            history_a_view: None,
            history_b: None,
            history_b_view: None,
            write_b: false,
            sampler: None,
            motion_fallback: None,
            motion_fallback_view: None,
            width: 0,
            height: 0,
            frame_index: 0,
            prev_vp: identity,
            enabled: true,
            initialized: false,
            cached_resolve_bg: None,
        }
    }
}

impl TaaPass {
    /// Compute the sub-pixel jitter for this frame.
    pub fn jitter(&self) -> (f32, f32) {
        let idx = (self.frame_index % 16) + 1;
        let jx = halton(idx, 2) * 2.0 - 1.0;
        let jy = halton(idx, 3) * 2.0 - 1.0;
        (jx, jy)
    }

    /// Apply jitter to a projection matrix. Call before rendering the main pass.
    pub fn apply_jitter(&self, proj: &mut glam::Mat4, width: f32, height: f32) {
        if !self.enabled {
            return;
        }
        let (jx, jy) = self.jitter();
        let mut m = *proj;
        // Add sub-pixel offset to projection matrix
        *m.col_mut(2) += glam::Vec4::new(jx / width, jy / height, 0.0, 0.0);
        *proj = m;
    }

    /// Initialize pipelines and allocate history buffers.
    pub fn init(&mut self, device: &wgpu::Device, width: u32, height: u32) {
        if self.initialized && self.width == width && self.height == height {
            return;
        }

        let format = wgpu::TextureFormat::Rgba16Float;

        let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("taa_bgl"),
            entries: &[
                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,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Float { filterable: true },
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 2,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Float { filterable: true },
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 3,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Depth,
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 4,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 5,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::StorageTexture {
                        access: wgpu::StorageTextureAccess::WriteOnly,
                        format,
                        view_dimension: wgpu::TextureViewDimension::D2,
                    },
                    count: None,
                },
                // TD-007: Motion vector texture for per-pixel reprojection
                wgpu::BindGroupLayoutEntry {
                    binding: 6,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Float { filterable: true },
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
            ],
        });

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

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

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

        // Allocate ping-pong history buffers
        let make_history = |label: &str| -> (wgpu::Texture, wgpu::TextureView) {
            let tex = 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,
                usage: wgpu::TextureUsages::STORAGE_BINDING
                    | wgpu::TextureUsages::TEXTURE_BINDING
                    | wgpu::TextureUsages::COPY_SRC,
                view_formats: &[],
            });
            let view = tex.create_view(&wgpu::TextureViewDescriptor::default());
            (tex, view)
        };
        let (ha, hav) = make_history("taa_history_a");
        let (hb, hbv) = make_history("taa_history_b");
        self.history_a = Some(ha);
        self.history_a_view = Some(hav);
        self.history_b = Some(hb);
        self.history_b_view = Some(hbv);

        if self.sampler.is_none() {
            self.sampler = Some(device.create_sampler(&wgpu::SamplerDescriptor {
                label: Some("taa_sampler"),
                mag_filter: wgpu::FilterMode::Linear,
                min_filter: wgpu::FilterMode::Linear,
                ..Default::default()
            }));
        }
        if self.params_buf.is_none() {
            self.params_buf = Some(device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("taa_params"),
                size: std::mem::size_of::<TaaParams>() as u64,
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            }));
        }

        // 1x1 motion vector fallback (shader detects via dimensions > 1)
        if self.motion_fallback.is_none() {
            let tex = device.create_texture(&wgpu::TextureDescriptor {
                label: Some("taa_motion_fallback"),
                size: wgpu::Extent3d {
                    width: 1,
                    height: 1,
                    depth_or_array_layers: 1,
                },
                mip_level_count: 1,
                sample_count: 1,
                dimension: wgpu::TextureDimension::D2,
                format: wgpu::TextureFormat::Rgba8Unorm, // universally filterable
                usage: wgpu::TextureUsages::TEXTURE_BINDING,
                view_formats: &[],
            });
            self.motion_fallback_view = Some(tex.create_view(&wgpu::TextureViewDescriptor::default()));
            self.motion_fallback = Some(tex);
        }

        self.bgl = Some(bgl);
        self.width = width;
        self.height = height;
        self.initialized = true;
        self.cached_resolve_bg = None;

        log::info!(
            "TAA initialized ({}x{}, Halton 2,3 jitter, 16-sample sequence)",
            width,
            height
        );
    }

    /// Dispatch TAA resolve. Reads current HDR + depth, writes to history buffer.
    /// Returns the resolved output view.
    pub fn dispatch(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
        current_view: &wgpu::TextureView,
        depth_view: &wgpu::TextureView,
        camera: &crate::camera::Camera,
    ) -> Option<&wgpu::TextureView> {
        if !self.enabled || !self.initialized {
            return None;
        }
        let Some(ref pipeline) = self.pipeline else { return None };
        let Some(ref bgl) = self.bgl else { return None };
        let Some(ref params_buf) = self.params_buf else {
            return None;
        };
        let Some(ref sampler) = self.sampler else { return None };

        // Select read/write history buffers (ping-pong)
        let (read_view, write_view) = if self.write_b {
            (self.history_a_view.as_ref()?, self.history_b_view.as_ref()?)
        } else {
            (self.history_b_view.as_ref()?, self.history_a_view.as_ref()?)
        };

        let (jx, jy) = self.jitter();
        let curr_vp = (camera.projection_matrix * camera.view_matrix).to_cols_array();
        let curr_inv_vp = (camera.projection_matrix * camera.view_matrix)
            .inverse()
            .to_cols_array();

        let params = TaaParams {
            screen_size: [self.width as f32, self.height as f32],
            jitter: [jx, jy],
            prev_vp: self.prev_vp,
            curr_inv_vp,
        };
        queue.write_buffer(params_buf, 0, bytemuck::bytes_of(&params));

        let bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("taa_bg"),
            layout: bgl,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: params_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(current_view),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::TextureView(read_view),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: wgpu::BindingResource::TextureView(depth_view),
                },
                wgpu::BindGroupEntry {
                    binding: 4,
                    resource: wgpu::BindingResource::Sampler(sampler),
                },
                wgpu::BindGroupEntry {
                    binding: 5,
                    resource: wgpu::BindingResource::TextureView(write_view),
                },
                // TD-007: Motion vector fallback (1x1 = shader uses matrix reprojection)
                wgpu::BindGroupEntry {
                    binding: 6,
                    resource: wgpu::BindingResource::TextureView(
                        self.motion_fallback_view.as_ref().unwrap_or(write_view),
                    ),
                },
            ],
        });

        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("taa_resolve"),
                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);
        }

        // Store current VP for next frame's reprojection
        self.prev_vp = curr_vp;
        self.frame_index = self.frame_index.wrapping_add(1);
        self.write_b = !self.write_b;

        // Return the just-written history as the resolved output
        if self.write_b {
            self.history_b_view.as_ref()
        } else {
            self.history_a_view.as_ref()
        }
    }

    /// TD-006: Get the last-written history texture for HDR swap.
    /// Returns the texture that was just written by the most recent dispatch().
    pub fn output_texture(&self) -> Option<&wgpu::Texture> {
        // After dispatch, write_b is flipped. So the just-written texture is:
        if self.write_b {
            self.history_b.as_ref()
        } else {
            self.history_a.as_ref()
        }
    }

    pub fn destroy(&mut self) {
        if let Some(t) = self.history_a.take() {
            t.destroy();
        }
        if let Some(t) = self.history_b.take() {
            t.destroy();
        }
        if let Some(b) = self.params_buf.take() {
            b.destroy();
        }
        self.history_a_view = None;
        self.history_b_view = None;
        self.initialized = false;
    }
}

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

    #[test]
    fn taa_params_size() {
        // 2 floats + 2 floats + 16 floats + 16 floats = 36 × 4 = 144 bytes
        assert_eq!(std::mem::size_of::<TaaParams>(), 144);
    }

    #[test]
    fn halton_sequence_unique() {
        let mut values = Vec::new();
        for i in 1..=16 {
            values.push(halton(i, 2));
        }
        // All values should be unique
        for i in 0..values.len() {
            for j in (i + 1)..values.len() {
                assert!(
                    (values[i] - values[j]).abs() > 0.001,
                    "Halton values not unique: {} and {}",
                    values[i],
                    values[j]
                );
            }
        }
    }

    #[test]
    fn halton_range() {
        for i in 1..=32 {
            let v2 = halton(i, 2);
            let v3 = halton(i, 3);
            assert!(v2 >= 0.0 && v2 < 1.0, "Halton(2) out of range: {}", v2);
            assert!(v3 >= 0.0 && v3 < 1.0, "Halton(3) out of range: {}", v3);
        }
    }

    #[test]
    fn taa_default_enabled() {
        let pass = TaaPass::default();
        assert!(pass.enabled);
        assert!(!pass.initialized);
    }
}
