//! Deterministic CBOR (RFC 8949 §4.2.1) restricted to the profile of spec §5.1:
//! unsigned integers, byte strings, text strings, arrays, maps with unsigned
//! integer keys in increasing order, booleans, and null. Everything else is
//! rejected on decode and cannot be expressed on encode.

use std::collections::BTreeMap;

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Value {
    Uint(u64),
    Bytes(Vec<u8>),
    Text(String),
    Array(Vec<Value>),
    Map(BTreeMap<u64, Value>),
    Bool(bool),
    Null,
}

#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum CborError {
    #[error("unexpected end of input")]
    UnexpectedEnd,
    #[error("non-canonical encoding: {0}")]
    NonCanonical(&'static str),
    #[error("unsupported CBOR item: {0}")]
    Unsupported(&'static str),
    #[error("trailing bytes after value")]
    TrailingBytes,
    #[error("invalid UTF-8 in text string")]
    InvalidUtf8,
    #[error("map keys must be strictly increasing unsigned integers")]
    BadMapKey,
    #[error("length exceeds input")]
    LengthOverflow,
}

impl Value {
    pub fn map(pairs: impl IntoIterator<Item = (u64, Value)>) -> Value {
        Value::Map(pairs.into_iter().collect())
    }
    pub fn as_uint(&self) -> Option<u64> {
        if let Value::Uint(n) = self { Some(*n) } else { None }
    }
    pub fn as_bytes(&self) -> Option<&[u8]> {
        if let Value::Bytes(b) = self { Some(b) } else { None }
    }
    pub fn as_text(&self) -> Option<&str> {
        if let Value::Text(t) = self { Some(t) } else { None }
    }
    pub fn as_array(&self) -> Option<&[Value]> {
        if let Value::Array(a) = self { Some(a) } else { None }
    }
    pub fn as_map(&self) -> Option<&BTreeMap<u64, Value>> {
        if let Value::Map(m) = self { Some(m) } else { None }
    }
    pub fn as_bool(&self) -> Option<bool> {
        if let Value::Bool(b) = self { Some(*b) } else { None }
    }
    pub fn is_null(&self) -> bool {
        matches!(self, Value::Null)
    }
}

pub fn encode(v: &Value) -> Vec<u8> {
    let mut out = Vec::new();
    encode_into(v, &mut out);
    out
}

pub fn encode_into(v: &Value, out: &mut Vec<u8>) {
    match v {
        Value::Uint(n) => head(out, 0, *n),
        Value::Bytes(b) => {
            head(out, 2, b.len() as u64);
            out.extend_from_slice(b);
        }
        Value::Text(t) => {
            head(out, 3, t.len() as u64);
            out.extend_from_slice(t.as_bytes());
        }
        Value::Array(a) => {
            head(out, 4, a.len() as u64);
            for x in a {
                encode_into(x, out);
            }
        }
        Value::Map(m) => {
            head(out, 5, m.len() as u64);
            for (k, x) in m {
                head(out, 0, *k);
                encode_into(x, out);
            }
        }
        Value::Bool(false) => out.push(0xf4),
        Value::Bool(true) => out.push(0xf5),
        Value::Null => out.push(0xf6),
    }
}

fn head(out: &mut Vec<u8>, major: u8, n: u64) {
    let m = major << 5;
    if n < 24 {
        out.push(m | n as u8);
    } else if n <= 0xff {
        out.push(m | 24);
        out.push(n as u8);
    } else if n <= 0xffff {
        out.push(m | 25);
        out.extend_from_slice(&(n as u16).to_be_bytes());
    } else if n <= 0xffff_ffff {
        out.push(m | 26);
        out.extend_from_slice(&(n as u32).to_be_bytes());
    } else {
        out.push(m | 27);
        out.extend_from_slice(&n.to_be_bytes());
    }
}

pub fn decode(bytes: &[u8]) -> Result<Value, CborError> {
    let mut d = Decoder { b: bytes, pos: 0, depth: 0 };
    let v = d.value()?;
    if d.pos != bytes.len() {
        return Err(CborError::TrailingBytes);
    }
    Ok(v)
}

const MAX_DEPTH: u32 = 256;

struct Decoder<'a> {
    b: &'a [u8],
    pos: usize,
    depth: u32,
}

impl<'a> Decoder<'a> {
    fn byte(&mut self) -> Result<u8, CborError> {
        let x = *self.b.get(self.pos).ok_or(CborError::UnexpectedEnd)?;
        self.pos += 1;
        Ok(x)
    }

    fn take(&mut self, n: u64) -> Result<&'a [u8], CborError> {
        let n = usize::try_from(n).map_err(|_| CborError::LengthOverflow)?;
        let end = self.pos.checked_add(n).ok_or(CborError::LengthOverflow)?;
        if end > self.b.len() {
            return Err(CborError::UnexpectedEnd);
        }
        let s = &self.b[self.pos..end];
        self.pos = end;
        Ok(s)
    }

    /// Reads a head for majors 0..=5 and returns (major, argument), enforcing
    /// shortest-form arguments.
    fn head(&mut self) -> Result<(u8, u64), CborError> {
        let ib = self.byte()?;
        let major = ib >> 5;
        let ai = ib & 0x1f;
        if major == 7 {
            return Err(CborError::Unsupported("simple value or float where a length or integer was expected"));
        }
        let n = match ai {
            0..=23 => ai as u64,
            24 => {
                let n = self.byte()? as u64;
                if n < 24 {
                    return Err(CborError::NonCanonical("argument fits in the initial byte"));
                }
                n
            }
            25 => {
                let n = u16::from_be_bytes(self.take(2)?.try_into().unwrap()) as u64;
                if n <= 0xff {
                    return Err(CborError::NonCanonical("argument fits in one byte"));
                }
                n
            }
            26 => {
                let n = u32::from_be_bytes(self.take(4)?.try_into().unwrap()) as u64;
                if n <= 0xffff {
                    return Err(CborError::NonCanonical("argument fits in two bytes"));
                }
                n
            }
            27 => {
                let n = u64::from_be_bytes(self.take(8)?.try_into().unwrap());
                if n <= 0xffff_ffff {
                    return Err(CborError::NonCanonical("argument fits in four bytes"));
                }
                n
            }
            28..=30 => return Err(CborError::Unsupported("reserved additional information")),
            _ => return Err(CborError::Unsupported("indefinite-length item")),
        };
        Ok((major, n))
    }

    fn value(&mut self) -> Result<Value, CborError> {
        self.depth += 1;
        if self.depth > MAX_DEPTH {
            return Err(CborError::Unsupported("nesting deeper than 256"));
        }
        let ib = *self.b.get(self.pos).ok_or(CborError::UnexpectedEnd)?;
        let v = if ib >> 5 == 7 {
            self.pos += 1;
            match ib {
                0xf4 => Value::Bool(false),
                0xf5 => Value::Bool(true),
                0xf6 => Value::Null,
                _ => return Err(CborError::Unsupported("simple value, float, undefined, or break")),
            }
        } else {
            let (major, n) = self.head()?;
            match major {
                0 => Value::Uint(n),
                1 => return Err(CborError::Unsupported("negative integer")),
                2 => Value::Bytes(self.take(n)?.to_vec()),
                3 => Value::Text(
                    std::str::from_utf8(self.take(n)?)
                        .map_err(|_| CborError::InvalidUtf8)?
                        .to_owned(),
                ),
                4 => {
                    let mut a = Vec::new();
                    for _ in 0..n {
                        a.push(self.value()?);
                    }
                    Value::Array(a)
                }
                5 => {
                    let mut m = BTreeMap::new();
                    let mut last: Option<u64> = None;
                    for _ in 0..n {
                        let (km, k) = self.head()?;
                        if km != 0 {
                            return Err(CborError::BadMapKey);
                        }
                        if let Some(l) = last {
                            if k <= l {
                                return Err(CborError::BadMapKey);
                            }
                        }
                        last = Some(k);
                        let v = self.value()?;
                        m.insert(k, v);
                    }
                    Value::Map(m)
                }
                _ => return Err(CborError::Unsupported("tag")),
            }
        };
        self.depth -= 1;
        Ok(v)
    }
}

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

    fn rt(v: Value) -> Vec<u8> {
        let b = encode(&v);
        assert_eq!(decode(&b).unwrap(), v);
        b
    }

    #[test]
    fn integers_use_shortest_form() {
        assert_eq!(rt(Value::Uint(0)), [0x00]);
        assert_eq!(rt(Value::Uint(23)), [0x17]);
        assert_eq!(rt(Value::Uint(24)), [0x18, 24]);
        assert_eq!(rt(Value::Uint(255)), [0x18, 255]);
        assert_eq!(rt(Value::Uint(256)), [0x19, 1, 0]);
        assert_eq!(rt(Value::Uint(65536)), [0x1a, 0, 1, 0, 0]);
        assert_eq!(rt(Value::Uint(1 << 32)), [0x1b, 0, 0, 0, 1, 0, 0, 0, 0]);
        assert_eq!(rt(Value::Uint(u64::MAX)).len(), 9);
    }

    #[test]
    fn rejects_non_shortest_integers() {
        assert_eq!(decode(&[0x18, 5]), Err(CborError::NonCanonical("argument fits in the initial byte")));
        assert_eq!(decode(&[0x19, 0, 5]), Err(CborError::NonCanonical("argument fits in one byte")));
        assert_eq!(decode(&[0x1a, 0, 0, 0, 5]), Err(CborError::NonCanonical("argument fits in two bytes")));
        assert_eq!(decode(&[0x1b, 0, 0, 0, 0, 0, 0, 0, 5]), Err(CborError::NonCanonical("argument fits in four bytes")));
    }

    #[test]
    fn rejects_unsupported_items() {
        assert!(matches!(decode(&[0x20]), Err(CborError::Unsupported(_)))); // -1
        assert!(matches!(decode(&[0xc0, 0x00]), Err(CborError::Unsupported(_)))); // tag
        assert!(matches!(decode(&[0xf9, 0, 0]), Err(CborError::Unsupported(_)))); // float16
        assert!(matches!(decode(&[0x9f, 0xff]), Err(CborError::Unsupported(_)))); // indefinite array
        assert!(matches!(decode(&[0xf7]), Err(CborError::Unsupported(_)))); // undefined
    }

    #[test]
    fn maps_are_sorted_and_unique() {
        let v = Value::map([(2, Value::Uint(1)), (0, Value::Null), (1, Value::Bool(true))]);
        let b = rt(v);
        assert_eq!(b, [0xa3, 0x00, 0xf6, 0x01, 0xf5, 0x02, 0x01]);
        assert_eq!(decode(&[0xa2, 0x01, 0x00, 0x00, 0x00]), Err(CborError::BadMapKey));
        assert_eq!(decode(&[0xa2, 0x01, 0x00, 0x01, 0x00]), Err(CborError::BadMapKey));
        assert_eq!(decode(&[0xa1, 0x61, 0x61, 0x00]), Err(CborError::BadMapKey)); // text key
    }

    #[test]
    fn strings_arrays_and_trailing() {
        rt(Value::Bytes(vec![1, 2, 3]));
        rt(Value::Text("héllo".into()));
        rt(Value::Array(vec![Value::Uint(1), Value::Text("x".into()), Value::Array(vec![])]));
        assert_eq!(decode(&[0x00, 0x00]), Err(CborError::TrailingBytes));
        assert_eq!(decode(&[0x62, 0xff, 0xfe]), Err(CborError::InvalidUtf8));
        assert_eq!(decode(&[0x43, 1, 2]), Err(CborError::UnexpectedEnd));
    }

    #[test]
    fn huge_length_does_not_allocate() {
        assert_eq!(decode(&[0x5b, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff]), Err(CborError::LengthOverflow));
        assert!(decode(&[0x9b, 0x00, 0x00, 0x00, 0x01, 0, 0, 0, 0]).is_err());
    }
}