//! CBOR support.
//!
//! We can test this as follows:
//!
//! ~~~ text
//! echo 0x826161a161626163 | xxd -r | cargo run -- --from cbor
//! ~~~
//!
//! The [examples](https://www.rfc-editor.org/rfc/rfc8949.html#section-appendix.a)
//! from the CBOR specification are quite helpful here.
//! They can be pasted directly into the command above.
use crate::{invalid_data, BoxError};
use alloc::{string::String, vec::Vec};
use ciborium_io::Read;
use ciborium_ll::{simple, tag, Decoder, Error, Header};
use core::fmt::{self, Formatter};
use jaq_json::{Num, Val};
use num_bigint::{BigInt, BigUint};
use std::io;

const BUF_SIZE: usize = 4096;

/// Limit pre-allocated container capacity to this number of elements.
///
/// CBOR allows users to supply the number of elements in a container,
/// e.g. a string or an array.
/// However, when we allocate memory for the user-given number of elements,
/// we could quickly run out of memory, if the user gives a number too high.
/// Therefore, we limit the amount of pre-allocated capacity.
///
/// The worst case is an (invalid) value of the shape `[[[...`, where
/// each array is specified to have the maximal capacity.
/// Suppose that:
///
/// - 16 bytes is the size of a [`jaq_json::Val`] (as should be the case on x86_64),
/// - 1024 is the maximal capacity that we pre-allocate, and
/// - 8192 is the maximal stack size.
///
/// Then the maximum amount of memory that can be consumed by such an array is
/// 16 bytes/value * 1024 values/array * 8192 arrays = 134217728 bytes = 128MB.
const MAX_CAP: usize = 1024;

/// Error that may indicate end of file.
trait IsEof {
    /// Is the error caused by EOF?
    fn is_eof(&self) -> bool;
}

// TODO: implement IsEof for EndOfFile, which is only visible when
// `ciborium_io` is compiled without the `std` feature
impl IsEof for std::io::Error {
    fn is_eof(&self) -> bool {
        matches!(self.kind(), std::io::ErrorKind::UnexpectedEof)
    }
}

/// Parse error.
#[derive(Debug)]
enum PError<E> {
    /// Lex error, coming from ciborium
    Lex(Error<E>),
    /// unsupported simple value
    Simple(u8),
    /// unsupported tag
    Tag(u64),
    /// unexpected break
    Break,
}

impl<E: fmt::Debug + fmt::Display> std::error::Error for PError<E> {}

impl<E> From<Error<E>> for PError<E> {
    fn from(e: Error<E>) -> Self {
        Self::Lex(e)
    }
}

impl<E: fmt::Display> fmt::Display for PError<E> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        match self {
            Self::Lex(Error::Io(e)) => e.fmt(f),
            Self::Lex(Error::Syntax(offset)) => write!(f, "lexing failed at byte {offset}"),
            Self::Simple(s) => write!(f, "unsupported simple value: {s}"),
            Self::Tag(t) => write!(f, "unsupported tag: {t}"),
            Self::Break => write!(f, "unexpected break"),
        }
    }
}

impl From<PError<io::Error>> for io::Error {
    fn from(e: PError<io::Error>) -> Self {
        match e {
            PError::Lex(Error::Io(e)) => e,
            e => invalid_data(e),
        }
    }
}

/// Parse a sequence of CBOR values.
pub fn parse_many(b: &[u8]) -> impl Iterator<Item = Result<Val, BoxError>> + '_ {
    decode_many(b).map(|r| r.map_err(Into::into))
}

/// Read a sequence of CBOR values.
pub fn read_many(read: impl io::Read) -> impl Iterator<Item = io::Result<Val>> {
    decode_many(read).map(|r| r.map_err(Into::into))
}

/// Decode a sequence of CBOR values.
fn decode_many<R: Read>(read: R) -> impl Iterator<Item = Result<Val, PError<R::Error>>>
where
    R::Error: IsEof,
{
    let mut decoder = Decoder::from(read);
    core::iter::from_fn(move || {
        let offset = decoder.offset();
        match decoder.pull() {
            Ok(header) => Some(parse(header, &mut decoder)),
            // The following is a hack which is unfortunately necessary because
            // ciborium does not currently give us any means of checking for EOF
            // (<https://github.com/enarx/ciborium/issues/159>).
            // I found that an impartial header, such as triggered by 0x18,
            // advances the offset, so by negation, I conclude that
            // if the offset has remained the same after header decoding
            // AND the error is EOF, that means that there was no input.
            // Note that we cannot omit the EOF check, because a different
            // I/O error might have been triggered, and we should propagate that.
            Err(Error::Io(e)) if e.is_eof() && offset == decoder.offset() => None,
            Err(e) => Some(Err(e.into())),
        }
    })
}

fn with_size<R: Read, T>(
    size: Option<usize>,
    decoder: &mut Decoder<R>,
    f: impl Fn(Header, &mut Decoder<R>) -> Result<T, PError<R::Error>>,
) -> Result<Vec<T>, PError<R::Error>> {
    let mut a;
    if let Some(size) = size {
        a = Vec::with_capacity(size.min(MAX_CAP));
        for _ in 0..size {
            a.push(f(decoder.pull()?, decoder)?);
        }
    } else {
        a = Vec::new();
        loop {
            match decoder.pull()? {
                Header::Break => break,
                header => a.push(f(header, decoder)?),
            }
        }
    }
    a.shrink_to_fit();
    Ok(a)
}

fn parse_str<R: Read>(
    len: Option<usize>,
    decoder: &mut Decoder<R>,
) -> Result<String, Error<R::Error>> {
    let mut s = len.map_or_else(String::new, |c| String::with_capacity(c.min(MAX_CAP)));
    let mut segments = decoder.text(len);
    while let Some(mut segment) = segments.pull()? {
        let mut buffer = [0; BUF_SIZE];
        while let Some(chunk) = segment.pull(&mut buffer[..])? {
            s.push_str(chunk);
        }
    }
    s.shrink_to_fit();
    Ok(s)
}

fn parse_bytes<R: Read>(
    len: Option<usize>,
    decoder: &mut Decoder<R>,
) -> Result<Vec<u8>, Error<R::Error>> {
    let mut b = len.map_or_else(Vec::new, |c| Vec::with_capacity(c.min(MAX_CAP)));
    let mut segments = decoder.bytes(len);
    while let Some(mut segment) = segments.pull()? {
        let mut buffer = [0; BUF_SIZE];
        while let Some(chunk) = segment.pull(&mut buffer)? {
            b.extend(chunk);
        }
    }
    b.shrink_to_fit();
    Ok(b)
}

fn biguint<R: Read>(decoder: &mut Decoder<R>) -> Result<BigUint, Error<R::Error>> {
    match decoder.pull()? {
        Header::Bytes(len) => Ok(BigUint::from_bytes_be(&parse_bytes(len, decoder)?)),
        _ => Err(Error::Syntax(decoder.offset())),
    }
}

fn parse<R: Read>(header: Header, decoder: &mut Decoder<R>) -> Result<Val, PError<R::Error>> {
    match header {
        Header::Text(len) => Ok(Val::from(parse_str(len, decoder)?)),
        Header::Bytes(len) => Ok(Val::byte_str(parse_bytes(len, decoder)?)),
        Header::Simple(simple::NULL) => Ok(Val::Null),
        Header::Simple(simple::FALSE) => Ok(Val::Bool(false)),
        Header::Simple(simple::TRUE) => Ok(Val::Bool(true)),
        Header::Simple(simple) => Err(PError::Simple(simple)),
        Header::Tag(tag::BIGNEG) => {
            Ok(biguint(decoder).map(|u| Val::Num(Num::big_int(-BigInt::from(u) - 1)))?)
        }
        Header::Tag(tag::BIGPOS) => {
            Ok(biguint(decoder).map(|u| Val::Num(Num::big_int(BigInt::from(u))))?)
        }
        Header::Tag(tag) => Err(PError::Tag(tag)),
        Header::Positive(pos) => Ok(Val::Num(Num::from_integral(pos))),
        Header::Negative(neg) => Ok(Val::Num(Num::from_integral(neg as i128 ^ !0))),
        Header::Float(f) => Ok(Val::from(f)),
        Header::Array(size) => Ok(Val::Arr(
            with_size(size, decoder, |h, d| parse(h, d))?.into(),
        )),
        Header::Map(size) => {
            let o = with_size(size, decoder, |h, d| {
                Ok((parse(h, d)?, parse(d.pull()?, d)?))
            })?;
            Ok(Val::obj(o.into_iter().collect()))
        }
        Header::Break => Err(PError::Break),
    }
}
