"""Decode Spectra websocket replay payloads.

Spectra historically sent zstd-compressed JSON bytes. Newer frames may be
zstd-compressed msgpack bytes. Keep the format check on the decompressed bytes
so both rollout states are accepted without guessing from websocket frame type.
"""
from __future__ import annotations

import json
from typing import Any

import msgpack
import zstandard as zstd

_ZSTD_DECOMPRESSOR = zstd.ZstdDecompressor()
_MAX_OUTPUT_SIZE = 64 * 1024 * 1024


class SpectraPayloadError(ValueError):
    """Raised when a Spectra websocket payload cannot be decoded."""


def _frame_bytes(message: str | bytes | bytearray | memoryview) -> bytes:
    if isinstance(message, str):
        return message.encode("utf-8")
    return bytes(message)


def _decompress_if_needed(raw: bytes) -> bytes:
    try:
        return _ZSTD_DECOMPRESSOR.decompress(raw, max_output_size=_MAX_OUTPUT_SIZE)
    except zstd.ZstdError:
        return raw


def _first_non_ws_byte(payload: bytes) -> int | None:
    for byte in payload:
        if byte not in b" \t\r\n":
            return byte
    return None


def _looks_like_json(payload: bytes) -> bool:
    first = _first_non_ws_byte(payload)
    return first in (ord("{"), ord("["))


def _looks_like_msgpack_container(payload: bytes) -> bool:
    first = _first_non_ws_byte(payload)
    if first is None:
        return False
    return (
        0x80 <= first <= 0x9F  # fixmap/fixarray
        or first in (0xDC, 0xDD, 0xDE, 0xDF)  # array16/32, map16/32
    )


def decode_spectra_ws_payload(message: str | bytes | bytearray | memoryview) -> Any:
    """Decode a Spectra WS frame as zstd-json, zstd-msgpack, or plain JSON.

    Returns the decoded Python object. Raises ``SpectraPayloadError`` when the
    decompressed payload is neither a JSON object/array nor a msgpack map/array.
    """
    payload = _decompress_if_needed(_frame_bytes(message))

    if _looks_like_json(payload):
        try:
            return json.loads(payload.decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError) as exc:
            raise SpectraPayloadError("Spectra WS payload looked like JSON but was invalid") from exc

    if _looks_like_msgpack_container(payload):
        try:
            return msgpack.unpackb(payload, raw=False)
        except (msgpack.ExtraData, msgpack.FormatError, msgpack.StackError, ValueError) as exc:
            raise SpectraPayloadError("Spectra WS payload looked like msgpack but was invalid") from exc

    raise SpectraPayloadError("Spectra WS payload was neither JSON nor msgpack")