"""Serialization utilities using cloudpickle with compression."""

from __future__ import annotations

import zlib
from typing import Any, Final

import cloudpickle
from loguru import logger

log = logger.bind(component="serialization")

COMPRESSED_MAGIC: Final = b"\x00CZ"
COMPRESSION_LEVEL: Final = 6


def serialize(obj: Any, compress: bool = True) -> bytes:
    """Serialize an object to bytes using cloudpickle with optional compression.

    Args:
        obj: Any Python object to serialize.
        compress: Whether to compress the output (default True).

    Returns:
        Serialized (and optionally compressed) bytes.
    """
    pickled: bytes = cloudpickle.dumps(obj)

    if compress:
        compressed = zlib.compress(pickled, level=COMPRESSION_LEVEL)
        if len(compressed) < len(pickled):
            log.debug(
                "Serialized {raw} -> {compressed} bytes",
                raw=len(pickled), compressed=len(compressed),
            )
            return COMPRESSED_MAGIC + compressed

    log.debug("Serialized {size} bytes (uncompressed)", size=len(pickled))
    return pickled


def deserialize(data: bytes) -> Any:
    """Deserialize bytes back to a Python object.

    Automatically detects and handles compressed data.

    Args:
        data: Bytes to deserialize (may be compressed).

    Returns:
        The deserialized Python object.
    """
    if data.startswith(COMPRESSED_MAGIC):
        data = zlib.decompress(data[len(COMPRESSED_MAGIC) :])

    return cloudpickle.loads(data)
