use crate::*;
use std::fmt;
use std::io::{self, Read, Write};
use std::mem::size_of;

/// Application policy limits, not maxima inferred from the legacy format.
#[derive(Clone, Copy, Debug)]
pub struct Limits {
    /// Maximum complete input or encoded output, including the header.
    pub max_bytes: usize,
    /// Maximum bytes in any individual ASCII name or value, excluding NUL.
    pub max_string_bytes: usize,
    /// Aggregate names, brushes, planes, properties, patches and vec3 elements.
    pub max_items: usize,
    /// Decoder's aggregate requested owned storage, excluding allocator overhead
    /// and the input buffer. Includes record slots and string payloads.
    pub max_owned_bytes: usize,
}

impl Default for Limits {
    fn default() -> Self {
        Self {
            max_bytes: 64 * 1024 * 1024,
            max_string_bytes: 1024 * 1024,
            max_items: 1_000_000,
            max_owned_bytes: 128 * 1024 * 1024,
        }
    }
}

#[derive(Debug)]
#[non_exhaustive]
pub enum ErrorKind {
    Io(io::Error),
    UnexpectedEof,
    InvalidMagic,
    UnsupportedVersion(u32),
    InvalidFlags(u8),
    ReservedWidth,
    NegativeCount,
    CountOverflow,
    LimitExceeded(&'static str),
    Allocation,
    NonAscii,
    EmbeddedNul,
    UnterminatedString,
    NonFinite,
    PropertyIndex(u32),
    DuplicateProperty(u32),
    UnsortedProperties,
    InconsistentComputedData,
    InvalidPatchLength,
    LegacyProperties,
    TrailingData,
}

#[derive(Debug)]
pub struct Error {
    /// Byte offset in the input, encoded output, or completed I/O prefix.
    pub offset: usize,
    pub kind: ErrorKind,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "MB3 DAT at byte {}: {:?}", self.offset, self.kind)
    }
}

impl std::error::Error for Error {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match &self.kind {
            ErrorKind::Io(e) => Some(e),
            _ => None,
        }
    }
}

type Result<T> = std::result::Result<T, Error>;

fn err(offset: usize, kind: ErrorKind) -> Error {
    Error { offset, kind }
}

#[derive(Default)]
struct Budget {
    items: usize,
    owned: usize,
}

impl Budget {
    fn items(&mut self, n: usize, limits: &Limits, at: usize) -> Result<()> {
        self.items = self
            .items
            .checked_add(n)
            .ok_or_else(|| err(at, ErrorKind::CountOverflow))?;
        if self.items > limits.max_items {
            return Err(err(at, ErrorKind::LimitExceeded("aggregate items")));
        }
        Ok(())
    }

    fn bytes(&mut self, n: usize, limits: &Limits, at: usize) -> Result<()> {
        self.owned = self
            .owned
            .checked_add(n)
            .ok_or_else(|| err(at, ErrorKind::CountOverflow))?;
        if self.owned > limits.max_owned_bytes {
            return Err(err(at, ErrorKind::LimitExceeded("owned storage")));
        }
        Ok(())
    }

    fn claim<T>(&mut self, n: usize, limits: &Limits, at: usize) -> Result<()> {
        self.items(n, limits, at)?;
        self.bytes(
            n.checked_mul(size_of::<T>())
                .ok_or_else(|| err(at, ErrorKind::CountOverflow))?,
            limits,
            at,
        )
    }
}

#[derive(Clone, Copy)]
struct Width(u8);

impl Width {
    fn from_bits(bits: u8, at: usize) -> Result<Self> {
        if bits == 3 {
            Err(err(at, ErrorKind::ReservedWidth))
        } else {
            Ok(Self(bits))
        }
    }

    fn for_max(n: usize, version: Version, at: usize) -> Result<Self> {
        if n > i32::MAX as usize {
            return Err(err(at, ErrorKind::CountOverflow));
        }
        Ok(Self(if n <= 255 {
            0
        } else if n
            <= if version == Version::V201501280 {
                32767
            } else {
                65535
            }
        {
            1
        } else {
            2
        }))
    }

    fn bytes(self) -> usize {
        [1, 2, 4][self.0 as usize]
    }

    fn sentinel(self) -> u32 {
        [255, 65535, u32::MAX][self.0 as usize]
    }
}

struct Decoder<'a> {
    data: &'a [u8],
    pos: usize,
    limits: &'a Limits,
    budget: Budget,
    version: Version,
}

impl<'a> Decoder<'a> {
    fn take(&mut self, n: usize) -> Result<&'a [u8]> {
        if n > self.data.len() - self.pos {
            return Err(err(self.pos, ErrorKind::UnexpectedEof));
        }
        let start = self.pos;
        self.pos += n;
        Ok(&self.data[start..self.pos])
    }

    fn byte(&mut self) -> Result<u8> {
        Ok(self.take(1)?[0])
    }

    fn raw(&mut self, width: Width) -> Result<u32> {
        let mut bytes = [0; 4];
        bytes[..width.bytes()].copy_from_slice(self.take(width.bytes())?);
        Ok(u32::from_le_bytes(bytes))
    }

    fn count(&mut self, width: Width) -> Result<usize> {
        let at = self.pos;
        let n = self.raw(width)?;
        if (width.0 == 2 && n > i32::MAX as u32)
            || (width.0 == 1 && self.version == Version::V201501280 && n > i16::MAX as u32)
        {
            return Err(err(at, ErrorKind::NegativeCount));
        }
        Ok(n as usize)
    }

    fn reserve<T>(&mut self, n: usize, minimum_wire_bytes: usize) -> Result<Vec<T>> {
        if n > (self.data.len() - self.pos) / minimum_wire_bytes.max(1) {
            return Err(err(self.pos, ErrorKind::UnexpectedEof));
        }
        self.budget.claim::<T>(n, self.limits, self.pos)?;
        let mut out = Vec::new();
        out.try_reserve_exact(n)
            .map_err(|_| err(self.pos, ErrorKind::Allocation))?;
        Ok(out)
    }

    fn push<T>(&mut self, out: &mut Vec<T>, item: T, maximum: usize) -> Result<()> {
        self.budget.items(1, self.limits, self.pos)?;
        if out.len() == out.capacity() {
            let capacity = out.capacity().saturating_mul(2).max(1).min(maximum);
            let bytes = (capacity - out.capacity())
                .checked_mul(size_of::<T>())
                .ok_or_else(|| err(self.pos, ErrorKind::CountOverflow))?;
            self.budget.bytes(bytes, self.limits, self.pos)?;
            out.try_reserve_exact(capacity - out.len())
                .map_err(|_| err(self.pos, ErrorKind::Allocation))?;
        }
        out.push(item);
        Ok(())
    }

    fn string(&mut self) -> Result<String> {
        let at = self.pos;
        let remaining = &self.data[at..];
        let n = remaining
            .iter()
            .take(self.limits.max_string_bytes.saturating_add(1))
            .position(|b| *b == 0)
            .ok_or_else(|| {
                err(
                    at,
                    if remaining.len() > self.limits.max_string_bytes {
                        ErrorKind::LimitExceeded("string bytes")
                    } else {
                        ErrorKind::UnterminatedString
                    },
                )
            })?;
        let raw = self.take(n + 1)?;
        if !raw[..n].is_ascii() {
            return Err(err(at, ErrorKind::NonAscii));
        }
        self.budget.bytes(n, self.limits, at)?;
        let mut value = String::new();
        value
            .try_reserve_exact(n)
            .map_err(|_| err(at, ErrorKind::Allocation))?;
        value.push_str(std::str::from_utf8(&raw[..n]).map_err(|_| err(at, ErrorKind::NonAscii))?);
        Ok(value)
    }

    fn table(&mut self, width: Width) -> Result<Vec<String>> {
        let n = self.count(width)?;
        let mut out = self.reserve(n, 1)?;
        for _ in 0..n {
            out.push(self.string()?);
        }
        Ok(out)
    }

    fn float(&mut self) -> Result<f32> {
        let at = self.pos;
        let value = f32::from_bits(self.raw(Width(2))?);
        if !value.is_finite() {
            return Err(err(at, ErrorKind::NonFinite));
        }
        Ok(value)
    }

    fn vector(&mut self) -> Result<Vec3> {
        Ok([self.float()?, self.float()?, self.float()?])
    }

    fn vectors(&mut self, n: usize) -> Result<Vec<Vec3>> {
        let mut out = self.reserve(n, 12)?;
        for _ in 0..n {
            out.push(self.vector()?);
        }
        Ok(out)
    }

    fn uv(&mut self) -> Result<Uv> {
        Ok(Uv {
            u: self.vector()?,
            v: self.vector()?,
            u_scale: self.float()?,
            u_shift: self.float()?,
            v_scale: self.float()?,
            v_shift: self.float()?,
        })
    }

    fn properties<T>(
        &mut self,
        names: usize,
        width: Width,
        mut value: impl FnMut(&mut Self) -> Result<T>,
    ) -> Result<Vec<Property<T>>> {
        let mut out = Vec::new();
        if names == 0 {
            return Ok(out);
        }
        loop {
            let at = self.pos;
            let index = self.raw(width)?;
            if index == width.sentinel() {
                break;
            }
            if index as usize >= names {
                return Err(err(at, ErrorKind::PropertyIndex(index)));
            }
            if out.len() >= names {
                return Err(err(at, ErrorKind::DuplicateProperty(index)));
            }
            let entry = Property {
                index,
                value: value(self)?,
            };
            self.push(&mut out, entry, names)?;
        }
        out.sort_unstable_by_key(|p| p.index);
        for pair in out.windows(2) {
            if pair[0].index == pair[1].index {
                return Err(err(self.pos, ErrorKind::DuplicateProperty(pair[0].index)));
            }
        }
        Ok(out)
    }

    fn brush(&mut self, world: &World, sw: Width, uw: Width) -> Result<Brush> {
        let at = self.pos;
        let flags = self.byte()?;
        if flags & 0xf0 != 0 {
            return Err(err(at, ErrorKind::InvalidFlags(flags)));
        }
        let pw = Width::from_bits(flags & 3, at)?;
        let vw = Width::from_bits((flags >> 2) & 3, at)?;
        let n = self.count(pw)?;
        let mut planes = self.reserve(n, 36)?;
        for _ in 0..n {
            let strings = self.properties(world.string_properties.len(), sw, Self::string)?;
            let uv = self.properties(world.uv_properties.len(), uw, Self::uv)?;
            let p0 = self.vector()?;
            let u = self.vector()?;
            let v = self.vector()?;
            let polygon = if world.computed {
                let n = self.count(vw)?;
                Some(self.vectors(n)?)
            } else {
                None
            };
            planes.push(Plane {
                p0,
                u,
                v,
                strings,
                uv,
                polygon,
            });
        }
        Ok(Brush { planes })
    }

    fn patch(&mut self, computed: bool) -> Result<Patch9> {
        let at = self.pos;
        let flags = self.byte()?;
        if flags & 0xf0 != 0 {
            return Err(err(at, ErrorKind::InvalidFlags(flags)));
        }
        let aw = Width::from_bits(flags & 3, at)?;
        let qw = Width::from_bits((flags >> 2) & 3, at)?;
        let au = self.count(aw)?;
        let av = self.count(aw)?;
        let qu = self.count(qw)?;
        let qv = self.count(qw)?;
        let n = product(au, av, at)?;
        let anchors = self.vectors(n)?;
        let computed = if computed {
            let n = grid_size(au, av, qu, qv, at)?;
            Some(PatchComputed {
                points: self.vectors(n)?,
                normals: self.vectors(n)?,
            })
        } else {
            None
        };
        Ok(Patch9 {
            anchors_u: au as u32,
            anchors_v: av as u32,
            quads_u: qu as u32,
            quads_v: qv as u32,
            anchors,
            computed,
        })
    }
}

fn product(a: usize, b: usize, at: usize) -> Result<usize> {
    a.checked_mul(b)
        .ok_or_else(|| err(at, ErrorKind::CountOverflow))
}

fn grid_size(au: usize, av: usize, qu: usize, qv: usize, at: usize) -> Result<usize> {
    let u = product(au / 2, qu, at)?
        .checked_add(1)
        .ok_or_else(|| err(at, ErrorKind::CountOverflow))?;
    let v = product(av / 2, qv, at)?
        .checked_add(1)
        .ok_or_else(|| err(at, ErrorKind::CountOverflow))?;
    product(u, v, at)
}

/// Decode one complete document. Unknown bits, selector 3, non-ASCII strings,
/// duplicate slots, non-finite floats and trailing bytes are rejected.
pub fn decode(data: &[u8], limits: &Limits) -> Result<Document> {
    if data.len() > limits.max_bytes {
        return Err(err(0, ErrorKind::LimitExceeded("file bytes")));
    }
    let mut r = Decoder {
        data,
        pos: 0,
        limits,
        budget: Budget::default(),
        version: Version::V201512200,
    };
    if r.take(8)? != b"Mb3World" {
        return Err(err(0, ErrorKind::InvalidMagic));
    }
    r.version = match r.raw(Width(2))? {
        201501280 => Version::V201501280,
        201512200 => Version::V201512200,
        n => return Err(err(8, ErrorKind::UnsupportedVersion(n))),
    };
    let flags = r.byte()?;
    let old = r.version == Version::V201501280;
    if (old && flags > 1) || (!old && flags & 0xe0 != 0) {
        return Err(err(12, ErrorKind::InvalidFlags(flags)));
    }
    let sw = Width::from_bits(if old { 0 } else { (flags >> 1) & 3 }, 12)?;
    let uw = Width::from_bits(if old { 0 } else { (flags >> 3) & 3 }, 12)?;
    let mut world = World {
        computed: flags & 1 != 0,
        ..World::default()
    };
    if !old {
        world.string_properties = r.table(sw)?;
        world.uv_properties = r.table(uw)?;
    }
    let brushes = r.count(Width(2))?;
    let patches = if old { Some(r.count(Width(2))?) } else { None };
    world.brushes = r.reserve(brushes, 2)?;
    for _ in 0..brushes {
        let brush = r.brush(&world, sw, uw)?;
        world.brushes.push(brush);
    }
    let patches = match patches {
        Some(n) => n,
        None => r.count(Width(2))?,
    };
    world.patches = r.reserve(patches, 5)?;
    for _ in 0..patches {
        world.patches.push(r.patch(world.computed)?);
    }
    if r.pos != data.len() {
        return Err(err(r.pos, ErrorKind::TrailingData));
    }
    Ok(Document {
        version: r.version,
        world,
    })
}

// Growth stays within the wire-byte limit and uses fallible allocation.
fn append(out: &mut Vec<u8>, bytes: &[u8], limit: usize) -> Result<()> {
    let at = out.len();
    let n = at
        .checked_add(bytes.len())
        .ok_or_else(|| err(at, ErrorKind::CountOverflow))?;
    if n > limit {
        return Err(err(at, ErrorKind::LimitExceeded("file bytes")));
    }
    if n > out.capacity() {
        let capacity = n.max(out.capacity().saturating_mul(2)).min(limit);
        out.try_reserve_exact(capacity - at)
            .map_err(|_| err(at, ErrorKind::Allocation))?;
    }
    out.extend_from_slice(bytes);
    Ok(())
}

struct Encoder<'a> {
    data: Vec<u8>,
    limits: &'a Limits,
    budget: Budget,
    version: Version,
}

impl Encoder<'_> {
    fn bytes(&mut self, bytes: &[u8]) -> Result<()> {
        append(&mut self.data, bytes, self.limits.max_bytes)
    }

    fn byte(&mut self, value: u8) -> Result<()> {
        self.bytes(&[value])
    }

    fn raw(&mut self, value: u32, width: Width) -> Result<()> {
        self.bytes(&value.to_le_bytes()[..width.bytes()])
    }

    fn count(&mut self, value: usize, width: Width) -> Result<()> {
        let max = match width.0 {
            0 => 255,
            1 if self.version == Version::V201501280 => 32767,
            1 => 65535,
            _ => i32::MAX as usize,
        };
        if value > max {
            return Err(err(self.data.len(), ErrorKind::CountOverflow));
        }
        self.raw(value as u32, width)
    }

    fn items(&mut self, n: usize) -> Result<()> {
        self.budget.items(n, self.limits, self.data.len())
    }

    fn string(&mut self, s: &str) -> Result<()> {
        let at = self.data.len();
        if s.len() > self.limits.max_string_bytes {
            return Err(err(at, ErrorKind::LimitExceeded("string bytes")));
        }
        if !s.is_ascii() {
            return Err(err(at, ErrorKind::NonAscii));
        }
        if s.as_bytes().contains(&0) {
            return Err(err(at, ErrorKind::EmbeddedNul));
        }
        self.bytes(s.as_bytes())?;
        self.byte(0)
    }

    fn table(&mut self, names: &[String], width: Width) -> Result<()> {
        self.items(names.len())?;
        self.count(names.len(), width)?;
        for name in names {
            self.string(name)?;
        }
        Ok(())
    }

    fn float(&mut self, value: f32) -> Result<()> {
        if !value.is_finite() {
            return Err(err(self.data.len(), ErrorKind::NonFinite));
        }
        self.bytes(&value.to_le_bytes())
    }

    fn vector(&mut self, value: &Vec3) -> Result<()> {
        for f in value {
            self.float(*f)?;
        }
        Ok(())
    }

    fn vectors(&mut self, values: &[Vec3]) -> Result<()> {
        self.items(values.len())?;
        for value in values {
            self.vector(value)?;
        }
        Ok(())
    }

    fn uv(&mut self, value: &Uv) -> Result<()> {
        self.vector(&value.u)?;
        self.vector(&value.v)?;
        for f in [value.u_scale, value.u_shift, value.v_scale, value.v_shift] {
            self.float(f)?;
        }
        Ok(())
    }

    fn properties<T>(
        &mut self,
        values: &[Property<T>],
        names: usize,
        width: Width,
        mut write: impl FnMut(&mut Self, &T) -> Result<()>,
    ) -> Result<()> {
        self.items(values.len())?;
        let mut previous = None;
        for p in values {
            if p.index as usize >= names {
                return Err(err(self.data.len(), ErrorKind::PropertyIndex(p.index)));
            }
            if previous == Some(p.index) {
                return Err(err(self.data.len(), ErrorKind::DuplicateProperty(p.index)));
            }
            if previous.is_some_and(|n| n > p.index) {
                return Err(err(self.data.len(), ErrorKind::UnsortedProperties));
            }
            previous = Some(p.index);
            self.raw(p.index, width)?;
            write(self, &p.value)?;
        }
        if names != 0 {
            self.raw(width.sentinel(), width)?;
        }
        Ok(())
    }

    fn brush(&mut self, brush: &Brush, world: &World, sw: Width, uw: Width) -> Result<()> {
        let at = self.data.len();
        self.items(brush.planes.len())?;
        let pw = Width::for_max(brush.planes.len(), self.version, at)?;
        let mut largest = 0;
        for plane in &brush.planes {
            if plane.polygon.is_some() != world.computed {
                return Err(err(at, ErrorKind::InconsistentComputedData));
            }
            largest = largest.max(plane.polygon.as_ref().map_or(0, Vec::len));
        }
        let vw = Width::for_max(largest, self.version, at)?;
        self.byte(pw.0 | vw.0 << 2)?;
        self.count(brush.planes.len(), pw)?;
        for plane in &brush.planes {
            if self.version == Version::V201501280
                && (!plane.strings.is_empty() || !plane.uv.is_empty())
            {
                return Err(err(self.data.len(), ErrorKind::LegacyProperties));
            }
            self.properties(&plane.strings, world.string_properties.len(), sw, |e, s| {
                e.string(s)
            })?;
            self.properties(&plane.uv, world.uv_properties.len(), uw, Self::uv)?;
            self.vector(&plane.p0)?;
            self.vector(&plane.u)?;
            self.vector(&plane.v)?;
            if let Some(polygon) = &plane.polygon {
                self.count(polygon.len(), vw)?;
                self.vectors(polygon)?;
            }
        }
        Ok(())
    }

    fn patch(&mut self, patch: &Patch9, computed: bool) -> Result<()> {
        let at = self.data.len();
        if patch.computed.is_some() != computed {
            return Err(err(at, ErrorKind::InconsistentComputedData));
        }
        let au = patch.anchors_u as usize;
        let av = patch.anchors_v as usize;
        let qu = patch.quads_u as usize;
        let qv = patch.quads_v as usize;
        if patch.anchors.len() != product(au, av, at)? {
            return Err(err(at, ErrorKind::InvalidPatchLength));
        }
        if let Some(data) = &patch.computed {
            let n = grid_size(au, av, qu, qv, at)?;
            if data.points.len() != n || data.normals.len() != n {
                return Err(err(at, ErrorKind::InvalidPatchLength));
            }
        }
        let aw = Width::for_max(au.max(av), self.version, at)?;
        let qw = Width::for_max(qu.max(qv), self.version, at)?;
        self.byte(aw.0 | qw.0 << 2)?;
        self.count(au, aw)?;
        self.count(av, aw)?;
        self.count(qu, qw)?;
        self.count(qv, qw)?;
        self.vectors(&patch.anchors)?;
        if let Some(data) = &patch.computed {
            self.vectors(&data.points)?;
            self.vectors(&data.normals)?;
        }
        Ok(())
    }
}

/// Encode a document without optimizing properties or dropping data. Legacy
/// output refuses properties and uses four-byte counts for 32768..=65535.
/// All validation completes before bytes are returned or any I/O is attempted.
pub fn encode(document: &Document, limits: &Limits) -> Result<Vec<u8>> {
    let world = &document.world;
    let old = document.version == Version::V201501280;
    if old && (!world.string_properties.is_empty() || !world.uv_properties.is_empty()) {
        return Err(err(0, ErrorKind::LegacyProperties));
    }
    let mut e = Encoder {
        data: Vec::new(),
        limits,
        budget: Budget::default(),
        version: document.version,
    };
    let sw = Width::for_max(world.string_properties.len(), document.version, 12)?;
    let uw = Width::for_max(world.uv_properties.len(), document.version, 12)?;
    e.bytes(b"Mb3World")?;
    e.raw(document.version as u32, Width(2))?;
    e.byte(u8::from(world.computed) | sw.0 << 1 | uw.0 << 3)?;
    if !old {
        e.table(&world.string_properties, sw)?;
        e.table(&world.uv_properties, uw)?;
    }
    e.items(world.brushes.len())?;
    e.count(world.brushes.len(), Width(2))?;
    if old {
        e.count(world.patches.len(), Width(2))?;
    }
    for brush in &world.brushes {
        e.brush(brush, world, sw, uw)?;
    }
    e.items(world.patches.len())?;
    if !old {
        e.count(world.patches.len(), Width(2))?;
    }
    for patch in &world.patches {
        e.patch(patch, world.computed)?;
    }
    Ok(e.data)
}

/// Read until EOF, bounded to `max_bytes + 1` observed bytes. Short reads and
/// Interrupted are supported. The buffered input is then strictly decoded.
pub fn read(mut source: impl Read, limits: &Limits) -> Result<Document> {
    let mut data = Vec::new();
    let mut chunk = [0; 8192];
    loop {
        let n = limits
            .max_bytes
            .saturating_sub(data.len())
            .saturating_add(1)
            .min(chunk.len());
        let n = match source.read(&mut chunk[..n]) {
            Ok(0) => break,
            Ok(n) => n,
            Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
            Err(e) => return Err(err(data.len(), ErrorKind::Io(e))),
        };
        append(&mut data, &chunk[..n], limits.max_bytes)?;
    }
    decode(&data, limits)
}

/// Validate and encode before touching the destination. I/O failure can still
/// leave a written prefix; callers needing atomic files must save via a temp file.
pub fn write(mut destination: impl Write, document: &Document, limits: &Limits) -> Result<()> {
    let data = encode(document, limits)?;
    let mut written = 0;
    while written != data.len() {
        match destination.write(&data[written..]) {
            Ok(0) => return Err(err(written, ErrorKind::Io(io::ErrorKind::WriteZero.into()))),
            Ok(n) => written += n,
            Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
            Err(e) => return Err(err(written, ErrorKind::Io(e))),
        }
    }
    Ok(())
}