//! Heat2 decoder.
//!
//! Input is bytes off a socket, so every read is bounds-checked and every
//! failure is an `Error`, never a panic. The oracle (`heat2.py`) indexes
//! optimistically and would raise on malformed input; matching its *bytes* is
//! required, matching its *crash behaviour* is not.

use super::tag::Tag;
use super::value::{Struct, TypeId, Value, UNION_UNSET};
use super::varint;
use crate::error::{Error, Result};

/// Decode a top-level payload body (unterminated, delimited by `buf`).
pub fn decode(buf: &[u8]) -> Result<Struct> {
    let (s, _) = decode_members(buf, 0, buf.len(), false)?;
    Ok(s)
}

/// Read members until `end` (top level) or a `0x00` terminator (nested).
fn decode_members(
    buf: &[u8],
    mut i: usize,
    end: usize,
    terminated: bool,
) -> Result<(Struct, usize)> {
    let mut fields = Vec::new();
    while i < end {
        if terminated && buf[i] == 0x00 {
            i += 1;
            break;
        }
        if i + 4 > end {
            return Err(Error::Truncated {
                what: "field header",
                need: 4,
                have: end - i,
            });
        }
        let tag = Tag([buf[i], buf[i + 1], buf[i + 2]]);
        let type_byte = buf[i + 3];
        let ty = TypeId::from_byte(type_byte).ok_or(Error::UnknownType {
            type_byte,
            at: i + 3,
        })?;
        i += 4;
        let (value, next) = decode_value(buf, i, end, ty)?;
        i = next;
        fields.push((tag, value));
    }
    Ok((Struct { fields }, i))
}

fn decode_value(buf: &[u8], i: usize, end: usize, ty: TypeId) -> Result<(Value, usize)> {
    match ty {
        TypeId::Int => {
            let (v, i) = varint::decode(buf, i)?;
            Ok((Value::Int(v), i))
        }

        TypeId::String => {
            let (len, i) = varint::decode(buf, i)?;
            let len = checked_len(len, i, end, "string")?;
            let raw = &buf[i..i + len];
            // The declared length includes the NUL; strip any trailing NULs so
            // the value round-trips through the encoder unchanged.
            let cut = raw.iter().rposition(|&b| b != 0).map_or(0, |p| p + 1);
            let s = String::from_utf8_lossy(&raw[..cut]).into_owned();
            Ok((Value::String(s), i + len))
        }

        TypeId::Blob => {
            let (len, i) = varint::decode(buf, i)?;
            let len = checked_len(len, i, end, "blob")?;
            Ok((Value::Blob(buf[i..i + len].to_vec()), i + len))
        }

        TypeId::Struct => {
            let (s, i) = decode_members(buf, i, end, true)?;
            Ok((Value::Struct(s), i))
        }

        TypeId::List => {
            let elem_byte = *byte_at(buf, i, end, "list element type")?;
            let elem = TypeId::from_byte(elem_byte).ok_or(Error::UnknownType {
                type_byte: elem_byte,
                at: i,
            })?;
            let (count, mut i) = varint::decode(buf, i + 1)?;
            let count = checked_count(count, end - i.min(end), "list")?;
            let mut items = Vec::with_capacity(count.min(1024));
            for _ in 0..count {
                let (v, next) = decode_value(buf, i, end, elem)?;
                i = next;
                items.push(v);
            }
            Ok((Value::List { elem, items }, i))
        }

        TypeId::Map => {
            let key_byte = *byte_at(buf, i, end, "map key type")?;
            let val_byte = *byte_at(buf, i + 1, end, "map value type")?;
            let key = TypeId::from_byte(key_byte).ok_or(Error::UnknownType {
                type_byte: key_byte,
                at: i,
            })?;
            let val = TypeId::from_byte(val_byte).ok_or(Error::UnknownType {
                type_byte: val_byte,
                at: i + 1,
            })?;
            let (count, mut i) = varint::decode(buf, i + 2)?;
            let count = checked_count(count, end - i.min(end), "map")?;
            let mut entries = Vec::with_capacity(count.min(1024));
            for _ in 0..count {
                let (k, next) = decode_value(buf, i, end, key)?;
                let (v, next) = decode_value(buf, next, end, val)?;
                i = next;
                entries.push((k, v));
            }
            Ok((Value::Map { key, val, entries }, i))
        }

        TypeId::Union => {
            let key = *byte_at(buf, i, end, "union discriminator")?;
            let i = i + 1;
            if key == UNION_UNSET {
                return Ok((Value::Union { key, member: None }, i));
            }
            if i + 4 > end {
                return Err(Error::Truncated {
                    what: "union member header",
                    need: 4,
                    have: end.saturating_sub(i),
                });
            }
            let tag = Tag([buf[i], buf[i + 1], buf[i + 2]]);
            let mtype_byte = buf[i + 3];
            let mtype = TypeId::from_byte(mtype_byte).ok_or(Error::UnknownType {
                type_byte: mtype_byte,
                at: i + 3,
            })?;
            let (value, i) = decode_value(buf, i + 4, end, mtype)?;
            Ok((
                Value::Union {
                    key,
                    member: Some(Box::new((tag, value))),
                },
                i,
            ))
        }

        TypeId::VarList => {
            let (count, mut i) = varint::decode(buf, i)?;
            let count = checked_count(count, end - i.min(end), "varlist")?;
            let mut items = Vec::with_capacity(count.min(1024));
            for _ in 0..count {
                let (v, next) = varint::decode(buf, i)?;
                i = next;
                items.push(v);
            }
            Ok((Value::VarList(items), i))
        }

        TypeId::ObjType => {
            let (component, i) = varint::decode(buf, i)?;
            let (ty, i) = varint::decode(buf, i)?;
            Ok((Value::ObjType { component, ty }, i))
        }

        TypeId::ObjId => {
            let (component, i) = varint::decode(buf, i)?;
            let (ty, i) = varint::decode(buf, i)?;
            let (id, i) = varint::decode(buf, i)?;
            Ok((Value::ObjId { component, ty, id }, i))
        }

        TypeId::Float => {
            if i + 4 > end {
                return Err(Error::Truncated {
                    what: "float",
                    need: 4,
                    have: end.saturating_sub(i),
                });
            }
            let f = f32::from_be_bytes([buf[i], buf[i + 1], buf[i + 2], buf[i + 3]]);
            Ok((Value::Float(f), i + 4))
        }
    }
}

fn byte_at<'a>(buf: &'a [u8], i: usize, end: usize, what: &'static str) -> Result<&'a u8> {
    if i >= end {
        return Err(Error::Truncated {
            what,
            need: 1,
            have: 0,
        });
    }
    buf.get(i).ok_or(Error::Truncated {
        what,
        need: 1,
        have: 0,
    })
}

fn checked_len(len: i64, i: usize, end: usize, what: &'static str) -> Result<usize> {
    let available = end.saturating_sub(i);
    if len < 0 || len as u64 > available as u64 {
        return Err(Error::Truncated {
            what,
            need: len.max(0) as usize,
            have: available,
        });
    }
    Ok(len as usize)
}

/// Reject a declared element count that cannot fit in the remaining bytes.
///
/// Without this a two-byte varint can ask for a billion elements and the
/// allocation, not the parse, becomes the failure.
fn checked_count(count: i64, remaining: usize, what: &'static str) -> Result<usize> {
    if count < 0 || count as u64 > remaining as u64 {
        return Err(Error::Truncated {
            what,
            need: count.max(0) as usize,
            have: remaining,
        });
    }
    Ok(count as usize)
}

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

    fn round_trip(s: Struct) {
        let bytes = encode(&s);
        let back = decode(&bytes).expect("decodes");
        assert_eq!(encode(&back), bytes, "re-encode must be byte-identical");
    }

    #[test]
    fn round_trips_scalars() {
        round_trip(
            Struct::new()
                .with("INTV", Value::Int(0x2000))
                .with("STRV", Value::String("hello".into()))
                .with("BLBV", Value::Blob(vec![0, 1, 2, 255]))
                .with("FLTV", Value::Float(-0.25)),
        );
    }

    #[test]
    fn round_trips_nesting() {
        round_trip(
            Struct::new()
                .with(
                    "OUTR",
                    Value::Struct(
                        Struct::new()
                            .with(
                                "INNR",
                                Value::Struct(Struct::new().with("LEAF", Value::Int(42))),
                            )
                            .with("SIBL", Value::String("s".into())),
                    ),
                )
                .with("TAIL", Value::Int(9)),
        );
    }

    #[test]
    fn nested_terminator_does_not_swallow_following_members() {
        let s = Struct::new()
            .with(
                "AAAA",
                Value::Struct(Struct::new().with("X", Value::Int(1))),
            )
            .with("BBBB", Value::Int(7));
        let back = decode(&encode(&s)).unwrap();
        assert_eq!(back.len(), 2);
        assert_eq!(back.get("BBBB").and_then(Value::as_int), Some(7));
    }

    #[test]
    fn rejects_an_unknown_type_byte() {
        // tag "AAAA" then type 0x0B, which is not a Heat2 type.
        let mut bytes = Tag::from_label("AAAA").as_bytes().to_vec();
        bytes.push(0x0B);
        assert!(matches!(
            decode(&bytes),
            Err(Error::UnknownType {
                type_byte: 0x0B,
                ..
            })
        ));
    }

    #[test]
    fn rejects_a_string_longer_than_the_buffer() {
        let mut bytes = Tag::from_label("S").as_bytes().to_vec();
        bytes.push(TypeId::String.as_byte());
        bytes.push(0x3F); // claims 63 bytes
        bytes.extend_from_slice(b"short");
        assert!(matches!(decode(&bytes), Err(Error::Truncated { .. })));
    }

    #[test]
    fn rejects_an_absurd_list_count_without_allocating() {
        let mut bytes = Tag::from_label("L").as_bytes().to_vec();
        bytes.push(TypeId::List.as_byte());
        bytes.push(TypeId::Int.as_byte());
        bytes.extend_from_slice(&[0xFF, 0xFF, 0xFF, 0x7F]); // huge count
        assert!(matches!(decode(&bytes), Err(Error::Truncated { .. })));
    }

    #[test]
    fn rejects_a_truncated_field_header() {
        let bytes = vec![0x96, 0xed]; // two bytes of a three-byte tag
        assert!(decode(&bytes).is_err());
    }

    #[test]
    fn never_panics_on_arbitrary_bytes() {
        // Cheap structured fuzz: every 3-byte prefix followed by each type byte.
        for type_byte in 0x00u8..=0x0C {
            for pattern in [0x00u8, 0x01, 0x7F, 0x80, 0xFF] {
                let bytes = vec![pattern, pattern, pattern, type_byte, pattern, pattern];
                let _ = decode(&bytes); // must return, not panic
            }
        }
    }
}