from __future__ import annotations

import socket
import threading
from dataclasses import dataclass
from typing import Callable

from .state import RuntimeStats

TZSP_VERSION = 1
TZSP_TYPE_RECEIVED = 0
TZSP_TYPE_TRANSMIT = 1
TZSP_ENCAP_ETHERNET = 1
TAG_PADDING = 0
TAG_END = 1


class TZSPError(ValueError):
    pass


@dataclass(frozen=True)
class TZSPPacket:
    packet_type: int
    encapsulation: int
    frame: bytes


def decode_tzsp(data: bytes) -> TZSPPacket:
    if len(data) < 5:
        raise TZSPError("datagram too short")

    version = data[0]
    packet_type = data[1]
    encapsulation = int.from_bytes(data[2:4], "big")

    if version != TZSP_VERSION:
        raise TZSPError(f"unsupported TZSP version {version}")
    if packet_type not in {TZSP_TYPE_RECEIVED, TZSP_TYPE_TRANSMIT}:
        raise TZSPError(f"TZSP packet type {packet_type} has no packet payload")

    offset = 4
    found_end = False
    while offset < len(data):
        tag_type = data[offset]
        offset += 1
        if tag_type == TAG_PADDING:
            continue
        if tag_type == TAG_END:
            found_end = True
            break
        if offset >= len(data):
            raise TZSPError("truncated TZSP tag length")
        tag_len = data[offset]
        offset += 1
        if offset + tag_len > len(data):
            raise TZSPError("truncated TZSP tag value")
        offset += tag_len

    if not found_end:
        raise TZSPError("missing TZSP END tag")
    if offset >= len(data):
        raise TZSPError("TZSP datagram contains no encapsulated frame")

    return TZSPPacket(packet_type=packet_type, encapsulation=encapsulation, frame=data[offset:])


class TZSPReceiver(threading.Thread):
    def __init__(
        self,
        bind_host: str,
        port: int,
        frame_writer: Callable[[bytes], int],
        stats: RuntimeStats,
        stop_event: threading.Event,
    ) -> None:
        super().__init__(name="tzsp-receiver", daemon=True)
        self.bind_host = bind_host
        self.port = port
        self.frame_writer = frame_writer
        self.stats = stats
        self.stop_event = stop_event
        self.sock: socket.socket | None = None

    def run(self) -> None:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind((self.bind_host, self.port))
        sock.settimeout(1.0)
        self.sock = sock
        print(f"[tzsp] listening on udp://{self.bind_host}:{self.port}", flush=True)

        try:
            while not self.stop_event.is_set():
                try:
                    data, _addr = sock.recvfrom(65535)
                except socket.timeout:
                    continue
                except OSError:
                    if self.stop_event.is_set():
                        break
                    raise

                self.stats.inc("tzsp_datagrams")
                self.stats.stamp("last_packet_at")
                try:
                    packet = decode_tzsp(data)
                except TZSPError:
                    self.stats.inc("tzsp_decode_errors")
                    continue

                if packet.encapsulation != TZSP_ENCAP_ETHERNET:
                    self.stats.inc("tzsp_unsupported")
                    continue

                try:
                    self.frame_writer(packet.frame)
                    self.stats.inc("frames_injected")
                except OSError:
                    self.stats.inc("inject_errors")
        finally:
            sock.close()

    def close(self) -> None:
        if self.sock is not None:
            try:
                self.sock.close()
            except OSError:
                pass