#!/usr/bin/env python

from __future__ import print_function

import collections
import re

from gen_utils import Arg, project_dir, read_file


def parse_rfc7541_huffman_table():
    HuffmanCode = collections.namedtuple("HuffmanCode", ['length', 'char', 'bits', 'code'])
    code_regex = re.compile(
        "[(] *(?P<char>\\d+)[)] +[|](?P<bits>[01|]+) +(?P<code>[a-f0-9]+) +\\[ *(?P<length>\\d+)\\]$"
    )

    result = []
    for line in read_file(project_dir("huffman.txt")).split("\n"):
        line = line.strip()
        if not line:
            continue
        match = code_regex.search(line)
        if not match:
            continue
        bits = match.group("bits").split("|")
        result.append(HuffmanCode(
            length=int(match.group("length")),
            char=int(match.group("char")),
            bits=bits,
            code=match.group("code"),
        ))

    if len(result) != 257:
        raise Exception("bad parse: {}".format(len(result)))

    return result


def _main():
    huffman = parse_rfc7541_huffman_table()

    @Arg()
    def check_huffman_canonical():
        if sorted(huffman, key=lambda c: (c.bit_len, c.char)) != sorted(huffman, key=lambda c: c.bits_be):
            raise Exception("not canonical")

    @Arg()
    def print_huffman_code_length_table():
        buf = []
        cnt = 0
        for item in huffman:
            cnt += 1
            buf += ["{", "{:14s}".format("{:11s} {}".format("0x{},".format(item.code), item.length)), "},"]
            if not (cnt % 4):
                print(" ".join(buf))
                buf = []
        if buf:
            print (" ".join(buf))


    Arg.run()


if __name__ == "__main__":
    _main()
