"""Emergency-command acknowledgement carried by unchanged Lab033 packets."""

from __future__ import annotations

from dataclasses import dataclass
import struct

from protocol.link_packet import Direction, LinkPacket, TrafficClass


ACK_MAGIC = b"ACK1"
ACK_PAYLOAD_FORMAT = "!4sHI"
ACK_PAYLOAD_SIZE = struct.calcsize(ACK_PAYLOAD_FORMAT)
STREAM_EMERGENCY_ACK = 5


class EmergencyAckError(ValueError):
    """An emergency acknowledgement is malformed."""


@dataclass(frozen=True)
class EmergencyIdentity:
    stream_id: int
    sequence_number: int


def encode_ack_payload(identity: EmergencyIdentity) -> bytes:
    if not 0 <= identity.stream_id <= 0xFFFF:
        raise EmergencyAckError("emergency stream_id is outside uint16")
    if not 0 <= identity.sequence_number <= 0xFFFFFFFF:
        raise EmergencyAckError("emergency sequence_number is outside uint32")
    return struct.pack(ACK_PAYLOAD_FORMAT, ACK_MAGIC, identity.stream_id, identity.sequence_number)


def decode_ack_payload(payload: bytes) -> EmergencyIdentity:
    if not isinstance(payload, (bytes, bytearray)):
        raise TypeError("payload must be bytes or bytearray")
    if len(payload) != ACK_PAYLOAD_SIZE:
        raise EmergencyAckError("invalid acknowledgement payload length")
    magic, stream_id, sequence_number = struct.unpack(ACK_PAYLOAD_FORMAT, bytes(payload))
    if magic != ACK_MAGIC:
        raise EmergencyAckError("invalid acknowledgement payload magic")
    return EmergencyIdentity(stream_id, sequence_number)


def build_emergency_ack(command: LinkPacket, ack_sequence: int, generation_time_us: int) -> LinkPacket:
    """Build a high-priority rover-to-ground ACK for one emergency command."""

    if command.traffic_class is not TrafficClass.EMERGENCY:
        raise ValueError("acknowledged packet must be an emergency command")
    return LinkPacket(
        traffic_class=TrafficClass.EMERGENCY,
        direction=Direction.ROVER_TO_GROUND,
        stream_id=STREAM_EMERGENCY_ACK,
        sequence_number=ack_sequence,
        generation_time_us=generation_time_us,
        deadline_ms=50,
        payload=encode_ack_payload(EmergencyIdentity(command.stream_id, command.sequence_number)),
    )


def acknowledged_identity(packet: LinkPacket) -> EmergencyIdentity:
    """Validate ACK envelope fields and return the referenced command."""

    if packet.traffic_class is not TrafficClass.EMERGENCY:
        raise EmergencyAckError("ack must use high-priority traffic class")
    if packet.direction is not Direction.ROVER_TO_GROUND:
        raise EmergencyAckError("ack must travel rover to ground")
    if packet.stream_id != STREAM_EMERGENCY_ACK:
        raise EmergencyAckError("unexpected acknowledgement stream")
    return decode_ack_payload(packet.payload)


class EmergencyAckReceiver:
    """Recognize first and late acknowledgements independently per command."""

    def __init__(self) -> None:
        self._received: set[EmergencyIdentity] = set()
        self.unique = 0
        self.duplicates = 0

    def accept(self, packet: LinkPacket) -> bool:
        identity = acknowledged_identity(packet)
        if identity in self._received:
            self.duplicates += 1
            return False
        self._received.add(identity)
        self.unique += 1
        return True

    def has_received(self, identity: EmergencyIdentity) -> bool:
        return identity in self._received