#!/usr/bin/env python3
"""Bidirectional EDGE pipe probe over WG+gretap — never injects to POG422.

Purpose: prove spoke↔hub eth path is live before (and while) gameplay runs.
Addresses first-game aborts where spoke eth_tx existed but hub saw no eth_rx
for ~15–20 s.

Wire format (0x88B5 payload; dual-channel seats prefix FIXTURE_ETH_TAG_PIPE):
  b'EGPIPE1' + kind_u8 + seq_u32be + t_tx_ns_u64be + role_u8

  kind: 1=PING, 2=PONG
  role: 1=hub, 2=spoke, 3=peer

Env:
  PIPE_PROBE=1              enable (default on when EDGE_MODE=1)
  PIPE_PROBE_TIMEOUT_MS=250 per-attempt wait
  PIPE_PROBE_INTERVAL_S=30  ongoing RTT sample after ready (0=off)
  PIPE_PROBE_MAX_ATTEMPTS=0 unlimited until success (startup phase)
"""
from __future__ import annotations

import os
import struct
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Callable

MAGIC = b"EGPIPE1"
KIND_PING = 1
KIND_PONG = 2
ROLE_HUB = 1
ROLE_SPOKE = 2
ROLE_PEER = 3

# Dual-channel eth tag — not machine/OIS; eth path must not UART-inject.
FIXTURE_ETH_TAG_PIPE = 0x50  # 'P'


def _env_truthy(name: str, default: str = "0") -> bool:
    return os.environ.get(name, default).strip().lower() in ("1", "true", "yes", "on")


def pipe_probe_enabled() -> bool:
    if os.environ.get("PIPE_PROBE", "").strip() != "":
        return _env_truthy("PIPE_PROBE", "1")
    return _env_truthy("EDGE_MODE", "0")


def pipe_probe_timeout_s() -> float:
    try:
        return max(0.05, float(os.environ.get("PIPE_PROBE_TIMEOUT_MS", "250")) / 1000.0)
    except ValueError:
        return 0.25


def pipe_probe_interval_s() -> float:
    try:
        return max(0.0, float(os.environ.get("PIPE_PROBE_INTERVAL_S", "30")))
    except ValueError:
        return 30.0


def pipe_probe_max_attempts() -> int:
    try:
        return max(0, int(os.environ.get("PIPE_PROBE_MAX_ATTEMPTS", "0")))
    except ValueError:
        return 0


def resolve_pipe_role() -> int:
    topo = (os.environ.get("EDGE_TOPOLOGY") or "").strip().lower()
    role = (os.environ.get("EDGE_ROLE") or os.environ.get("POG_EDGE_ROLE") or "").strip().lower()
    if topo == "hub" or role in ("hub", "peer-a", "office", "a", "1"):
        return ROLE_HUB
    if topo == "spoke" or role in ("spoke", "peer-b", "pos4", "b", "4"):
        return ROLE_SPOKE
    if role.startswith("peer"):
        return ROLE_PEER
    return ROLE_SPOKE


def build_pipe_payload(*, kind: int, seq: int, t_tx_ns: int, role: int) -> bytes:
    return MAGIC + struct.pack("!BIQB", kind & 0xFF, seq & 0xFFFFFFFF, t_tx_ns & 0xFFFFFFFFFFFFFFFF, role & 0xFF)


def parse_pipe_payload(raw: bytes) -> dict[str, Any] | None:
    if not raw or not raw.startswith(MAGIC):
        return None
    body = raw[len(MAGIC) :]
    if len(body) < 1 + 4 + 8 + 1:
        return None
    kind, seq, t_tx_ns, role = struct.unpack_from("!BIQB", body, 0)
    return {"kind": kind, "seq": seq, "t_tx_ns": t_tx_ns, "role": role}


def is_pipe_wire(wire: bytes, *, dual: bool) -> tuple[bool, bytes]:
    """Return (is_pipe, payload_without_dual_tag)."""
    if not wire:
        return False, b""
    if dual and wire[0] == FIXTURE_ETH_TAG_PIPE:
        return True, wire[1:]
    if wire.startswith(MAGIC):
        return True, wire
    if dual and len(wire) > 1 and wire[1:].startswith(MAGIC):
        # Mis-tagged but recognizable
        return True, wire[1:]
    return False, b""


def tag_pipe_wire(payload: bytes, *, dual: bool) -> bytes:
    if dual:
        return bytes([FIXTURE_ETH_TAG_PIPE]) + payload
    return payload


@dataclass
class PipeProbeState:
    ready: bool = False
    ready_mono: float = 0.0
    ready_delay_s: float | None = None
    ready_attempts: int = 0
    ready_rtt_ms: float | None = None
    last_rtt_ms: float | None = None
    last_ok_mono: float = 0.0
    last_fail_mono: float = 0.0
    attempts: int = 0
    fails: int = 0
    pending_seq: int | None = None
    pending_t0_ns: int | None = None
    _pong_event: threading.Event = field(default_factory=threading.Event)
    _lock: threading.Lock = field(default_factory=threading.Lock)
    _start_mono: float = field(default_factory=time.monotonic)

    def note_pong(self, seq: int, rtt_ms: float) -> bool:
        with self._lock:
            if self.pending_seq is None or seq != self.pending_seq:
                return False
            self.last_rtt_ms = rtt_ms
            self.last_ok_mono = time.monotonic()
            self.pending_seq = None
            self.pending_t0_ns = None
            self._pong_event.set()
            if not self.ready:
                self.ready = True
                self.ready_mono = self.last_ok_mono
                self.ready_delay_s = self.ready_mono - self._start_mono
                self.ready_attempts = self.attempts
                self.ready_rtt_ms = rtt_ms
            return True

    def begin_ping(self, seq: int, t0_ns: int) -> None:
        with self._lock:
            self.attempts += 1
            self.pending_seq = seq
            self.pending_t0_ns = t0_ns
            self._pong_event.clear()

    def note_timeout(self) -> None:
        with self._lock:
            self.fails += 1
            self.last_fail_mono = time.monotonic()
            self.pending_seq = None
            self.pending_t0_ns = None

    def snapshot(self) -> dict[str, Any]:
        with self._lock:
            return {
                "pipe_ready": self.ready,
                "pipe_ready_delay_s": (
                    round(self.ready_delay_s, 3) if self.ready_delay_s is not None else None
                ),
                "pipe_ready_attempts": self.ready_attempts if self.ready else self.attempts,
                "pipe_ready_rtt_ms": (
                    round(self.ready_rtt_ms, 2) if self.ready_rtt_ms is not None else None
                ),
                "pipe_last_rtt_ms": (
                    round(self.last_rtt_ms, 2) if self.last_rtt_ms is not None else None
                ),
                "pipe_attempts": self.attempts,
                "pipe_fails": self.fails,
            }


# Process-wide state for eth RX thread ↔ probe thread.
_STATE: PipeProbeState | None = None
_SEQ = 0
_SEQ_LOCK = threading.Lock()


def get_pipe_probe_state() -> PipeProbeState | None:
    return _STATE


def _next_seq() -> int:
    global _SEQ
    with _SEQ_LOCK:
        _SEQ = (_SEQ + 1) & 0xFFFFFFFF
        if _SEQ == 0:
            _SEQ = 1
        return _SEQ


def handle_pipe_eth_rx(
    *,
    payload: bytes,
    eth_port: str | None,
    router: Any,
    emit: Callable[..., None],
    dual: bool,
) -> bool:
    """Handle pipe wire on eth RX. Returns True if consumed (never inject)."""
    parsed = parse_pipe_payload(payload)
    if parsed is None:
        emit("pipe_probe_rx_skip", reason="bad_payload", eth_port=eth_port, len=len(payload))
        return True
    kind = int(parsed["kind"])
    seq = int(parsed["seq"])
    role = int(parsed["role"])
    emit(
        "pipe_probe_rx",
        kind="ping" if kind == KIND_PING else "pong" if kind == KIND_PONG else kind,
        seq=seq,
        peer_role=role,
        eth_port=eth_port,
    )
    if kind == KIND_PING:
        # Reply on ingress port only — proves return path to that peer.
        pong = build_pipe_payload(
            kind=KIND_PONG,
            seq=seq,
            t_tx_ns=int(parsed["t_tx_ns"]),
            role=resolve_pipe_role(),
        )
        wire = tag_pipe_wire(pong, dual=dual)
        try:
            only = [eth_port] if eth_port else None
            sends = router.send_uart_to_eth(wire, only_ports=only)
            emit(
                "pipe_probe_tx",
                kind="pong",
                seq=seq,
                eth_port=eth_port,
                ports=[p for p, _ in sends],
            )
        except Exception as exc:
            emit("pipe_probe_error", where="pong_tx", message=str(exc), seq=seq)
        return True
    if kind == KIND_PONG:
        st = _STATE
        if st is not None:
            t0 = st.pending_t0_ns
            rtt_ms = None
            if t0 is not None:
                rtt_ms = (time.perf_counter_ns() - t0) / 1_000_000.0
            elif parsed.get("t_tx_ns"):
                # Fallback if state raced — still mark approx from peer stamp unused
                rtt_ms = 0.0
            if rtt_ms is not None and st.note_pong(seq, float(rtt_ms)):
                emit(
                    "pipe_probe_pong",
                    seq=seq,
                    rtt_ms=round(float(rtt_ms), 2),
                    eth_port=eth_port,
                    **st.snapshot(),
                )
        return True
    return True


def _send_ping(router: Any, dual: bool, emit: Callable[..., None], st: PipeProbeState) -> int:
    seq = _next_seq()
    t0 = time.perf_counter_ns()
    st.begin_ping(seq, t0)
    payload = build_pipe_payload(
        kind=KIND_PING,
        seq=seq,
        t_tx_ns=t0,
        role=resolve_pipe_role(),
    )
    wire = tag_pipe_wire(payload, dual=dual)
    sends = router.send_uart_to_eth(wire)
    emit(
        "pipe_probe_tx",
        kind="ping",
        seq=seq,
        attempt=st.attempts,
        ports=[p for p, _ in sends],
    )
    return seq


def start_pipe_probe_for_bridge(
    router: Any,
    emit: Callable[..., None],
    stop: threading.Event,
    *,
    dual: bool = False,
) -> PipeProbeState | None:
    """Start reply-capable state + initiator thread when PIPE_PROBE enabled."""
    global _STATE
    if not pipe_probe_enabled():
        return None
    if router is None:
        return None
    st = PipeProbeState()
    _STATE = st
    emit("pipe_probe_start", role=resolve_pipe_role(), dual=dual, **st.snapshot())

    def _loop() -> None:
        timeout = pipe_probe_timeout_s()
        max_attempts = pipe_probe_max_attempts()
        interval = pipe_probe_interval_s()
        # Phase 1: retry until first success
        while not stop.is_set() and not st.ready:
            if max_attempts and st.attempts >= max_attempts:
                emit("pipe_probe_give_up", **st.snapshot())
                break
            try:
                _send_ping(router, dual, emit, st)
            except Exception as exc:
                st.note_timeout()
                emit("pipe_probe_error", where="ping_tx", message=str(exc), **st.snapshot())
                if stop.wait(timeout):
                    return
                continue
            if st._pong_event.wait(timeout):
                emit("pipe_probe_ready", **st.snapshot())
                break
            st.note_timeout()
            emit(
                "pipe_probe_retry",
                reason="timeout",
                timeout_s=timeout,
                **st.snapshot(),
            )
        # Phase 2: optional ongoing RTT samples (single attempt each interval)
        while not stop.is_set() and interval > 0 and st.ready:
            if stop.wait(interval):
                return
            try:
                _send_ping(router, dual, emit, st)
            except Exception as exc:
                emit("pipe_probe_error", where="ping_tx_watch", message=str(exc))
                continue
            if st._pong_event.wait(timeout):
                emit("pipe_probe_rtt", **st.snapshot())
            else:
                st.note_timeout()
                emit("pipe_probe_miss", reason="timeout", **st.snapshot())

    threading.Thread(target=_loop, name="pipe-probe", daemon=True).start()
    return st