def _chr(x):
    return bytes([x])

def lzw_compress(bstr):
    # type TrieNode = (Code, Map Byte TrieNode)
    # type Trie = Map Byte TrieNode
    d = {x: (x, {}) for x in range(0, 256)}
    d_counter = 256
    def _gensym():
        nonlocal d_counter
        res = d_counter
        # maximum code 12-bit
        if res >= 4095:
            return None
        d_counter += 1
        return res
    res = []
    p = b''
    def _chkpc(pc):
        nonlocal d
        d_pointer = d
        res = True
        code = -1
        for c in pc:
            if d_pointer.get(c):
                code = d_pointer[c][0]
                d_pointer = d_pointer[c][1]
            else:
                res = False
                break
        return res, code, d_pointer

    for i, c in enumerate(bstr):
        cc = _chr(c)
        chkres, pcode, d_parent = _chkpc(p + cc)
        if chkres:
            p = p + cc
        else:
            res.append(pcode)
            sym = _gensym()
            if sym is not None:
                d_parent[c] = (sym, {})
            p = cc
        
    chkres, pcode, d_parent = _chkpc(p)
    res.append(pcode)
    return res, d

def lzw_decompress(code_list):
    d = {x: _chr(x) for x in range(0, 256)}
    d_counter = 256
    def _gensym():
        nonlocal d_counter
        res = d_counter
        d_counter += 1
        return res
    cw = code_list[0]
    res = [d[cw]]
    pw = cw
    p = b''
    for i, cw in enumerate(code_list[1:]):
        if d.get(cw):
            res.append(d[cw])
            p = d[pw]
            c = _chr(d[cw][0])
            d[_gensym()] = p+c
        else:
            p = d[pw]
            c = _chr(d[pw][0])
            res.append(p+c)
            d[_gensym()] = p+c
        pw = cw
    return b''.join(res)

def _pack(code_list):
    if len(code_list) % 2:
        code_list.append(4095)
    i = 0
    lenx = len(code_list)
    res = []
    while i < lenx:
        code1 = code_list[i]
        code2 = code_list[i+1]
        byte1 = (code1 >> 4)
        byte2 = ((code1 & 0xf) << 4) | (code2 >> 8)
        byte3 = (code2 & 0xff)
        res.append(byte1)
        res.append(byte2)
        res.append(byte3)
        i += 2
    return bytes(res)

def _unpack(bstr):
    res = []
    i = 0
    lenx = len(bstr)
    while i < lenx:
        piece = bstr[i:i+3]
        code1 = (piece[0] << 4) | (piece[1] >> 4)
        code2 = ((piece[1] & 0xf) << 8) | piece[2]
        res.append(code1)
        res.append(code2)
        i += 3
    return res

import sys

if len(sys.argv) < 3:
    print('Usage: lzw_demo [compress|decompress] [filename]')
elif sys.argv[1] == 'compress':
    with open(sys.argv[2], 'rb') as f:
        data = f.read()
    res, _ = lzw_compress(data)
    with open(f'{sys.argv[2]}.lzw', 'wb') as f:
        f.write(_pack(res))
elif sys.argv[1] == 'decompress':
    with open(sys.argv[2], 'rb') as f:
        data = f.read()
    unpacked = _unpack(data)
    if (unpacked[len(unpacked) - 1] == 4095):
        unpacked.pop()
    res = lzw_decompress(unpacked)
    with open(f'{sys.argv[2]}.lzw_dec', 'wb') as f:
        f.write(res)