"""BLE fragment reassembly and acknowledgement helpers for TRK1 frames."""

from __future__ import annotations

import struct

BLE_FRAGMENT_HEADER = struct.Struct("<IHH")
BLE_ACK = struct.Struct("<4sI")
BLE_BEGIN_SESSION = struct.Struct("<4sQ")
BLE_MIN_FRAME_SIZE = 36
BLE_MAX_FRAME_SIZE = 196


class BleFrameReassembler:
    def __init__(self) -> None:
        self.rejected_fragment_count = 0
        self._sequence: int | None = None
        self._total_size = 0
        self._frame = bytearray()

    def reset(self) -> None:
        self._sequence = None
        self._total_size = 0
        self._frame.clear()

    def feed(self, fragment: bytes) -> bytes | None:
        if len(fragment) <= BLE_FRAGMENT_HEADER.size:
            self.rejected_fragment_count += 1
            self.reset()
            return None
        sequence, offset, total_size = BLE_FRAGMENT_HEADER.unpack_from(fragment)
        data = fragment[BLE_FRAGMENT_HEADER.size :]
        if (
            total_size < BLE_MIN_FRAME_SIZE
            or total_size > BLE_MAX_FRAME_SIZE
            or offset >= total_size
            or offset + len(data) > total_size
        ):
            self.rejected_fragment_count += 1
            self.reset()
            return None

        # Offset zero is an explicit replay boundary, including when the same
        # packet restarts after an ACK timeout or reconnect.
        if offset == 0:
            self._sequence = sequence
            self._total_size = total_size
            self._frame = bytearray()
        if (
            self._sequence != sequence
            or self._total_size != total_size
            or offset != len(self._frame)
        ):
            self.rejected_fragment_count += 1
            self.reset()
            return None

        self._frame.extend(data)
        if len(self._frame) != self._total_size:
            return None
        frame = bytes(self._frame)
        self.reset()
        return frame


def encode_ack(packet_sequence: int) -> bytes:
    return BLE_ACK.pack(b"ACK1", packet_sequence & 0xFFFFFFFF)


def encode_begin_session(session_token: int) -> bytes:
    return BLE_BEGIN_SESSION.pack(b"BGN1", session_token & 0xFFFFFFFFFFFFFFFF)