// GPU-accelerated tweening — evaluates animation tracks on the GPU via compute shaders.
// Uses the tween.wgsl compute shader for parallel interpolation.

use crate::shader;

/// Tween input data uploaded to GPU. Must match tween.wgsl TweenInput layout.
#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
pub struct TweenInput {
    pub time_ms: u32,
    pub start_time_ms: u32,
    pub end_time_ms: u32,
    pub tween_mode: u32,
    pub start_value: [f32; 4],
    pub end_value: [f32; 4],
}

/// Tween output data read back from GPU.
#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
pub struct TweenOutput {
    pub value: [f32; 4],
}

/// Cached compute pipeline and layout for the tween shader.
/// Created once and reused across buffer reallocations.
struct TweenPipelineCache {
    pipeline: wgpu::ComputePipeline,
    bind_group_layout: wgpu::BindGroupLayout,
}

/// GPU tweening context.
pub struct GpuTweenContext {
    pub shader_module: wgpu::ShaderModule,
    pipeline_cache: Option<TweenPipelineCache>,
    bind_group: Option<wgpu::BindGroup>,
    input_buffer: Option<wgpu::Buffer>,
    output_buffer: Option<wgpu::Buffer>,
    capacity: u32,
}

impl GpuTweenContext {
    /// Create the tween context with the embedded tween shader.
    pub fn new(device: &wgpu::Device) -> Self {
        let shader_module = shader::load_wgsl_embedded(device, "tween_compute", include_str!("../shaders/tween.wgsl"));

        Self {
            shader_module,
            pipeline_cache: None,
            bind_group: None,
            input_buffer: None,
            output_buffer: None,
            capacity: 0,
        }
    }

    /// Ensure the pipeline cache is created. Called once; the pipeline and layout
    /// are reused across buffer reallocations.
    fn ensure_pipeline(&mut self, device: &wgpu::Device) {
        if self.pipeline_cache.is_some() {
            return;
        }

        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("tween_main"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    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,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: false },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
            ],
        });

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("tween_main"),
            bind_group_layouts: &[Some(&bind_group_layout)],
            immediate_size: 0,
        });

        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("tween_main"),
            layout: Some(&pipeline_layout),
            module: &self.shader_module,
            entry_point: Some("tween_main"),
            compilation_options: Default::default(),
            cache: None,
        });

        self.pipeline_cache = Some(TweenPipelineCache {
            pipeline,
            bind_group_layout,
        });
    }

    /// Ensure buffers are large enough for the given count.
    /// Rounds capacity up to the next power of two to avoid repeated reallocations.
    /// The pipeline is created once and cached; only the bind group is recreated
    /// when buffers change.
    pub fn ensure_capacity(&mut self, device: &wgpu::Device, count: u32) {
        if count <= self.capacity && self.bind_group.is_some() {
            return;
        }

        self.ensure_pipeline(device);

        let new_capacity = count.next_power_of_two();
        let input_size = (new_capacity as u64) * std::mem::size_of::<TweenInput>() as u64;
        let output_size = (new_capacity as u64) * std::mem::size_of::<TweenOutput>() as u64;

        let input_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("tween_input"),
            size: input_size,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let output_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("tween_output"),
            size: output_size,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });

        let Some(cache) = self.pipeline_cache.as_ref() else {
            return;
        };
        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("tween_main"),
            layout: &cache.bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: input_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: output_buffer.as_entire_binding(),
                },
            ],
        });

        self.bind_group = Some(bind_group);
        self.input_buffer = Some(input_buffer);
        self.output_buffer = Some(output_buffer);
        self.capacity = new_capacity;
    }

    /// Upload tween inputs and dispatch the compute shader.
    pub fn dispatch(&self, queue: &wgpu::Queue, encoder: &mut wgpu::CommandEncoder, inputs: &[TweenInput]) {
        let Some(input_buffer) = &self.input_buffer else {
            return;
        };
        let Some(cache) = &self.pipeline_cache else {
            return;
        };
        let Some(bind_group) = &self.bind_group else {
            return;
        };

        if inputs.len() > self.capacity as usize {
            log::error!(
                "tween_dispatch:inputs {} exceed capacity {}",
                inputs.len(),
                self.capacity
            );
            return;
        }

        queue.write_buffer(input_buffer, 0, bytemuck::cast_slice(inputs));

        let workgroups = (inputs.len() as u32).div_ceil(64);
        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
            label: None,
            timestamp_writes: None,
        });
        pass.set_pipeline(&cache.pipeline);
        pass.set_bind_group(0, bind_group, &[]);
        pass.dispatch_workgroups(workgroups, 1, 1);
    }

    /// Explicitly destroy GPU resources.
    pub fn destroy(&mut self) {
        if let Some(buf) = self.input_buffer.take() {
            buf.destroy();
        }
        if let Some(buf) = self.output_buffer.take() {
            buf.destroy();
        }
        self.bind_group = None;
        self.pipeline_cache = None;
        self.capacity = 0;
    }
}

/// Tween mode constants matching the WGSL shader.
pub const TWEEN_LINEAR: u32 = 0;
pub const TWEEN_EASE_IN_OUT: u32 = 1;
pub const TWEEN_EASE_IN: u32 = 2;
pub const TWEEN_EASE_OUT: u32 = 3;
pub const TWEEN_STEP: u32 = 4;
pub const TWEEN_CATMULL_ROM: u32 = 5;

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

    #[test]
    fn tween_input_size() {
        assert_eq!(std::mem::size_of::<TweenInput>(), 48);
    }

    #[test]
    fn tween_output_size() {
        assert_eq!(std::mem::size_of::<TweenOutput>(), 16);
    }

    #[test]
    fn tween_mode_constants() {
        assert_eq!(TWEEN_LINEAR, 0);
        assert_eq!(TWEEN_EASE_IN_OUT, 1);
        assert_eq!(TWEEN_STEP, 4);
    }
}
