// Animation clips + batch dispatch via GpuTweenContext.
// Provides a higher-level animation API on top of the raw tween compute shader.

use crate::tween::{GpuTweenContext, TweenInput};

/// A single animation clip with named property tracks.
pub struct AnimationClip {
    pub name: String,
    pub duration_ms: u32,
    pub tracks: Vec<AnimationTrack>,
    pub looping: bool,
}

/// A property track within an animation clip.
pub struct AnimationTrack {
    pub property: String,
    pub keyframes: Vec<Keyframe>,
    pub tween_mode: u32,
}

/// A keyframe at a specific time with a vec4 value.
#[derive(Debug, Clone, Copy)]
pub struct Keyframe {
    pub time_ms: u32,
    pub value: [f32; 4],
}

/// An active animation instance referencing a clip.
pub struct AnimationInstance {
    pub clip_index: usize,
    pub start_time_ms: u32,
    pub speed: f32,
    pub playing: bool,
}

/// Maximum concurrent GPU animation instances.
const MAX_INSTANCES: usize = 1024;

/// Animation batch processor — evaluates multiple animation instances via GPU compute.
#[derive(Default)]
pub struct AnimationBatch {
    clips: Vec<AnimationClip>,
    instances: Vec<AnimationInstance>,
}

impl AnimationBatch {
    /// Register a clip and return its index.
    pub fn add_clip(&mut self, clip: AnimationClip) -> usize {
        let idx = self.clips.len();
        self.clips.push(clip);
        idx
    }

    /// Start playing a clip. Returns the instance index, or `None` if at capacity.
    pub fn play(&mut self, clip_index: usize, current_time_ms: u32) -> Option<usize> {
        if self.instances.len() >= MAX_INSTANCES {
            self.compact(); // Try to reclaim stopped instances first
            if self.instances.len() >= MAX_INSTANCES {
                return None;
            }
        }
        let idx = self.instances.len();
        self.instances.push(AnimationInstance {
            clip_index,
            start_time_ms: current_time_ms,
            speed: 1.0,
            playing: true,
        });
        Some(idx)
    }

    /// Stop an animation instance.
    pub fn stop(&mut self, instance_index: usize) {
        if let Some(inst) = self.instances.get_mut(instance_index) {
            inst.playing = false;
        }
    }

    /// Build tween inputs for all active instances at the given time.
    /// Appends to the caller-provided Vec to allow allocation reuse across frames.
    pub fn build_inputs(&self, current_time_ms: u32, inputs: &mut Vec<TweenInput>) {
        inputs.clear();

        for inst in &self.instances {
            if !inst.playing {
                continue;
            }

            let Some(clip) = self.clips.get(inst.clip_index) else {
                continue;
            };

            let elapsed = ((current_time_ms.saturating_sub(inst.start_time_ms)) as f32 * inst.speed) as u32;

            let local_time = if clip.looping && clip.duration_ms > 0 {
                elapsed % clip.duration_ms
            } else {
                elapsed.min(clip.duration_ms)
            };

            for track in &clip.tracks {
                // Find surrounding keyframes
                if track.keyframes.len() < 2 {
                    continue;
                }

                let mut k0 = &track.keyframes[0];
                let mut k1 = &track.keyframes[track.keyframes.len() - 1];

                for window in track.keyframes.windows(2) {
                    if local_time >= window[0].time_ms && local_time <= window[1].time_ms {
                        k0 = &window[0];
                        k1 = &window[1];
                        break;
                    }
                }

                inputs.push(TweenInput {
                    time_ms: local_time,
                    start_time_ms: k0.time_ms,
                    end_time_ms: k1.time_ms,
                    tween_mode: track.tween_mode,
                    start_value: k0.value,
                    end_value: k1.value,
                });
            }
        }
    }

    /// Dispatch all active animations via GPU compute.
    /// Compacts stopped instances before dispatch to bound memory growth.
    pub fn dispatch(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
        tween_ctx: &mut GpuTweenContext,
        current_time_ms: u32,
    ) {
        self.compact();

        let mut inputs = Vec::new();
        self.build_inputs(current_time_ms, &mut inputs);
        if inputs.is_empty() {
            return;
        }

        tween_ctx.ensure_capacity(device, inputs.len() as u32);
        tween_ctx.dispatch(queue, encoder, &inputs);
    }

    /// Number of registered clips.
    pub fn clip_count(&self) -> usize {
        self.clips.len()
    }

    /// Number of animation instances (active + stopped).
    pub fn instance_count(&self) -> usize {
        self.instances.len()
    }

    /// Remove stopped instances to reclaim memory.
    pub fn compact(&mut self) {
        self.instances.retain(|i| i.playing);
    }
}

// ── Skeletal animation ────────────────────────────────────────────────

/// Maximum joints per skeleton. Matches typical glTF limits.
pub const MAX_JOINTS: usize = 128;

/// Joint-space keyframe for skeletal animation.
#[derive(Debug, Clone, Copy)]
pub struct JointKeyframe {
    pub time_ms: u32,
    /// 4x4 column-major joint-space matrix.
    pub matrix: [f32; 16],
}

/// A skeletal animation clip from glTF.
pub struct SkeletalClip {
    pub name: String,
    pub duration_ms: u32,
    pub looping: bool,
    /// Per-joint keyframe tracks. Outer vec indexed by joint index.
    pub joint_tracks: Vec<Vec<JointKeyframe>>,
}

/// Active skeletal animation instance.
pub struct SkeletalInstance {
    pub clip_index: usize,
    pub start_time_ms: u32,
    pub speed: f32,
    pub playing: bool,
}

/// Skeletal animation evaluator. Evaluates joint matrices per frame.
#[derive(Default)]
pub struct SkeletalBatch {
    clips: Vec<SkeletalClip>,
    instances: Vec<SkeletalInstance>,
}

impl SkeletalBatch {
    /// Register a skeletal clip. Returns its index.
    pub fn add_clip(&mut self, clip: SkeletalClip) -> usize {
        let idx = self.clips.len();
        self.clips.push(clip);
        idx
    }

    /// Start playing a skeletal clip. Returns None if at capacity.
    pub fn play(&mut self, clip_index: usize, current_time_ms: u32) -> Option<usize> {
        if self.instances.len() >= MAX_INSTANCES {
            self.compact();
            if self.instances.len() >= MAX_INSTANCES {
                return None;
            }
        }
        let idx = self.instances.len();
        self.instances.push(SkeletalInstance {
            clip_index,
            start_time_ms: current_time_ms,
            speed: 1.0,
            playing: true,
        });
        Some(idx)
    }

    /// Stop a skeletal animation instance.
    pub fn stop(&mut self, instance_index: usize) {
        if let Some(inst) = self.instances.get_mut(instance_index) {
            inst.playing = false;
        }
    }

    /// Evaluate joint matrices for all active instances.
    /// Returns a flat array: [instance_0_joint_0..joint_N, instance_1_joint_0..joint_N, ...].
    /// Each entry is a 4x4 column-major matrix.
    pub fn evaluate(&self, current_time_ms: u32) -> Vec<[f32; 16]> {
        let mut result = Vec::new();
        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,
        ];

        for inst in &self.instances {
            if !inst.playing {
                continue;
            }
            let Some(clip) = self.clips.get(inst.clip_index) else {
                continue;
            };

            let elapsed = ((current_time_ms.saturating_sub(inst.start_time_ms)) as f32 * inst.speed) as u32;
            let local_time = if clip.looping && clip.duration_ms > 0 {
                elapsed % clip.duration_ms
            } else {
                elapsed.min(clip.duration_ms)
            };

            for track in &clip.joint_tracks {
                if track.len() < 2 {
                    result.push(if let Some(kf) = track.first() {
                        kf.matrix
                    } else {
                        identity
                    });
                    continue;
                }

                // Find surrounding keyframes and lerp.
                let mut k0 = &track[0];
                let mut k1 = &track[track.len() - 1];
                for window in track.windows(2) {
                    if local_time >= window[0].time_ms && local_time <= window[1].time_ms {
                        k0 = &window[0];
                        k1 = &window[1];
                        break;
                    }
                }

                let span = k1.time_ms.saturating_sub(k0.time_ms);
                let t = if span > 0 {
                    (local_time.saturating_sub(k0.time_ms)) as f32 / span as f32
                } else {
                    0.0
                };

                // Component-wise lerp of matrices.
                let mut m = [0.0f32; 16];
                for i in 0..16 {
                    m[i] = k0.matrix[i] * (1.0 - t) + k1.matrix[i] * t;
                }
                result.push(m);
            }
        }

        result
    }

    /// Number of clips.
    pub fn clip_count(&self) -> usize {
        self.clips.len()
    }

    /// Number of instances.
    pub fn instance_count(&self) -> usize {
        self.instances.len()
    }

    /// Remove stopped instances.
    pub fn compact(&mut self) {
        self.instances.retain(|i| i.playing);
    }
}

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

    #[test]
    fn add_clip_and_play() {
        let mut batch = AnimationBatch::default();
        let clip = AnimationClip {
            name: "test".into(),
            duration_ms: 1000,
            tracks: vec![AnimationTrack {
                property: "position".into(),
                keyframes: vec![
                    Keyframe {
                        time_ms: 0,
                        value: [0.0, 0.0, 0.0, 0.0],
                    },
                    Keyframe {
                        time_ms: 1000,
                        value: [1.0, 1.0, 1.0, 1.0],
                    },
                ],
                tween_mode: TWEEN_LINEAR,
            }],
            looping: false,
        };
        let clip_idx = batch.add_clip(clip);
        let _inst_idx = batch.play(clip_idx, 0).expect("under capacity");
        assert_eq!(batch.clip_count(), 1);
        assert_eq!(batch.instance_count(), 1);

        let mut inputs = Vec::new();
        batch.build_inputs(500, &mut inputs);
        assert_eq!(inputs.len(), 1);
        assert_eq!(inputs[0].time_ms, 500);
    }

    #[test]
    fn looping_wraps_time() {
        let mut batch = AnimationBatch::default();
        let clip = AnimationClip {
            name: "loop".into(),
            duration_ms: 100,
            tracks: vec![AnimationTrack {
                property: "alpha".into(),
                keyframes: vec![
                    Keyframe {
                        time_ms: 0,
                        value: [0.0; 4],
                    },
                    Keyframe {
                        time_ms: 100,
                        value: [1.0; 4],
                    },
                ],
                tween_mode: TWEEN_LINEAR,
            }],
            looping: true,
        };
        batch.add_clip(clip);
        batch.play(0, 0).expect("under capacity");

        let mut inputs = Vec::new();
        batch.build_inputs(250, &mut inputs);
        assert_eq!(inputs.len(), 1);
        assert_eq!(inputs[0].time_ms, 50); // 250 % 100 = 50
    }

    #[test]
    fn stop_and_compact() {
        let mut batch = AnimationBatch::default();
        let clip = AnimationClip {
            name: "t".into(),
            duration_ms: 100,
            tracks: vec![],
            looping: false,
        };
        batch.add_clip(clip);
        let idx = batch.play(0, 0).expect("under capacity");
        batch.stop(idx);
        assert_eq!(batch.instance_count(), 1);
        batch.compact();
        assert_eq!(batch.instance_count(), 0);
    }

    // ── Skeletal animation tests ──────────────────────────────────────

    #[test]
    fn skeletal_add_and_play() {
        let mut batch = SkeletalBatch::default();
        #[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,
        ];
        let clip = SkeletalClip {
            name: "walk".into(),
            duration_ms: 1000,
            looping: true,
            joint_tracks: vec![vec![
                JointKeyframe {
                    time_ms: 0,
                    matrix: identity,
                },
                JointKeyframe {
                    time_ms: 1000,
                    matrix: identity,
                },
            ]],
        };
        batch.add_clip(clip);
        batch.play(0, 0);
        assert_eq!(batch.clip_count(), 1);
        assert_eq!(batch.instance_count(), 1);

        let matrices = batch.evaluate(500);
        assert_eq!(matrices.len(), 1);
    }

    #[test]
    fn skeletal_lerp_midpoint() {
        let mut batch = SkeletalBatch::default();
        let mut m0 = [0.0f32; 16];
        m0[0] = 1.0;
        let mut m1 = [0.0f32; 16];
        m1[0] = 3.0;
        let clip = SkeletalClip {
            name: "test".into(),
            duration_ms: 100,
            looping: false,
            joint_tracks: vec![vec![
                JointKeyframe { time_ms: 0, matrix: m0 },
                JointKeyframe {
                    time_ms: 100,
                    matrix: m1,
                },
            ]],
        };
        batch.add_clip(clip);
        batch.play(0, 0);
        let matrices = batch.evaluate(50);
        // Midpoint: lerp(1.0, 3.0, 0.5) = 2.0
        assert!((matrices[0][0] - 2.0).abs() < 0.01);
    }

    #[test]
    fn skeletal_compact() {
        let mut batch = SkeletalBatch::default();
        let clip = SkeletalClip {
            name: "t".into(),
            duration_ms: 100,
            looping: false,
            joint_tracks: vec![],
        };
        batch.add_clip(clip);
        let idx = batch.play(0, 0).unwrap();
        batch.stop(idx);
        batch.compact();
        assert_eq!(batch.instance_count(), 0);
    }
}
