//! Text/general-binary compression wrappers. zstd (universal balance)
//! and XZ/LZMA (max classical ratio). Both lossless by construction.

use std::io::Read;
#[cfg(feature = "codecs")]
use std::io::Write;

#[derive(Debug)]
pub struct TextError(pub String);

impl<E: std::fmt::Display> From<E> for TextError {
    fn from(e: E) -> Self { TextError(e.to_string()) }
}

pub type Result<T> = std::result::Result<T, TextError>;

/// zstd at the given compression level (1..=22; 22 is "ultra", slow). Refuses more
/// than [`TEXT_DECODED_MAX_LEN`] bytes (`text_decoded_too_long`): a conforming encoder
/// writes no payload a conforming reader may refuse (SPEC §2, *Reader limits*). No
/// level asks for a window past [`TEXT_WINDOW_MAX_LEN`].
pub fn zstd_encode(data: &[u8], level: i32) -> Result<Vec<u8>> {
    if data.len() > TEXT_DECODED_MAX_LEN {
        return Err(too_long(TEXT_DECODED_MAX_LEN));
    }
    Ok(zstd::stream::encode_all(data, level)?)
}

/// The most bytes a `wai.text.zstd` or `wai.text.xz` payload may decode to that a reader
/// must accept (SPEC §2, *Reader limits*): 2^27. The decoders here write no more than
/// this, whatever the payload asks for: one that decodes to more is refused
/// (`text_decoded_too_long`) when its output would pass the limit, so the output buffer
/// never holds more. The reference writers refuse to encode more.
pub const TEXT_DECODED_MAX_LEN: usize = 1 << 27;

/// The base-2 logarithm of [`TEXT_WINDOW_MAX_LEN`].
pub const TEXT_WINDOW_MAX_LOG: u32 = 27;

/// The most history a `wai.text.zstd` or `wai.text.xz` stream may ask its decoder to
/// keep that a reader must accept (SPEC §2, *Reader limits*): 2^27 bytes, as the zstd
/// frame's window size or the LZMA2 dictionary size. A stream that declares more is
/// refused (`text_window_too_large`) from the declaration, before the decoder allocates
/// it. Together with [`TEXT_DECODED_MAX_LEN`] it bounds the memory a decode takes: the
/// length limit bounds the output, and this bounds what the decoder holds beside it.
pub const TEXT_WINDOW_MAX_LEN: usize = 1 << TEXT_WINDOW_MAX_LOG;

#[cfg(feature = "codecs")]
/// What liblzma may allocate to decode one stream: the dictionary, at most
/// [`TEXT_WINDOW_MAX_LEN`], and the decoder's own fixed state beside it.
const XZ_DECODER_MEMLIMIT: u64 = TEXT_WINDOW_MAX_LEN as u64 + (1 << 20);

fn too_long(max_out: usize) -> TextError {
    TextError(format!(
        "text_decoded_too_long: the payload decodes to more than {max_out} bytes, the most this reader accepts (SPEC §2)"
    ))
}

fn window_too_large() -> TextError {
    TextError(format!(
        "text_window_too_large: the stream asks its decoder to keep more than {TEXT_WINDOW_MAX_LEN} bytes of history, the most a reader need give it (SPEC §2)"
    ))
}

/// Reads `r` to its end into a buffer that never holds, and is never sized for, more
/// than `max_out` bytes. `Ok(None)` when the reader has more than that to give.
fn read_within<R: Read>(mut r: R, max_out: usize) -> std::io::Result<Option<Vec<u8>>> {
    let mut out: Vec<u8> = Vec::new();
    let mut chunk = vec![0u8; 64 * 1024];
    loop {
        let n = match r.read(&mut chunk) {
            Ok(0) => return Ok(Some(out)),
            Ok(n) => n,
            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
            Err(e) => return Err(e),
        };
        let room = max_out - out.len();
        if n > room {
            return Ok(None);
        }
        if out.capacity() - out.len() < n {
            // Grow as a `Vec` would, by about its length, but never past the limit.
            out.reserve_exact(out.len().max(chunk.len()).min(room).max(n));
        }
        out.extend_from_slice(&chunk[..n]);
    }
}

/// Decodes a `wai.text.zstd` payload within the reader limits of SPEC §2:
/// [`TEXT_DECODED_MAX_LEN`] bytes of output and [`TEXT_WINDOW_MAX_LEN`] bytes of window.
pub fn zstd_decode(bytes: &[u8]) -> Result<Vec<u8>> {
    zstd_decode_within(bytes, TEXT_DECODED_MAX_LEN)
}

/// [`zstd_decode`] with the output limit the caller gives: for a caller that knows what
/// its payload may hold, and for tests. `max_out` is capped at
/// [`TEXT_DECODED_MAX_LEN`]. The window limit is not the caller's to change.
pub fn zstd_decode_within(bytes: &[u8], max_out: usize) -> Result<Vec<u8>> {
    let max_out = max_out.min(TEXT_DECODED_MAX_LEN);
    // A frame that declares its content size past the limit is refused from the
    // declaration, before anything is decoded. A frame need not declare it, and a
    // payload may hold more frames than the first, so the output is counted too.
    if let Ok(Some(declared)) = zstd::zstd_safe::get_frame_content_size(bytes) {
        if declared > max_out as u64 {
            return Err(too_long(max_out));
        }
    }
    let mut dec = zstd::stream::read::Decoder::with_buffer(bytes)?;
    dec.window_log_max(TEXT_WINDOW_MAX_LOG)?;
    match read_within(dec, max_out) {
        Ok(Some(out)) => Ok(out),
        Ok(None) => Err(too_long(max_out)),
        // libzstd's name for a frame whose window is past the decoder's limit.
        Err(e) if e.to_string().contains("too much memory") => Err(window_too_large()),
        Err(e) => Err(e.into()),
    }
}

#[cfg(feature = "codecs")]
/// XZ/LZMA2 at the given preset (0..=9; 6 is the default, 9 is "extreme"). Refuses more
/// than [`TEXT_DECODED_MAX_LEN`] bytes (`text_decoded_too_long`), as [`zstd_encode`]
/// does. No preset asks for a dictionary past [`TEXT_WINDOW_MAX_LEN`].
pub fn xz_encode(data: &[u8], level: u32) -> Result<Vec<u8>> {
    if data.len() > TEXT_DECODED_MAX_LEN {
        return Err(too_long(TEXT_DECODED_MAX_LEN));
    }
    let mut out = Vec::with_capacity(data.len() / 4);
    let mut enc = xz2::write::XzEncoder::new(&mut out, level);
    enc.write_all(data)?;
    enc.finish()?;
    Ok(out)
}

#[cfg(feature = "codecs")]
/// Decodes a `wai.text.xz` payload within the reader limits of SPEC §2:
/// [`TEXT_DECODED_MAX_LEN`] bytes of output and [`TEXT_WINDOW_MAX_LEN`] bytes of
/// dictionary.
pub fn xz_decode(bytes: &[u8]) -> Result<Vec<u8>> {
    xz_decode_within(bytes, TEXT_DECODED_MAX_LEN)
}

#[cfg(feature = "codecs")]
/// [`xz_decode`] with the output limit the caller gives: for a caller that knows what
/// its payload may hold, and for tests. `max_out` is capped at
/// [`TEXT_DECODED_MAX_LEN`]. The dictionary limit is not the caller's to change.
pub fn xz_decode_within(bytes: &[u8], max_out: usize) -> Result<Vec<u8>> {
    let max_out = max_out.min(TEXT_DECODED_MAX_LEN);
    // The decoder is given a memory limit, so a block that declares a dictionary past
    // it is refused from its header, before the dictionary is allocated. An xz stream
    // states its decoded length only at its end, so the output is counted as it comes.
    let stream = xz2::stream::Stream::new_stream_decoder(XZ_DECODER_MEMLIMIT, 0)?;
    let dec = xz2::read::XzDecoder::new_stream(bytes, stream);
    match read_within(dec, max_out) {
        Ok(Some(out)) => Ok(out),
        Ok(None) => Err(too_long(max_out)),
        Err(e) if matches!(
            e.get_ref().and_then(|i| i.downcast_ref::<xz2::stream::Error>()),
            Some(xz2::stream::Error::MemLimit)
        ) => Err(window_too_large()),
        Err(e) => Err(e.into()),
    }
}

// The tests exercise every codec here: a build with `codecs` runs them; the floor's are
// tests/floor.rs.
#[cfg(all(test, feature = "codecs"))]
mod tests {
    use super::*;

    #[test]
    fn zstd_lossless_round_trip() {
        // mix of structure and PRNG — zstd should compress the structured
        // part well. Pure PRNG never compresses (high entropy) so use a
        // partly-repetitive payload to exercise both paths.
        let mut data = b"WAI standard zeroth text path test payload. ".repeat(500);
        data.extend((0..10_000u32).flat_map(|i| {
            (i.wrapping_mul(1_103_515_245).wrapping_add(12_345)).to_le_bytes()
        }));
        let e = zstd_encode(&data, 19).unwrap();
        let d = zstd_decode(&e).unwrap();
        assert_eq!(d, data, "zstd must be lossless");
        assert!(e.len() < data.len(), "zstd should compress mixed-entropy data");
    }

    #[test]
    fn xz_lossless_round_trip() {
        let data = b"the quick brown fox ".repeat(2_000);
        let e = xz_encode(&data, 6).unwrap();
        let d = xz_decode(&e).unwrap();
        assert_eq!(d, data);
        assert!(e.len() < data.len() / 5, "XZ should crush highly repetitive data");
    }
}
