//! [STAR](https://dasl.ing/star.html), DASL's streaming archive:
//!
//! ```text
//! 2a 6c 01 00 | varint len | meta | ( 01 | varint len | key | varint len | value )* | 00
//! ```
//!
//! the writers are free functions over any [`Write`]. the header and the end byte can go in zstd
//! frames of their own.

use std::io::{self, BufRead, Write};

use ipld_core::ipld::Ipld;
use serde::Serialize;

pub const MAGIC: [u8; 4] = [0x2a, 0x6c, 0x01, 0x00];
pub const PAIR: u8 = 0x01;
pub const END: u8 = 0x00;
/// a varint holds 63 bits, 7 per byte
const VARINT_MAX_BYTES: u32 = 9;
const SHORT_KEY: usize = 64;

pub type Pair = (Vec<u8>, Vec<u8>);

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error(transparent)]
    Io(io::Error),
    #[error("not a STAR")]
    Magic,
    #[error("a STAR cut short")]
    Truncated,
    #[error("a varint {0}")]
    Varint(&'static str),
    #[error("a header {0}")]
    Meta(String),
    #[error("control byte {0:#04x}, neither a pair nor the end")]
    Control(u8),
    #[error("bytes after the end of a STAR")]
    TrailingData,
}

impl From<io::Error> for Error {
    fn from(e: io::Error) -> Self {
        match e.kind() {
            io::ErrorKind::UnexpectedEof => Self::Truncated,
            _ => Self::Io(e),
        }
    }
}

/// refuses a `meta` that isn't a DRISL map with a text `$type`
pub fn write_header(out: &mut impl Write, meta: &impl Serialize) -> Result<(), Error> {
    let meta = serde_ipld_dagcbor::to_vec(meta)
        .map_err(|e| Error::Meta(format!("that can't be encoded: {e}")))?;
    check_meta(&meta)?;
    out.write_all(&MAGIC)?;
    write_varint(out, meta.len() as u64)?;
    out.write_all(&meta)?;
    Ok(())
}

/// a short key goes out in one write with the control byte and both lengths, because a zstd writer
/// pays for every call
pub fn write_pair(out: &mut impl Write, key: &[u8], value: &[u8]) -> io::Result<()> {
    if key.len() > SHORT_KEY {
        out.write_all(&[PAIR])?;
        write_varint(out, key.len() as u64)?;
        out.write_all(key)?;
        write_varint(out, value.len() as u64)?;
        return out.write_all(value);
    }
    let mut head = [0; 1 + 2 * VARINT_MAX_BYTES as usize + SHORT_KEY];
    head[0] = PAIR;
    let mut at = 1;
    at += put_varint(&mut head[at..], key.len() as u64);
    head[at..at + key.len()].copy_from_slice(key);
    at += key.len();
    at += put_varint(&mut head[at..], value.len() as u64);
    out.write_all(&head[..at])?;
    out.write_all(value)
}

/// nothing may be written after it
pub fn write_end(out: &mut impl Write) -> io::Result<()> {
    out.write_all(&[END])
}

fn write_varint(out: &mut impl Write, n: u64) -> io::Result<()> {
    let mut buf = [0; VARINT_MAX_BYTES as usize];
    let len = put_varint(&mut buf, n);
    out.write_all(&buf[..len])
}

/// lengths are at most `isize::MAX`, which fits in the 63 bits a varint holds
fn put_varint(buf: &mut [u8], mut n: u64) -> usize {
    let mut len = 0;
    while n >= 0x80 {
        buf[len] = n as u8 | 0x80;
        n >>= 7;
        len += 1;
    }
    buf[len] = n as u8;
    len + 1
}

/// the first error ends it. pairs before the error come from an invalid archive.
/// collect them all first if you can't take them back
pub struct Reader<R> {
    inner: R,
    meta: Vec<u8>,
    /// the current pair, key then value
    buf: Vec<u8>,
    done: bool,
}

impl<R: BufRead> Reader<R> {
    pub fn new(mut inner: R) -> Result<Self, Error> {
        let mut magic = [0; 4];
        inner.read_exact(&mut magic)?;
        if magic != MAGIC {
            return Err(Error::Magic);
        }
        let len = read_varint(&mut inner)?;
        let mut meta = Vec::new();
        read_into(&mut inner, &mut meta, len)?;
        check_meta(&meta)?;
        Ok(Self {
            inner,
            meta,
            buf: Vec::new(),
            done: false,
        })
    }

    /// the header's DRISL bytes, already checked to be a map with a text `$type`
    pub fn meta(&self) -> &[u8] {
        &self.meta
    }

    // llvm stops inlining the read path once a program uses both this and the iterator, and short
    // pairs then read twice as slow in `benches/star.rs`
    #[inline(always)]
    #[allow(clippy::type_complexity)]
    pub fn next_pair(&mut self) -> Result<Option<(&[u8], &[u8])>, Error> {
        if self.done {
            return Ok(None);
        }
        let key_len = self.pair();
        self.done = !matches!(key_len, Ok(Some(_)));
        Ok(key_len?.map(|at| self.buf.split_at(at)))
    }

    /// the key's length, with the key and value in `buf`
    #[inline(always)]
    fn pair(&mut self) -> Result<Option<usize>, Error> {
        match read_byte(&mut self.inner)? {
            PAIR => {
                self.buf.clear();
                let len = read_varint(&mut self.inner)?;
                read_into(&mut self.inner, &mut self.buf, len)?;
                let key_len = self.buf.len();
                let len = read_varint(&mut self.inner)?;
                read_into(&mut self.inner, &mut self.buf, len)?;
                Ok(Some(key_len))
            }
            END if self.inner.fill_buf()?.is_empty() => Ok(None),
            END => Err(Error::TrailingData),
            other => Err(Error::Control(other)),
        }
    }
}

impl<R: BufRead> Iterator for Reader<R> {
    type Item = Result<Pair, Error>;

    fn next(&mut self) -> Option<Self::Item> {
        let pair = self.next_pair().transpose()?;
        Some(pair.map(|(key, value)| (key.to_vec(), value.to_vec())))
    }
}

#[inline(always)]
fn read_byte(r: &mut impl BufRead) -> Result<u8, Error> {
    let byte = *r.fill_buf()?.first().ok_or(Error::Truncated)?;
    r.consume(1);
    Ok(byte)
}

#[inline(always)]
fn read_varint(r: &mut impl BufRead) -> Result<u64, Error> {
    let mut n = 0;
    for i in 0..VARINT_MAX_BYTES {
        let byte = read_byte(r)?;
        n |= u64::from(byte & 0x7f) << (7 * i);
        if byte & 0x80 == 0 {
            if byte == 0 && i > 0 {
                return Err(Error::Varint("that isn't minimally encoded"));
            }
            return Ok(n);
        }
    }
    Err(Error::Varint("longer than 9 bytes"))
}

/// grows `buf` as bytes arrive and never ahead of them, because a huge length on a short input
/// would otherwise allocate it all up front
#[inline(always)]
fn read_into(r: &mut impl BufRead, buf: &mut Vec<u8>, len: u64) -> Result<(), Error> {
    let mut left = usize::try_from(len).unwrap_or(usize::MAX);
    while left > 0 {
        let available = r.fill_buf()?;
        if available.is_empty() {
            return Err(Error::Truncated);
        }
        let n = available.len().min(left);
        if buf.capacity() - buf.len() < n {
            buf.reserve_exact(left.min(buf.len().max(n)));
        }
        buf.extend_from_slice(&available[..n]);
        r.consume(n);
        left -= n;
    }
    Ok(())
}

/// a value that changes when decoded and re-encoded isn't DRISL, which has one encoding per value
fn check_meta(meta: &[u8]) -> Result<(), Error> {
    let value: Ipld = serde_ipld_dagcbor::from_slice(meta)
        .map_err(|e| Error::Meta(format!("that isn't DRISL: {e}")))?;
    let canonical = serde_ipld_dagcbor::to_vec(&value)
        .map_err(|e| Error::Meta(format!("that isn't DRISL: {e}")))?;
    if canonical != meta {
        return Err(Error::Meta(
            "that isn't deterministically encoded".to_owned(),
        ));
    }
    let Ipld::Map(map) = value else {
        return Err(Error::Meta("that isn't a map".to_owned()));
    };
    match map.get("$type") {
        Some(Ipld::String(_)) => Ok(()),
        _ => Err(Error::Meta("without a text $type".to_owned())),
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use super::*;

    fn varint(n: u64) -> Vec<u8> {
        let mut out = Vec::new();
        write_varint(&mut out, n).unwrap();
        out
    }

    fn read(bytes: &[u8]) -> Result<u64, Error> {
        read_varint(&mut &bytes[..])
    }

    #[test]
    fn varints_round_trip_at_every_byte_boundary() {
        for bytes in 1..=VARINT_MAX_BYTES {
            let low = if bytes == 1 {
                0
            } else {
                1u64 << (7 * (bytes - 1))
            };
            let high = (1u64 << (7 * bytes)) - 1;
            for n in [low, high] {
                let encoded = varint(n);
                assert_eq!(encoded.len(), bytes as usize, "{n}");
                assert_eq!(read(&encoded).unwrap(), n);
            }
        }
    }

    #[test]
    fn varints_the_spec_forbids_are_refused() {
        assert!(matches!(read(&[0x80, 0x00]), Err(Error::Varint(_))));
        assert!(matches!(read(&[0xff, 0x80, 0x00]), Err(Error::Varint(_))));
        assert!(matches!(read(&[0xff; 9]), Err(Error::Varint(_))));
        assert!(matches!(read(&[0x80]), Err(Error::Truncated)));
        assert!(matches!(read(&[]), Err(Error::Truncated)));
    }

    #[derive(Serialize)]
    struct Meta {
        #[serde(rename = "$type")]
        kind: &'static str,
        stream: &'static str,
        first: u64,
    }

    #[test]
    fn what_is_written_reads_back() {
        let mut star = Vec::new();
        let meta = Meta {
            kind: "a.b.c",
            stream: "s",
            first: 300,
        };
        // keys either side of `SHORT_KEY`, values either side of the 127/128 varint boundary
        let pairs = [
            (b"b".to_vec(), vec![1; 200]),
            (vec![], vec![]),
            (vec![2; SHORT_KEY], vec![3; 127]),
            (vec![4; SHORT_KEY + 1], vec![5; 128]),
            (b"b".to_vec(), b"x".to_vec()),
        ];
        write_header(&mut star, &meta).unwrap();
        for (key, value) in &pairs {
            write_pair(&mut star, key, value).unwrap();
        }
        write_end(&mut star).unwrap();

        let reader = Reader::new(&star[..]).unwrap();
        assert_eq!(reader.meta(), serde_ipld_dagcbor::to_vec(&meta).unwrap());
        let read: Vec<_> = reader.collect::<Result<_, _>>().unwrap();
        assert_eq!(read, pairs);

        let mut reader = Reader::new(std::io::BufReader::with_capacity(1, &star[..])).unwrap();
        for (key, value) in &pairs {
            assert_eq!(reader.next_pair().unwrap(), Some((&key[..], &value[..])));
        }
        assert_eq!(reader.next_pair().unwrap(), None);
    }

    #[test]
    fn a_header_without_a_text_type_is_not_written() {
        let no_type = BTreeMap::from([("kind", "a.b.c")]);
        assert!(matches!(
            write_header(&mut Vec::new(), &no_type),
            Err(Error::Meta(_))
        ));
        let int_type = BTreeMap::from([("$type", 1)]);
        assert!(matches!(
            write_header(&mut Vec::new(), &int_type),
            Err(Error::Meta(_))
        ));
        assert!(matches!(
            write_header(&mut Vec::new(), &"a.b.c"),
            Err(Error::Meta(_))
        ));
    }

    #[test]
    fn the_reader_stops_at_the_first_error() {
        let mut star = Vec::new();
        write_header(&mut star, &BTreeMap::from([("$type", "")])).unwrap();
        write_pair(&mut star, b"k", b"v").unwrap();
        star.push(0x02);
        let mut reader = Reader::new(&star[..]).unwrap();
        assert!(matches!(reader.next(), Some(Ok(_))));
        assert!(matches!(reader.next(), Some(Err(Error::Control(0x02)))));
        assert!(reader.next().is_none());
    }
}