"""PoE Bundles2 container header + 256 KiB Oodle chunks."""
from __future__ import annotations

import struct
from dataclasses import dataclass

from .oodle import decompress_chunk


@dataclass(frozen=True)
class BundleHeader:
    uncompressed_size: int
    total_payload_size: int
    head_payload_size: int
    encoder: int
    block_count: int
    granularity: int
    block_sizes: tuple[int, ...]
    payload_offset: int


def parse_bundle_header(blob: bytes) -> BundleHeader:
    if len(blob) < 12:
        raise ValueError("bundle too small for header")
    uncompressed_size, total_payload_size, head_payload_size = struct.unpack_from("<III", blob, 0)
    inner = blob[12 : 12 + head_payload_size]
    if len(inner) < 48:
        raise ValueError("bundle head payload truncated")
    encoder = struct.unpack_from("<I", inner, 0)[0]
    block_count = struct.unpack_from("<I", inner, 24)[0]
    granularity = struct.unpack_from("<I", inner, 28)[0]
    sizes = struct.unpack_from("<" + "I" * block_count, inner, 48)
    if len(sizes) != block_count:
        raise ValueError("bundle block_sizes truncated")
    return BundleHeader(
        uncompressed_size=uncompressed_size,
        total_payload_size=total_payload_size,
        head_payload_size=head_payload_size,
        encoder=encoder,
        block_count=block_count,
        granularity=granularity,
        block_sizes=tuple(sizes),
        payload_offset=12 + head_payload_size,
    )


def decompress_bundle(blob: bytes) -> bytes:
    hdr = parse_bundle_header(blob)
    src_pos = hdr.payload_offset
    out = bytearray()
    remaining = hdr.uncompressed_size
    for i, comp_size in enumerate(hdr.block_sizes):
        chunk = blob[src_pos : src_pos + comp_size]
        src_pos += comp_size
        raw_len = hdr.granularity if remaining > hdr.granularity else remaining
        if len(chunk) == raw_len and (not chunk or chunk[:2] != b"\x8c"):
            out.extend(chunk)
        else:
            out.extend(decompress_chunk(chunk, raw_len))
        remaining -= raw_len
        if remaining <= 0:
            break
    if len(out) != hdr.uncompressed_size:
        raise ValueError(
            f"bundle decompress size mismatch: got {len(out)} expected {hdr.uncompressed_size}"
        )
    return bytes(out)