module deflate

const deflate_hash_bits = 15
const deflate_hash_size = 1 << deflate_hash_bits
const deflate_max_chain = 64
const deflate_min_match = 3
const deflate_max_match = 258
const deflate_window = 32768

// fixed_litlen_encode returns (reversed_codes, code_lengths) for fixed Huffman lit/len.
fn fixed_litlen_encode() ([]u32, []int) {
    lens := fixed_litlen_lengths()
    mut max_bits := 0
    for l in lens {
        if l > max_bits {
            max_bits = l
        }
    }
    mut bl_count := []int{len: max_bits + 1}
    for l in lens {
        if l > 0 {
            bl_count[l]++
        }
    }
    mut next_code := []u32{len: max_bits + 1}
    mut c := u32(0)
    for bits in 1 .. max_bits + 1 {
        c = (c + u32(bl_count[bits - 1])) << 1
        next_code[bits] = c
    }
    mut codes := []u32{len: 288}
    for sym in 0 .. 288 {
        l := lens[sym]
        if l == 0 {
            continue
        }
        codes[sym] = bit_reverse(next_code[l], l)
        next_code[l]++
    }
    return codes, lens
}

// fixed_dist_encode returns (reversed_codes, code_lengths) for fixed Huffman distance.
fn fixed_dist_encode() ([]u32, []int) {
    mut codes := []u32{len: 30}
    for i in 0 .. 30 {
        codes[i] = bit_reverse(u32(i), 5)
    }
    return codes, []int{len: 30, init: 5}
}

fn length_code_info(length int) (int, int, int) {
    for i := length_bases.len - 1; i >= 0; i-- {
        if length >= length_bases[i] {
            return i, length - length_bases[i], length_extra_bits[i]
        }
    }
    return 0, 0, 0
}

fn dist_code_info(distance int) (int, int, int) {
    for i := dist_bases.len - 1; i >= 0; i-- {
        if distance >= dist_bases[i] {
            return i, distance - dist_bases[i], dist_extra_bits[i]
        }
    }
    return 0, 0, 0
}

fn hash3(data []u8, pos int) int {
    v := u32(data[pos]) | (u32(data[pos + 1]) << 8) | (u32(data[pos + 2]) << 16)
    return int((v * u32(2654435761)) >> u32(32 - deflate_hash_bits))
}

@[direct_array_access]
fn find_lz_match(data []u8, pos int, last []int, prev []int) (int, int) {
    if pos + deflate_min_match > data.len {
        return 0, 0
    }
    max_len := if pos + deflate_max_match < data.len { deflate_max_match } else { data.len - pos }
    mut best_len := 0
    mut best_off := 0
    mut i := last[hash3(data, pos)]
    mut chain := 0
    for i >= 0 && chain < deflate_max_chain {
        off := pos - i
        if off > deflate_window {
            break
        }
        mut l := 0
        for l < max_len && data[i + l] == data[pos + l] {
            l++
        }
        if l > best_len {
            best_len = l
            best_off = off
            if best_len == max_len {
                break
            }
        }
        i = prev[i]
        chain++
    }
    return best_off, best_len
}

struct BitWriter {
mut:
    buf   []u8
    bits  u32
    nbits int
}

@[direct_array_access; inline]
fn (mut w BitWriter) write_bits(value u32, nbits int) {
    if nbits == 0 {
        return
    }
    w.bits |= value << w.nbits
    w.nbits += nbits
    for w.nbits >= 8 {
        w.buf << u8(w.bits & 0xff)
        w.bits >>= 8
        w.nbits -= 8
    }
}

fn (mut w BitWriter) flush() {
    if w.nbits > 0 {
        w.buf << u8(w.bits & 0xff)
        w.bits = 0
        w.nbits = 0
    }
}

// deflate_compress_fixed compresses data to RFC 1951 DEFLATE using fixed Huffman codes.
@[direct_array_access]
fn deflate_compress_fixed(data []u8) []u8 {
    ll_codes, ll_lens := fixed_litlen_encode()
    d_codes, d_lens := fixed_dist_encode()
    mut w := BitWriter{}
    // BFINAL=1, BTYPE=01 (fixed Huffman)
    w.write_bits(1, 1)
    w.write_bits(1, 2)
    if data.len == 0 {
        w.write_bits(ll_codes[256], ll_lens[256])
        w.flush()
        return w.buf
    }
    mut last := []int{len: deflate_hash_size, init: -1}
    mut prev := []int{len: data.len, init: -1}
    mut pos := 0
    for pos < data.len {
        off, match_len := find_lz_match(data, pos, last, prev)
        if match_len >= deflate_min_match {
            li, lext, lext_bits := length_code_info(match_len)
            sym := 257 + li
            w.write_bits(ll_codes[sym], ll_lens[sym])
            w.write_bits(u32(lext), lext_bits)
            di, dext, dext_bits := dist_code_info(off)
            w.write_bits(d_codes[di], d_lens[di])
            w.write_bits(u32(dext), dext_bits)
            for i in pos .. pos + match_len {
                if i + deflate_min_match < data.len {
                    h := hash3(data, i)
                    prev[i] = last[h]
                    last[h] = i
                }
            }
            pos += match_len
        } else {
            w.write_bits(ll_codes[int(data[pos])], ll_lens[int(data[pos])])
            if pos + deflate_min_match < data.len {
                h := hash3(data, pos)
                prev[pos] = last[h]
                last[h] = pos
            }
            pos++
        }
    }
    w.write_bits(ll_codes[256], ll_lens[256])
    w.flush()
    return w.buf
}