#!/usr/bin/env python3
"""Prometheus exporter for the labsim connectivity matrix.

Runs the same sweep as labsim-matrix.py on an interval and exposes it as
metrics, so Grafana can show the mesh as a heatmap and — more usefully — a
history of exactly when a cell flipped after a firewall change.

  labsim_reachable{src,dst,proto}   1 = reachable, 0 = blocked
  labsim_sweep_seconds              how long the last sweep took
  labsim_sweep_total                sweeps completed since start
  labsim_up                         1 while the exporter is alive

Deliberately stdlib-only (http.server + threads): this runs on the workstation
next to libvirt, and adding a dependency to watch a lab network is silly.

  ./labsim-exporter.py --port 9101 --interval 15
"""
from __future__ import annotations

import argparse
import http.server
import json
import os
import threading
import time

import labsim_matrix_lib as m  # thin import shim, see below


class Collector:
    def __init__(self, interval: int, timeout: int) -> None:
        self.interval = interval
        self.timeout = timeout
        self.vlans = m.load_vlans()
        self.lock = threading.Lock()
        self.results: dict = {}
        self.duration = 0.0
        self.sweeps = 0

    def loop(self) -> None:
        while True:
            started = time.time()
            try:
                results = m.sweep(self.vlans, self.timeout)
                with self.lock:
                    self.results = results
                    self.duration = time.time() - started
                    self.sweeps += 1
            except Exception:  # noqa: BLE001 - never let the loop die
                pass
            time.sleep(max(1.0, self.interval - (time.time() - started)))

    def snapshot(self) -> dict:
        """Everything the topology page needs, in one JSON payload."""
        with self.lock:
            results, duration = dict(self.results), self.duration
        reach = total = 0
        for data in results.values():
            if "__error__" in data:
                continue
            for protos in data.values():
                for proto, ok in protos.items():
                    if proto == "rtt_ms":
                        continue
                    total += 1
                    if ok:
                        reach += 1
        return {"vlans": self.vlans, "results": results, "reachable": reach,
                "total": total, "sweep_seconds": duration}

    def render(self) -> str:
        with self.lock:
            results, duration, sweeps = dict(self.results), self.duration, self.sweeps

        out = [
            "# HELP labsim_reachable 1 if dst is reachable from src over proto",
            "# TYPE labsim_reachable gauge",
        ]
        rtts = []
        for src, data in results.items():
            if "__error__" in data:
                continue
            for dst, protos in data.items():
                for proto, ok in protos.items():
                    if proto == "rtt_ms":
                        if isinstance(ok, (int, float)):
                            rtts.append((src, dst, ok))
                        continue
                    out.append(
                        f'labsim_reachable{{src="{src}",dst="{dst}",proto="{proto}"}} {1 if ok else 0}')
        out += ["# HELP labsim_rtt_ms ICMP round-trip time",
                "# TYPE labsim_rtt_ms gauge"]
        for src, dst, val in rtts:
            out.append(f'labsim_rtt_ms{{src="{src}",dst="{dst}"}} {val}')
        out += [
            "# HELP labsim_sweep_seconds duration of the last sweep",
            "# TYPE labsim_sweep_seconds gauge",
            f"labsim_sweep_seconds {duration:.3f}",
            "# HELP labsim_sweep_total sweeps completed",
            "# TYPE labsim_sweep_total counter",
            f"labsim_sweep_total {sweeps}",
            "# HELP labsim_up exporter liveness",
            "# TYPE labsim_up gauge",
            "labsim_up 1",
        ]
        return "\n".join(out) + "\n"


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=9101)
    ap.add_argument("--interval", type=int, default=15)
    ap.add_argument("--timeout", type=int, default=30)
    args = ap.parse_args()

    collector = Collector(args.interval, args.timeout)
    threading.Thread(target=collector.loop, daemon=True).start()

    here = os.path.dirname(os.path.abspath(__file__))

    class Handler(http.server.BaseHTTPRequestHandler):
        def _send(self, body: bytes, ctype: str) -> None:
            self.send_response(200)
            self.send_header("Content-Type", ctype)
            self.send_header("Content-Length", str(len(body)))
            self.send_header("Cache-Control", "no-store")
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self) -> None:  # noqa: N802 - stdlib API
            path = self.path.split("?")[0].rstrip("/")
            if path in ("", "/topology"):
                # Live topology view — the thing you actually watch.
                try:
                    with open(os.path.join(here, "topology.html"), "rb") as fh:
                        self._send(fh.read(), "text/html; charset=utf-8")
                except OSError:
                    self.send_error(500, "topology.html missing")
            elif path == "/api/matrix":
                self._send(json.dumps(collector.snapshot()).encode(), "application/json")
            elif path == "/metrics":
                self._send(collector.render().encode(), "text/plain; version=0.0.4")
            else:
                self.send_error(404)

        def log_message(self, *_args) -> None:  # keep the console quiet
            return

    srv = http.server.ThreadingHTTPServer(("0.0.0.0", args.port), Handler)
    print(f"labsim topology  http://localhost:{args.port}/")
    print(f"labsim metrics   http://localhost:{args.port}/metrics  (sweep every {args.interval}s)")
    srv.serve_forever()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())