// Copyright (c) 2019-2024 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
module http

import hash.huffman

// h2_huffman_eos is the index of the EOS (end-of-string) symbol in the
// HPACK Huffman table (RFC 7541 Appendix B).
const h2_huffman_eos = 256

// h2_max_code_bits is the longest HPACK Huffman code length (RFC 7541
// Appendix B). 30 bits is too wide for a flat decode table, so decoding goes
// bit-at-a-time against h2_huffman_decode_map.
const h2_max_code_bits = 30

// h2_huffman_table holds the canonical HPACK Huffman codes, rebuilt at startup
// from the per-symbol bit lengths via the shared hash.huffman builder. The
// codes are MSB-first (the order they appear on the wire).
const h2_huffman_table = build_h2_huffman_table()

// h2_huffman_decode_map maps a (bit_length, code) pair, packed as
// `(bit_length << 32) | code`, to its symbol. It is built once at startup
// from the canonical code table and is read-only afterwards.
const h2_huffman_decode_map = h2_huffman_table.decode_map() or { panic('hpack: ${err}') }

fn build_h2_huffman_table() huffman.Table {
	return huffman.build(
		lengths:   h2_huffman_code_lens[..].map(int(it))
		max_bits:  h2_max_code_bits
		bit_order: .msb_first
	) or { panic('hpack: ${err}') }
}

// h2_huffman_encode returns the HPACK Huffman encoding of `input`
// (RFC 7541 Section 5.2). The final byte is padded with the most significant
// bits of the EOS code, i.e. all ones.
fn h2_huffman_encode(input []u8) []u8 {
	mut out := []u8{cap: input.len}
	mut acc := u64(0)
	mut nbits := 0
	for b in input {
		code := u64(h2_huffman_table.codes[b])
		clen := h2_huffman_table.lengths[b]
		acc = (acc << clen) | code
		nbits += clen
		for nbits >= 8 {
			nbits -= 8
			out << u8((acc >> nbits) & 0xff)
		}
		// Keep only the still-buffered low bits, so acc stays bounded.
		acc = if nbits == 0 { u64(0) } else { acc & ((u64(1) << nbits) - 1) }
	}
	if nbits > 0 {
		pad := 8 - nbits
		out << u8(((acc << pad) | ((u64(1) << pad) - 1)) & 0xff)
	}
	return out
}

// h2_huffman_decode reverses h2_huffman_encode. It returns an error for the
// invalid cases called out by RFC 7541 Section 5.2: an explicit EOS symbol,
// padding longer than 7 bits, or padding that is not all ones.
fn h2_huffman_decode(input []u8) ![]u8 {
	mut out := []u8{cap: input.len + input.len / 2}
	mut cur := u64(0)
	mut cur_len := 0
	for b in input {
		for bit := 7; bit >= 0; bit-- {
			cur = (cur << 1) | u64((b >> u8(bit)) & 1)
			cur_len++
			if cur_len > h2_max_code_bits {
				return error('hpack: invalid huffman code (no symbol within ${h2_max_code_bits} bits)')
			}
			if sym := h2_huffman_decode_map[(u64(cur_len) << 32) | cur] {
				if sym == h2_huffman_eos {
					return error('hpack: EOS symbol encountered in huffman string')
				}
				out << u8(sym)
				cur = 0
				cur_len = 0
			}
		}
	}
	if cur_len >= 8 {
		return error('hpack: invalid huffman padding (more than 7 bits left)')
	}
	if cur_len > 0 {
		mask := (u64(1) << cur_len) - 1
		if cur & mask != mask {
			return error('hpack: invalid huffman padding (not all ones)')
		}
	}
	return out
}