"""CRDT gossip for node liveness heartbeats.

Uses a last-write-wins (LWW) register per inference node: each tracker node
keeps its own copy of {node_id → last_seen_monotonic_timestamp} and merges
incoming gossip by taking the max per key.  This is eventually consistent —
a heartbeat received by one tracker propagates to all others within a few
gossip intervals.

Monotonic timestamps are local-clock-relative; for cross-machine gossip the
caller should use wall-clock seconds (time.time()).  The tracker converts
monotonic to wall-clock when recording and back when comparing.
"""

from __future__ import annotations

import json
import random
import threading
import time
import urllib.request


class NodeGossip:
    """LWW gossip table for inference-node heartbeat timestamps.

    ``record(node_id)`` is called when a node sends a heartbeat to *this*
    tracker.  ``merge(remote)`` is called when gossip arrives from a peer
    tracker.  The table is periodically pushed to one random peer.
    """

    PUSH_INTERVAL = 3.0  # seconds between gossip pushes to a random peer

    def __init__(self, peers: list[str]) -> None:
        self.peers = list(peers)
        # Maps node_id → wall-clock seconds of last known heartbeat.
        self._table: dict[str, float] = {}
        self._lock = threading.Lock()
        self._running = False

    def start(self) -> None:
        self._running = True
        threading.Thread(target=self._push_loop, daemon=True, name="gossip").start()

    def stop(self) -> None:
        self._running = False

    def record(self, node_id: str, wall_ts: float | None = None) -> None:
        """Record a heartbeat for *node_id* at *wall_ts* (default: now)."""
        ts = wall_ts if wall_ts is not None else time.time()
        with self._lock:
            if ts > self._table.get(node_id, 0.0):
                self._table[node_id] = ts

    def merge(self, remote: dict[str, float]) -> None:
        """Merge a gossip snapshot from a peer tracker (LWW per key)."""
        with self._lock:
            for node_id, ts in remote.items():
                if ts > self._table.get(node_id, 0.0):
                    self._table[node_id] = ts

    def last_seen(self, node_id: str) -> float | None:
        """Return wall-clock timestamp of last known heartbeat, or None."""
        with self._lock:
            return self._table.get(node_id)

    def snapshot(self) -> dict[str, float]:
        with self._lock:
            return dict(self._table)

    def _push_loop(self) -> None:
        while self._running:
            time.sleep(self.PUSH_INTERVAL)
            if not self.peers:
                continue
            peer = random.choice(self.peers)
            try:
                snap = self.snapshot()
                body = json.dumps(snap).encode()
                req = urllib.request.Request(
                    f"{peer}/v1/gossip",
                    data=body,
                    headers={"Content-Type": "application/json"},
                    method="POST",
                )
                with urllib.request.urlopen(req, timeout=2.0) as r:
                    r.read()
            except Exception:
                pass