"""Optional reliability layer for datagram transport."""

from __future__ import annotations

import asyncio
import struct
import weakref
from collections import deque
from dataclasses import dataclass, field
from types import TracebackType
from typing import TYPE_CHECKING, Self

from pywebtransport.exceptions import DatagramError, TimeoutError
from pywebtransport.types import Data, EventType
from pywebtransport.utils import ensure_bytes, get_logger, get_timestamp

if TYPE_CHECKING:
    from pywebtransport.events import Event
    from pywebtransport.session.session import WebTransportSession


__all__: list[str] = ["DatagramReliabilityLayer"]

logger = get_logger(name=__name__)


class DatagramReliabilityLayer:
    """Add a TCP-like reliability layer over an unreliable datagram transport."""

    def __init__(self, session: WebTransportSession, *, ack_timeout: float = 2.0, max_retries: int = 5) -> None:
        """Initialize the datagram reliability layer."""
        self._session = weakref.ref(session)
        self._ack_timeout = ack_timeout
        self._max_retries = max_retries
        self._closed = False
        self._send_sequence = 0
        self._receive_sequence = 0
        self._pending_acks: dict[int, _ReliableDatagram] = {}
        self._received_sequences: deque[int] = deque(maxlen=1024)
        self._incoming_queue: asyncio.Queue[bytes] | None = None
        self._lock: asyncio.Lock | None = None
        self._retry_task: asyncio.Task[None] | None = None

    async def __aenter__(self) -> Self:
        """Enter the async context and start background tasks."""
        self._lock = asyncio.Lock()
        self._incoming_queue = asyncio.Queue()
        if session := self._session():
            session.events.on(event_type=EventType.DATAGRAM_RECEIVED, handler=self._on_datagram_received)
        self._start_background_tasks()
        return self

    async def __aexit__(
        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
    ) -> None:
        """Exit the async context and close the reliability layer."""
        await self.close()

    async def close(self) -> None:
        """Gracefully close the reliability layer and clean up resources."""
        if self._closed:
            return

        self._closed = True

        if session := self._session():
            session.events.off(event_type=EventType.DATAGRAM_RECEIVED, handler=self._on_datagram_received)

        if self._retry_task:
            self._retry_task.cancel()
            try:
                await self._retry_task
            except asyncio.CancelledError:
                pass

        if self._lock:
            async with self._lock:
                self._pending_acks.clear()

        self._received_sequences.clear()
        logger.debug("Reliability layer closed")

    async def receive(self, *, timeout: float | None = None) -> bytes:
        """Receive a reliable datagram, waiting if necessary."""
        if self._incoming_queue is None:
            raise DatagramError(
                message=(
                    "DatagramReliabilityLayer has not been activated. It must be used as an "
                    "asynchronous context manager (`async with ...`)."
                )
            )
        if self._closed:
            raise DatagramError(message="Reliability layer is closed.")
        try:
            async with asyncio.timeout(delay=timeout):
                return await self._incoming_queue.get()
        except asyncio.TimeoutError:
            raise TimeoutError(message=f"Receive timeout after {timeout}s") from None

    async def send(self, *, data: Data) -> None:
        """Send a datagram with guaranteed delivery."""
        if self._lock is None:
            raise DatagramError(message="Reliability layer has not been activated.")

        session = self._get_session()
        data_bytes = ensure_bytes(data=data)

        async with self._lock:
            seq = self._send_sequence
            self._send_sequence += 1
            data_payload = struct.pack("!I", seq) + data_bytes
            frame = self._pack_frame(message_type="DATA", payload=data_payload)
            datagram = _ReliableDatagram(data=frame, sequence=seq)
            self._pending_acks[seq] = datagram

        await session.send_datagram(data=frame)
        logger.debug("Sent reliable datagram with sequence %d", seq)

    def _get_session(self) -> WebTransportSession:
        """Get the underlying session or raise an error if it is gone or closed."""
        session = self._session()
        if self._closed or not session or session.is_closed:
            raise DatagramError(message="Reliability layer or underlying session is closed.")
        return session

    async def _handle_ack_message(self, *, payload: bytes) -> None:
        """Handle an incoming ACK message."""
        if self._lock is None:
            return
        try:
            seq = int(payload.decode("utf-8"))
            async with self._lock:
                if seq in self._pending_acks:
                    del self._pending_acks[seq]
                    logger.debug("Received ACK for sequence %d", seq)
        except (ValueError, UnicodeDecodeError):
            logger.warning("Received malformed ACK: %r", payload)

    async def _handle_data_message(self, *, payload: bytes) -> None:
        """Handle an incoming DATA message."""
        if self._incoming_queue is None:
            return
        if len(payload) < 4:
            return

        seq = struct.unpack("!I", payload[:4])[0]
        data = payload[4:]

        try:
            session = self._get_session()
            ack_payload = str(seq).encode("utf-8")
            frame = self._pack_frame(message_type="ACK", payload=ack_payload)
            await session.send_datagram(data=frame)
        except DatagramError as e:
            logger.warning("Failed to send ACK for sequence %d: %s", seq, e, exc_info=True)
            return

        if seq in self._received_sequences:
            logger.debug("Ignoring duplicate reliable datagram with sequence %d", seq)
            return

        self._received_sequences.append(seq)
        await self._incoming_queue.put(item=data)

    async def _on_datagram_received(self, event: Event) -> None:
        """Handle all incoming datagrams from the underlying transport."""
        if self._incoming_queue is None:
            return
        if not isinstance(event.data, dict):
            return

        raw_data = event.data.get("data")
        if not isinstance(raw_data, bytes):
            return

        unpacked = self._unpack_frame(raw_data=raw_data)
        if not unpacked:
            return

        message_type, payload = unpacked
        try:
            match message_type:
                case "ACK":
                    await self._handle_ack_message(payload=payload)
                case "DATA":
                    await self._handle_data_message(payload=payload)
        except Exception as e:
            logger.error("Error processing received datagram for reliability: %s", e, exc_info=e)

    def _pack_frame(self, *, message_type: str, payload: bytes) -> bytes:
        """Pack a message type and payload into a single bytes frame."""
        type_bytes = message_type.encode("utf-8")
        if len(type_bytes) > 255:
            raise DatagramError(message="Message type too long (max 255 bytes)")

        return struct.pack("!B", len(type_bytes)) + type_bytes + payload

    async def _retry_loop(self) -> None:
        """Periodically check for and retry unacknowledged datagrams."""
        if self._lock is None:
            return
        try:
            while not self._closed:
                await asyncio.sleep(delay=self._ack_timeout)
                current_time = get_timestamp()
                to_retry: list[_ReliableDatagram] = []

                async with self._lock:
                    for datagram in list(self._pending_acks.values()):
                        if current_time - datagram.timestamp > self._ack_timeout:
                            if datagram.retry_count >= self._max_retries:
                                if datagram.sequence is not None and datagram.sequence in self._pending_acks:
                                    del self._pending_acks[datagram.sequence]
                                logger.warning(
                                    "Gave up on sequence %s after %d retries.", datagram.sequence, datagram.retry_count
                                )
                            else:
                                to_retry.append(datagram)

                if not to_retry:
                    continue

                try:
                    session = self._get_session()
                except DatagramError:
                    logger.warning("Could not retry datagrams, transport is closed.")
                    self._closed = True
                    return

                try:
                    async with asyncio.TaskGroup() as tg:
                        for datagram in to_retry:
                            async with self._lock:
                                datagram.retry_count += 1
                                datagram.timestamp = get_timestamp()
                            tg.create_task(coro=session.send_datagram(data=datagram.data))
                            logger.debug("Retrying sequence %s, attempt %d", datagram.sequence, datagram.retry_count)
                except* Exception as eg:
                    logger.warning("Errors occurred during datagram retry: %s", eg.exceptions, exc_info=eg)

        except asyncio.CancelledError:
            pass
        except Exception as e:
            logger.error("Reliability retry loop crashed: %s", e, exc_info=e)

    def _start_background_tasks(self) -> None:
        """Start the background retry task if it is not already running."""
        if self._retry_task is None:
            try:
                self._retry_task = asyncio.create_task(coro=self._retry_loop())
            except RuntimeError:
                logger.warning("Could not start reliability layer tasks: No running event loop.")

    def _unpack_frame(self, *, raw_data: bytes) -> tuple[str, bytes] | None:
        """Unpack a raw datagram into a message type and payload."""
        try:
            type_len = raw_data[0]
            if len(raw_data) < 1 + type_len:
                return None
            message_type = raw_data[1 : 1 + type_len].decode("utf-8")
            payload = raw_data[1 + type_len :]
            return message_type, payload
        except (IndexError, UnicodeDecodeError):
            return None


@dataclass(kw_only=True)
class _ReliableDatagram:
    """An internal datagram message with added reliability metadata."""

    data: bytes
    sequence: int
    timestamp: float = field(default_factory=get_timestamp)
    retry_count: int = 0
