// Copyright (C) 2026 alto555
// SPDX-License-Identifier: GPL-2.0-only

//! Minimal PNG writer: adaptive filtering + fixed-Huffman deflate.
//!
//! Written out rather than pulled from a crate because the whole tree is
//! dependency-free and a plot image is not worth breaking that for.  Fixed
//! Huffman with greedy LZ77 is a fraction of a real compressor, but plot
//! images are mostly flat background and compress ~50x anyway.

const CRC_POLY: u32 = 0xEDB88320;

fn crc32(data: &[u8]) -> u32 {
    let mut table = [0u32; 256];
    for (i, e) in table.iter_mut().enumerate() {
        let mut c = i as u32;
        for _ in 0..8 {
            c = if c & 1 != 0 { CRC_POLY ^ (c >> 1) } else { c >> 1 };
        }
        *e = c;
    }
    let mut c = 0xFFFF_FFFFu32;
    for b in data {
        c = table[((c ^ *b as u32) & 0xFF) as usize] ^ (c >> 8);
    }
    c ^ 0xFFFF_FFFF
}

fn adler32(data: &[u8]) -> u32 {
    let (mut a, mut b) = (1u32, 0u32);
    for chunk in data.chunks(5552) {
        for x in chunk {
            a += *x as u32;
            b += a;
        }
        a %= 65521;
        b %= 65521;
    }
    (b << 16) | a
}

struct BitWriter {
    out: Vec<u8>,
    buf: u32,
    n: u32,
}

impl BitWriter {
    fn new() -> BitWriter {
        BitWriter {
            out: Vec::new(),
            buf: 0,
            n: 0,
        }
    }
    /// Deflate packs bits starting at the least significant bit of each byte.
    fn bits(&mut self, v: u32, n: u32) {
        self.buf |= (v & ((1 << n) - 1)) << self.n;
        self.n += n;
        while self.n >= 8 {
            self.out.push((self.buf & 0xFF) as u8);
            self.buf >>= 8;
            self.n -= 8;
        }
    }
    /// Huffman codes are stored most-significant-bit first, so reverse.
    fn code(&mut self, code: u32, len: u32) {
        let mut r = 0u32;
        for i in 0..len {
            r |= ((code >> i) & 1) << (len - 1 - i);
        }
        self.bits(r, len);
    }
    fn finish(mut self) -> Vec<u8> {
        if self.n > 0 {
            self.out.push((self.buf & 0xFF) as u8);
        }
        self.out
    }
}

const LEN_BASE: [u16; 29] = [
    3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 15, 17, 19, 23, 27, 31, 35, 43, 51, 59, 67, 83, 99, 115, 131,
    163, 195, 227, 258,
];
const LEN_EXTRA: [u8; 29] = [
    0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 0,
];
const DIST_BASE: [u16; 30] = [
    1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 33, 49, 65, 97, 129, 193, 257, 385, 513, 769, 1025, 1537,
    2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577,
];
const DIST_EXTRA: [u8; 30] = [
    0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13,
    13,
];

fn write_literal(bw: &mut BitWriter, sym: u32) {
    if sym <= 143 {
        bw.code(0x30 + sym, 8);
    } else if sym <= 255 {
        bw.code(0x190 + (sym - 144), 9);
    } else if sym <= 279 {
        bw.code(sym - 256, 7);
    } else {
        bw.code(0xC0 + (sym - 280), 8);
    }
}

const WIN: usize = 32768;
const MIN_MATCH: usize = 3;
const MAX_MATCH: usize = 258;
const HASH_BITS: usize = 15;

/// Single fixed-Huffman block with greedy matching.
fn deflate(data: &[u8]) -> Vec<u8> {
    let mut bw = BitWriter::new();
    bw.bits(1, 1); // BFINAL
    bw.bits(1, 2); // BTYPE = fixed Huffman

    let mut head = vec![usize::MAX; 1 << HASH_BITS];
    let mut prev = vec![usize::MAX; data.len().max(1)];
    let hash = |d: &[u8], i: usize| -> usize {
        ((d[i] as usize) << 10 ^ (d[i + 1] as usize) << 5 ^ (d[i + 2] as usize)) & ((1 << HASH_BITS) - 1)
    };

    let mut i = 0usize;
    while i < data.len() {
        let mut best_len = 0usize;
        let mut best_dist = 0usize;
        if i + MIN_MATCH <= data.len() {
            let h = hash(data, i);
            let mut cand = head[h];
            let mut tries = 0;
            while cand != usize::MAX && tries < 32 {
                if i - cand > WIN {
                    break;
                }
                let max = MAX_MATCH.min(data.len() - i);
                let mut l = 0usize;
                while l < max && data[cand + l] == data[i + l] {
                    l += 1;
                }
                if l > best_len {
                    best_len = l;
                    best_dist = i - cand;
                    if l >= MAX_MATCH {
                        break;
                    }
                }
                cand = prev[cand];
                tries += 1;
            }
        }

        if best_len >= MIN_MATCH {
            let mut lc = 0usize;
            while lc + 1 < 29 && LEN_BASE[lc + 1] as usize <= best_len {
                lc += 1;
            }
            write_literal(&mut bw, 257 + lc as u32);
            let e = LEN_EXTRA[lc] as u32;
            if e > 0 {
                bw.bits(best_len as u32 - LEN_BASE[lc] as u32, e);
            }
            let mut dc = 0usize;
            while dc + 1 < 30 && DIST_BASE[dc + 1] as usize <= best_dist {
                dc += 1;
            }
            bw.code(dc as u32, 5);
            let de = DIST_EXTRA[dc] as u32;
            if de > 0 {
                bw.bits(best_dist as u32 - DIST_BASE[dc] as u32, de);
            }
            for k in 0..best_len {
                if i + k + MIN_MATCH <= data.len() {
                    let h = hash(data, i + k);
                    prev[i + k] = head[h];
                    head[h] = i + k;
                }
            }
            i += best_len;
        } else {
            write_literal(&mut bw, data[i] as u32);
            if i + MIN_MATCH <= data.len() {
                let h = hash(data, i);
                prev[i] = head[h];
                head[h] = i;
            }
            i += 1;
        }
    }
    write_literal(&mut bw, 256); // end of block
    bw.finish()
}

fn chunk(out: &mut Vec<u8>, tag: &[u8; 4], body: &[u8]) {
    out.extend_from_slice(&(body.len() as u32).to_be_bytes());
    let mut full = Vec::with_capacity(4 + body.len());
    full.extend_from_slice(tag);
    full.extend_from_slice(body);
    out.extend_from_slice(&full);
    out.extend_from_slice(&crc32(&full).to_be_bytes());
}

/// Encode an RGB8 buffer (`w * h * 3` bytes) as a PNG.
pub fn encode_rgb(w: usize, h: usize, rgb: &[u8]) -> Vec<u8> {
    assert_eq!(rgb.len(), w * h * 3);
    let stride = w * 3;

    // Per-scanline adaptive filter: None / Sub / Up, whichever has the
    // smallest sum of absolute deviations.  Plots are mostly flat, so Up
    // usually wins by a wide margin.
    let mut raw = Vec::with_capacity((stride + 1) * h);
    let mut cand = vec![0u8; stride];
    for y in 0..h {
        let row = &rgb[y * stride..(y + 1) * stride];
        let up = if y > 0 {
            &rgb[(y - 1) * stride..y * stride]
        } else {
            &[][..]
        };
        let mut best = (0u8, u64::MAX, Vec::new());
        for f in 0..3u8 {
            let mut score = 0u64;
            for x in 0..stride {
                let a = if x >= 3 { row[x - 3] } else { 0 };
                let b = if y > 0 { up[x] } else { 0 };
                let v = match f {
                    0 => row[x],
                    1 => row[x].wrapping_sub(a),
                    _ => row[x].wrapping_sub(b),
                };
                cand[x] = v;
                score += (v as i8).unsigned_abs() as u64;
            }
            if score < best.1 {
                best = (f, score, cand.clone());
            }
        }
        raw.push(best.0);
        raw.extend_from_slice(&best.2);
    }

    let mut zlib = vec![0x78, 0x01];
    zlib.extend_from_slice(&deflate(&raw));
    zlib.extend_from_slice(&adler32(&raw).to_be_bytes());

    let mut out = vec![0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A];
    let mut ihdr = Vec::new();
    ihdr.extend_from_slice(&(w as u32).to_be_bytes());
    ihdr.extend_from_slice(&(h as u32).to_be_bytes());
    ihdr.extend_from_slice(&[8, 2, 0, 0, 0]); // 8-bit, truecolour
    chunk(&mut out, b"IHDR", &ihdr);
    chunk(&mut out, b"IDAT", &zlib);
    chunk(&mut out, b"IEND", &[]);
    out
}
