enum Token {
    Lit(u8),
    Match { len: u16, dist: u16 },
}

pub fn deflate(input: &[u8]) -> Vec<u8> {
    let mut best: Option<Vec<u8>> = None;
    let consider = |block: Vec<u8>, best: &mut Option<Vec<u8>>| {
        if best.as_ref().map_or(true, |b| block.len() < b.len()) {
            *best = Some(block);
        }
    };

    let lazy = lz77(input);
    consider(emit_dynamic_block(&lazy), &mut best);
    consider(emit_fixed_block(&lazy), &mut best);

    let matches = build_match_table(input);
    let default_parse = lz77_optimal(input, &matches, &default_costs());
    consider(emit_dynamic_block(&default_parse), &mut best);
    consider(emit_fixed_block(&default_parse), &mut best);

    for seed in [&default_parse, &lazy] {
        let mut last = lz77_optimal(input, &matches, &costs_from_tokens(seed));
        for _ in 0..30 {
            consider(emit_dynamic_block(&last), &mut best);
            consider(emit_fixed_block(&last), &mut best);
            let next = lz77_optimal(input, &matches, &costs_from_tokens(&last));
            last = next;
        }
    }

    best.unwrap()
}

struct Costs {
    lit: [u32; 288],
    dist: [u32; 30],
}

fn default_costs() -> Costs {
    let mut lit = [9u32; 288];
    for i in 0..144 {
        lit[i] = 8;
    }
    for i in 256..288 {
        lit[i] = 8;
    }
    Costs { lit, dist: [5u32; 30] }
}

fn costs_from_tokens(tokens: &[Token]) -> Costs {
    let (lit_freq, dist_freq) = tally(tokens);
    let lit_lens = huffman_lengths(&lit_freq, 15);
    let dist_lens = huffman_lengths(&dist_freq, 15);
    let mut lit = [0u32; 288];
    for i in 0..288 {
        lit[i] = if lit_lens[i] > 0 { lit_lens[i] as u32 } else { 15 };
    }
    let mut dist = [0u32; 30];
    for i in 0..30 {
        dist[i] = if dist_lens[i] > 0 { dist_lens[i] as u32 } else { 15 };
    }
    Costs { lit, dist }
}

fn build_match_table(input: &[u8]) -> (Vec<u16>, Vec<u16>) {
    let n = input.len();
    let mut best_len = vec![0u16; n];
    let mut best_dist = vec![0u16; n];
    if n == 0 {
        return (best_len, best_dist);
    }
    let mut head = vec![usize::MAX; 1 << 15];
    let mut prev = vec![usize::MAX; n];
    let hash = |a: u8, b: u8, c: u8| -> usize {
        (((a as usize) << 10) ^ ((b as usize) << 5) ^ (c as usize)) & ((1 << 15) - 1)
    };
    for i in 0..n {
        if i + 3 <= n {
            let h = hash(input[i], input[i + 1], input[i + 2]);
            let mut j = head[h];
            let mut chain = 4096;
            let max_len = (n - i).min(258);
            let mut bl = 0;
            let mut bd = 0;
            while j != usize::MAX && chain > 0 {
                if i - j > 32768 {
                    break;
                }
                let mut l = 0;
                while l < max_len && input[j + l] == input[i + l] {
                    l += 1;
                }
                if l > bl {
                    bl = l;
                    bd = i - j;
                    if l >= max_len {
                        break;
                    }
                }
                j = prev[j];
                chain -= 1;
            }
            best_len[i] = bl as u16;
            best_dist[i] = bd as u16;
            prev[i] = head[h];
            head[h] = i;
        }
    }
    (best_len, best_dist)
}

fn lz77_optimal(input: &[u8], matches: &(Vec<u16>, Vec<u16>), costs: &Costs) -> Vec<Token> {
    let n = input.len();
    if n == 0 {
        return Vec::new();
    }
    let (best_len, best_dist) = matches;

    let lit_cost = |b: u8| -> u32 { costs.lit[b as usize] };
    let match_cost = |len: u16, dist: u16| -> u32 {
        let (lsym, _, lex) = length_code(len);
        let (dsym, _, dex) = dist_code(dist);
        costs.lit[lsym as usize] + lex + costs.dist[dsym as usize] + dex
    };

    let mut cost = vec![u32::MAX; n + 1];
    let mut choice_len = vec![0u16; n + 1];
    cost[n] = 0;
    for i in (0..n).rev() {
        let mut bc = lit_cost(input[i]).saturating_add(cost[i + 1]);
        let mut bl = 0u16;
        let ml = best_len[i];
        if ml >= 3 {
            let d = best_dist[i];
            let mut l = 3u16;
            while l <= ml {
                let nxt = cost[i + l as usize];
                if nxt != u32::MAX {
                    let c = match_cost(l, d).saturating_add(nxt);
                    if c < bc {
                        bc = c;
                        bl = l;
                    }
                }
                l += 1;
            }
        }
        cost[i] = bc;
        choice_len[i] = bl;
    }

    let mut tokens = Vec::new();
    let mut i = 0;
    while i < n {
        let l = choice_len[i];
        if l >= 3 {
            tokens.push(Token::Match { len: l, dist: best_dist[i] });
            i += l as usize;
        } else {
            tokens.push(Token::Lit(input[i]));
            i += 1;
        }
    }
    tokens
}

fn lz77(input: &[u8]) -> Vec<Token> {
    let n = input.len();
    let mut tokens = Vec::new();
    let mut head = vec![usize::MAX; 1 << 15];
    let mut prev = vec![usize::MAX; n.max(1)];

    let hash = |a: u8, b: u8, c: u8| -> usize {
        (((a as usize) << 10) ^ ((b as usize) << 5) ^ (c as usize)) & ((1 << 15) - 1)
    };

    let find = |p: usize, head: &[usize], prev: &[usize]| -> (usize, usize) {
        if p + 3 > n {
            return (0, 0);
        }
        let h = hash(input[p], input[p + 1], input[p + 2]);
        let mut j = head[h];
        let mut best_len = 0;
        let mut best_dist = 0;
        let mut chain = 4096;
        let max_len = (n - p).min(258);
        while j != usize::MAX && chain > 0 {
            if p - j > 32768 {
                break;
            }
            let mut l = 0;
            while l < max_len && input[j + l] == input[p + l] {
                l += 1;
            }
            if l > best_len {
                best_len = l;
                best_dist = p - j;
                if l >= max_len {
                    break;
                }
            }
            j = prev[j];
            chain -= 1;
        }
        (best_len, best_dist)
    };
    let insert = |p: usize, head: &mut Vec<usize>, prev: &mut Vec<usize>| {
        if p + 3 <= n {
            let h = hash(input[p], input[p + 1], input[p + 2]);
            prev[p] = head[h];
            head[h] = p;
        }
    };

    let mut i = 0;
    let mut inserted = 0usize;
    macro_rules! insert_before {
        ($before:expr) => {{
            let before = $before;
            while inserted < before && inserted < n {
                insert(inserted, &mut head, &mut prev);
                inserted += 1;
            }
        }};
    }
    while i < n {
        insert_before!(i);
        let (mlen, mdist) = find(i, &head, &prev);

        if mlen >= 3 {
            if i + 1 < n {
                insert_before!(i + 1);
                let (next_len, _) = find(i + 1, &head, &prev);
                if next_len > mlen {
                    tokens.push(Token::Lit(input[i]));
                    i += 1;
                    continue;
                }
            }
            tokens.push(Token::Match { len: mlen as u16, dist: mdist as u16 });
            i += mlen;
        } else {
            tokens.push(Token::Lit(input[i]));
            i += 1;
        }
    }
    tokens
}

fn length_code(len: u16) -> (u16, u32, u32) {
    const BASE: [(u16, u16, u32); 29] = [
        (257, 3, 0), (258, 4, 0), (259, 5, 0), (260, 6, 0), (261, 7, 0),
        (262, 8, 0), (263, 9, 0), (264, 10, 0), (265, 11, 1), (266, 13, 1),
        (267, 15, 1), (268, 17, 1), (269, 19, 2), (270, 23, 2), (271, 27, 2),
        (272, 31, 2), (273, 35, 3), (274, 43, 3), (275, 51, 3), (276, 59, 3),
        (277, 67, 4), (278, 83, 4), (279, 99, 4), (280, 115, 4), (281, 131, 5),
        (282, 163, 5), (283, 195, 5), (284, 227, 5), (285, 258, 0),
    ];
    let mut chosen = BASE[0];
    for &e in BASE.iter() {
        if e.1 <= len {
            chosen = e;
        } else {
            break;
        }
    }
    let (sym, base, ebits) = chosen;
    (sym, (len - base) as u32, ebits)
}

fn dist_code(dist: u16) -> (u16, u32, u32) {
    const BASE: [(u16, u16, u32); 30] = [
        (0, 1, 0), (1, 2, 0), (2, 3, 0), (3, 4, 0), (4, 5, 1), (5, 7, 1),
        (6, 9, 2), (7, 13, 2), (8, 17, 3), (9, 25, 3), (10, 33, 4), (11, 49, 4),
        (12, 65, 5), (13, 97, 5), (14, 129, 6), (15, 193, 6), (16, 257, 7),
        (17, 385, 7), (18, 513, 8), (19, 769, 8), (20, 1025, 9), (21, 1537, 9),
        (22, 2049, 10), (23, 3073, 10), (24, 4097, 11), (25, 6145, 11),
        (26, 8193, 12), (27, 12289, 12), (28, 16385, 13), (29, 24577, 13),
    ];
    let mut chosen = BASE[0];
    for &e in BASE.iter() {
        if e.1 <= dist {
            chosen = e;
        } else {
            break;
        }
    }
    let (sym, base, ebits) = chosen;
    (sym, (dist - base) as u32, ebits)
}

fn length_extra_bits(sym: u16) -> u32 {
    length_code_extra(sym)
}
fn length_code_extra(sym: u16) -> u32 {
    match sym {
        257..=264 => 0,
        265..=268 => 1,
        269..=272 => 2,
        273..=276 => 3,
        277..=280 => 4,
        281..=284 => 5,
        285 => 0,
        _ => 0,
    }
}
fn dist_extra_bits(sym: u16) -> u32 {
    match sym {
        0..=3 => 0,
        s => ((s / 2) - 1) as u32,
    }
}

fn huffman_lengths(freqs: &[u32], max_bits: usize) -> Vec<u8> {
    let n = freqs.len();
    let mut lens = vec![0u8; n];

    let used: Vec<usize> = (0..n).filter(|&i| freqs[i] > 0).collect();
    if used.is_empty() {
        return lens;
    }
    if used.len() == 1 {
        lens[used[0]] = 1;
        return lens;
    }

    #[derive(Clone)]
    struct Node {
        freq: u64,
        sym: i32,
        left: i32,
        right: i32,
    }
    let mut nodes: Vec<Node> = Vec::new();
    let mut heap: Vec<usize> = Vec::new();
    for &s in &used {
        nodes.push(Node { freq: freqs[s] as u64, sym: s as i32, left: -1, right: -1 });
        heap.push(nodes.len() - 1);
    }
    let pop_min = |heap: &mut Vec<usize>, nodes: &Vec<Node>| -> usize {
        let mut mi = 0;
        for k in 1..heap.len() {
            if nodes[heap[k]].freq < nodes[heap[mi]].freq {
                mi = k;
            }
        }
        heap.swap_remove(mi)
    };
    while heap.len() > 1 {
        let a = pop_min(&mut heap, &nodes);
        let b = pop_min(&mut heap, &nodes);
        let f = nodes[a].freq + nodes[b].freq;
        nodes.push(Node { freq: f, sym: -1, left: a as i32, right: b as i32 });
        heap.push(nodes.len() - 1);
    }
    let root = heap[0];

    fn depth(nodes: &[Node], idx: i32, d: u8, lens: &mut Vec<u8>) {
        let node = &nodes[idx as usize];
        if node.sym >= 0 {
            lens[node.sym as usize] = d.max(1);
        } else {
            depth(nodes, node.left, d + 1, lens);
            depth(nodes, node.right, d + 1, lens);
        }
    }
    depth(&nodes, root as i32, 0, &mut lens);

    clamp_lengths(&mut lens, max_bits);
    lens
}

fn clamp_lengths(lens: &mut [u8], max_bits: usize) {
    let maxb = max_bits as u8;
    loop {
        let mut over = false;
        for l in lens.iter() {
            if *l > maxb {
                over = true;
                break;
            }
        }
        if !over {
            break;
        }
        let mut long_idx = None;
        for (i, l) in lens.iter().enumerate() {
            if *l > maxb {
                long_idx = Some(i);
                break;
            }
        }
        let li = long_idx.unwrap();
        let mut donor = None;
        for (i, l) in lens.iter().enumerate() {
            if *l > 0 && *l < maxb && i != li {
                donor = Some(i);
                break;
            }
        }
        match donor {
            Some(di) => {
                lens[li] = maxb;
                lens[di] += 1;
            }
            None => {
                lens[li] = maxb;
            }
        }
    }
}

fn canonical_codes(lens: &[u8]) -> Vec<(u32, u32)> {
    let max_len = *lens.iter().max().unwrap_or(&0) as usize;
    let mut bl_count = vec![0u32; max_len + 1];
    for &l in lens {
        if l > 0 {
            bl_count[l as usize] += 1;
        }
    }
    let mut next_code = vec![0u32; max_len + 2];
    let mut code = 0u32;
    for bits in 1..=max_len {
        code = (code + bl_count[bits - 1]) << 1;
        next_code[bits] = code;
    }
    let mut codes = vec![(0u32, 0u32); lens.len()];
    for (i, &l) in lens.iter().enumerate() {
        if l > 0 {
            codes[i] = (next_code[l as usize], l as u32);
            next_code[l as usize] += 1;
        }
    }
    codes
}

fn tally(tokens: &[Token]) -> ([u32; 288], [u32; 30]) {
    let mut lit = [0u32; 288];
    let mut dst = [0u32; 30];
    for t in tokens {
        match t {
            Token::Lit(b) => lit[*b as usize] += 1,
            Token::Match { len, dist } => {
                let (lc, _, _) = length_code(*len);
                lit[lc as usize] += 1;
                let (dc, _, _) = dist_code(*dist);
                dst[dc as usize] += 1;
            }
        }
    }
    lit[256] += 1;
    (lit, dst)
}

fn emit_dynamic_block(tokens: &[Token]) -> Vec<u8> {
    let (lit_freq, dist_freq) = tally(tokens);

    let mut lit_lens = huffman_lengths(&lit_freq, 15);
    let mut dist_lens = huffman_lengths(&dist_freq, 15);

    let mut hlit = 286;
    while hlit > 257 && lit_lens[hlit - 1] == 0 {
        hlit -= 1;
    }
    let mut hdist = 30;
    while hdist > 1 && dist_lens[hdist - 1] == 0 {
        hdist -= 1;
    }
    lit_lens.truncate(hlit);
    dist_lens.truncate(hdist);

    let lit_codes = canonical_codes(&lit_lens);
    let dist_codes = canonical_codes(&dist_lens);

    let mut cl_seq: Vec<u8> = Vec::new();
    cl_seq.extend_from_slice(&lit_lens);
    cl_seq.extend_from_slice(&dist_lens);
    let (cl_syms, cl_extra) = rle_code_lengths(&cl_seq);

    let mut cl_freq = [0u32; 19];
    for &s in &cl_syms {
        cl_freq[s as usize] += 1;
    }
    let cl_lens_full = huffman_lengths(&cl_freq, 7);
    let cl_codes = canonical_codes(&cl_lens_full);

    const ORDER: [usize; 19] = [
        16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15,
    ];
    let mut hclen = 19;
    while hclen > 4 && cl_lens_full[ORDER[hclen - 1]] == 0 {
        hclen -= 1;
    }

    let mut bw = BitWriter::new();
    bw.write_bits(1, 1);
    bw.write_bits(0b10, 2);
    bw.write_bits((hlit - 257) as u32, 5);
    bw.write_bits((hdist - 1) as u32, 5);
    bw.write_bits((hclen - 4) as u32, 4);
    for k in 0..hclen {
        bw.write_bits(cl_lens_full[ORDER[k]] as u32, 3);
    }
    for (idx, &s) in cl_syms.iter().enumerate() {
        let (code, len) = cl_codes[s as usize];
        bw.write_bits_msb(code, len);
        match s {
            16 => bw.write_bits(cl_extra[idx], 2),
            17 => bw.write_bits(cl_extra[idx], 3),
            18 => bw.write_bits(cl_extra[idx], 7),
            _ => {}
        }
    }

    emit_tokens(&mut bw, tokens, &lit_codes, &dist_codes);
    bw.finish()
}

fn emit_fixed_block(tokens: &[Token]) -> Vec<u8> {
    let mut bw = BitWriter::new();
    bw.write_bits(1, 1);
    bw.write_bits(0b01, 2);

    let lit_codes: Vec<(u32, u32)> = (0..288).map(fixed_lit_code_u).collect();
    let dist_codes: Vec<(u32, u32)> = (0..30).map(|s| (s as u32, 5u32)).collect();
    emit_tokens(&mut bw, tokens, &lit_codes, &dist_codes);
    bw.finish()
}

fn fixed_lit_code_u(sym: usize) -> (u32, u32) {
    match sym {
        0..=143 => (0x30 + sym as u32, 8),
        144..=255 => (0x190 + (sym as u32 - 144), 9),
        256..=279 => (sym as u32 - 256, 7),
        280..=287 => (0xc0 + (sym as u32 - 280), 8),
        _ => unreachable!(),
    }
}

fn emit_tokens(
    bw: &mut BitWriter,
    tokens: &[Token],
    lit_codes: &[(u32, u32)],
    dist_codes: &[(u32, u32)],
) {
    for t in tokens {
        match t {
            Token::Lit(b) => {
                let (c, l) = lit_codes[*b as usize];
                bw.write_bits_msb(c, l);
            }
            Token::Match { len, dist } => {
                let (lc, lextra, lebits) = length_code(*len);
                let (c, l) = lit_codes[lc as usize];
                bw.write_bits_msb(c, l);
                if lebits > 0 {
                    bw.write_bits(lextra, lebits);
                }
                let (dc, dextra, debits) = dist_code(*dist);
                let (c2, l2) = dist_codes[dc as usize];
                bw.write_bits_msb(c2, l2);
                if debits > 0 {
                    bw.write_bits(dextra, debits);
                }
            }
        }
    }
    let (c, l) = lit_codes[256];
    bw.write_bits_msb(c, l);
    let _ = (length_extra_bits, length_code_extra, dist_extra_bits);
}

fn rle_code_lengths(lens: &[u8]) -> (Vec<u8>, Vec<u32>) {
    let mut syms = Vec::new();
    let mut extra = Vec::new();
    let n = lens.len();
    let mut i = 0;
    while i < n {
        let cur = lens[i];
        let mut run = 1;
        while i + run < n && lens[i + run] == cur {
            run += 1;
        }
        if cur == 0 {
            while run >= 11 {
                let take = run.min(138);
                syms.push(18);
                extra.push((take - 11) as u32);
                run -= take;
                i += take;
            }
            while run >= 3 {
                let take = run.min(10);
                syms.push(17);
                extra.push((take - 3) as u32);
                run -= take;
                i += take;
            }
            for _ in 0..run {
                syms.push(0);
                extra.push(0);
                i += 1;
            }
        } else {
            syms.push(cur);
            extra.push(0);
            i += 1;
            let mut rem = run - 1;
            while rem >= 3 {
                let take = rem.min(6);
                syms.push(16);
                extra.push((take - 3) as u32);
                rem -= take;
                i += take;
            }
            for _ in 0..rem {
                syms.push(cur);
                extra.push(0);
                i += 1;
            }
        }
    }
    (syms, extra)
}

struct BitWriter {
    out: Vec<u8>,
    bitbuf: u32,
    bitcnt: u32,
}
impl BitWriter {
    fn new() -> Self {
        BitWriter { out: Vec::new(), bitbuf: 0, bitcnt: 0 }
    }
    fn write_bits(&mut self, val: u32, n: u32) {
        self.bitbuf |= (val & ((1u32 << n) - 1)) << self.bitcnt;
        self.bitcnt += n;
        while self.bitcnt >= 8 {
            self.out.push((self.bitbuf & 0xff) as u8);
            self.bitbuf >>= 8;
            self.bitcnt -= 8;
        }
    }
    fn write_bits_msb(&mut self, code: u32, n: u32) {
        for k in (0..n).rev() {
            self.write_bits((code >> k) & 1, 1);
        }
    }
    fn finish(mut self) -> Vec<u8> {
        if self.bitcnt > 0 {
            self.out.push((self.bitbuf & 0xff) as u8);
        }
        self.out
    }
}