from abc import ABC, abstractmethod
from os.path import getsize
from heapq import heappush, heappop
from struct import pack, unpack
from time import time

class CompressionAlgorithm(ABC):
    """压缩算法基类，定义通用接口"""
    
    @abstractmethod
    def compress(self, data):
        """压缩数据，返回压缩后的完整数据块（包含元数据）"""
        pass
    
    @abstractmethod
    def decompress(self, compressed_block):
        """解压完整数据块（包含元数据）"""
        pass
    
    @property
    @abstractmethod
    def marker(self):
        """返回算法标识字节（用于文件头）"""
        pass

class LZWCompression(CompressionAlgorithm):
    """LZW压缩算法实现"""
    
    def __init__(self):
        self.marker_char = b'L'
    
    @property
    def marker(self):
        return self.marker_char
    
    def compress(self, data):
        """实现LZW压缩接口，返回完整数据块（元数据+压缩数据）"""
        compressed_block = bytearray()
        
        if not data:
            # 空数据情况
            compressed_block.extend(pack('<I', 0))  # 原始长度=0
            return bytes(compressed_block)
        
        # 初始化字典：0-255为单个字节
        dictionary = {bytes([i]): i for i in range(256)}
        next_code = 256  # 下一个可用的字典码
        max_code = 65535  # 最大字典大小（16位）
        
        w = b""
        compressed = bytearray()
        bit_buffer = 0
        bit_count = 0
        
        for c in data:
            c = bytes([c])
            wc = w + c
            if wc in dictionary:
                w = wc
            else:
                # 写入w的编码（16位）
                code = dictionary[w]
                bit_buffer = (bit_buffer << 16) | code
                bit_count += 16
                
                # 处理位缓冲区
                while bit_count >= 8:
                    # 提取高8位
                    byte_val = (bit_buffer >> (bit_count - 8)) & 0xFF
                    compressed.append(byte_val)
                    bit_count -= 8
                    bit_buffer = bit_buffer & ((1 << bit_count) - 1)
                
                # 添加新序列到字典
                if next_code <= max_code:
                    dictionary[wc] = next_code
                    next_code += 1
                
                # 重置w为当前字符
                w = c
        
        # 处理最后一个序列
        if w:
            code = dictionary[w]
            bit_buffer = (bit_buffer << 16) | code
            bit_count += 16
            
            while bit_count >= 8:
                byte_val = (bit_buffer >> (bit_count - 8)) & 0xFF
                compressed.append(byte_val)
                bit_count -= 8
                bit_buffer = bit_buffer & ((1 << bit_count) - 1)
        
        # 处理剩余的位
        if bit_count > 0:
            # 左移并填充0
            bit_buffer <<= (8 - bit_count)
            compressed.append(bit_buffer & 0xFF)
        
        # 添加元数据（原始数据长度）
        compressed_block.extend(pack('<I', len(data)))
        # 添加压缩数据
        compressed_block.extend(compressed)
        
        return bytes(compressed_block)
    
    def decompress(self, compressed_block):
        """实现LZW解压接口，处理完整数据块（元数据+压缩数据）"""
        # 解析元数据（原始数据长度）
        original_length = unpack('<I', compressed_block[:4])[0]
        
        if original_length == 0:
            return b""
        
        # 获取压缩数据部分
        compressed = compressed_block[4:]
        
        # 初始化字典
        dictionary = {i: bytes([i]) for i in range(256)}
        next_code = 256  # 下一个可用的字典码
        
        # 准备位缓冲区
        bit_buffer = 0
        bit_count = 0
        data_index = 0
        decompressed = bytearray()
        
        # 读取第一个16位编码
        while bit_count < 16 and data_index < len(compressed):
            bit_buffer = (bit_buffer << 8) | compressed[data_index]
            data_index += 1
            bit_count += 8
        
        if bit_count < 16:
            return b""
        
        # 提取第一个编码
        code = (bit_buffer >> (bit_count - 16)) & 0xFFFF
        bit_count -= 16
        bit_buffer = bit_buffer & ((1 << bit_count) - 1)
        
        w = dictionary[code]
        decompressed.extend(w)
        
        # 处理后续编码
        while len(decompressed) < original_length:
            # 读取下一个16位编码
            while bit_count < 16 and data_index < len(compressed):
                bit_buffer = (bit_buffer << 8) | compressed[data_index]
                data_index += 1
                bit_count += 8
            
            if bit_count < 16:
                break
            
            code = (bit_buffer >> (bit_count - 16)) & 0xFFFF
            bit_count -= 16
            bit_buffer = bit_buffer & ((1 << bit_count) - 1)
            
            if code in dictionary:
                entry = dictionary[code]
            elif code == next_code:
                entry = w + w[:1]
            else:
                raise ValueError("无效的LZW编码")
            
            decompressed.extend(entry)
            
            # 添加到字典
            if next_code < 65536:
                dictionary[next_code] = w + entry[:1]
                next_code += 1
            
            w = entry
        
        return bytes(decompressed[:original_length])

提高这个lzw的压缩率（我测试使用的图片压缩后反而更大了）