avatar-anim 0.1.2

use binrw::io::{Read, Seek, Write};
use binrw::{BinRead, BinResult, BinWrite, Endian, binrw};
use glam::{EulerRot, Quat, Vec3};
use llsd_rs::Llsd;
use std::collections::HashSet;
use thiserror::Error;

pub mod io;
pub mod skeleton;

use crate::io::{
    AnimReadContext, AnimVersion, KEYFRAME_MOTION_SUBVERSION, KEYFRAME_MOTION_VERSION,
    classify_anim_version, read_fixed_length_string, read_null_terminated_string, read_pos_vec3,
    read_position_keys, read_rot_quat, read_rotation_keys, write_fixed_length_string,
    write_null_terminated_string, write_pos_vec3, write_rot_quat,
};

pub use AnimError as Error;
pub use skeleton::{SkeletonBone, SkeletonDefinition};
pub type Result<T> = std::result::Result<T, AnimError>;

#[derive(Debug, Error)]
pub enum AnimError {
    #[error("I/O error: {0}")]
    Io(#[from] std::io::Error),
    #[error("Binary parsing error: {0}")]
    BinRw(#[from] binrw::Error),
    #[error("Invalid structure: {0}")]
    InvalidStructure(String),
    #[error("LLSD parse error: {0}")]
    Llsd(String),
}

#[binrw]
#[brw(little)]
#[derive(Clone, Debug, PartialEq)]
pub struct AnimationHeader {
    pub version: u16,
    pub sub_version: u16,
    pub base_priority: i32,
    pub duration: f32,
    #[br(parse_with = read_null_terminated_string)]
    #[bw(write_with = write_null_terminated_string)]
    pub emote_name: String,
    pub loop_in_point: f32,
    pub loop_out_point: f32,
    pub looped: i32,
    pub ease_in_duration: f32,
    pub ease_out_duration: f32,
    pub hand_pose: u32,
}

impl Default for AnimationHeader {
    fn default() -> Self {
        Self {
            version: 1,
            sub_version: 0,
            base_priority: 6,
            duration: 0.017,
            emote_name: String::new(),
            loop_in_point: 0.0,
            loop_out_point: 0.017,
            looped: 1,
            ease_in_duration: 1.0,
            ease_out_duration: 1.0,
            hand_pose: 0,
        }
    }
}

#[binrw]
#[brw(little)]
#[derive(Clone, Debug, Default, PartialEq)]
pub struct RotationKey {
    pub time: u16,
    #[br(parse_with = read_rot_quat)]
    #[bw(write_with = write_rot_quat)]
    pub rot: Quat,
}

impl From<Quat> for RotationKey {
    fn from(rot: Quat) -> Self {
        Self { time: 0, rot }
    }
}

#[binrw]
#[brw(little)]
#[derive(Clone, Debug, Default, PartialEq)]
pub struct PositionKey {
    pub time: u16,
    #[br(parse_with = read_pos_vec3)]
    #[bw(write_with = write_pos_vec3)]
    pub pos: Vec3,
}

impl From<Vec3> for PositionKey {
    fn from(pos: Vec3) -> Self {
        Self { time: 0, pos }
    }
}

#[derive(Clone, Debug, Default, PartialEq)]
pub struct JointData {
    pub name: String,
    pub priority: i32,
    pub rotation_keys: Vec<RotationKey>,
    pub position_keys: Vec<PositionKey>,
}

impl BinRead for JointData {
    type Args<'a> = AnimReadContext;

    fn read_options<R: Read + Seek>(
        reader: &mut R,
        endian: Endian,
        ctx: Self::Args<'_>,
    ) -> BinResult<Self> {
        let name = read_null_terminated_string(reader, endian, ())?;
        let priority = i32::read_options(reader, endian, ())?;

        let num_rot_keys = i32::read_options(reader, endian, ())?;
        let rotation_keys = read_rotation_keys(reader, endian, num_rot_keys, ctx)?;

        let num_pos_keys = i32::read_options(reader, endian, ())?;
        let position_keys = read_position_keys(reader, endian, num_pos_keys, ctx)?;

        Ok(Self {
            name,
            priority,
            rotation_keys,
            position_keys,
        })
    }
}

impl BinWrite for JointData {
    type Args<'a> = ();

    fn write_options<W: Write + Seek>(
        &self,
        writer: &mut W,
        endian: Endian,
        _args: Self::Args<'_>,
    ) -> BinResult<()> {
        write_null_terminated_string(&self.name, writer, endian, ())?;
        self.priority.write_options(writer, endian, ())?;
        (self.rotation_keys.len() as i32).write_options(writer, endian, ())?;
        for key in &self.rotation_keys {
            key.write_options(writer, endian, ())?;
        }

        (self.position_keys.len() as i32).write_options(writer, endian, ())?;
        for key in &self.position_keys {
            key.write_options(writer, endian, ())?;
        }
        Ok(())
    }
}

#[binrw]
#[brw(little)]
#[derive(Clone, Debug, Default, PartialEq)]
pub struct Constraint {
    pub chain_length: u8,
    pub constraint_type: u8,

    #[br(parse_with = read_fixed_length_string, args(16usize))]
    #[bw(write_with = write_fixed_length_string, args(16usize))]
    pub source_volume: String,

    pub source_offset: [f32; 3],

    #[br(parse_with = read_fixed_length_string, args(16usize))]
    #[bw(write_with = write_fixed_length_string, args(16usize))]
    pub target_volume: String,

    pub target_offset: [f32; 3],
    pub target_dir: [f32; 3],
    pub ease_in_start: f32,
    pub ease_in_stop: f32,
    pub ease_out_start: f32,
    pub ease_out_stop: f32,
}

#[derive(Clone, Debug, Default, PartialEq)]
pub struct Animation {
    pub header: AnimationHeader,
    pub joints: Vec<JointData>,
    pub constraints: Vec<Constraint>,
}

impl BinRead for Animation {
    type Args<'a> = ();

    fn read_options<R: Read + Seek>(
        reader: &mut R,
        endian: Endian,
        _args: Self::Args<'_>,
    ) -> BinResult<Self> {
        let mut header = AnimationHeader::read_options(reader, endian, ())?;
        let version = classify_anim_version(header.version, header.sub_version)?;
        if version == AnimVersion::Old {
            header.version = KEYFRAME_MOTION_VERSION;
            header.sub_version = KEYFRAME_MOTION_SUBVERSION;
        }
        let ctx = AnimReadContext {
            version,
            duration: header.duration,
        };

        let num_joints = u32::read_options(reader, endian, ())?;
        let mut joints = Vec::with_capacity(num_joints as usize);
        for _ in 0..num_joints {
            joints.push(JointData::read_options(reader, endian, ctx)?);
        }

        let num_constraints = i32::read_options(reader, endian, ())?;
        if num_constraints < 0 {
            return Err(binrw::Error::AssertFail {
                pos: 0,
                message: "num_constraints must be non-negative".into(),
            });
        }
        let mut constraints = Vec::with_capacity(num_constraints as usize);
        for _ in 0..num_constraints {
            constraints.push(Constraint::read_options(reader, endian, ())?);
        }

        Ok(Self {
            header,
            joints,
            constraints,
        })
    }
}

impl BinWrite for Animation {
    type Args<'a> = ();

    fn write_options<W: Write + Seek>(
        &self,
        writer: &mut W,
        endian: Endian,
        _args: Self::Args<'_>,
    ) -> BinResult<()> {
        let header = AnimationHeader {
            version: KEYFRAME_MOTION_VERSION,
            sub_version: KEYFRAME_MOTION_SUBVERSION,
            ..self.header.clone()
        };
        header.write_options(writer, endian, ())?;
        (self.joints.len() as u32).write_options(writer, endian, ())?;
        for joint in &self.joints {
            joint.write_options(writer, endian, ())?;
        }

        (self.constraints.len() as i32).write_options(writer, endian, ())?;
        for constraint in &self.constraints {
            constraint.write_options(writer, endian, ())?;
        }
        Ok(())
    }
}

/// Strategy for handling duplicate keyframe times when cleaning up keys.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum DuplicateKeyStrategy {
    /// Keep the first encountered key (time-stable).
    KeepFirst,
    /// Keep the last encountered key.
    KeepLast,
    /// Average all keys with the same timestamp (rotation via progressive slerp, position via arithmetic mean).
    Average,
}

fn group_average_rot(keys: &[RotationKey]) -> Vec<RotationKey> {
    if keys.is_empty() {
        return Vec::new();
    }
    let mut out = Vec::new();
    let mut i = 0usize;
    while i < keys.len() {
        let t = keys[i].time;
        let mut acc = glam::Quat::IDENTITY;
        let mut count = 0f32;
        let mut j = i;
        while j < keys.len() && keys[j].time == t {
            acc = if count == 0.0 {
                keys[j].rot
            } else {
                acc.slerp(keys[j].rot, 1.0 / (count + 1.0))
            };
            count += 1.0;
            j += 1;
        }
        out.push(RotationKey {
            time: t,
            rot: acc.normalize(),
        });
        i = j;
    }
    out
}

fn group_average_pos(keys: &[PositionKey]) -> Vec<PositionKey> {
    if keys.is_empty() {
        return Vec::new();
    }
    let mut out = Vec::new();
    let mut i = 0usize;
    while i < keys.len() {
        let t = keys[i].time;
        let mut acc = glam::Vec3::ZERO;
        let mut count = 0.0f32;
        let mut j = i;
        while j < keys.len() && keys[j].time == t {
            acc += keys[j].pos;
            count += 1.0;
            j += 1;
        }
        out.push(PositionKey {
            time: t,
            pos: acc / count,
        });
        i = j;
    }
    out
}

impl Animation {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn set_priority(&mut self, priority: i32) -> &mut Self {
        self.header.base_priority = priority;
        for joint in &mut self.joints {
            joint.priority = priority;
        }
        self
    }

    pub fn set_duration(&mut self, duration: f32) -> &mut Self {
        let duration = duration.max(0.0);
        self.header.duration = duration;
        self.header.loop_in_point = 0.0;
        self.header.loop_out_point = duration;
        self
    }

    pub fn set_joint_priority(&mut self, priority: i32) -> &mut Self {
        for joint in &mut self.joints {
            joint.priority = priority;
        }
        self
    }

    pub fn drop_empty_joints(&mut self) -> &mut Self {
        self.joints
            .retain(|joint| !joint.position_keys.is_empty() || !joint.rotation_keys.is_empty());
        self
    }

    pub fn drop_position_keys(&mut self) -> &mut Self {
        for joint in &mut self.joints {
            joint.position_keys.clear();
        }
        self
    }

    pub fn drop_position(&mut self, joints: impl Fn(&JointData) -> bool) -> &mut Self {
        for joint in &mut self.joints {
            if joints(joint) {
                joint.position_keys.clear();
            }
        }
        self
    }

    pub fn drop_zero_position_keys(&mut self, epsilon: f32) -> &mut Self {
        let epsilon_sq = epsilon.max(0.0) * epsilon.max(0.0);
        for joint in &mut self.joints {
            joint
                .position_keys
                .retain(|key| key.pos.length_squared() > epsilon_sq);
        }
        self
    }

    pub fn drop_rotation_keys(&mut self) -> &mut Self {
        for joint in &mut self.joints {
            joint.rotation_keys.clear();
        }
        self
    }

    pub fn drop_rotation(&mut self, joints: impl Fn(&JointData) -> bool) -> &mut Self {
        for joint in &mut self.joints {
            if joints(joint) {
                joint.rotation_keys.clear();
            }
        }
        self
    }

    pub fn cleanup_keys(&mut self) -> &mut Self {
        for joint in &mut self.joints {
            let mut seen_times = HashSet::new();
            joint.rotation_keys.reverse();
            joint
                .rotation_keys
                .retain(|key| seen_times.insert(key.time));
            joint.rotation_keys.sort_by_key(|key| key.time);
            seen_times.clear();
            joint.position_keys.reverse();
            joint
                .position_keys
                .retain(|key| seen_times.insert(key.time));
            joint.position_keys.sort_by_key(|key| key.time);
        }
        self
    }

    /// Cleanup duplicate keyframe times with a customizable strategy.
    pub fn cleanup_keys_with(&mut self, strategy: DuplicateKeyStrategy) -> &mut Self {
        for joint in &mut self.joints {
            match strategy {
                DuplicateKeyStrategy::KeepFirst => {
                    let mut seen = HashSet::new();
                    joint.rotation_keys.retain(|k| seen.insert(k.time));
                    seen.clear();
                    joint.position_keys.retain(|k| seen.insert(k.time));
                }
                DuplicateKeyStrategy::KeepLast => {
                    // Retain last: iterate reverse, keep first occurrence in reverse order.
                    let mut seen = HashSet::new();
                    joint.rotation_keys.reverse();
                    joint.rotation_keys.retain(|k| seen.insert(k.time));
                    joint.rotation_keys.reverse();
                    seen.clear();
                    joint.position_keys.reverse();
                    joint.position_keys.retain(|k| seen.insert(k.time));
                    joint.position_keys.reverse();
                }
                DuplicateKeyStrategy::Average => {
                    // Group by time then average.
                    joint.rotation_keys.sort_by_key(|k| k.time);
                    joint.position_keys.sort_by_key(|k| k.time);
                    joint.rotation_keys = group_average_rot(&joint.rotation_keys);
                    joint.position_keys = group_average_pos(&joint.position_keys);
                }
            }
            joint.rotation_keys.sort_by_key(|k| k.time);
            joint.position_keys.sort_by_key(|k| k.time);
        }
        self
    }

    pub fn joint(&self, name: &str) -> Option<&JointData> {
        self.joints.iter().find(|joint| joint.name == name)
    }

    pub fn joint_mut(&mut self, name: &str) -> Option<&mut JointData> {
        self.joints.iter_mut().find(|joint| joint.name == name)
    }

    pub fn position_reset_from_skeleton<'a>(
        skeleton: &SkeletonDefinition,
        joint_names: impl IntoIterator<Item = &'a str>,
        priority: i32,
    ) -> Result<Self> {
        let mut bones = Vec::new();
        for name in joint_names {
            let bone = skeleton.bone(name).ok_or_else(|| {
                AnimError::InvalidStructure(format!("Skeleton is missing required bone '{name}'"))
            })?;
            bones.push(bone);
        }
        Self::position_reset_from_bones(bones, priority)
    }

    pub fn position_reset_from_bones<'a>(
        bones: impl IntoIterator<Item = &'a SkeletonBone>,
        priority: i32,
    ) -> Result<Self> {
        let priority = priority.clamp(0, 7);
        let mut animation = Self::default();
        animation.header.base_priority = priority;

        for bone in bones {
            animation.joints.push(JointData {
                name: bone.name.clone(),
                priority,
                rotation_keys: Vec::new(),
                position_keys: vec![
                    PositionKey {
                        time: 0,
                        pos: bone.pos,
                    },
                    PositionKey {
                        time: u16::MAX,
                        pos: bone.pos,
                    },
                ],
            });
        }

        Ok(animation)
    }

    pub fn add_skeleton_positions(&mut self, skeleton: &SkeletonDefinition) -> Result<&mut Self> {
        for joint in &mut self.joints {
            if joint.position_keys.is_empty() {
                continue;
            }
            let base_pos = skeleton.position(&joint.name).ok_or_else(|| {
                AnimError::InvalidStructure(format!(
                    "Skeleton is missing position for joint '{}'",
                    joint.name
                ))
            })?;
            for key in &mut joint.position_keys {
                key.pos += base_pos;
            }
        }
        Ok(self)
    }

    /// Creates an animation from LLSD data, typically from Firestorm poser files.
    ///
    /// This function parses LLSD-XML data exported by Firestorm's poser system and converts
    /// it into an animation structure. Firestorm stores pose data in LLSD-XML format in
    /// the user's configuration directory.
    ///
    /// # Arguments
    ///
    /// * `llsd` - The parsed LLSD data containing joint poses
    /// * `check_enabled` - If true, only includes joints where the "enabled" field is true
    ///
    /// # File Locations
    ///
    /// Firestorm poser files are typically found at:
    /// * **Linux**: `~/.firestorm_x64/user_settings/poses/`
    /// * **Windows**: `%APPDATA%/Firestorm_x64/user_settings/poses/`
    /// * **macOS**: `~/Library/Application Support/Firestorm_x64/user_settings/poses/`
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use std::fs::File;
    /// use std::io::BufReader;
    /// use avatar_anim::Animation;
    ///
    /// # fn main() -> avatar_anim::Result<()> {
    /// // Load LLSD-XML file from Firestorm poses directory
    /// let file = BufReader::new(File::open("my_pose.xml")?);
    /// let llsd = llsd_rs::xml::from_reader(file).map_err(|e| avatar_anim::AnimError::Llsd(e.to_string()))?;
    ///
    /// // Convert to animation, including only enabled joints
    /// let animation = Animation::from_llsd(&llsd, true)?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn from_llsd(llsd: &Llsd, check_enabled: bool) -> Result<Self> {
        let Some(joints) = llsd.as_map() else {
            return Err(AnimError::InvalidStructure("LLSD must be a map".into()));
        };
        let mut animation = Self::default();
        for (key, value) in joints {
            let Some(inner) = value.as_map() else {
                continue;
            };
            if check_enabled
                && inner
                    .get("enabled")
                    .is_none_or(|e| e.as_boolean() != Some(&true))
            {
                continue;
            }
            let extract_key = |key: &str| -> Option<(f32, f32, f32)> {
                inner.get(key).and_then(|e| e.as_array()).map(|arr| {
                    (
                        *arr.first().and_then(|e| e.as_real()).unwrap_or(&0.0f64) as f32,
                        *arr.get(1).and_then(|e| e.as_real()).unwrap_or(&0.0f64) as f32,
                        *arr.get(2).and_then(|e| e.as_real()).unwrap_or(&0.0f64) as f32,
                    )
                })
            };
            let rotation = extract_key("rotation").map(|(roll, pitch, yaw)| RotationKey {
                time: u16::MAX,
                rot: Quat::from_euler(EulerRot::XYZ, roll, pitch, yaw).normalize(),
            });
            let position = extract_key("position").map(|(x, y, z)| PositionKey {
                time: u16::MAX,
                pos: Vec3::new(x, y, z),
            });
            animation.joints.push(JointData {
                name: key.clone(),
                rotation_keys: rotation.into_iter().collect(),
                position_keys: position.into_iter().collect(),
                ..Default::default()
            });
        }
        Ok(animation)
    }

    /// Load an animation from a .anim file
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use avatar_anim::Animation;
    ///
    /// # fn main() -> avatar_anim::Result<()> {
    /// let animation = Animation::from_file("my_animation.anim")?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn from_file<P: AsRef<std::path::Path>>(path: P) -> Result<Self> {
        use binrw::BinRead;
        use std::fs::File;
        use std::io::BufReader;
        let file = File::open(path).map_err(AnimError::Io)?;
        let mut reader = BufReader::new(file);
        Self::read_options(&mut reader, Endian::Little, ()).map_err(AnimError::BinRw)
    }

    /// Save an animation to a .anim file
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use avatar_anim::Animation;
    ///
    /// # fn main() -> avatar_anim::Result<()> {
    /// let animation = Animation::default();
    /// animation.to_file("output.anim")?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn to_file<P: AsRef<std::path::Path>>(&self, path: P) -> Result<()> {
        use binrw::BinWrite;
        use std::fs::File;
        use std::io::BufWriter;
        let file = File::create(path).map_err(AnimError::Io)?;
        let mut writer = BufWriter::new(file);
        self.write_options(&mut writer, Endian::Little, ())
            .map_err(AnimError::BinRw)
    }

    /// Load LLSD-XML data from a Firestorm pose file
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use avatar_anim::Animation;
    ///
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let animation = Animation::from_llsd_file("pose.xml", true)?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn from_llsd_file<P: AsRef<std::path::Path>>(path: P, check_enabled: bool) -> Result<Self> {
        use std::fs::File;
        use std::io::BufReader;
        let file = File::open(path).map_err(AnimError::Io)?;
        let reader = BufReader::new(file);
        let llsd = llsd_rs::xml::from_reader(reader).map_err(|e| AnimError::Llsd(e.to_string()))?;
        Self::from_llsd(&llsd, check_enabled)
    }
}