"""Response payload encoding for the TCP serve <-> client leg.

The wire frame stays ``4-byte LE length + payload``. Payloads are plain JSON
for legacy peers, or a 1-byte codec tag followed by serialized/compressed data
when the peer advertised support for it.
"""
from __future__ import annotations

import json

from browser_cli.constants import COMP_GZIP, COMP_NONE, COMP_ZLIB, COMP_ZSTD, DEFAULT_TRANSPORT_THRESHOLD, SER_JSON, SER_MSGPACK
from browser_cli.transport.binary import hoist_screenshot as _hoist_screenshot, unhoist_binary as _unhoist_binary
from browser_cli.transport.codecs import (
  JSON_FIRST_BYTES as _JSON_FIRST_BYTES,
  _msgpack,
  choose_codec as _choose,
  client_accept_encoding,
  compress_payload as _compress,
  decompress_payload as _decompress,
  msgpack_available,
  supported_compression,
  supported_serialization,
  zstd_available,
)

def encode_response(
  obj,
  accept: dict | None = None,
  command: str | None = None,
  threshold: int = DEFAULT_TRANSPORT_THRESHOLD,
) -> bytes:
  """Encode a response object for the chosen/accepted codec.

  Returns bare JSON bytes when no encoding is negotiated, which is byte-for-byte
  what an old server would have sent.
  """
  ser, comp = _choose(accept)

  if ser == SER_MSGPACK:
    body = _msgpack.packb(_hoist_screenshot(obj, command), use_bin_type=True)
  else:
    body = json.dumps(obj).encode("utf-8")

  if comp != COMP_NONE and len(body) >= threshold:
    body = _compress(comp, body)
  else:
    comp = COMP_NONE

  if ser == SER_JSON and comp == COMP_NONE:
    return body  # plain JSON — historical wire format, no tag byte

  return bytes([(ser << 4) | comp]) + body

def decode_response(raw: bytes | None):
  """Decode a payload produced by :func:`encode_response` (or plain JSON)."""
  if raw is None:
    return None
  if not raw:
    raise ValueError("empty response payload")
  if raw[0] in _JSON_FIRST_BYTES:
    return json.loads(raw)

  tag = raw[0]
  ser, comp = tag >> 4, tag & 0x0F
  body = _decompress(comp, raw[1:])
  if ser == SER_MSGPACK:
    if _msgpack is None:
      raise ValueError("msgpack payload received but msgpack is not installed")
    return _unhoist_binary(_msgpack.unpackb(body, raw=False))
  if ser == SER_JSON:
    return json.loads(body)
  raise ValueError(f"unknown serialization id {ser}")