//! 3D Gaussian splat codec — the `F4:Photoreal` cargo codec.
//!
//! A compact, quantized, gzip-compressed container for 3D Gaussian splats, in the
//! spirit of Niantic SPZ / the `KHR_gaussian_splatting_compression` family:
//! positions as fixed-point, scales log-quantized, rotation quaternion and colour
//! as bytes. This is the reader **and** writer half a conforming F4 tile needs —
//! the codec the photoreal tier was missing (raster covers F2, MVT covers F3).
//!
//! Spherical-harmonic degree is recorded in the header; degree 0 (DC colour) is
//! encoded here, higher degrees are a forward-compatible extension.

use std::io::{Read, Write};

use flate2::read::GzDecoder;
use flate2::write::GzEncoder;
use flate2::Compression;

use crate::TileError;

/// One 3D Gaussian splat: mean, per-axis scale, rotation quaternion `(x,y,z,w)`,
/// linear RGB DC colour, opacity, and higher-order **spherical-harmonic**
/// coefficients for view-dependent colour. `sh` holds the bands above degree 0
/// (the DC term is `color`), RGB per coefficient, length `(degree+1)² − 1`.
#[derive(Debug, Clone, PartialEq)]
pub struct Splat {
    pub position: [f32; 3],
    pub scale: [f32; 3],
    pub rotation: [f32; 4],
    pub color: [f32; 3],
    pub opacity: f32,
    pub sh: Vec<[f32; 3]>,
}

const MAGIC: &[u8; 4] = b"MSPL";
const VERSION: u8 = 1;
/// Fixed-point fractional bits for positions (±~8 units at ~0.24 mm on a unit globe).
const POS_FRAC: i32 = 12;
/// Log-scale quantization: `q = (ln(s) + LOG_BIAS) * LOG_SCALE`, clamped to u8.
const LOG_BIAS: f32 = 9.5;
const LOG_SCALE: f32 = 255.0 / 9.5;
/// Bytes per splat for the degree-0 base (pos 12 + scale 3 + rot 4 + colour 3 + opacity 1).
const BASE_STRIDE: usize = 23;
/// Fixed-point scale for SH coefficients (range ±~4, precision 1/32).
const SH_QUANT: f32 = 32.0;

/// SH coefficient count above degree 0 for a given degree: `(degree+1)² − 1`.
fn sh_count(degree: u8) -> usize {
    let d = degree as usize + 1;
    d * d - 1
}

/// Encode splats into a compact gzip-compressed F4 payload. The SH degree is
/// taken from the first splat's `sh` length (all splats must match).
pub fn encode_splats(splats: &[Splat]) -> Result<Vec<u8>, TileError> {
    let coeffs = splats.first().map(|s| s.sh.len()).unwrap_or(0);
    // degree d satisfies (d+1)² − 1 == coeffs.
    let degree = ((coeffs + 1) as f64).sqrt().round() as u8 - 1;
    if sh_count(degree) != coeffs {
        return Err(TileError::GeometryError(format!("invalid SH coeff count {coeffs}")));
    }
    let stride = BASE_STRIDE + 3 * coeffs;

    let mut out = Vec::with_capacity(16 + splats.len() * stride);
    out.extend_from_slice(MAGIC);
    out.push(VERSION);
    out.push(degree);
    out.push(POS_FRAC as u8);
    out.push(0); // reserved
    out.extend_from_slice(&(splats.len() as u32).to_le_bytes());

    let mut body = Vec::with_capacity(splats.len() * stride);
    for s in splats {
        for &p in &s.position {
            let q = (p * (1 << POS_FRAC) as f32).round() as i32;
            body.extend_from_slice(&q.to_le_bytes());
        }
        for &sc in &s.scale {
            body.push(quantize_log(sc));
        }
        for &r in &s.rotation {
            body.push((r.clamp(-1.0, 1.0) * 127.0).round() as i8 as u8);
        }
        for &c in &s.color {
            body.push((c.clamp(0.0, 1.0) * 255.0).round() as u8);
        }
        body.push((s.opacity.clamp(0.0, 1.0) * 255.0).round() as u8);
        for coeff in &s.sh {
            for &c in coeff {
                body.push((c * SH_QUANT).round().clamp(-127.0, 127.0) as i8 as u8);
            }
        }
    }

    let mut gz = GzEncoder::new(Vec::new(), Compression::fast());
    gz.write_all(&body).map_err(|e| TileError::CompressionError(e.to_string()))?;
    let compressed = gz.finish().map_err(|e| TileError::CompressionError(e.to_string()))?;
    out.extend_from_slice(&compressed);
    Ok(out)
}

/// Decode an F4 splat payload (inverse of [`encode_splats`]).
pub fn decode_splats(bytes: &[u8]) -> Result<Vec<Splat>, TileError> {
    if bytes.len() < 12 || &bytes[0..4] != MAGIC {
        return Err(TileError::GeometryError("bad splat magic".into()));
    }
    if bytes[4] != VERSION {
        return Err(TileError::GeometryError(format!("unsupported splat version {}", bytes[4])));
    }
    let degree = bytes[5];
    let frac = bytes[6] as i32;
    let count = u32::from_le_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]) as usize;
    let coeffs = sh_count(degree);
    let stride = BASE_STRIDE + 3 * coeffs;

    let mut body = Vec::new();
    GzDecoder::new(&bytes[12..])
        .read_to_end(&mut body)
        .map_err(|e| TileError::CompressionError(e.to_string()))?;

    if body.len() < count * stride {
        return Err(TileError::GeometryError("truncated splat body".into()));
    }
    let scale = (1 << frac) as f32;
    let mut splats = Vec::with_capacity(count);
    for i in 0..count {
        let b = &body[i * stride..];
        let pos = [
            i32::from_le_bytes([b[0], b[1], b[2], b[3]]) as f32 / scale,
            i32::from_le_bytes([b[4], b[5], b[6], b[7]]) as f32 / scale,
            i32::from_le_bytes([b[8], b[9], b[10], b[11]]) as f32 / scale,
        ];
        let sc = [dequantize_log(b[12]), dequantize_log(b[13]), dequantize_log(b[14])];
        let rot = [
            b[15] as i8 as f32 / 127.0,
            b[16] as i8 as f32 / 127.0,
            b[17] as i8 as f32 / 127.0,
            b[18] as i8 as f32 / 127.0,
        ];
        let color = [b[19] as f32 / 255.0, b[20] as f32 / 255.0, b[21] as f32 / 255.0];
        let opacity = b[22] as f32 / 255.0;
        let sh = (0..coeffs)
            .map(|k| {
                let o = BASE_STRIDE + k * 3;
                [
                    b[o] as i8 as f32 / SH_QUANT,
                    b[o + 1] as i8 as f32 / SH_QUANT,
                    b[o + 2] as i8 as f32 / SH_QUANT,
                ]
            })
            .collect();
        splats.push(Splat { position: pos, scale: sc, rotation: rot, color, opacity, sh });
    }
    Ok(splats)
}

fn quantize_log(s: f32) -> u8 {
    let v = (s.max(1e-5).ln() + LOG_BIAS) * LOG_SCALE;
    v.round().clamp(0.0, 255.0) as u8
}
fn dequantize_log(q: u8) -> f32 {
    (q as f32 / LOG_SCALE - LOG_BIAS).exp()
}

/// Zeroth-band SH normalization constant (the 3DGS DC convention).
const SH_C0: f32 = 0.282_094_79;

fn sigmoid(x: f32) -> f32 {
    1.0 / (1.0 + (-x).exp())
}

/// Ingest a captured 3D Gaussian scene in the de-facto INRIA `.ply` interchange
/// format (the output of the reference 3DGS trainer and most tooling) into the
/// F4 [`Splat`] model — so a real capture can be re-encoded with [`encode_splats`]
/// into a conforming tile. Supports ASCII and binary-little-endian PLY.
///
/// The trainer's storage conventions are inverted here: scales are `exp`'d (stored
/// as log), opacity is `sigmoid`'d (stored as logit), the rotation `(w,x,y,z)` is
/// reordered to `(x,y,z,w)` and normalized, the DC colour is `0.5 + C₀·f_dc`, and
/// the `f_rest` higher SH coefficients (stored channel-major) are interleaved to
/// per-coefficient RGB.
pub fn read_ply(bytes: &[u8]) -> Result<Vec<Splat>, TileError> {
    let marker = b"end_header";
    let hpos = bytes
        .windows(marker.len())
        .position(|w| w == marker)
        .ok_or_else(|| TileError::GeometryError("PLY: no end_header".into()))?;
    // Body starts after the newline following the marker.
    let body_start = bytes[hpos..]
        .iter()
        .position(|&b| b == b'\n')
        .map(|n| hpos + n + 1)
        .ok_or_else(|| TileError::GeometryError("PLY: header not terminated".into()))?;
    let header = std::str::from_utf8(&bytes[..hpos])
        .map_err(|_| TileError::GeometryError("PLY: non-UTF8 header".into()))?;

    let mut ascii = false;
    let mut count = 0usize;
    let mut props: Vec<(String, PlyType)> = Vec::new();
    let mut in_vertex = false;
    for line in header.lines() {
        let mut tok = line.split_whitespace();
        match tok.next() {
            Some("format") => {
                ascii = matches!(tok.next(), Some("ascii"));
            }
            Some("element") => {
                in_vertex = tok.next() == Some("vertex");
                if in_vertex {
                    count = tok.next().and_then(|n| n.parse().ok()).unwrap_or(0);
                }
            }
            Some("property") if in_vertex => {
                let ty = tok.next().unwrap_or("");
                if ty == "list" {
                    return Err(TileError::GeometryError("PLY: list props unsupported".into()));
                }
                let name = tok.next().unwrap_or("").to_string();
                props.push((name, PlyType::parse(ty)?));
            }
            _ => {}
        }
    }
    let is_le = header.contains("binary_little_endian");
    if !ascii && !is_le {
        return Err(TileError::GeometryError("PLY: unsupported format".into()));
    }

    let idx = |n: &str| props.iter().position(|(name, _)| name == n);
    let rest_count = props.iter().filter(|(n, _)| n.starts_with("f_rest_")).count();
    let coeffs = rest_count / 3;

    // Read all vertices into a flat [vertex][prop] f32 table.
    let mut table: Vec<f32> = Vec::with_capacity(count * props.len());
    if ascii {
        let body = std::str::from_utf8(&bytes[body_start..])
            .map_err(|_| TileError::GeometryError("PLY: non-UTF8 body".into()))?;
        for v in body.split_whitespace().take(count * props.len()) {
            table.push(v.parse().map_err(|_| TileError::GeometryError("PLY: bad number".into()))?);
        }
    } else {
        let stride: usize = props.iter().map(|(_, t)| t.size()).sum();
        let body = &bytes[body_start..];
        if body.len() < count * stride {
            return Err(TileError::GeometryError("PLY: truncated body".into()));
        }
        for i in 0..count {
            let mut off = i * stride;
            for (_, t) in &props {
                table.push(t.read_le(&body[off..]));
                off += t.size();
            }
        }
    }
    if table.len() < count * props.len() {
        return Err(TileError::GeometryError("PLY: truncated body".into()));
    }

    let get = |row: usize, name: &str| idx(name).map(|c| table[row * props.len() + c]);
    let mut splats = Vec::with_capacity(count);
    for r in 0..count {
        let g = |n: &str| get(r, n).unwrap_or(0.0);
        let rot_raw = [g("rot_0"), g("rot_1"), g("rot_2"), g("rot_3")];
        let norm = (rot_raw.iter().map(|v| v * v).sum::<f32>()).sqrt().max(1e-12);
        let rotation = [rot_raw[1] / norm, rot_raw[2] / norm, rot_raw[3] / norm, rot_raw[0] / norm];
        let color = [
            (0.5 + SH_C0 * g("f_dc_0")).clamp(0.0, 1.0),
            (0.5 + SH_C0 * g("f_dc_1")).clamp(0.0, 1.0),
            (0.5 + SH_C0 * g("f_dc_2")).clamp(0.0, 1.0),
        ];
        // f_rest is channel-major: [R₀..R_{c-1}, G₀.., B₀..].
        let sh = (0..coeffs)
            .map(|k| {
                [
                    g(&format!("f_rest_{k}")),
                    g(&format!("f_rest_{}", coeffs + k)),
                    g(&format!("f_rest_{}", 2 * coeffs + k)),
                ]
            })
            .collect();
        splats.push(Splat {
            position: [g("x"), g("y"), g("z")],
            scale: [g("scale_0").exp(), g("scale_1").exp(), g("scale_2").exp()],
            rotation,
            color,
            opacity: sigmoid(g("opacity")),
            sh,
        });
    }
    Ok(splats)
}

/// PLY scalar property types, enough to compute strides and read values as f32.
#[derive(Clone, Copy)]
enum PlyType {
    F32,
    F64,
    I8,
    U8,
    I16,
    U16,
    I32,
    U32,
}
impl PlyType {
    fn parse(s: &str) -> Result<Self, TileError> {
        Ok(match s {
            "float" | "float32" => PlyType::F32,
            "double" | "float64" => PlyType::F64,
            "char" | "int8" => PlyType::I8,
            "uchar" | "uint8" => PlyType::U8,
            "short" | "int16" => PlyType::I16,
            "ushort" | "uint16" => PlyType::U16,
            "int" | "int32" => PlyType::I32,
            "uint" | "uint32" => PlyType::U32,
            other => return Err(TileError::GeometryError(format!("PLY: bad type {other}"))),
        })
    }
    fn size(self) -> usize {
        match self {
            PlyType::I8 | PlyType::U8 => 1,
            PlyType::I16 | PlyType::U16 => 2,
            PlyType::F32 | PlyType::I32 | PlyType::U32 => 4,
            PlyType::F64 => 8,
        }
    }
    fn read_le(self, b: &[u8]) -> f32 {
        match self {
            PlyType::F32 => f32::from_le_bytes([b[0], b[1], b[2], b[3]]),
            PlyType::F64 => f64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]) as f32,
            PlyType::I8 => b[0] as i8 as f32,
            PlyType::U8 => b[0] as f32,
            PlyType::I16 => i16::from_le_bytes([b[0], b[1]]) as f32,
            PlyType::U16 => u16::from_le_bytes([b[0], b[1]]) as f32,
            PlyType::I32 => i32::from_le_bytes([b[0], b[1], b[2], b[3]]) as f32,
            PlyType::U32 => u32::from_le_bytes([b[0], b[1], b[2], b[3]]) as f32,
        }
    }
}

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

    fn sample() -> Vec<Splat> {
        (0..256)
            .map(|i| {
                let t = i as f32 / 256.0;
                Splat {
                    position: [t - 0.5, (t * std::f32::consts::TAU).sin() * 0.5, (t * std::f32::consts::PI).cos() * 0.5],
                    scale: [0.01 + t * 0.02, 0.008, 0.003],
                    rotation: [0.0, 0.0, (t * std::f32::consts::PI).sin(), (t * std::f32::consts::PI).cos()],
                    color: [t, 1.0 - t, 0.5],
                    opacity: 0.9,
                    sh: vec![],
                }
            })
            .collect()
    }

    #[test]
    fn roundtrips_within_quantization() {
        let splats = sample();
        let bytes = encode_splats(&splats).unwrap();
        let back = decode_splats(&bytes).unwrap();
        assert_eq!(back.len(), splats.len());
        for (a, b) in splats.iter().zip(&back) {
            for k in 0..3 {
                assert!((a.position[k] - b.position[k]).abs() < 1e-3, "pos");
                assert!((a.color[k] - b.color[k]).abs() < 0.01, "color");
                assert!((a.scale[k] - b.scale[k]).abs() / a.scale[k] < 0.05, "scale");
            }
            assert!((a.opacity - b.opacity).abs() < 0.01);
        }
    }

    #[test]
    fn compresses_below_raw() {
        let splats = sample();
        let encoded = encode_splats(&splats).unwrap();
        // Raw would be 13 floats * 4 bytes = 52 B/splat; quantized+gzip is far less.
        assert!(encoded.len() < splats.len() * 52);
    }

    #[test]
    fn rejects_bad_magic() {
        assert!(decode_splats(b"nope............").is_err());
    }

    #[test]
    fn reads_ascii_ply_with_trainer_conventions() {
        // Degree-1 (9 f_rest) ASCII PLY, one vertex, INRIA property layout.
        let header = "ply\nformat ascii 1.0\nelement vertex 1\n\
            property float x\nproperty float y\nproperty float z\n\
            property float scale_0\nproperty float scale_1\nproperty float scale_2\n\
            property float rot_0\nproperty float rot_1\nproperty float rot_2\nproperty float rot_3\n\
            property float opacity\n\
            property float f_dc_0\nproperty float f_dc_1\nproperty float f_dc_2\n\
            property float f_rest_0\nproperty float f_rest_1\nproperty float f_rest_2\n\
            property float f_rest_3\nproperty float f_rest_4\nproperty float f_rest_5\n\
            property float f_rest_6\nproperty float f_rest_7\nproperty float f_rest_8\n\
            end_header\n";
        // pos | log-scales | quat(w,x,y,z, unnormalized) | logit-opacity | f_dc | f_rest(R0..2,G0..2,B0..2)
        let row = "1.0 2.0 3.0  0.0 -0.6931472 0.6931472  2.0 0.0 0.0 0.0  0.0  \
            1.7724539 0.0 -1.7724539  0.1 0.2 0.3  0.4 0.5 0.6  0.7 0.8 0.9\n";
        let ply = format!("{header}{row}");
        let s = &read_ply(ply.as_bytes()).unwrap()[0];

        assert_eq!(s.position, [1.0, 2.0, 3.0]);
        // exp(log-scale): exp(0)=1, exp(-ln2)=0.5, exp(ln2)=2.
        assert!((s.scale[0] - 1.0).abs() < 1e-4);
        assert!((s.scale[1] - 0.5).abs() < 1e-4);
        assert!((s.scale[2] - 2.0).abs() < 1e-4);
        // rot (w,x,y,z)=(2,0,0,0) → normalized (x,y,z,w)=(0,0,0,1).
        assert!((s.rotation[3] - 1.0).abs() < 1e-5 && s.rotation[0].abs() < 1e-5);
        // opacity sigmoid(0)=0.5.
        assert!((s.opacity - 0.5).abs() < 1e-5);
        // DC: 0.5 + C0*f_dc. f_dc_0=1.7724539 → 0.5 + 0.2820948*1.7724539 ≈ 1.0 (clamped).
        assert!((s.color[0] - 1.0).abs() < 1e-3);
        assert!((s.color[1] - 0.5).abs() < 1e-3);
        assert!((s.color[2] - 0.0).abs() < 1e-3);
        // 3 SH coeffs, channel-major de-interleave: coeff0 = (R0,G0,B0)=(0.1,0.4,0.7).
        assert_eq!(s.sh.len(), 3);
        assert!((s.sh[0][0] - 0.1).abs() < 1e-5);
        assert!((s.sh[0][1] - 0.4).abs() < 1e-5);
        assert!((s.sh[0][2] - 0.7).abs() < 1e-5);
        assert!((s.sh[2][2] - 0.9).abs() < 1e-5);

        // Ingested capture re-encodes into a conforming F4 tile.
        let bytes = encode_splats(&read_ply(ply.as_bytes()).unwrap()).unwrap();
        assert_eq!(decode_splats(&bytes).unwrap().len(), 1);
    }

    #[test]
    fn reads_binary_le_ply() {
        // Minimal degree-0 binary PLY: x,y,z,scale_0..2,rot_0..3,opacity,f_dc_0..2.
        let header = "ply\nformat binary_little_endian 1.0\nelement vertex 1\n\
            property float x\nproperty float y\nproperty float z\n\
            property float scale_0\nproperty float scale_1\nproperty float scale_2\n\
            property float rot_0\nproperty float rot_1\nproperty float rot_2\nproperty float rot_3\n\
            property float opacity\n\
            property float f_dc_0\nproperty float f_dc_1\nproperty float f_dc_2\n\
            end_header\n";
        let vals: [f32; 14] = [
            4.0, 5.0, 6.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
        ];
        let mut bytes = header.as_bytes().to_vec();
        for v in vals {
            bytes.extend_from_slice(&v.to_le_bytes());
        }
        let s = &read_ply(&bytes).unwrap()[0];
        assert_eq!(s.position, [4.0, 5.0, 6.0]);
        assert!((s.scale[0] - 1.0).abs() < 1e-4); // exp(0)
        assert!((s.opacity - 0.5).abs() < 1e-5); // sigmoid(0)
        assert!(s.sh.is_empty()); // no f_rest → degree 0
    }

    #[test]
    fn rejects_headerless_ply() {
        assert!(read_ply(b"not a ply at all").is_err());
    }

    #[test]
    fn sh_degree1_roundtrips() {
        // Degree 1 → 3 SH coefficients (RGB each) above the DC term.
        let splats: Vec<Splat> = (0..64)
            .map(|i| {
                let t = i as f32 / 64.0;
                Splat {
                    position: [t, -t, 0.5],
                    scale: [0.01, 0.01, 0.004],
                    rotation: [0.0, 0.0, 0.0, 1.0],
                    color: [0.5, 0.5, 0.5],
                    opacity: 0.8,
                    sh: vec![[t - 0.5, 0.2, -0.3], [0.1, t, 0.0], [-0.2, 0.0, t * 0.5]],
                }
            })
            .collect();
        let bytes = encode_splats(&splats).unwrap();
        // Header records degree 1.
        assert_eq!(bytes[5], 1);
        let back = decode_splats(&bytes).unwrap();
        assert_eq!(back.len(), splats.len());
        for (a, b) in splats.iter().zip(&back) {
            assert_eq!(a.sh.len(), b.sh.len());
            for (ca, cb) in a.sh.iter().zip(&b.sh) {
                for k in 0..3 {
                    assert!((ca[k] - cb[k]).abs() < 1.0 / 32.0 + 1e-6, "sh coeff");
                }
            }
        }
    }
}