// Motion Vectors — per-pixel velocity buffer for TAA ghosting reduction.
//
// Renders a fullscreen motion vector texture (RG16Float) by comparing
// current and previous frame clip positions. The TAA resolve shader uses
// this for accurate history reprojection instead of depth-only reprojection.

/// Motion vector pass uniform parameters.
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct MotionVectorParams {
    pub curr_vp: [f32; 16],
    pub prev_vp: [f32; 16],
    pub screen_size: [f32; 2],
    pub _pad: [f32; 2],
}

/// Motion vector pass manager.
pub struct MotionVectorPass {
    pipeline: Option<wgpu::ComputePipeline>,
    bgl: Option<wgpu::BindGroupLayout>,
    params_buf: Option<wgpu::Buffer>,
    pub motion_texture: Option<wgpu::Texture>,
    pub motion_view: Option<wgpu::TextureView>,
    prev_vp: [f32; 16],
    width: u32,
    height: u32,
    pub enabled: bool,
    initialized: bool,
}

impl Default for MotionVectorPass {
    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,
            motion_texture: None,
            motion_view: None,
            prev_vp: identity,
            width: 0,
            height: 0,
            enabled: true,
            initialized: false,
        }
    }
}

impl MotionVectorPass {
    /// Initialize pipeline and allocate motion vector texture.
    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("motion_vector_bgl"),
            entries: &[
                // binding 0: params (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: depth texture
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Depth,
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
                // binding 2: motion vector output
                wgpu::BindGroupLayoutEntry {
                    binding: 2,
                    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("motion_vectors.wgsl"),
            source: wgpu::ShaderSource::Wgsl(include_str!("../shaders/motion_vectors.wgsl").into()),
        });

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

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

        let tex = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("motion_vectors"),
            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,
            view_formats: &[],
        });
        self.motion_view = Some(tex.create_view(&wgpu::TextureViewDescriptor::default()));
        self.motion_texture = Some(tex);

        if self.params_buf.is_none() {
            self.params_buf = Some(device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("motion_vector_params"),
                size: std::mem::size_of::<MotionVectorParams>() as u64,
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            }));
        }

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

        log::info!("Motion vectors initialized ({}×{})", width, height);
    }

    /// Dispatch motion vector computation.
    pub fn dispatch(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
        depth_view: &wgpu::TextureView,
        camera: &crate::camera::Camera,
    ) {
        if !self.enabled || !self.initialized {
            return;
        }
        let Some(ref pipeline) = self.pipeline else { return };
        let Some(ref bgl) = self.bgl else { return };
        let Some(ref params_buf) = self.params_buf else { return };
        let Some(ref motion_view) = self.motion_view else {
            return;
        };

        let curr_vp = (camera.projection_matrix * camera.view_matrix).to_cols_array();

        let params = MotionVectorParams {
            curr_vp,
            prev_vp: self.prev_vp,
            screen_size: [self.width as f32, self.height as f32],
            _pad: [0.0; 2],
        };
        queue.write_buffer(params_buf, 0, bytemuck::bytes_of(&params));

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

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

        self.prev_vp = curr_vp;
    }

    pub fn destroy(&mut self) {
        if let Some(t) = self.motion_texture.take() {
            t.destroy();
        }
        if let Some(b) = self.params_buf.take() {
            b.destroy();
        }
        self.motion_view = None;
        self.initialized = false;
    }
}

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

    #[test]
    fn motion_vector_params_size() {
        assert_eq!(std::mem::size_of::<MotionVectorParams>(), 144);
    }

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