"""Common Lab033 link packet shared by commands, telemetry, and video.

The fixed 32-byte header uses network byte order and no implicit padding::

    !4sBBBBHIQHII

    Offset  Size  Field
       0      4   magic (b"SLP1")
       4      1   version
       5      1   traffic_class
       6      1   direction
       7      1   flags
       8      2   stream_id
      10      4   sequence_number
      14      8   generation_time_us
      22      2   deadline_ms
      24      4   payload_length
      28      4   packet_crc32

``packet_crc32`` is IEEE CRC32 over the complete header with that field set
to zero, followed by the complete payload.
"""

from __future__ import annotations

from dataclasses import dataclass, replace
from enum import IntEnum
import struct
import zlib


MAGIC = b"SLP1"
VERSION = 1
HEADER_FORMAT = "!4sBBBBHIQHII"
HEADER_SIZE = struct.calcsize(HEADER_FORMAT)
SUPPORTED_FLAGS_MASK = 0


class TrafficClass(IntEnum):
    """Traffic classes in descending scheduling priority."""

    EMERGENCY = 1
    CONTROL = 2
    TELEMETRY = 3
    VIDEO = 4


class Direction(IntEnum):
    """Logical direction of a packet on the shared modelled resource."""

    GROUND_TO_ROVER = 1
    ROVER_TO_GROUND = 2


class LinkPacketError(ValueError):
    """Base class for malformed Lab033 link packets."""


class LinkPacketCRCError(LinkPacketError):
    """The Lab033 packet header or payload failed CRC32 validation."""


@dataclass(frozen=True)
class LinkPacket:
    """Decoded common packet, including its CRC-validated payload."""

    traffic_class: TrafficClass
    direction: Direction
    stream_id: int
    sequence_number: int
    generation_time_us: int
    deadline_ms: int
    payload: bytes
    flags: int = 0
    version: int = VERSION
    packet_crc32: int = 0


def crc32(data: bytes) -> int:
    """Return unsigned IEEE CRC32."""

    return zlib.crc32(data) & 0xFFFFFFFF


def _validated(packet: LinkPacket) -> LinkPacket:
    if not isinstance(packet, LinkPacket):
        raise TypeError("packet must be LinkPacket")
    try:
        traffic_class = TrafficClass(packet.traffic_class)
        direction = Direction(packet.direction)
    except ValueError as error:
        raise LinkPacketError("unsupported enumeration value") from error
    if packet.version != VERSION:
        raise LinkPacketError("unsupported link packet version")
    if packet.flags & ~SUPPORTED_FLAGS_MASK:
        raise LinkPacketError("unsupported link packet flags")
    if not 0 <= packet.stream_id <= 0xFFFF:
        raise LinkPacketError("stream_id is outside uint16")
    if not 0 <= packet.sequence_number <= 0xFFFFFFFF:
        raise LinkPacketError("sequence_number is outside uint32")
    if not 0 <= packet.generation_time_us <= 0xFFFFFFFFFFFFFFFF:
        raise LinkPacketError("generation_time_us is outside uint64")
    if not 0 <= packet.deadline_ms <= 0xFFFF:
        raise LinkPacketError("deadline_ms is outside uint16")
    if not isinstance(packet.payload, (bytes, bytearray)):
        raise TypeError("payload must be bytes or bytearray")
    payload = bytes(packet.payload)
    if len(payload) > 0xFFFFFFFF:
        raise LinkPacketError("payload is outside uint32 length")
    return replace(
        packet,
        traffic_class=traffic_class,
        direction=direction,
        payload=payload,
    )


def _pack_header(packet: LinkPacket, crc_value: int) -> bytes:
    return struct.pack(
        HEADER_FORMAT,
        MAGIC,
        packet.version,
        int(packet.traffic_class),
        int(packet.direction),
        packet.flags,
        packet.stream_id,
        packet.sequence_number,
        packet.generation_time_us,
        packet.deadline_ms,
        len(packet.payload),
        crc_value,
    )


def encode_link_packet(packet: LinkPacket) -> bytes:
    """Serialize a common packet and calculate its CRC32."""

    packet = _validated(packet)
    header_without_crc = _pack_header(packet, 0)
    packet_crc32 = crc32(header_without_crc + packet.payload)
    return _pack_header(packet, packet_crc32) + packet.payload


def decode_link_packet(wire_packet: bytes) -> LinkPacket:
    """Deserialize a common packet and validate its format and CRC32."""

    if not isinstance(wire_packet, (bytes, bytearray)):
        raise TypeError("wire_packet must be bytes or bytearray")
    wire_packet = bytes(wire_packet)
    if len(wire_packet) < HEADER_SIZE:
        raise LinkPacketError("link packet is shorter than its header")
    (
        magic,
        version,
        traffic_class,
        direction,
        flags,
        stream_id,
        sequence_number,
        generation_time_us,
        deadline_ms,
        payload_length,
        received_crc32,
    ) = struct.unpack(HEADER_FORMAT, wire_packet[:HEADER_SIZE])
    if magic != MAGIC:
        raise LinkPacketError("invalid link packet magic")
    if len(wire_packet) != HEADER_SIZE + payload_length:
        raise LinkPacketError("payload_length does not match packet length")
    try:
        packet = LinkPacket(
            traffic_class=TrafficClass(traffic_class),
            direction=Direction(direction),
            stream_id=stream_id,
            sequence_number=sequence_number,
            generation_time_us=generation_time_us,
            deadline_ms=deadline_ms,
            payload=wire_packet[HEADER_SIZE:],
            flags=flags,
            version=version,
            packet_crc32=received_crc32,
        )
    except ValueError as error:
        raise LinkPacketError("unsupported enumeration value") from error
    packet = _validated(packet)
    calculated_crc32 = crc32(_pack_header(packet, 0) + packet.payload)
    if calculated_crc32 != received_crc32:
        raise LinkPacketCRCError(
            "link packet CRC mismatch: "
            f"received 0x{received_crc32:08X}, "
            f"calculated 0x{calculated_crc32:08X}"
        )
    return packet


assert HEADER_SIZE == 32