#! python3

'''
解压被huffman编码后的文件
'''

import sys
from huffman import *
from constant import * # 导入常量

def read_weights(fin):
    '''
    读取频率表
    fin: 文件对象
    '''
    weights = [ int.from_bytes(fin.read(FOUR), 'little') for i in range(CHAR_NUM)]

    return weights


def decoding(fin, fout, root, total, num):
    '''
    解码文件
    fin: 输入文件对象
    fout: 输出文件对象
    root: huffman树的根节点
    total: 要解码的串的长度(字节为单位)
    num: 最后一个字节的有效位数
    '''
    cur = root # cur指向当前节点
    for byte in range(total):
        code = int.from_bytes(fin.read(1), 'little') # 读入一个字节的编码串, 并转换为整数
        # debug{
        #print('code=%s' % code)
        #sys.exit()
        # }
        n = num if byte==total-1 else BYTE_SIZE
        for bit in range(n):
            if code & MASKS[bit]:
                # 此位为1, 向右子树移动
                cur = cur.right
                #print('cur = cur.right') # debug
            else:
                # 此为为0, 向左子树移动
                cur = cur.left
                #print('cur = cur.left') # debug
            
            if not (cur.left or cur.right):
                # 左右子树都为空, 表示到达叶节点
                char = cur.char
                #print(char)
                #sys.exit()
                fout.write(char.to_bytes(1, 'little'))
                cur = root # 重新从根节点开始
    

def decompress(input_file, output_file):
    '''
    input_file: 输入文件
    output_file: 输出文件
    '''
    fin = open(input_file, 'rb')
    fout = open(output_file, 'wb')

    fin.seek(-(FOUR*2),2) # 文件最后的8字节是编码串的长度(字节为单位),和最后一个字节编码的有效位数
    total = int.from_bytes(fin.read(FOUR), 'little')
    num = int.from_bytes(fin.read(FOUR), 'little')

    #print('total=%s, num=%s' % (total, num)) # debug

    # 读取字符的频率列表
    fin.seek(total, 0)
    weights = read_weights(fin)

    # debug{
    #logfile = input_file + '.log1'
    #with open(logfile, 'w') as fd:
    #    fd.write(str(weights))
    # }

    # 构建huffman树
    huff_tree = HuffTree()
    huff_tree.build_huffman_tree(weights)

    #print('huff_tree.root.right.char=%s' % huff_tree.root.right.char) # debug

    #print('huff_tree.root.right(%s,%s)' % (huff_tree.root.right.left, huff_tree.root.right.right) ) # debug

    # 解码后输出到输出文件
    fin.seek(0, 0)
    decoding(fin, fout, huff_tree.root, total, num)

    fin.close()
    fout.close()

    print('%s decompress as %s' % (input_file, output_file) )


if __name__ == '__main__':
    if len(sys.argv) != 2:
        print('usage: python %s file_name' % sys.argv[0])
        sys.exit()
    
    input_file = sys.argv[1]
    output_file = input_file.rsplit('.', 1)[0] # 
    decompress(input_file, output_file)