#!/usr/bin/env python3
# Montana explorer data.json collector.
# Polls the three Genesis nodes for their montana-node status and merges in
# any auto-discovered peer (any peer connected with label=unknown that has
# emitted heartbeat OK in the recent journal window). The result is written
# to /var/www/efir/explorer/data.json once per minute.

import json
import re
import subprocess
import time
import os
from datetime import datetime, timezone

OUT = "/var/www/efir/explorer/data.json"
DISCOVERY_WINDOW_SECONDS = 600  # consider a peer "live" if heartbeat OK within 10 min

GENESIS_NODES = [
    ("Moscow",    "local"),
    ("Helsinki",  "91.132.142.42"),
    ("Frankfurt", "89.19.208.158"),
]


def run_local(cmd, timeout=10):
    return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout).stdout


def run_ssh(host, remote_cmd, timeout=10):
    cmd = [
        "ssh", "-o", "ConnectTimeout=5", "-o", "StrictHostKeyChecking=no",
        f"root@{host}", remote_cmd,
    ]
    return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout).stdout


def parse_status(text):
    if not text:
        return None
    def grep(pattern, default=None, group=1):
        m = re.search(pattern, text, re.MULTILINE)
        return m.group(group) if m else default
    return {
        "current_window": int(grep(r"^current_window\s*:\s*(\d+)", "0") or 0),
        "phase": grep(r"^phase\s*:\s*(\S+)", "unknown"),
        "D": int(grep(r"^D \(current\)\s*:\s*(\d+)", "0") or 0),
        "account_id": grep(r"^account_id\s*:\s*([0-9a-f]+)", ""),
        "node_id": grep(r"^node_id\s*:\s*([0-9a-f]+)", ""),
        "balance_nj": int(grep(r"^balance\s*:\s*(\d+)\s*nɈ", "0") or 0),
        "supply_nj": int(grep(r"supply \(closed-form\)\s*:\s*(\d+)", "0") or 0),
        "account_table": int(grep(r"^AccountTable\s*:\s*(\d+)", "0") or 0),
        "node_table": int(grep(r"^NodeTable\s*:\s*(\d+)", "0") or 0),
    }


def fetch_genesis(label, host):
    try:
        if host == "local":
            out = run_local(
                ["/usr/local/bin/montana-node", "status",
                 "--data-dir", "/var/lib/montana"]
            )
        else:
            out = run_ssh(
                host,
                "/usr/local/bin/montana-node status --data-dir /var/lib/montana"
            )
        st = parse_status(out)
        if st:
            return {"label": label, "host": host, "status": "active", **st}
        return {"label": label, "host": host, "status": "no_data"}
    except Exception as e:
        return {"label": label, "host": host, "status": "unreachable",
                "error": str(e)[:100]}


# Patterns:
# [network] CONNECTION ESTABLISHED peer=<XX> label=<L> remote=/ip4/<IP>/tcp/<PORT>...
# [network] heartbeat OK peer=<XX> request_id=N
# [network] connection closed peer=<XX> label=<L> cause=...
CONN_ESTABLISHED_RE = re.compile(
    r"CONNECTION ESTABLISHED peer=(\S+) label=(\S+) remote=/ip4/(\d+\.\d+\.\d+\.\d+)/tcp/(\d+)"
)
HEARTBEAT_RE = re.compile(r"heartbeat OK peer=(\S+)")
CONN_CLOSED_RE = re.compile(r"connection closed peer=(\S+)")


def collect_discovery(witness_label, witness_host):
    """Scan the witness node's recent journal for discovered (unknown-label) peers.

    Returns: dict mapping peer_id → {label, remote_ip, last_heartbeat_seconds_ago}.
    """
    try:
        cmd_str = (
            f"journalctl -u montana-node --since '{DISCOVERY_WINDOW_SECONDS} seconds ago' "
            "--no-pager -o short-iso"
        )
        if witness_host == "local":
            raw = run_local(["bash", "-lc", cmd_str], timeout=15)
        else:
            raw = run_ssh(witness_host, cmd_str, timeout=15)
    except Exception:
        return {}

    discovered = {}
    closed = set()
    # First pass: collect ConnectionEstablished events (peer_id → remote_ip)
    for line in raw.splitlines():
        m = CONN_ESTABLISHED_RE.search(line)
        if m:
            peer_id, label, remote_ip, _port = m.group(1), m.group(2), m.group(3), m.group(4)
            if label == "unknown":
                discovered[peer_id] = {
                    "peer_id": peer_id,
                    "remote_ip": remote_ip,
                    "witness": witness_label,
                    "first_seen_via": witness_label,
                }
        m = CONN_CLOSED_RE.search(line)
        if m:
            closed.add(m.group(1))

    # Second pass: track latest heartbeat for each discovered peer
    last_hb = {}
    for line in raw.splitlines():
        m = HEARTBEAT_RE.search(line)
        if m:
            peer_id = m.group(1)
            ts_match = re.match(r"(\S+)", line)
            if ts_match:
                last_hb[peer_id] = ts_match.group(1)

    now = datetime.now(timezone.utc)
    result = {}
    for peer_id, info in discovered.items():
        # Skip peers that disconnected after their last ConnectionEstablished
        # (heuristic: if peer present in `closed` set and no heartbeat after).
        if peer_id in closed and peer_id not in last_hb:
            continue
        if peer_id in last_hb:
            try:
                last_ts = datetime.fromisoformat(last_hb[peer_id])
                age = (now - last_ts).total_seconds()
                info["last_heartbeat_seconds_ago"] = int(age)
            except Exception:
                info["last_heartbeat_seconds_ago"] = None
        result[peer_id] = info
    return result


def merge_discoveries(*discovery_maps):
    merged = {}
    for dmap in discovery_maps:
        for peer_id, info in dmap.items():
            if peer_id not in merged:
                merged[peer_id] = {**info, "witnessed_by": [info["witness"]]}
            else:
                if info["witness"] not in merged[peer_id]["witnessed_by"]:
                    merged[peer_id]["witnessed_by"].append(info["witness"])
                # keep smallest heartbeat age
                cur = merged[peer_id].get("last_heartbeat_seconds_ago")
                new = info.get("last_heartbeat_seconds_ago")
                if cur is None or (new is not None and new < cur):
                    merged[peer_id]["last_heartbeat_seconds_ago"] = new
    return [
        {
            "peer_id": p["peer_id"],
            "remote_ip": p["remote_ip"],
            "witnessed_by": p["witnessed_by"],
            "last_heartbeat_seconds_ago": p.get("last_heartbeat_seconds_ago"),
            "status": ("active" if (p.get("last_heartbeat_seconds_ago") or 999999) < 60 else "stale"),
        }
        for p in merged.values()
    ]


# --- Build the document ---
nodes = [fetch_genesis(label, host) for (label, host) in GENESIS_NODES]

discoveries = []
for (label, host) in GENESIS_NODES:
    dmap = collect_discovery(label, host)
    discoveries.append(dmap)
discovered_peers = merge_discoveries(*discoveries)

doc = {
    "updated": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
    "updated_unix": int(time.time()),
    "nodes": nodes,
    "discovered_peers": discovered_peers,
    "network_summary": {
        "active_nodes": sum(1 for n in nodes if n["status"] == "active"),
        "total_nodes": len(nodes),
        "discovered_peer_count": len(discovered_peers),
        "max_window": max((n.get("current_window", 0) for n in nodes), default=0),
        "total_supply_nj": sum(
            n.get("supply_nj", 0) for n in nodes if n["status"] == "active"
        ),
    },
}

os.makedirs(os.path.dirname(OUT), exist_ok=True)
with open(OUT, "w") as f:
    json.dump(doc, f, indent=2, ensure_ascii=False)
os.chmod(OUT, 0o644)
print(json.dumps(doc, indent=2, ensure_ascii=False)[:800])