# Copyright 2009 Wayne See
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import struct

_NUM_CHARS = 256

_LZX_PRETREE_MAXSYMBOLS = 20
_LZX_PRETREE_NUM_ELEMENTS_BITS = 4
_LZX_PRETREE_TABLEBITS = 6

_LZX_MAINTREE_TABLEBITS = 12
_LZX_MAINTREE_MAXSYMBOLS = _NUM_CHARS + 50 * 8

_NUM_SECONDARY_LENGTHS = 249
_LZX_LENGTH_TABLEBITS = 12
_LZX_LENGTH_MAXSYMBOLS = _NUM_SECONDARY_LENGTHS + 1

_LZX_NUM_PRIMARY_LENGTHS = 7

_MIN_MATCH = 2

_EXTRA_BITS = [ 0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4,
            5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13, 13, 14,
            14, 15, 15, 16, 16, 17, 17, 17, 17, 17, 17, 17, 17, 17, 17, 17, 17,
            17, 17, 17 ]

_POSITION_BASE = [ 0, 1, 2, 3, 4, 6, 8, 12, 16, 24,
            32, 48, 64, 96, 128, 192, 256, 384, 512, 768, 1024, 1536, 2048,
            3072, 4096, 6144, 8192, 12288, 16384, 24576, 32768, 49152, 65536,
            98304, 131072, 196608, 262144, 393216, 524288, 655360, 786432,
            917504, 1048576, 1179648, 1310720, 1441792, 1572864, 1703936,
            1835008, 1966080, 2097152 ]

class _LzxBlock:
    
    def __init__(self):
        self.content_length = 0
        self.lzx_state = _LzxState()
    
class _LzxState:
    
    def __init__(self):
        self.block_length = 0
        self.block_remaining = 0
        self.type = None
        self.intel_file_size = 0
        self.header_read = False
        self.intel_started = False
        self.R0 = 1
        self.R1 = 1
        self.R2 = 1
        
class _BitBuffer:
    
    def __init__(self, bytes):
        self.left = 0
        self.value = 0
        self.bytes = bytes
        self.length = len(bytes)
        self.traversed = 0
        
    def read_bits(self, get_bits, remove_bits=None):
        if remove_bits is None:
            remove_bits = get_bits
        while self.left < 16:
            self.value = (self.value << 16) + self.pop() + (self.pop() << 8)
            self.left += 16
        temp = self.value >> (self.left - get_bits)
        self.left -= remove_bits
        t = self.value >> self.left
        self.value -= (t << self.left)
        return temp
    
    def pop(self):
        if self.traversed < self.length:
            val, = struct.unpack("B", self.bytes[self.traversed])
            self.traversed += 1
            return val
        else:
            return 0
    
    def check_bit(self, i):
        n = 1 << (self.left - i)
        if (self.value & n) == 0:
            return 0
        return 1
        

def create_lzx_block(block_no, window, bytes, block_length, prev_block=None):
    block = _create_empty_block(window)
    block.block_no = block_no
    if prev_block is None:
        prev_block = _create_empty_block(window)
    lzx_state = prev_block.lzx_state
    block.lzx_state = lzx_state
    if lzx_state.block_length > lzx_state.block_remaining:
        prev_content = prev_block.content
    else:
        prev_content = []
    buf = _BitBuffer(bytes)
    block.content = [0 for i in xrange(block_length)]
    if not lzx_state.header_read:
        lzx_state.header_read = True
        if buf.read_bits(1) == 1:
            lzx_state.intel_file_size = (buf.read_bits(16) << 16) + buf.read_bits(16) 
    while block.content_length < block_length:
        if lzx_state.block_remaining == 0:
            lzx_state.type = buf.read_bits(3)
            assert lzx_state.type == 1 #can't handle anything but verbatim
            lzx_state.block_length = (buf.read_bits(16) << 8) + buf.read_bits(8)
            lzx_state.block_remaining = lzx_state.block_length
            lzx_state._main_tree_table = _create_main_tree_table(lzx_state, buf)
            lzx_state._length_tree_table = _create_length_tree_table(lzx_state, buf)
            if lzx_state._main_tree_length_table[0xe8] != 0:
                lzx_state.intel_started = True
        if block.content_length + lzx_state.block_remaining > block_length:
            lzx_state.block_remaining = block.content_length + lzx_state.block_remaining - block_length
            length = block_length
        else:
            length = block.content_length + lzx_state.block_remaining
            lzx_state.block_remaining = 0
        _decompress_verbatim_block(block.content, lzx_state, block.content_length, buf, length, block_length, prev_content)
        block.content_length = length
    return block

def _get_main_tree_index(buf, main_bits, tree_table, main_max_symbol):
    f = buf.read_bits(main_bits, 0)
    z = tree_table[f]
    if z >= main_max_symbol:
        x = main_bits
        while True:
            x += 1
            z <<= 1
            z += buf.check_bit(x)
            z = tree_table[z]
            if z < main_max_symbol:
                break
    return z

def _decompress_verbatim_block(content, lzx_state, content_length, buf, length, block_length, prev_content):
    main_tree = lzx_state._main_tree_table
    main_tree_length = lzx_state._main_tree_length_table
    length_tree = lzx_state._length_tree_table
    length_tree_length = lzx_state._length_tree_length_table
    R0 = lzx_state.R0
    R1 = lzx_state.R1
    R2 = lzx_state.R2
    i = content_length
    while i < length:
        s = _get_main_tree_index(buf, _LZX_MAINTREE_TABLEBITS, main_tree, lzx_state._main_tree_elements) 
        buf.read_bits(main_tree_length[s])
        if s < _NUM_CHARS:
            content[i] = s
        else:
            s -= _NUM_CHARS
            match_length = s & _LZX_NUM_PRIMARY_LENGTHS
            if match_length == _LZX_NUM_PRIMARY_LENGTHS:
                match_footer = _get_main_tree_index(buf, _LZX_LENGTH_TABLEBITS, length_tree, _NUM_SECONDARY_LENGTHS)
                buf.read_bits(length_tree_length[match_footer])
                match_length += match_footer
            match_length += _MIN_MATCH
            match_offset = s >> 3
            if match_offset > 2:
                if match_offset != 3:
                    extra = _EXTRA_BITS[match_offset]
                    l = buf.read_bits(extra)
                    match_offset = _POSITION_BASE[match_offset] - 2 + l
                else:
                    match_offset = 1
                R2 = R1
                R1 = R0
                R0 = match_offset
            elif match_offset == 0:
                match_offset = R0
            elif match_offset == 1:
                match_offset = R1
                R1 = R0
                R0 = match_offset
            else:
                match_offset = R2
                R2 = R0
                R0 = match_offset
            run_dest = i
            run_src = run_dest - match_offset
            i += (match_length - 1)
            if i > length:
                break
            if run_src < 0:
                if match_length + run_src <= 0:
                    run_src += len(prev_content)
                    while match_length > 0:
                        match_length -= 1
                        content[run_dest] = prev_content[run_src]
                        run_dest += 1
                        run_src += 1
                else:
                    prev_content_length = len(prev_content) 
                    run_src += prev_content_length
                    while run_src < prev_content_length:
                        content[run_dest] = prev_content[run_src]
                        run_dest += 1
                        run_src += 1
                    match_length = match_length + run_src - prev_content_length
                    run_src = 0
                    while match_length > 0:
                        match_length -= 1
                        content[run_dest] = content[run_src]
                        run_dest += 1
                        run_src += 1
            else:
                while (run_src < 0) and (match_length > 0):
                    content[run_dest] = content[run_src + block_length]
                    run_dest += 1