"""
AmunPy — Compression handling for AMUN snapshot payloads.

This module provides a small, centralized decompression layer used by snapshot
parsers. It is intentionally independent of XML parsing and AMR logic.
"""
from __future__ import annotations
from typing import Optional

class CompressionError(RuntimeError):
    pass

_ALIASES = {
    "": "none",
    "none": "none",
    "no": "none",
    "off": "none",
    "zstd": "zstd",
    "zstandard": "zstd",
    "lz4": "lz4",
    "lz4frame": "lz4",
    "lz4-frame": "lz4",
    "lzma": "lzma",
}


def decompress(
    data: bytes,
    compression_format: Optional[str],
    *,
    max_output_size: Optional[int] = None,
) -> bytes:
    """
    Decompress `data` according to AMUN XML compression_format.

    Supported:
      - None / "" / "none": no-op
      - "zstd" / "zstandard"
      - "lz4"
      - "lzma"

    Parameters
    ----------
    max_output_size:
        If set, raises CompressionError if decompressed output exceeds this many bytes.
        For zstd this is enforced during decompression; for others it is checked after.
    """
    if compression_format is None:
        return data

    key = str(compression_format).strip().lower()
    fmt = _ALIASES.get(key)
    if fmt is None:
        raise CompressionError(f"Unsupported compression_format={compression_format!r}")

    if fmt == "none":
        return data

    try:
        if fmt == "zstd":
            try:
                import zstandard as zstd
            except ImportError as e:
                raise CompressionError("zstd compression requires the 'zstandard' package") from e

            dctx = zstd.ZstdDecompressor()
            if max_output_size is None:
                out = dctx.decompress(data)
            else:
                out = dctx.decompress(data, max_output_size=max_output_size)
            return out

        if fmt == "lz4":
            try:
                import lz4.frame
            except ImportError as e:
                raise CompressionError("lz4 compression requires the 'lz4' package") from e

            out = lz4.frame.decompress(data)
            if max_output_size is not None and len(out) > max_output_size:
                raise CompressionError(
                    f"Decompressed data exceeds max_output_size={max_output_size} bytes (got {len(out)})"
                )
            return out

        if fmt == "lzma":
            import lzma
            # memlimit constrains memory usage in the decoder (not exactly output size),
            # but it's still a useful guardrail when max_output_size is set.
            if max_output_size is None:
                return lzma.decompress(data)
            out = lzma.decompress(data, memlimit=max_output_size)
            if len(out) > max_output_size:
                raise CompressionError(
                    f"Decompressed data exceeds max_output_size={max_output_size} bytes (got {len(out)})"
                )
            return out

        # Should be unreachable due to alias check
        raise CompressionError(f"Unsupported compression_format={compression_format!r}")

    except CompressionError:
        raise
    except Exception as e:
        # Normalize all codec failures to CompressionError
        raise CompressionError(f"Failed to decompress using format '{fmt}': {e}") from e