//! The `W4DS` container + the clip contract — mirrors the `WHAP`
//! (haptics) and `WSCN` (spatial-audio) section layout: a 4-byte magic,
//! a `u16` section count, a `(kind u8, off u32, len u32)` table, then
//! 8-byte-aligned section payloads. Two sections: the canonical-JSON
//! contract and the per-Gaussian keyframe tracks.

use super::{GaussKey, Splat4dClipData};
use crate::world::Fx;

pub const MAGIC_W4DS: &[u8; 4] = b"W4DS";
const SECTION_CONTRACT: u8 = 0x01;
const SECTION_TRACKS: u8 = 0x02;

#[derive(Debug)]
pub enum Splat4dError {
    BadMagic,
    Truncated,
    MissingSection(u8),
    BadContract(String),
}

impl std::fmt::Display for Splat4dError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::BadMagic => write!(f, "payload is not a W4DS container"),
            Self::Truncated => write!(f, "byte stream truncated"),
            Self::MissingSection(k) => write!(f, "required section 0x{k:02x} missing"),
            Self::BadContract(why) => write!(f, "bad clip contract: {why}"),
        }
    }
}
impl std::error::Error for Splat4dError {}

/// Clip contract — canonical JSON (sorted keys, no whitespace), bound by
/// the clip hash and carried verbatim on the wire.
#[derive(Clone, Debug, PartialEq)]
pub struct Splat4dContract {
    pub frame_rate: u32,
    pub numeric: String,
    pub n_gaussians: u32,
    pub n_frames: u64,
}

impl Splat4dContract {
    pub fn to_canonical_json(&self) -> Vec<u8> {
        format!(
            "{{\"frame_rate\":{},\"n_frames\":{},\"n_gaussians\":{},\"numeric\":\"{}\"}}",
            self.frame_rate, self.n_frames, self.n_gaussians, self.numeric
        )
        .into_bytes()
    }

    pub fn from_json(b: &[u8]) -> Result<Splat4dContract, Splat4dError> {
        let v: serde_json::Value =
            serde_json::from_slice(b).map_err(|e| Splat4dError::BadContract(e.to_string()))?;
        let o = v.as_object().ok_or_else(|| Splat4dError::BadContract("not an object".into()))?;
        let u = |k: &str| {
            o.get(k)
                .and_then(|x| x.as_u64())
                .ok_or_else(|| Splat4dError::BadContract(format!("missing `{k}`")))
        };
        let s = |k: &str| {
            o.get(k)
                .and_then(|x| x.as_str())
                .map(str::to_owned)
                .ok_or_else(|| Splat4dError::BadContract(format!("missing `{k}`")))
        };
        let fr = u("frame_rate")?;
        if fr == 0 || fr > u32::MAX as u64 {
            return Err(Splat4dError::BadContract("frame_rate out of range".into()));
        }
        let ng = u("n_gaussians")?;
        if ng > u32::MAX as u64 {
            return Err(Splat4dError::BadContract("n_gaussians out of range".into()));
        }
        let n_frames = u("n_frames")?;
        if n_frames > crate::limits::RENDER_MAX_FRAMES {
            return Err(Splat4dError::BadContract(crate::limits::duration_too_long(
                "frames", n_frames, crate::limits::RENDER_MAX_FRAMES)));
        }
        Ok(Splat4dContract {
            frame_rate: fr as u32,
            numeric: s("numeric")?,
            n_gaussians: ng as u32,
            n_frames,
        })
    }
}

/// Tracks section bytes: `n_gaussians u32`, then per Gaussian
/// `n_keys u16` then `n × ( frame u64 | 14×i64 Fx )` in the field order
/// `pos[3], scale[3], rot[4], color[3], opacity`.
pub fn tracks_bytes(gaussians: &[Vec<GaussKey>]) -> Vec<u8> {
    let mut out = Vec::new();
    out.extend_from_slice(&(gaussians.len() as u32).to_le_bytes());
    for track in gaussians {
        out.extend_from_slice(&(track.len() as u16).to_le_bytes());
        for k in track {
            out.extend_from_slice(&k.frame.to_le_bytes());
            for v in k
                .pos
                .iter()
                .chain(k.scale.iter())
                .chain(k.rot.iter())
                .chain(k.color.iter())
                .chain(std::iter::once(&k.opacity))
            {
                out.extend_from_slice(&v.to_bits().to_le_bytes());
            }
        }
    }
    out
}

fn parse_tracks(b: &[u8]) -> Result<Vec<Vec<GaussKey>>, Splat4dError> {
    let rd = |p: usize, n: usize| -> Result<&[u8], Splat4dError> {
        if p + n > b.len() {
            Err(Splat4dError::Truncated)
        } else {
            Ok(&b[p..p + n])
        }
    };
    let ng = u32::from_le_bytes(rd(0, 4)?.try_into().unwrap()) as usize;
    let mut p = 4usize;
    // Reserve no more entries than the bytes left could hold: the count is the input's,
    // and a reservation sized by it alone let a few bytes ask for gigabytes.
    let mut gaussians = Vec::with_capacity(ng.min(b.len() / 2));
    for _ in 0..ng {
        let nk = u16::from_le_bytes(rd(p, 2)?.try_into().unwrap()) as usize;
        p += 2;
        let mut track = Vec::with_capacity(nk);
        for _ in 0..nk {
            let frame = u64::from_le_bytes(rd(p, 8)?.try_into().unwrap());
            p += 8;
            let mut fx = [Fx::ZERO; 14];
            for slot in &mut fx {
                *slot = Fx::from_bits(i64::from_le_bytes(rd(p, 8)?.try_into().unwrap()));
                p += 8;
            }
            track.push(GaussKey {
                frame,
                pos: [fx[0], fx[1], fx[2]],
                scale: [fx[3], fx[4], fx[5]],
                rot: [fx[6], fx[7], fx[8], fx[9]],
                color: [fx[10], fx[11], fx[12]],
                opacity: fx[13],
            });
        }
        gaussians.push(track);
    }
    Ok(gaussians)
}

impl Splat4dClipData {
    pub fn to_bytes(&self) -> Vec<u8> {
        let tracks = tracks_bytes(&self.gaussians);
        let sections = [(SECTION_CONTRACT, &self.contract_bytes[..]), (SECTION_TRACKS, &tracks[..])];
        let mut out = Vec::new();
        out.extend_from_slice(MAGIC_W4DS);
        out.extend_from_slice(&(sections.len() as u16).to_le_bytes());
        let mut off = 0u32;
        let mut offs = Vec::new();
        for (_, b) in &sections {
            offs.push(off);
            off += b.len().div_ceil(8) as u32 * 8;
        }
        for ((kind, b), o) in sections.iter().zip(&offs) {
            out.push(*kind);
            out.extend_from_slice(&o.to_le_bytes());
            out.extend_from_slice(&(b.len() as u32).to_le_bytes());
        }
        for (_, b) in &sections {
            out.extend_from_slice(b);
            out.extend_from_slice(&[0u8; 8][..b.len().div_ceil(8) * 8 - b.len()]);
        }
        out
    }

    pub fn from_bytes(bytes: &[u8]) -> Result<Splat4dClipData, Splat4dError> {
        if bytes.len() < 6 {
            return Err(Splat4dError::Truncated);
        }
        if &bytes[..4] != MAGIC_W4DS {
            return Err(Splat4dError::BadMagic);
        }
        let n = u16::from_le_bytes(bytes[4..6].try_into().unwrap()) as usize;
        let table_end = 6 + n * 9;
        if bytes.len() < table_end {
            return Err(Splat4dError::Truncated);
        }
        let pay = &bytes[table_end..];
        let (mut contract_bytes, mut trackbytes) = (None, None);
        for i in 0..n {
            let e = &bytes[6 + i * 9..6 + (i + 1) * 9];
            let off = u32::from_le_bytes(e[1..5].try_into().unwrap()) as usize;
            let len = u32::from_le_bytes(e[5..9].try_into().unwrap()) as usize;
            let end = off.checked_add(len).ok_or(Splat4dError::Truncated)?;
            if end > pay.len() {
                return Err(Splat4dError::Truncated);
            }
            match e[0] {
                SECTION_CONTRACT => contract_bytes = Some(pay[off..end].to_vec()),
                SECTION_TRACKS => trackbytes = Some(&pay[off..end]),
                _ => {}
            }
        }
        let contract_bytes = contract_bytes.ok_or(Splat4dError::MissingSection(SECTION_CONTRACT))?;
        let trackbytes = trackbytes.ok_or(Splat4dError::MissingSection(SECTION_TRACKS))?;
        let contract = Splat4dContract::from_json(&contract_bytes)?;
        let gaussians = parse_tracks(trackbytes)?;
        Ok(Splat4dClipData { contract, contract_bytes, gaussians })
    }
}
