#!/usr/bin/env python3
"""Clear-WLAN bank path: signed 0x0A meters, jackpot announces, Master discovery.

Gameplay stays on WG+gretap (TID-unicast). Soft bank-wide traffic rides the
non-VPN uplink as UDP broadcast so N spokes do not pay sequential WG fan-out.

Env:
  EDGE_BANK_WLAN=1          enable publish/listen
  EDGE_TOPOLOGY=hub|spoke   hub publishes; spoke listens (also peer-a aliases → hub)
  EDGE_BANK=N               bank scope (must match envelope)
  EDGE_FLEET_HMAC_KEY       hex or raw secret (or path via EDGE_FLEET_HMAC_KEY_FILE)
  EDGE_BANK_WLAN_PORT=41234 UDP port (default)
  EDGE_BANK_WLAN_IF=wlan0   bind/broadcast iface (default EDGE_UPLINK or wlan0)
  EDGE_HUB_WG_PUB=...       hub WireGuard public key (discovery field)
  EDGE_HUB_WG_PORT=51820
"""

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import os
import socket
import struct
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Callable

from pog_edge_identity import latch_cabinet_identity

KIND_METER = "meter_0a"
KIND_JACKPOT = "jackpot"
KIND_DISCOVERY = "discovery"
KIND_KEEPALIVE = "keepalive"

DEFAULT_PORT = 41234
DEFAULT_IFACE = "wlan0"
REPLAY_WINDOW_S = 120.0
# Hub announce for spoke enroll / EDGE_UPLINK_IP_LEARNED — must beat Master
# service return budget (~5s). Override with EDGE_BANK_ANNOUNCE_S.
def _bank_announce_interval_s() -> float:
    raw = (os.environ.get("EDGE_BANK_ANNOUNCE_S") or "").strip()
    if raw:
        try:
            return max(1.0, min(30.0, float(raw)))
        except ValueError:
            pass
    return 3.0


KEEPALIVE_AFTER_QUIET_S = 3.0  # default; maybe_keepalive reads live via helper


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


def bank_wlan_enabled() -> bool:
    """Clear-WLAN path when explicitly requested or EDGE_TOPOLOGY=hub|spoke."""
    if _env_truthy("EDGE_BANK_WLAN", "0"):
        return True
    return resolve_edge_topology() in ("hub", "spoke")


def resolve_edge_topology() -> str:
    """Explicit activation only.

    EDGE_TOPOLOGY=hub|spoke turns star services on. GAMESTATS identity
    (MACHINE_ROLE=M/S) is a *suggestion* surfaced in session_start
    (edge_topology_suggested) — it must never auto-start hub/spoke services,
    or copper Masters and the single-hop bench would grow listeners/publishers
    they did not ask for.
    """
    raw = os.environ.get("EDGE_TOPOLOGY", "").strip().lower()
    if raw in ("hub", "spoke"):
        return raw
    # Back-compat: peer-a/office → hub, peer-b/pos4 → spoke only when bank-wlan forced.
    if _env_truthy("EDGE_BANK_WLAN", "0"):
        role = os.environ.get("EDGE_ROLE", os.environ.get("POG_EDGE_ROLE", "")).strip().lower()
        if role in ("office", "peer-a", "a", "1", "hub"):
            return "hub"
        if role in ("pos4", "peer-b", "b", "4", "spoke"):
            return "spoke"
    return ""


def resolve_edge_bank() -> int:
    raw = os.environ.get("EDGE_BANK", "").strip()
    if raw.isdigit():
        return int(raw)
    try:
        return int(latch_cabinet_identity().bank_number)
    except Exception:
        return 0


def resolve_wlan_port() -> int:
    raw = os.environ.get("EDGE_BANK_WLAN_PORT", str(DEFAULT_PORT)).strip()
    try:
        return int(raw)
    except ValueError:
        return DEFAULT_PORT


def resolve_wlan_iface() -> str:
    return (
        os.environ.get("EDGE_BANK_WLAN_IF", "").strip()
        or os.environ.get("EDGE_UPLINK", "").strip()
        or DEFAULT_IFACE
    )


def load_fleet_hmac_key() -> bytes:
    path = os.environ.get("EDGE_FLEET_HMAC_KEY_FILE", "").strip()
    if path and os.path.isfile(path):
        data = open(path, "rb").read().strip()
        if data:
            return data
    raw = os.environ.get("EDGE_FLEET_HMAC_KEY", "").strip()
    if not raw:
        # Dev default — rotate in field via EDGE_FLEET_HMAC_KEY_FILE.
        raw = "pog-edge-dev-fleet-hmac-v1"
    try:
        return bytes.fromhex(raw)
    except ValueError:
        return raw.encode("utf-8")


def _sign_body(body: dict[str, Any], key: bytes) -> str:
    blob = json.dumps(body, sort_keys=True, separators=(",", ":")).encode("utf-8")
    return hmac.new(key, blob, hashlib.sha256).hexdigest()


def verify_envelope(msg: dict[str, Any], key: bytes) -> bool:
    if not isinstance(msg, dict):
        return False
    mac = msg.get("mac")
    if not isinstance(mac, str) or len(mac) < 32:
        return False
    body = {k: v for k, v in msg.items() if k != "mac"}
    expect = _sign_body(body, key)
    return hmac.compare_digest(expect, mac)


def build_envelope(
    *,
    kind: str,
    bank: int,
    counter: int,
    payload: bytes | None = None,
    extra: dict[str, Any] | None = None,
    key: bytes | None = None,
) -> dict[str, Any]:
    key = key if key is not None else load_fleet_hmac_key()
    body: dict[str, Any] = {
        "v": 1,
        "kind": kind,
        "bank": int(bank),
        "counter": int(counter),
        "ts": int(time.time()),
        "hub_uplink_ip": os.environ.get("EDGE_UPLINK_IP", "").strip() or _guess_uplink_ip(),
        "wg_port": int(os.environ.get("EDGE_HUB_WG_PORT", os.environ.get("POG_EDGE_WG_PORT", "51820"))),
        "hub_wg_pub": os.environ.get("EDGE_HUB_WG_PUB", "").strip(),
    }
    if payload is not None:
        body["payload_b64"] = base64.b64encode(payload).decode("ascii")
    if extra:
        body.update(extra)
    body["mac"] = _sign_body(body, key)
    return body


def _guess_uplink_ip() -> str:
    iface = resolve_wlan_iface()
    try:
        # Best-effort: parse `ip -4 -o addr show dev IF`
        import subprocess

        out = subprocess.check_output(
            ["ip", "-4", "-o", "addr", "show", "dev", iface],
            text=True,
            stderr=subprocess.DEVNULL,
        )
        # e.g. "2: wlan0    inet 10.1.10.140/24 ..."
        for part in out.split():
            if "/" in part and part[0].isdigit():
                return part.split("/", 1)[0]
    except Exception:
        pass
    return ""


def _broadcast_addr_for_iface(iface: str) -> str:
    """Prefer directed broadcast; fall back to limited broadcast."""
    try:
        import subprocess

        out = subprocess.check_output(
            ["ip", "-4", "-o", "addr", "show", "dev", iface],
            text=True,
            stderr=subprocess.DEVNULL,
        )
        # Parse inet A.B.C.D/NN — use host-bits-all broadcast when /24-like.
        for tok in out.split():
            if "/" in tok and tok[0].isdigit():
                ip_s, pfx_s = tok.split("/", 1)
                pfx = int(pfx_s)
                parts = [int(x) for x in ip_s.split(".")]
                if len(parts) == 4 and 0 < pfx <= 30:
                    mask = (0xFFFFFFFF << (32 - pfx)) & 0xFFFFFFFF
                    ip_i = (parts[0] << 24) | (parts[1] << 16) | (parts[2] << 8) | parts[3]
                    bcast = ip_i | (~mask & 0xFFFFFFFF)
                    return socket.inet_ntoa(struct.pack("!I", bcast))
    except Exception:
        pass
    return "255.255.255.255"


@dataclass
class BankWlanPublisher:
    """Hub-side UDP broadcast of signed meters / jackpots / discovery."""

    bank: int
    emit: Callable[..., None] | None = None
    port: int = field(default_factory=resolve_wlan_port)
    iface: str = field(default_factory=resolve_wlan_iface)
    key: bytes = field(default_factory=load_fleet_hmac_key)
    _counter: int = 0
    _lock: threading.Lock = field(default_factory=threading.Lock)
    _sock: socket.socket | None = None
    _last_meter_mono: float = 0.0
    # Last discovery/keepalive TX (mono). Meters alone must not gate announce —
    # spokes that start after hub boot need periodic hub_uplink_ip without 0x0A.
    _last_announce_mono: float = 0.0

    def _ensure_sock(self) -> socket.socket:
        if self._sock is not None:
            return self._sock
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        try:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_BINDTODEVICE, self.iface.encode() + b"\0")
        except (AttributeError, OSError):
            pass
        self._sock = sock
        return sock

    def next_counter(self) -> int:
        with self._lock:
            self._counter += 1
            return self._counter

    def publish(self, kind: str, payload: bytes | None = None, extra: dict[str, Any] | None = None) -> bool:
        env = build_envelope(
            kind=kind,
            bank=self.bank,
            counter=self.next_counter(),
            payload=payload,
            extra=extra,
            key=self.key,
        )
        data = json.dumps(env, separators=(",", ":")).encode("utf-8")
        try:
            sock = self._ensure_sock()
            dest = (_broadcast_addr_for_iface(self.iface), self.port)
            sock.sendto(data, dest)
            if self.emit:
                self.emit(
                    "bank_wlan_tx",
                    kind=kind,
                    bank=self.bank,
                    counter=env["counter"],
                    bytes=len(data),
                    dest=dest[0],
                    port=self.port,
                )
            if kind == KIND_METER:
                self._last_meter_mono = time.monotonic()
            if kind in (KIND_DISCOVERY, KIND_KEEPALIVE):
                self._last_announce_mono = time.monotonic()
            return True
        except OSError as exc:
            if self.emit:
                self.emit("bank_wlan_tx_error", kind=kind, message=str(exc))
            return False

    def publish_meter_0a(self, frame: bytes) -> bool:
        return self.publish(KIND_METER, payload=frame)

    def publish_jackpot(
        self,
        frame: bytes,
        *,
        tid_hex: str | None,
        progwin_cents: int,
        msg_type_hex: str | None,
    ) -> bool:
        return self.publish(
            KIND_JACKPOT,
            payload=frame,
            extra={
                "tid": (tid_hex or "").upper() or None,
                "progwin_cents": int(progwin_cents),
                "msg": msg_type_hex,
            },
        )

    def publish_discovery(self) -> bool:
        return self.publish(KIND_DISCOVERY)

    def publish_discovery_keepalive(self) -> bool:
        return self.publish(KIND_KEEPALIVE)

    def maybe_keepalive(self) -> None:
        """Periodic hub announce for spoke enroll (does not require prior meters)."""
        now = time.monotonic()
        interval = _bank_announce_interval_s()
        if self._last_announce_mono > 0 and (now - self._last_announce_mono) < interval:
            return
        self.publish_discovery_keepalive()


@dataclass
class BankWlanListener:
    """Spoke-side clear-WLAN listener → progressives.csv + fireworks."""

    bank: int
    emit: Callable[..., None] | None = None
    on_meter_frame: Callable[[bytes], None] | None = None
    on_jackpot: Callable[[dict[str, Any]], None] | None = None
    on_discovery: Callable[[dict[str, Any]], None] | None = None
    port: int = field(default_factory=resolve_wlan_port)
    iface: str = field(default_factory=resolve_wlan_iface)
    key: bytes = field(default_factory=load_fleet_hmac_key)
    _stop: threading.Event = field(default_factory=threading.Event)
    _thread: threading.Thread | None = None
    _seen: dict[int, float] = field(default_factory=dict)
    last_hub_uplink_ip: str = ""
    last_hub_wg_pub: str = ""
    last_counter: int = 0

    def start(self) -> None:
        if self._thread is not None and self._thread.is_alive():
            return
        self._stop.clear()
        self._thread = threading.Thread(target=self._run, name="bank-wlan-rx", daemon=True)
        self._thread.start()

    def stop(self) -> None:
        self._stop.set()
        if self._thread is not None:
            self._thread.join(timeout=2.0)

    def _replay_ok(self, counter: int, ts: int) -> bool:
        now = time.time()
        if abs(now - float(ts)) > REPLAY_WINDOW_S:
            return False
        last_ts = self._seen.get(counter)
        if last_ts is not None:
            return False
        self._seen[counter] = now
        # Bound memory
        if len(self._seen) > 4096:
            cutoff = now - REPLAY_WINDOW_S
            self._seen = {c: t for c, t in self._seen.items() if t >= cutoff}
        return True

    def _run(self) -> None:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        try:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_BINDTODEVICE, self.iface.encode() + b"\0")
        except (AttributeError, OSError):
            pass
        sock.bind(("", self.port))
        sock.settimeout(0.5)
        if self.emit:
            self.emit("bank_wlan_listen", iface=self.iface, port=self.port, bank=self.bank)
        while not self._stop.is_set():
            try:
                data, addr = sock.recvfrom(65535)
            except socket.timeout:
                continue
            except OSError:
                break
            try:
                msg = json.loads(data.decode("utf-8"))
            except (UnicodeDecodeError, json.JSONDecodeError):
                continue
            if not verify_envelope(msg, self.key):
                if self.emit:
                    self.emit("bank_wlan_rx_bad_mac", from_ip=addr[0])
                continue
            if int(msg.get("bank", -1)) != int(self.bank):
                continue
            counter = int(msg.get("counter", 0))
            ts = int(msg.get("ts", 0))
            if not self._replay_ok(counter, ts):
                continue
            self.last_counter = counter
            hub_ip = str(msg.get("hub_uplink_ip") or "")
            if hub_ip:
                self.last_hub_uplink_ip = hub_ip
                # Heal Endpoint without rewriting tip — optional env hint for tools.
                os.environ["EDGE_UPLINK_IP_LEARNED"] = hub_ip
            hub_pub = str(msg.get("hub_wg_pub") or "")
            if hub_pub:
                self.last_hub_wg_pub = hub_pub
            kind = str(msg.get("kind") or "")
            if self.emit:
                self.emit(
                    "bank_wlan_rx",
                    kind=kind,
                    bank=self.bank,
                    counter=counter,
                    from_ip=addr[0],
                    hub_uplink_ip=hub_ip or None,
                )
            if kind in (KIND_DISCOVERY, KIND_KEEPALIVE, KIND_METER) and self.on_discovery:
                self.on_discovery(msg)
            if kind == KIND_METER:
                b64 = msg.get("payload_b64")
                if isinstance(b64, str) and self.on_meter_frame:
                    try:
                        frame = base64.b64decode(b64)
                    except Exception:
                        continue
                    self.on_meter_frame(frame)
            elif kind == KIND_JACKPOT and self.on_jackpot:
                self.on_jackpot(msg)
        try:
            sock.close()
        except OSError:
            pass


# Process-wide publisher (hub) — set by start_bank_wlan_for_bridge.
_publisher: BankWlanPublisher | None = None
_listener: BankWlanListener | None = None
_keepalive_stop: threading.Event | None = None


def get_publisher() -> BankWlanPublisher | None:
    return _publisher


def maybe_publish_bank_wlan(bridge: Any, pkt: bytes, meta: Any) -> None:
    """Called from UART→ETH path on hub when meters / jackpot x9 are seen."""
    pub = _publisher
    if pub is None:
        return
    try:
        from pog422_core import (
            is_gameplay_response,
            is_progressive_broadcast,
            is_skill_jpot_grant,
            progressive_progwin_cents_be,
        )

        msg_type = getattr(meta, "msg_type", None)
        if is_progressive_broadcast(msg_type):
            pub.publish_meter_0a(pkt)
            return
        # Skill grants (0x5A/0x7A) carry the progressive amount for Draw Poker and
        # SpinJack, so they must announce like an outcome x9.
        if not (is_gameplay_response(msg_type) or is_skill_jpot_grant(msg_type)):
            return
        cents = progressive_progwin_cents_be(pkt)
        if cents <= 0:
            return
        pub.publish_jackpot(
            pkt,
            tid_hex=getattr(meta, "tid_hex", None),
            progwin_cents=cents,
            msg_type_hex=getattr(meta, "msg_type_hex", None),
        )
    except Exception as exc:
        emit = getattr(getattr(bridge, "event_log", None), "emit", None)
        if emit:
            emit("bank_wlan_publish_error", message=str(exc))


def _default_meter_handler(bridge: Any) -> Callable[[bytes], None]:
    def _on_meter(frame: bytes) -> None:
        try:
            bridge._try_progressive_csv_update(frame, source="bank_wlan")
        except Exception:
            pass

    return _on_meter


def _default_jackpot_handler(bridge: Any) -> Callable[[dict[str, Any]], None]:
    def _on_jpot(msg: dict[str, Any]) -> None:
        try:
            from fireworks_ipc import write_fireworks_request

            tid = str(msg.get("tid") or "")
            cents = int(msg.get("progwin_cents") or 0)
            if cents <= 0:
                return
            local = getattr(bridge, "_local_slave_tid", None) or ""
            write_fireworks_request(
                winner_tid=tid,
                progwin_cents=cents,
                game_id_hex=None,
                dedupe_key=f"bank_wlan:{msg.get('counter')}:{tid}:{cents}",
                local_wire_tid=local,
                is_progressive=True,
            )
            emit = getattr(getattr(bridge, "event_log", None), "emit", None)
            if emit:
                emit(
                    "bank_wlan_jackpot",
                    tid=tid,
                    progwin_cents=cents,
                    counter=msg.get("counter"),
                )
        except Exception as exc:
            emit = getattr(getattr(bridge, "event_log", None), "emit", None)
            if emit:
                emit("bank_wlan_jackpot_error", message=str(exc))

    return _on_jpot


def start_bank_wlan_for_bridge(
    bridge: Any,
    stop: threading.Event,
) -> tuple[BankWlanPublisher | None, BankWlanListener | None]:
    """Arm hub publisher and/or spoke listener when EDGE bank-WLAN is enabled."""
    global _publisher, _listener, _keepalive_stop
    if not bank_wlan_enabled():
        return None, None
    topo = resolve_edge_topology()
    if topo not in ("hub", "spoke"):
        return None, None
    bank = resolve_edge_bank()
    emit = bridge.event_log.emit

    if topo == "hub":
        # Publish a discoverable underlay IP for spokes (no manual EDGE_HUB_IP).
        if not (os.environ.get("EDGE_UPLINK_IP") or "").strip():
            guessed = _guess_uplink_ip()
            if guessed:
                os.environ["EDGE_UPLINK_IP"] = guessed
        _publisher = BankWlanPublisher(bank=bank, emit=emit)
        hub_ip = (os.environ.get("EDGE_UPLINK_IP") or "").strip() or None
        emit(
            "bank_wlan_hub",
            bank=bank,
            iface=resolve_wlan_iface(),
            port=resolve_wlan_port(),
            hub_uplink_ip=hub_ip,
        )
        # Immediate discovery so late spokes do not wait for the first 0x0A.
        _publisher.publish_discovery()
        _keepalive_stop = stop

        def _keepalive_loop() -> None:
            while not stop.is_set():
                stop.wait(5.0)
                pub = _publisher
                if pub is not None:
                    # Refresh IP if interface came up after bridge start.
                    if not (os.environ.get("EDGE_UPLINK_IP") or "").strip():
                        guessed = _guess_uplink_ip()
                        if guessed:
                            os.environ["EDGE_UPLINK_IP"] = guessed
                    pub.maybe_keepalive()

        threading.Thread(target=_keepalive_loop, name="bank-wlan-keepalive", daemon=True).start()

    if topo == "spoke":
        _listener = BankWlanListener(
            bank=bank,
            emit=emit,
            on_meter_frame=_default_meter_handler(bridge),
            on_jackpot=_default_jackpot_handler(bridge),
        )
        _listener.start()

    return _publisher, _listener


def stop_bank_wlan() -> None:
    global _publisher, _listener
    if _listener is not None:
        _listener.stop()
        _listener = None
    if _publisher is not None and _publisher._sock is not None:
        try:
            _publisher._sock.close()
        except OSError:
            pass
    _publisher = None