"""
RBSP Compression Utilities.

Provides transparent compression for stored fields and index data using
Python stdlib ``zlib``.  Two convenience classes are included:

    - ``CompressedStore``   Key-value store with size-based auto-compression.
    - ``BatchCompressor``   Bulk compress/decompress with length-prefixed framing.

All functions use zlib only (no external dependencies).
"""

from __future__ import annotations

import struct
import zlib
from typing import Dict, Optional

# ---------------------------------------------------------------------------
# Header magic — used to distinguish compressed from raw data in the store
# ---------------------------------------------------------------------------

_MAGIC = b"RBC\x01"  # "RB Compress" version 1
_HEADER_FMT = ">4sI"  # magic (4 bytes) + original size (4 bytes, big-endian)
_HEADER_SIZE = struct.calcsize(_HEADER_FMT)


# ---------------------------------------------------------------------------
# Low-level compress / decompress
# ---------------------------------------------------------------------------

def compress(data: bytes, level: int = 6) -> bytes:
    """Compress *data* with zlib and prepend a small header.

    The header stores a magic marker and the original uncompressed size so
    that ``decompress`` can validate and allocate efficiently.

    Args:
        data:  Raw bytes to compress.
        level: zlib compression level (0-9, default 6).

    Returns:
        Header + zlib-compressed payload.
    """
    header = struct.pack(_HEADER_FMT, _MAGIC, len(data))
    return header + zlib.compress(data, level)


def decompress(data: bytes) -> bytes:
    """Decompress bytes produced by ``compress``.

    Validates the magic header before decompressing.

    Raises:
        ValueError: If the header magic is missing or corrupt.
    """
    if len(data) < _HEADER_SIZE:
        raise ValueError("Data too short for compression header")
    magic, _orig_size = struct.unpack(_HEADER_FMT, data[:_HEADER_SIZE])
    if magic != _MAGIC:
        raise ValueError("Invalid compression header magic")
    return zlib.decompress(data[_HEADER_SIZE:])


# ---------------------------------------------------------------------------
# String convenience wrappers
# ---------------------------------------------------------------------------

def compress_string(text: str) -> bytes:
    """UTF-8 encode *text* then compress."""
    return compress(text.encode("utf-8"))


def decompress_string(data: bytes) -> str:
    """Decompress then UTF-8 decode back to a string."""
    return decompress(data).decode("utf-8")


# ---------------------------------------------------------------------------
# CompressedStore — key-value store with transparent compression
# ---------------------------------------------------------------------------

class CompressedStore:
    """Key-value store that transparently compresses values above a size
    threshold.

    Small values are stored raw to avoid the overhead of compression on
    tiny payloads.  Each entry tracks its original size so that
    ``compression_ratio`` can be reported accurately.

    Args:
        min_size: Only compress values whose length exceeds this many bytes
                  (default 128).
    """

    def __init__(self, min_size: int = 128) -> None:
        self._min_size = min_size
        # Stored entries: key -> (payload_bytes, original_size, is_compressed)
        self._data: Dict[str, tuple[bytes, int, bool]] = {}

    def put(self, key: str, value: bytes) -> None:
        """Store *value* under *key*, compressing if large enough."""
        original_size = len(value)
        if original_size > self._min_size:
            payload = compress(value)
            self._data[key] = (payload, original_size, True)
        else:
            self._data[key] = (value, original_size, False)

    def get(self, key: str) -> Optional[bytes]:
        """Retrieve and decompress the value for *key*, or ``None``."""
        entry = self._data.get(key)
        if entry is None:
            return None
        payload, _orig_size, is_compressed = entry
        if is_compressed:
            return decompress(payload)
        return payload

    def size_bytes(self) -> int:
        """Total compressed (on-disk) size of all stored values."""
        return sum(len(payload) for payload, _, _ in self._data.values())

    def compression_ratio(self) -> float:
        """Ratio of compressed size to original size.

        Returns a value in (0, 1] for typical data.  Returns 1.0 when the
        store is empty (no compression applied).
        """
        total_original = sum(orig for _, orig, _ in self._data.values())
        if total_original == 0:
            return 1.0
        return self.size_bytes() / total_original

    def keys(self) -> list[str]:
        """Return all stored keys."""
        return list(self._data.keys())


# ---------------------------------------------------------------------------
# BatchCompressor — bulk operations with length-prefixed framing
# ---------------------------------------------------------------------------

class BatchCompressor:
    """Compress and decompress batches of byte items into a single blob.

    Framing format (before compression)::

        [4-byte item count] [4-byte len_1][item_1] [4-byte len_2][item_2] ...

    The entire framed payload is then zlib-compressed as one unit for
    better compression ratio (shared dictionary across items).
    """

    @staticmethod
    def compress_batch(items: list[bytes]) -> bytes:
        """Compress *items* into a single length-prefixed blob."""
        parts: list[bytes] = [struct.pack(">I", len(items))]
        for item in items:
            parts.append(struct.pack(">I", len(item)))
            parts.append(item)
        return compress(b"".join(parts))

    @staticmethod
    def decompress_batch(data: bytes) -> list[bytes]:
        """Decompress a blob produced by ``compress_batch`` back into items."""
        raw = decompress(data)
        if len(raw) < 4:
            raise ValueError("Batch data too short for item count")
        (count,) = struct.unpack(">I", raw[:4])
        offset = 4
        items: list[bytes] = []
        for _ in range(count):
            if offset + 4 > len(raw):
                raise ValueError("Truncated batch: missing item length")
            (item_len,) = struct.unpack(">I", raw[offset : offset + 4])
            offset += 4
            if offset + item_len > len(raw):
                raise ValueError("Truncated batch: missing item data")
            items.append(raw[offset : offset + item_len])
            offset += item_len
        return items
