"""Zstd (with gzip fallback) compression helpers for otelctx."""

from __future__ import annotations

import gzip
import json
from typing import Any

try:
    import zstandard as zstd
    _HAS_ZSTD = True
except ImportError:
    _HAS_ZSTD = False


def compress_bytes(data: bytes, level: int = 3) -> bytes:
    if _HAS_ZSTD:
        cctx = zstd.ZstdCompressor(level=level)
        return cctx.compress(data)
    return gzip.compress(data, compresslevel=level)


def decompress_bytes(data: bytes) -> bytes:
    if _HAS_ZSTD:
        dctx = zstd.ZstdDecompressor()
        return dctx.decompress(data)
    return gzip.decompress(data)


def compress_json(obj: Any, level: int = 3) -> bytes:
    raw = json.dumps(obj, default=str).encode()
    return compress_bytes(raw, level=level)


def decompress_json(data: bytes) -> Any:
    raw = decompress_bytes(data)
    return json.loads(raw.decode())


def compression_backend() -> str:
    return "zstd" if _HAS_ZSTD else "gzip"
