from __future__ import annotations

import json
import os
import socket
import subprocess
import threading
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable

from .state import RuntimeStats


class RustTZSPReceiver:
    """Supervise the Rust TZSP data-plane and ingest its 1 Hz telemetry.

    Packet bytes never cross into Python. The Rust process owns UDP reception,
    TZSP decoding and TAP injection. Python receives only compact telemetry over
    a Unix datagram socket, so UI/Redis work cannot back-pressure packet capture.
    """

    def __init__(
        self,
        *,
        binary: str,
        telemetry_socket: str,
        stats: RuntimeStats,
        stop_event: threading.Event,
        throughput_sink: Callable[[dict[str, Any]], None] | None = None,
    ) -> None:
        self.binary = str(binary)
        self.telemetry_socket = str(telemetry_socket)
        self.stats = stats
        self.stop_event = stop_event
        self._throughput_sink = throughput_sink
        self._process: subprocess.Popen | None = None
        self._socket: socket.socket | None = None
        self._thread = threading.Thread(target=self._run_telemetry, name="tzsp-rust-telemetry", daemon=True)
        self._lock = threading.RLock()
        self._ready = threading.Event()
        self._last: dict[str, Any] = {}
        self._samples = 0
        self._telemetry_errors = 0
        self._started_at = time.monotonic()

    def start(self) -> None:
        if self._process is not None:
            return
        binary = Path(self.binary)
        if not binary.is_file():
            raise RuntimeError(f"Rust TZSP receiver binary not found: {self.binary}")

        path = Path(self.telemetry_socket)
        path.parent.mkdir(parents=True, exist_ok=True)
        try:
            path.unlink()
        except FileNotFoundError:
            pass

        sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
        sock.bind(self.telemetry_socket)
        sock.settimeout(0.5)
        self._socket = sock
        self._thread.start()

        env = os.environ.copy()
        env["TZSP_TELEMETRY_SOCKET"] = self.telemetry_socket
        self._process = subprocess.Popen([self.binary], env=env)
        print(f"[tzsp] Rust data-plane started, pid={self._process.pid}", flush=True)

    def wait_ready(self, timeout: float = 8.0) -> bool:
        deadline = time.monotonic() + max(0.1, float(timeout))
        while time.monotonic() < deadline:
            process = self._process
            if process is not None and process.poll() is not None:
                return False
            if self._ready.wait(timeout=min(0.1, max(0.0, deadline - time.monotonic()))):
                return True
        return False

    def set_throughput_sink(self, sink: Callable[[dict[str, Any]], None] | None) -> None:
        with self._lock:
            self._throughput_sink = sink

    def is_alive(self) -> bool:
        process = self._process
        return bool(process is not None and process.poll() is None and self._thread.is_alive())

    @property
    def pid(self) -> int | None:
        process = self._process
        return process.pid if process is not None and process.poll() is None else None

    def close(self, timeout: float = 3.0) -> None:
        process = self._process
        if process is not None and process.poll() is None:
            process.terminate()
            try:
                process.wait(timeout=timeout)
            except subprocess.TimeoutExpired:
                process.kill()
                try:
                    process.wait(timeout=1)
                except subprocess.TimeoutExpired:
                    pass
        sock = self._socket
        self._socket = None
        if sock is not None:
            try:
                sock.close()
            except OSError:
                pass
        if self._thread.is_alive():
            self._thread.join(timeout=1.0)
        try:
            Path(self.telemetry_socket).unlink()
        except FileNotFoundError:
            pass
        self._process = None

    def status(self) -> dict[str, Any]:
        with self._lock:
            data = dict(self._last)
        now_ms = int(time.time() * 1000)
        ts_ms = int(data.get("ts_ms") or 0)
        data.update(
            {
                "engine": "rust",
                "process_alive": self.is_alive(),
                "pid": self.pid or data.get("pid"),
                "ready": bool(self._ready.is_set() and self.is_alive()),
                "telemetry_age_ms": max(0, now_ms - ts_ms) if ts_ms else None,
                "throughput_samples": self._samples,
                "telemetry_errors": self._telemetry_errors,
            }
        )
        return data

    def current_throughput(self, window_seconds: int | None = None) -> dict[str, Any]:
        with self._lock:
            sample = dict(self._last)
        now_ms = int(time.time() * 1000)
        ts_ms = int(sample.get("ts_ms") or 0)
        interval = max(float(sample.get("interval_ms") or 1000) / 1000.0, 0.001)
        age_ms = max(0, now_ms - ts_ms) if ts_ms else 10**9
        fresh = bool(ts_ms and age_ms <= max(3000, int(interval * 3000)))
        if fresh:
            total = round(max(int(sample.get("bytes_total") or 0), 0) * 8 / interval)
            inbound = round(max(int(sample.get("bytes_in") or 0), 0) * 8 / interval)
            outbound = round(max(int(sample.get("bytes_out") or 0), 0) * 8 / interval)
            pps = round(max(int(sample.get("packets_total") or 0), 0) / interval, 2)
            ingress_bps = round(max(int(sample.get("rx_bytes_interval") or 0), 0) * 8 / interval)
            ingress_pps = round(max(int(sample.get("rx_datagrams_interval") or 0), 0) / interval, 2)
        else:
            total = inbound = outbound = ingress_bps = 0
            pps = ingress_pps = 0.0
        queue_depth = max(int(sample.get("queue_depth_batches") or 0), 0)
        queue_capacity = max(int(sample.get("queue_capacity_batches") or 0), 0)
        queue_fill_pct = round((queue_depth / queue_capacity) * 100.0, 1) if queue_capacity else 0.0
        inspection_ratio_pct = (
            round(min(100.0, (total / ingress_bps) * 100.0), 1)
            if ingress_bps > 0
            else (100.0 if total == 0 else 0.0)
        )
        loss_per_second = round(
            (
                max(int(sample.get("kernel_udp_drops_interval") or 0), 0)
                + max(int(sample.get("queue_drops_interval") or 0), 0)
                + max(int(sample.get("truncated_interval") or 0), 0)
            )
            / interval,
            2,
        ) if fresh else 0.0
        return {
            "window_seconds": int(window_seconds or 0),
            "current_bps": total,
            "current_in_bps": inbound,
            "current_out_bps": outbound,
            "current_other_bps": max(0, total - inbound - outbound),
            "current_pps": pps,
            "current_ingress_bps": ingress_bps,
            "current_ingress_pps": ingress_pps,
            "inspection_ratio_pct": inspection_ratio_pct,
            "capture_efficiency_pct": float(sample.get("capture_efficiency_pct") or 0.0),
            "loss_pps": loss_per_second,
            "current_sample_ts_ms": ts_ms,
            "current_sample_age_ms": age_ms if ts_ms else None,
            "current_sample_fresh": fresh,
            "receiver_engine": "rust",
            "receiver_pid": self.pid,
            "kernel_udp_drops": int(sample.get("kernel_udp_drops") or 0),
            "kernel_udp_drops_interval": int(sample.get("kernel_udp_drops_interval") or 0),
            "queue_dropped_datagrams": int(sample.get("queue_dropped_datagrams") or 0),
            "queue_drops_interval": int(sample.get("queue_drops_interval") or 0),
            "truncated_datagrams": int(sample.get("truncated_datagrams") or 0),
            "truncated_interval": int(sample.get("truncated_interval") or 0),
            "queue_depth_batches": queue_depth,
            "queue_capacity_batches": queue_capacity,
            "queue_capacity_bytes": int(sample.get("queue_capacity_bytes") or 0),
            "queue_high_water_batches": int(sample.get("queue_high_water_batches") or 0),
            "queue_fill_pct": queue_fill_pct,
            "rx_thread_alive": bool(sample.get("rx_thread_alive", False)),
            "worker_thread_alive": bool(sample.get("worker_thread_alive", False)),
            "rcvbuf_bytes": int(sample.get("rcvbuf_bytes") or 0),
            "batch_size": int(sample.get("batch_size") or 0),
            "datagram_bytes": int(sample.get("datagram_bytes") or 0),
        }

    def overlay_current(self, payload: dict[str, Any], window_seconds: int | None = None) -> dict[str, Any]:
        result = dict(payload)
        result.update(self.current_throughput(window_seconds or int(result.get("window_seconds") or 0)))
        return result

    def _run_telemetry(self) -> None:
        while not self.stop_event.is_set():
            sock = self._socket
            if sock is None:
                break
            try:
                raw = sock.recv(64 * 1024)
            except socket.timeout:
                continue
            except OSError:
                if self.stop_event.is_set() or self._socket is None:
                    break
                self._telemetry_errors += 1
                continue
            try:
                message = json.loads(raw.decode("utf-8"))
                if not isinstance(message, dict) or message.get("type") != "tzsp_sample":
                    continue
                self._ingest(message)
            except (UnicodeDecodeError, json.JSONDecodeError, TypeError, ValueError):
                self._telemetry_errors += 1

    def _ingest(self, message: dict[str, Any]) -> None:
        with self._lock:
            self._last = dict(message)
            sink = self._throughput_sink
        if message.get("ready"):
            self._ready.set()

        last_packet_ms = int(message.get("last_packet_ms") or 0)
        last_packet_at = None
        if last_packet_ms:
            last_packet_at = datetime.fromtimestamp(last_packet_ms / 1000.0, tz=timezone.utc).isoformat()
        self.stats.update_tzsp_receiver(
            {
                "tzsp_datagrams": int(message.get("tzsp_datagrams") or 0),
                "tzsp_decode_errors": int(message.get("tzsp_decode_errors") or 0),
                "tzsp_unsupported": int(message.get("tzsp_unsupported") or 0),
                "frames_injected": int(message.get("frames_injected") or 0),
                "inject_errors": int(message.get("inject_errors") or 0),
                "tzsp_kernel_udp_drops": int(message.get("kernel_udp_drops") or 0),
                "tzsp_queue_drops": int(message.get("queue_dropped_datagrams") or 0),
                "tzsp_truncated_datagrams": int(message.get("truncated_datagrams") or 0),
                "last_packet_at": last_packet_at,
            }
        )

        sample = {
            key: int(message.get(key) or 0)
            for key in (
                "ts_ms",
                "interval_ms",
                "bytes_total",
                "bytes_in",
                "bytes_out",
                "bytes_internal",
                "bytes_external",
                "packets_total",
                "packets_in",
                "packets_out",
                "packets_internal",
                "packets_external",
            )
        }
        self._samples += 1
        if sink is not None and sample["interval_ms"] > 0:
            try:
                sink(sample)
            except Exception:
                # Telemetry persistence is best-effort and is deliberately never
                # allowed to affect the independent Rust packet data-plane.
                self._telemetry_errors += 1