"""Serialization/compression primitives for TCP response payloads."""
from __future__ import annotations

import gzip
import zlib

from browser_cli.constants import COMP_GZIP, COMP_NONE, COMP_ZLIB, COMP_ZSTD, SER_JSON, SER_MSGPACK

try:  # optional: better ratio + speed than zlib/gzip
  import zstandard as _zstd
except Exception:  # pragma: no cover - depends on optional extra
  _zstd = None

try:  # optional: alternate serialization + raw binary for screenshots
  import msgpack as _msgpack
except Exception:  # pragma: no cover - depends on optional extra
  _msgpack = None

SERIALIZATION_NAME = {SER_JSON: "json", SER_MSGPACK: "msgpack"}
SERIALIZATION_ID = {value: key for key, value in SERIALIZATION_NAME.items()}
COMPRESSION_NAME = {COMP_NONE: "none", COMP_ZLIB: "zlib", COMP_GZIP: "gzip", COMP_ZSTD: "zstd"}
COMPRESSION_ID = {value: key for key, value in COMPRESSION_NAME.items()}
JSON_FIRST_BYTES = frozenset(b"{[")

def msgpack_available() -> bool:
  return _msgpack is not None

def zstd_available() -> bool:
  return _zstd is not None

def supported_serialization() -> list[str]:
  """Serializations this build can produce/consume, best first."""
  return (["msgpack"] if _msgpack is not None else []) + ["json"]

def supported_compression() -> list[str]:
  """Compression codecs this build can produce/consume, best first."""
  return (["zstd"] if _zstd is not None else []) + ["gzip", "zlib"]

def client_accept_encoding() -> dict:
  """What the local client advertises it can decode (sent with each request)."""
  return {"ser": supported_serialization(), "comp": supported_compression()}

def compress_payload(comp_id: int, data: bytes) -> bytes:
  if comp_id == COMP_NONE:
    return data
  if comp_id == COMP_ZLIB:
    return zlib.compress(data, 6)
  if comp_id == COMP_GZIP:
    return gzip.compress(data, compresslevel=6)
  if comp_id == COMP_ZSTD:
    if _zstd is None:
      raise ValueError("zstd compression requested but zstandard is not installed")
    return _zstd.ZstdCompressor(level=10).compress(data)
  raise ValueError(f"unknown compression id {comp_id}")

def decompress_payload(comp_id: int, data: bytes) -> bytes:
  if comp_id == COMP_NONE:
    return data
  if comp_id == COMP_ZLIB:
    return zlib.decompress(data)
  if comp_id == COMP_GZIP:
    return gzip.decompress(data)
  if comp_id == COMP_ZSTD:
    if _zstd is None:
      raise ValueError("zstd payload received but zstandard is not installed")
    return _zstd.ZstdDecompressor().decompress(data)
  raise ValueError(f"unknown compression id {comp_id}")

def choose_codec(accept: dict | None) -> tuple[int, int]:
  """Pick (serialization_id, compression_id) the peer accepts, server preference first."""
  accept = accept if isinstance(accept, dict) else {}
  accept_ser = accept.get("ser") or ["json"]
  accept_comp = accept.get("comp") or []

  serialization = SER_JSON
  if _msgpack is not None and "msgpack" in accept_ser:
    serialization = SER_MSGPACK

  compression = COMP_NONE
  for name in supported_compression():  # server preference: zstd > gzip > zlib
    if name in accept_comp:
      compression = COMPRESSION_ID[name]
      break
  return serialization, compression