from __future__ import annotations

import struct
from dataclasses import dataclass

from packet_codec import MAX_MESSAGE_BYTES

MAGIC = b"COG2"
VERSION = 1
HEADER = struct.Struct(">4sBBHI")
HEADER_SIZE = HEADER.size

_FAMILY_BY_PACKET_TYPE = {
    "transform": 1,
    "npcState": 2,
}
_PACKET_TYPE_BY_FAMILY = {value: key for key, value in _FAMILY_BY_PACKET_TYPE.items()}


class SnapshotEnvelopeError(ValueError):
    pass


@dataclass(frozen=True)
class SnapshotEnvelope:
    packet_type: str
    sequence: int
    payload: bytes


def is_snapshot_packet_type(packet_type: str) -> bool:
    return packet_type in _FAMILY_BY_PACKET_TYPE


def encode_snapshot(packet_type: str, payload: bytes, sequence: int) -> bytes:
    family = _FAMILY_BY_PACKET_TYPE.get(packet_type)
    if family is None:
        raise SnapshotEnvelopeError(f"packet type {packet_type!r} is not a GNS snapshot family")
    if not isinstance(sequence, int) or isinstance(sequence, bool) or not 1 <= sequence <= 0xFFFFFFFF:
        raise SnapshotEnvelopeError("snapshot sequence must be between 1 and 0xffffffff")
    raw = bytes(payload)
    if not raw:
        raise SnapshotEnvelopeError("snapshot payload cannot be empty")
    if HEADER_SIZE + len(raw) > MAX_MESSAGE_BYTES:
        raise SnapshotEnvelopeError("snapshot envelope exceeds maximum GNS message size")
    return HEADER.pack(MAGIC, VERSION, family, 0, sequence) + raw


def decode_snapshot(message: bytes | bytearray | memoryview) -> SnapshotEnvelope | None:
    raw = bytes(message)
    if not raw.startswith(MAGIC):
        return None
    if len(raw) < HEADER_SIZE:
        raise SnapshotEnvelopeError("truncated GNS snapshot envelope")
    magic, version, family, reserved, sequence = HEADER.unpack_from(raw)
    if magic != MAGIC or version != VERSION or reserved != 0:
        raise SnapshotEnvelopeError("invalid GNS snapshot envelope header")
    packet_type = _PACKET_TYPE_BY_FAMILY.get(family)
    if packet_type is None:
        raise SnapshotEnvelopeError("unknown GNS snapshot family")
    if sequence == 0:
        raise SnapshotEnvelopeError("snapshot sequence zero is reserved")
    payload = raw[HEADER_SIZE:]
    if not payload:
        raise SnapshotEnvelopeError("snapshot envelope payload is empty")
    return SnapshotEnvelope(packet_type, sequence, payload)