//! A parameter set as an RFC 9842 compression dictionary (extensions/http-delivery.md §8.8;
//! extensions/prior-carriage.md §16).
//!
//! The SHA-256 that RFC 9842's `Available-Dictionary` field carries is the digest a pin states
//! for a set in a one-file form. A sink that holds a set names it in that field; a server that
//! holds it MAY answer with the requested set coded against it, `Content-Encoding: dcz`
//! (Zstandard with the set as a raw dictionary, RFC 9842 §5). [`check_coded`] runs §8.8's
//! checks over such a response and decodes it; [`encode_dcz`] codes a set against another, as
//! an origin does. The reference decodes `dcz` only, and a sink built on it lists only `dcz`.

use crate::pin::sha256;

/// The `dcz` header's first 8 bytes: a Zstandard skippable frame of 32 bytes (RFC 9842 §5).
pub const DCZ_MAGIC: [u8; 8] = [0x5e, 0x2a, 0x4d, 0x18, 0x20, 0x00, 0x00, 0x00];
/// The `dcb` header's first 4 bytes (RFC 9842 §4).
pub const DCB_MAGIC: [u8; 4] = [0xff, 0x44, 0x43, 0x42];
/// The `dcz` window every client decodes: 8 MB, read as 2^23 bytes.
pub const DCZ_WINDOW_FLOOR: u64 = 1 << 23;
/// The largest `dcz` window: 128 MB, read as 2^27 bytes.
pub const DCZ_WINDOW_CEILING: u64 = 1 << 27;
/// The Zstandard frame magic (RFC 8878 §3.1.1).
const ZSTD_MAGIC: [u8; 4] = [0x28, 0xb5, 0x2f, 0xfd];

const B64: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

fn base64(b: &[u8]) -> String {
    let mut s = String::new();
    for c in b.chunks(3) {
        let n = (c[0] as u32) << 16 | (*c.get(1).unwrap_or(&0) as u32) << 8 | *c.get(2).unwrap_or(&0) as u32;
        for i in 0..4 {
            if i <= c.len() {
                s.push(B64[(n >> (18 - 6 * i) & 63) as usize] as char);
            } else {
                s.push('=');
            }
        }
    }
    s
}

/// The value of `Available-Dictionary` naming the set of digest `d`: a Structured Field Byte
/// Sequence (RFC 9651 §3.3.5), `:` base64 `:`.
pub fn available_dictionary(d: &[u8; 32]) -> String {
    format!(":{}:", base64(d))
}

/// The digest an `Available-Dictionary` value names: one Byte Sequence of 32 bytes, padded
/// base64 between colons, with spaces around it discarded; `None` otherwise.
pub fn parse_available_dictionary(v: &str) -> Option<[u8; 32]> {
    let inner = v.trim_matches(' ').strip_prefix(':')?.strip_suffix(':')?;
    if inner.len() != 44 || !inner.ends_with('=') || inner[..43].contains('=') {
        return None;
    }
    let mut bits: u32 = 0;
    let mut n = 0;
    let mut out = Vec::with_capacity(33);
    for c in inner[..43].bytes() {
        let x = B64.iter().position(|&b| b == c)? as u32;
        bits = bits << 6 | x;
        n += 6;
        if n >= 8 {
            n -= 8;
            out.push((bits >> n) as u8);
            bits &= (1 << n) - 1;
        }
    }
    // 43 characters carry 258 bits: 32 bytes and 2 pad bits, which RFC 4648 §3.5 has zero.
    (out.len() == 32 && bits == 0).then(|| out.try_into().ok()).flatten()
}

/// The value of `Use-As-Dictionary` a server sends with a set served below the root whose path
/// is `root_path`: every set below the root matches it (RFC 9842 §2.1).
pub fn use_as_dictionary(root_path: &str) -> String {
    format!("match=\"{root_path}*\"")
}

/// The window a Zstandard frame states (RFC 8878 §3.1.1.1): its Window_Size, or its
/// Frame_Content_Size when it is a single segment. `None` where the header does not read.
pub fn zstd_window(frame: &[u8]) -> Option<u64> {
    if frame.get(..4)? != ZSTD_MAGIC {
        return None;
    }
    let fhd = *frame.get(4)?;
    if fhd & 0x08 != 0 {
        return None;
    }
    let single = fhd & 0x20 != 0;
    let did = [0usize, 1, 2, 4][(fhd & 3) as usize];
    let fcs = match fhd >> 6 {
        0 if single => 1,
        0 => 0,
        1 => 2,
        2 => 4,
        _ => 8,
    };
    let mut at = 5;
    let window = if single {
        None
    } else {
        let wd = *frame.get(at)?;
        at += 1;
        let log = 10 + (wd >> 3) as u32;
        let base = 1u64 << log;
        Some(base + (base / 8) * (wd & 7) as u64)
    };
    at += did;
    let size = frame.get(at..at + fcs)?;
    let mut v = 0u64;
    for (i, b) in size.iter().enumerate() {
        v |= (*b as u64) << (8 * i);
    }
    if fcs == 2 {
        v += 256;
    }
    Some(window.unwrap_or(v))
}

/// Whether a `dcz` frame's window is one RFC 9842 §5 has a client decode against a dictionary
/// of `dict_len` bytes: at most 8 MB or 1.25 times the dictionary's length, whichever is
/// larger, and never above 128 MB.
pub fn dcz_window_admitted(window: u64, dict_len: u64) -> bool {
    window <= DCZ_WINDOW_CEILING && (window <= DCZ_WINDOW_FLOOR || window.saturating_mul(4) <= dict_len.saturating_mul(5))
}

/// Decodes a Zstandard frame against `dictionary` as raw content (RFC 8878 §5), reading at
/// most `limit` bytes of output. `Err("set_too_large")` past it, `Err("dictionary_coding_malformed")`
/// where it does not decode to exactly one whole frame.
fn zstd_decode(frame: &[u8], dictionary: &[u8], limit: u64) -> Result<Vec<u8>, &'static str> {
    use zstd_safe::{DCtx, DParameter, InBuffer, OutBuffer};
    const MALFORMED: &str = "dictionary_coding_malformed";
    let mut d = DCtx::create();
    d.set_parameter(DParameter::WindowLogMax(27)).map_err(|_| MALFORMED)?;
    d.ref_prefix(dictionary).map_err(|_| MALFORMED)?;
    let mut input = InBuffer::around(frame);
    let mut out = Vec::new();
    let mut chunk = vec![0u8; 1 << 16];
    loop {
        let mut o = OutBuffer::around(&mut chunk[..]);
        let hint = d.decompress_stream(&mut o, &mut input).map_err(|_| MALFORMED)?;
        let n = o.pos();
        out.extend_from_slice(&chunk[..n]);
        if out.len() as u64 > limit {
            return Err("set_too_large");
        }
        if hint == 0 {
            // A frame ended: nothing may follow it.
            return if input.pos() == frame.len() { Ok(out) } else { Err(MALFORMED) };
        }
        if input.pos() == frame.len() && n == 0 {
            return Err(MALFORMED);
        }
    }
}

/// Codes `target` against `dictionary` as a `dcz` body: the header, then a Zstandard frame at
/// level 19 with the dictionary as raw content, its content size and checksum stated.
pub fn encode_dcz(target: &[u8], dictionary: &[u8]) -> Result<Vec<u8>, String> {
    use zstd_safe::{CCtx, CParameter};
    let mut c = CCtx::create();
    for p in [CParameter::CompressionLevel(19), CParameter::ChecksumFlag(true), CParameter::ContentSizeFlag(true)] {
        c.set_parameter(p).map_err(|e| format!("zstd parameter: {e}"))?;
    }
    c.ref_prefix(dictionary).map_err(|e| format!("zstd prefix: {e}"))?;
    let mut frame = Vec::with_capacity(zstd_safe::compress_bound(target.len()));
    c.compress2(&mut frame, target).map_err(|e| format!("zstd: {e}"))?;
    let mut out = DCZ_MAGIC.to_vec();
    out.extend_from_slice(&sha256(dictionary));
    out.extend_from_slice(&frame);
    Ok(out)
}

/// A `dcz` body holding `target` in one raw block of a frame that states the window `2^log`
/// and no content size: a frame whose window a client checks before it decodes anything.
pub fn dcz_raw_with_window(target: &[u8], dictionary: &[u8], log: u8) -> Vec<u8> {
    assert!((10..=41).contains(&log) && target.len() < 1 << 17);
    let mut out = DCZ_MAGIC.to_vec();
    out.extend_from_slice(&sha256(dictionary));
    out.extend_from_slice(&ZSTD_MAGIC);
    out.push(0x00); // no content size, multi-segment, no checksum, no dictionary ID
    out.push((log - 10) << 3);
    let h = 1u32 | (target.len() as u32) << 3; // last block, raw
    out.extend_from_slice(&h.to_le_bytes()[..3]);
    out.extend_from_slice(target);
    out
}

/// What a sink sent with a request: the dictionary it named, if any. It lists `dcz` in
/// `Accept-Encoding` exactly when it names one.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Asked<'a> {
    pub dictionary: Option<(&'a [u8; 32], &'a [u8])>,
}

/// §8.8's checks over a response whose `Content-Encoding` is `coding`, in order, and its
/// decoding: the decoded bytes, or the reason. A coding other than `dcz` and `dcb` is not a
/// dictionary coding and is refused by §8.7 step 2, not here.
pub fn check_coded(coding: &str, body: &[u8], asked: Asked<'_>, limit: u64) -> Result<Vec<u8>, &'static str> {
    // 1. The request named a dictionary and listed this coding (the reference lists dcz only).
    let Some((named, dictionary)) = asked.dictionary.filter(|_| coding == "dcz") else {
        return Err("dictionary_unrequested");
    };
    // 2. The header: the magic, then the SHA-256 of the dictionary the request named.
    if body.len() < 40 || body[..8] != DCZ_MAGIC {
        return Err("dictionary_coding_malformed");
    }
    if body[8..40] != named[..] || sha256(dictionary) != *named {
        return Err("dictionary_mismatch");
    }
    // 3. The frame's window.
    let frame = &body[40..];
    let window = zstd_window(frame).ok_or("dictionary_coding_malformed")?;
    if !dcz_window_admitted(window, dictionary.len() as u64) {
        return Err("dictionary_window_too_large");
    }
    // 4. The decoding, within the bound.
    zstd_decode(frame, dictionary, limit)
}

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

    #[test]
    fn the_field_values() {
        let d = [0xa5u8; 32];
        let v = available_dictionary(&d);
        assert_eq!(v.len(), 46);
        assert_eq!(parse_available_dictionary(&v), Some(d));
        assert_eq!(parse_available_dictionary(&format!("  {v} ")), Some(d));
        // RFC 9842 §2.2's example value.
        assert!(parse_available_dictionary(":pZGm1Av0IEBKARczz7exkNYsZb8LzaMrV7J32a2fFG4=:").is_some());
        assert_eq!(parse_available_dictionary(&v[1..]), None);
        assert_eq!(parse_available_dictionary(":pZGm1Av0IEBKARczz7exkNYsZb8LzaMrV7J32a2fFG5=:"), None);
        assert_eq!(parse_available_dictionary(":pZGm1Av0IEBKARczz7exkNYsZb8LzaMrV7J32a2fFG4:"), None);
        assert_eq!(use_as_dictionary("/.well-known/wai-sets/"), "match=\"/.well-known/wai-sets/*\"");
    }

    #[test]
    fn windows() {
        assert!(dcz_window_admitted(1 << 23, 10));
        assert!(!dcz_window_admitted((1 << 23) + 1, 10));
        assert!(dcz_window_admitted(10 << 20, 8 << 20));
        assert!(!dcz_window_admitted((10 << 20) + 1, 8 << 20));
        assert!(!dcz_window_admitted((1 << 27) + 1, 1 << 30));
        let raw = dcz_raw_with_window(b"abc", b"dict", 24);
        assert_eq!(zstd_window(&raw[40..]), Some(1 << 24));
    }

    #[test]
    fn a_set_against_another_round_trips() {
        let dict: Vec<u8> = (0..4000u32).map(|i| (i * 7 % 251) as u8).collect();
        let mut target = dict.clone();
        target[100] ^= 0xff;
        target.extend_from_slice(b"tail");
        let body = encode_dcz(&target, &dict).expect("codes");
        assert!(body.len() < target.len() / 4, "{}", body.len());
        let named = sha256(&dict);
        let asked = Asked { dictionary: Some((&named, &dict)) };
        assert_eq!(check_coded("dcz", &body, asked, 1 << 20).as_deref(), Ok(&target[..]));
        assert_eq!(check_coded("dcb", &body, asked, 1 << 20), Err("dictionary_unrequested"));
        assert_eq!(check_coded("dcz", &body, Asked { dictionary: None }, 1 << 20), Err("dictionary_unrequested"));
        assert_eq!(check_coded("dcz", &body, asked, 100), Err("set_too_large"));
        assert_eq!(check_coded("dcz", &body[..body.len() - 3], asked, 1 << 20), Err("dictionary_coding_malformed"));
        let other = vec![0u8; 4000];
        let other_named = sha256(&other);
        assert_eq!(check_coded("dcz", &body, Asked { dictionary: Some((&other_named, &other)) }, 1 << 20), Err("dictionary_mismatch"));
        let wide = dcz_raw_with_window(&target, &dict, 24);
        assert_eq!(check_coded("dcz", &wide, asked, 1 << 20), Err("dictionary_window_too_large"));
        let fits = dcz_raw_with_window(&target, &dict, 13);
        assert_eq!(check_coded("dcz", &fits, asked, 1 << 20).as_deref(), Ok(&target[..]));
    }
}
