'''
    TRANSFER_MODEL:
    PACKAGE_ID(2) + CHUNK_NUM(2) + CHUNK_ID(2) + BODY_SIZE(4) + BODY_DATA(max: 1MB) + CHECKSUM(1)
'''
SIZE_FOR_PACKAGE_ID = 2 # 包索引大小
SIZE_FOR_CHUNK_NUM = 2 # 块数量大小
SIZE_FOR_CHUNK_ID = 2 # 块索引大小
SIZE_FOR_BODY_SIZE = 4 # 包体大小大小
SIZE_FOR_CHECKSUM = 1 # 校验和大小
BYTE_ORDER = 'big' # 字节顺序


def gen_checksum(data: bytes) -> int:
    """ 生成校验和 """
    checksum = 0
    for b in data:
        checksum ^= b
    return checksum
    

class Encoder:
    """ 编码器 """
    def __init__(self, chunk_size: int = 1024 * 1024) -> None:
        self.chunk_size = chunk_size
        self.package_id = 0
        
    def _gen_package_id(self) -> int:
        """ 生成包索引 """
        self.package_id += 1
        return self.package_id if self.package_id < 65536 else 0
        
    def _gen_chunk(self, package_id: int, chunk_num: int, chunk_id: int, chunk_data: bytes):
        """ 生成块 """
        return b''.join([
            package_id.to_bytes(SIZE_FOR_PACKAGE_ID, BYTE_ORDER),
            chunk_num.to_bytes(SIZE_FOR_CHUNK_NUM, BYTE_ORDER),
            chunk_id.to_bytes(SIZE_FOR_CHUNK_ID, BYTE_ORDER),
            len(chunk_data).to_bytes(SIZE_FOR_BODY_SIZE, BYTE_ORDER),
            chunk_data,
            gen_checksum(chunk_data).to_bytes(SIZE_FOR_CHECKSUM, BYTE_ORDER),
        ])

    def encode(self, data: bytes) -> list[bytes]:
        """ 编码数据 """
        package_id = self._gen_package_id()
        chunk_num = len(data) // self.chunk_size + 1
        return [
            self._gen_chunk(package_id, chunk_num, chunk_index, data[chunk_size: chunk_size + self.chunk_size])
            for chunk_index, chunk_size in enumerate(range(0, len(data), self.chunk_size), start=1)
        ]
        
        
class PackageCache:
    """ 包缓存 """
    def __init__(self, chunk_num: int) -> None:
        self.buffer: dict[int, bytes] = {}
        self.chunk_num = chunk_num
        
    def ok(self) -> bool:
        """ 是否完整 """
        return len(self.buffer) == self.chunk_num
    
    def add(self, chunk_id: int, chunk_data: bytes) -> None:
        """ 添加数据 """
        self.buffer[chunk_id] = chunk_data
        
    @property
    def value(self) -> bytes:
        """ 获取数据 """
        buffer = bytearray()
        for chunk_id in range(1, self.chunk_num + 1):
            buffer.extend(self.buffer[chunk_id])
        return bytes(buffer)


class Decoder:
    """ 解码器 """
    def __init__(self) -> None:
        self.package_cache: dict[int, PackageCache] = {}
        
    def decode(self, encoded_bytes: bytes) -> tuple[bool, bytes]:
        """ 解码数据 """
        cursor = 0
        package_id = int.from_bytes(encoded_bytes[cursor:SIZE_FOR_PACKAGE_ID], BYTE_ORDER)
        cursor += SIZE_FOR_PACKAGE_ID
        chunk_num = int.from_bytes(encoded_bytes[cursor:cursor + SIZE_FOR_CHUNK_NUM], BYTE_ORDER)
        cursor += SIZE_FOR_CHUNK_NUM
        chunk_id = int.from_bytes(encoded_bytes[cursor:cursor + SIZE_FOR_CHUNK_ID], BYTE_ORDER)
        cursor += SIZE_FOR_CHUNK_ID
        body_size = int.from_bytes(encoded_bytes[cursor:cursor + SIZE_FOR_BODY_SIZE], BYTE_ORDER)
        cursor += SIZE_FOR_BODY_SIZE
        body_data = encoded_bytes[cursor:cursor + body_size]
        cursor += body_size
        checksum = int.from_bytes(encoded_bytes[cursor:cursor + SIZE_FOR_CHECKSUM], BYTE_ORDER)
        if checksum != gen_checksum(body_data):
            # 校验失败
            return False, None
        if package_id not in self.package_cache:
            self.package_cache[package_id] = PackageCache(chunk_num)
        self.package_cache[package_id].add(chunk_id, body_data)
        if self.package_cache[package_id].ok():
            cache = self.package_cache.pop(package_id)
            return True, cache.value
        else:
            return False, None
