"""Payload compression using zstandard."""

from __future__ import annotations

import zstandard as zstd


class CompressionManager:
    """
    Manages compression and decompression of payloads using zstd.

    Zstandard provides excellent compression ratios with fast decompression,
    ideal for metric data which often contains repetitive patterns.
    """

    def __init__(self, level: int = 3) -> None:
        """
        Initialize compression manager.

        Args:
            level: Compression level (1-22, default 3 for balanced speed/ratio)
        """
        self._level = level
        self._compressor = zstd.ZstdCompressor(level=level)
        self._decompressor = zstd.ZstdDecompressor()

    @property
    def level(self) -> int:
        """Get the compression level."""
        return self._level

    def compress(self, data: bytes) -> bytes:
        """
        Compress data.

        Args:
            data: Raw bytes to compress

        Returns:
            Compressed bytes
        """
        return self._compressor.compress(data)

    def decompress(self, data: bytes) -> bytes:
        """
        Decompress data.

        Args:
            data: Compressed bytes

        Returns:
            Original uncompressed bytes
        """
        return self._decompressor.decompress(data)

    def compress_stream(self, data: bytes, chunk_size: int = 16384) -> bytes:
        """
        Compress data using streaming for large payloads.

        Args:
            data: Raw bytes to compress
            chunk_size: Size of chunks for streaming

        Returns:
            Compressed bytes
        """
        chunks: list[bytes] = []
        with self._compressor.stream_writer(
            ChunkWriter(chunks),  # type: ignore[arg-type]
            closefd=False,
        ) as writer:
            for i in range(0, len(data), chunk_size):
                writer.write(data[i : i + chunk_size])
        return b"".join(chunks)

    @staticmethod
    def estimate_compression_ratio(original_size: int, compressed_size: int) -> float:
        """
        Calculate compression ratio.

        Args:
            original_size: Size before compression
            compressed_size: Size after compression

        Returns:
            Compression ratio (e.g., 0.25 means 75% reduction)
        """
        if original_size == 0:
            return 1.0
        return compressed_size / original_size


class ChunkWriter:
    """Helper class for streaming compression."""

    def __init__(self, chunks: list[bytes]) -> None:
        self._chunks = chunks

    def write(self, data: bytes) -> int:
        self._chunks.append(data)
        return len(data)


# Module-level singleton for convenience
_default_manager: CompressionManager | None = None


def get_compression_manager(level: int = 3) -> CompressionManager:
    """Get or create the default compression manager."""
    global _default_manager
    if _default_manager is None or _default_manager.level != level:
        _default_manager = CompressionManager(level=level)
    return _default_manager


def compress(data: bytes, level: int = 3) -> bytes:
    """Convenience function to compress data."""
    return get_compression_manager(level).compress(data)


def decompress(data: bytes) -> bytes:
    """Convenience function to decompress data."""
    return get_compression_manager().decompress(data)
