#!/usr/bin/env python3
import csv
import datetime as dt
import json
import os
import re
import socket
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

RUN_ROOT = Path(os.environ.get("OURS_RUN_ROOT", "/root/ours-rp-continuous/portable-soak"))
PEER_ROOT = Path(os.environ.get("PEER_ROOT", "/root/ours-rp-continuous/portable-soak/inter-rp-peers"))
INSTANCE = os.environ.get("INTER_RP_INSTANCE", "remote200-inter-rp")
LISTEN = os.environ.get("INTER_RP_LISTEN", "0.0.0.0:9557")
SCAN_TTL = float(os.environ.get("INTER_RP_SCAN_TTL_SECONDS", "10"))

_cache_lock = threading.Lock()
_cache = {"deadline": 0.0, "metrics": "", "status": {}}
_count_cache = {}
_set_cache = {}

def unix_now():
    return time.time()

def parse_rfc3339(value):
    if not value:
        return None
    text = str(value).replace("Z", "+00:00")
    try:
        return dt.datetime.fromisoformat(text).timestamp()
    except Exception:
        return None

def read_json(path):
    with open(path, "r", encoding="utf-8") as handle:
        return json.load(handle)

def latest_ours_run():
    runs = RUN_ROOT / "runs"
    candidates = sorted(p for p in runs.glob("run_*") if (p / "run-summary.json").exists())
    return candidates[-1] if candidates else None

def unique_csv_set(path, cols):
    if not path.exists():
        return None
    stat = path.stat()
    key = (str(path), stat.st_mtime_ns, stat.st_size, cols)
    if key in _set_cache:
        return _set_cache[key]
    seen = set()
    with open(path, "r", encoding="utf-8", newline="") as handle:
        reader = csv.reader(handle)
        first = True
        for row in reader:
            if first:
                first = False
                continue
            if not row:
                continue
            if len(row) < cols:
                continue
            seen.add(tuple(cell.strip() for cell in row[:cols]))
    _set_cache.clear()
    _set_cache[key] = seen
    return seen

def count_unique_csv(path, cols):
    values = unique_csv_set(path, cols)
    if values is None:
        return None
    return len(values)

def metric_line(name, labels, value):
    label_text = ",".join(f'{k}="{str(v).replace(chr(92), chr(92)+chr(92)).replace(chr(34), chr(92)+chr(34))}"' for k, v in labels.items())
    return f"{name}{{{label_text}}} {value}\n"

def bool_num(value):
    return 1 if value else 0

def sample_ours(now, errors):
    sample = {"rp": "ours-rp", "present": False, "success": False, "max": {}, "errors": 0}
    run_dir = latest_ours_run()
    if run_dir is None:
        sample["errors"] += 1
        errors.append("ours-rp: no run-summary.json")
        return sample
    sample["present"] = True
    summary_path = run_dir / "run-summary.json"
    try:
        summary = read_json(summary_path)
    except Exception as exc:
        sample["errors"] += 1
        errors.append(f"ours-rp: read {summary_path}: {exc}")
        return sample
    sample["run_id"] = summary.get("runId") or run_dir.name
    sample["run_seq"] = summary.get("runSeq") or int(run_dir.name.split("_")[-1])
    sample["success"] = summary.get("status") == "success" and int(summary.get("exitCode", 0)) == 0
    if summary.get("wallMs") is not None:
        sample["wall"] = float(summary["wallMs"]) / 1000.0
    finished = parse_rfc3339(summary.get("finishedAtRfc3339Utc"))
    if finished is not None:
        sample["finished"] = finished
        sample["age"] = max(0.0, now - finished)
    rss = summary.get("processMetrics", {}).get("maxRssKb")
    if rss is not None:
        sample["max"]["parent"] = int(rss) * 1024
        sample["max"]["aggregate_peak"] = int(rss) * 1024
    sample["vrp_set"] = unique_csv_set(run_dir / "vrps.csv", 3)
    sample["vap_set"] = unique_csv_set(run_dir / "vaps.csv", 2)
    sample["vrps"] = len(sample["vrp_set"]) if sample["vrp_set"] is not None else None
    sample["vaps"] = len(sample["vap_set"]) if sample["vap_set"] is not None else None
    return sample

def sample_routinator(now, errors):
    rp = "routinator"
    latest = PEER_ROOT / rp / "latest"
    sample = {"rp": rp, "present": False, "success": False, "max": {}, "errors": 0}
    if not latest.exists():
        sample["errors"] += 1
        errors.append(f"routinator: missing latest directory: {latest}")
        return sample
    sample["present"] = True
    meta_path = latest / "run-meta.json"
    try:
        meta = read_json(meta_path)
    except Exception as exc:
        sample["errors"] += 1
        errors.append(f"routinator: read {meta_path}: {exc}")
        return sample
    sample["run_id"] = meta.get("runId")
    sample["run_seq"] = meta.get("runSeq")
    sample["success"] = bool(meta.get("success"))
    if meta.get("wallMs") is not None:
        sample["wall"] = float(meta["wallMs"]) / 1000.0
    finished = parse_rfc3339(meta.get("finishedAtRfc3339Utc"))
    if finished is not None:
        sample["finished"] = finished
        sample["age"] = max(0.0, now - finished)
    max_rss = meta.get("maxRssKb", {})
    for label, key in [("parent", "parent"), ("child_max", "childMax"), ("aggregate_peak", "aggregatePeak")]:
        if max_rss.get(key) is not None:
            sample["max"][label] = int(max_rss[key]) * 1024
    sample["vrp_set"] = unique_csv_set(latest / "vrps.csv", 3)
    sample["vap_set"] = unique_csv_set(latest / "vaps.csv", 2)
    sample["vrps"] = len(sample["vrp_set"]) if sample["vrp_set"] is not None else None
    sample["vaps"] = len(sample["vap_set"]) if sample["vap_set"] is not None else None
    return sample


def parse_prom_labels(label_text):
    labels = {}
    for key, value in re.findall(r'(\w+)="((?:[^"\\]|\\.)*)"', label_text):
        labels[key] = value.replace(r'\"', '"').replace(r'\\', '\\')
    return labels

def parse_prometheus_samples(text, metric_name):
    pattern = re.compile(r'^' + re.escape(metric_name) + r'\{([^}]*)\}\s+([-+0-9.eE]+)\s*$', re.MULTILINE)
    for match in pattern.finditer(text):
        try:
            value = float(match.group(2))
        except ValueError:
            continue
        yield parse_prom_labels(match.group(1)), value

def canonical_object_type(value):
    value = (value or "").strip().lower()
    if value in ("manifest", "mft"):
        return "mft"
    if value == "crl":
        return "crl"
    if value in ("certificate", "cert", "ca_cert", "router_cert", "ee_cert", "crt"):
        return "crt"
    if value == "roa":
        return "roa"
    if value == "aspa":
        return "aspa"
    return None

def empty_object_counts():
    return {"mft": 0, "crl": 0, "crt": 0, "roa": 0, "aspa": 0}

def add_object_count(counts_by_uri, uri, object_type, value):
    canonical = canonical_object_type(object_type)
    if not uri or canonical is None:
        return
    counts = counts_by_uri.setdefault(uri, empty_object_counts())
    counts[canonical] = counts.get(canonical, 0) + int(value)

def object_count_labels(counts_by_uri, uri):
    counts = counts_by_uri.get(uri, empty_object_counts())
    return {key: str(int(counts.get(key, 0))) for key in ["mft", "crl", "crt", "roa", "aspa"]}

def load_ours_repo_sets(errors):
    try:
        import urllib.request
        text = urllib.request.urlopen("http://127.0.0.1:9556/metrics", timeout=10).read().decode("utf-8", "replace")
    except Exception as exc:
        errors.append(f"ours-rp repo metrics: {exc}")
        return {"total": set(), "available": set(), "failed": set(), "info": {}, "object_counts": {}}
    info = {}
    states = {}
    object_counts = {}
    for labels, value in parse_prometheus_samples(text, "ours_rp_repository_info"):
        uri = labels.get("uri")
        if uri:
            info[uri] = labels
    for labels, value in parse_prometheus_samples(text, "ours_rp_repository_terminal_state_publication_points"):
        uri = None
        repo_id = labels.get("repo_id")
        for candidate_uri, candidate in info.items():
            if candidate.get("repo_id") == repo_id:
                uri = candidate_uri
                break
        if not uri:
            continue
        states.setdefault(uri, {})[labels.get("terminal_state", "unknown")] = value
    for labels, value in parse_prometheus_samples(text, "ours_rp_repository_objects_by_type"):
        uri = labels.get("uri")
        if not uri:
            repo_id = labels.get("repo_id")
            for candidate_uri, candidate in info.items():
                if candidate.get("repo_id") == repo_id:
                    uri = candidate_uri
                    break
        add_object_count(object_counts, uri, labels.get("object_type"), value)
    total = set(info)
    available = set()
    failed = set()
    for uri in total:
        repo_states = states.get(uri, {})
        non_failed = sum(value for state, value in repo_states.items() if state != "failed_no_cache")
        if non_failed > 0:
            available.add(uri)
        else:
            failed.add(uri)
    return {"total": total, "available": available, "failed": failed, "info": info, "object_counts": object_counts}

def load_routinator_repo_sets(errors):
    path = PEER_ROOT / "routinator" / "routinator-metrics.prom"
    if not path.exists():
        errors.append(f"routinator repo metrics: missing {path}")
        return {"total": set(), "success": set(), "failed": set(), "duration": {}, "object_counts": {}}
    try:
        text = path.read_text(encoding="utf-8", errors="replace")
    except Exception as exc:
        errors.append(f"routinator repo metrics: read {path}: {exc}")
        return {"total": set(), "success": set(), "failed": set(), "duration": {}, "object_counts": {}}
    total = set()
    success = set()
    failed = set()
    duration = {}
    object_counts = {}
    publication_point_states = {}
    for metric in ["routinator_rrdp_status", "routinator_rsync_status"]:
        for labels, value in parse_prometheus_samples(text, metric):
            uri = labels.get("uri")
            if not uri:
                continue
            total.add(uri)
            ok = value in (0.0, 200.0, 304.0)
            if ok:
                success.add(uri)
            else:
                failed.add(uri)
    for labels, value in parse_prometheus_samples(text, "routinator_repository_publication_points_total"):
        uri = labels.get("uri")
        state = labels.get("state", "unknown")
        if not uri:
            continue
        total.add(uri)
        publication_point_states.setdefault(uri, {})[state] = publication_point_states.setdefault(uri, {}).get(state, 0.0) + value
    for uri, states in publication_point_states.items():
        valid_count = states.get("valid", 0.0)
        non_valid_count = sum(value for state, value in states.items() if state != "valid")
        if valid_count > 0:
            success.add(uri)
        elif non_valid_count > 0:
            failed.add(uri)
    for metric in ["routinator_rrdp_duration", "routinator_rsync_duration"]:
        for labels, value in parse_prometheus_samples(text, metric):
            uri = labels.get("uri")
            if uri:
                duration[uri] = max(duration.get(uri, 0.0), value)
    for labels, value in parse_prometheus_samples(text, "routinator_repository_objects_total"):
        add_object_count(object_counts, labels.get("uri"), labels.get("type"), value)
    failed = failed - success
    return {"total": total, "success": success, "failed": failed, "duration": duration, "object_counts": object_counts}

def emit_repo_diff_metrics(out, errors):
    ours = load_ours_repo_sets(errors)
    routinator = load_routinator_repo_sets(errors)
    ours_ok = ours["available"]
    rout_ok = routinator["success"]
    only_ours = sorted(ours_ok - rout_ok)
    only_routinator = sorted(rout_ok - ours_ok)
    both = ours_ok & rout_ok
    neither = (ours["total"] | routinator["total"]) - (ours_ok | rout_ok)
    base = {"instance": INSTANCE, "left": "ours-rp", "right": "routinator"}
    out.append(metric_line("inter_rp_repo_sync_total", {"instance": INSTANCE, "rp": "ours-rp", "state": "total"}, len(ours["total"])))
    out.append(metric_line("inter_rp_repo_sync_total", {"instance": INSTANCE, "rp": "ours-rp", "state": "available"}, len(ours_ok)))
    out.append(metric_line("inter_rp_repo_sync_total", {"instance": INSTANCE, "rp": "ours-rp", "state": "failed"}, len(ours["failed"])))
    out.append(metric_line("inter_rp_repo_sync_total", {"instance": INSTANCE, "rp": "routinator", "state": "total"}, len(routinator["total"])))
    out.append(metric_line("inter_rp_repo_sync_total", {"instance": INSTANCE, "rp": "routinator", "state": "available"}, len(rout_ok)))
    out.append(metric_line("inter_rp_repo_sync_total", {"instance": INSTANCE, "rp": "routinator", "state": "failed"}, len(routinator["failed"])))
    out.append(metric_line("inter_rp_repo_sync_overlap_total", {**base, "class": "both_available"}, len(both)))
    out.append(metric_line("inter_rp_repo_sync_overlap_total", {**base, "class": "only_ours"}, len(only_ours)))
    out.append(metric_line("inter_rp_repo_sync_overlap_total", {**base, "class": "only_routinator"}, len(only_routinator)))
    out.append(metric_line("inter_rp_repo_sync_overlap_total", {**base, "class": "neither_available"}, len(neither)))
    for diff_class, uris, counts_by_uri in [
        ("only_ours", only_ours, ours["object_counts"]),
        ("only_routinator", only_routinator, routinator["object_counts"]),
    ]:
        for rank, uri in enumerate(uris[:50], start=1):
            labels = {
                **base,
                "class": diff_class,
                "rank": rank,
                "uri": uri,
                **object_count_labels(counts_by_uri, uri),
            }
            if uri in routinator["duration"]:
                labels["routinator_duration"] = f"{routinator['duration'][uri]:.3f}"
            out.append(metric_line("inter_rp_repo_sync_diff_info", labels, 1))
    return {
        "oursTotal": len(ours["total"]),
        "oursAvailable": len(ours_ok),
        "routinatorTotal": len(routinator["total"]),
        "routinatorAvailable": len(rout_ok),
        "onlyOurs": only_ours[:50],
        "onlyRoutinator": only_routinator[:50],
        "bothAvailable": len(both),
        "neitherAvailable": len(neither),
    }

def sync_metrics(now):
    path = PEER_ROOT / "sync-status.json"
    if not path.exists():
        return {"present": False, "success": False, "message": f"missing {path}"}
    try:
        value = read_json(path)
    except Exception as exc:
        return {"present": True, "success": False, "message": str(exc)}
    ts = parse_rfc3339(value.get("lastSyncAtRfc3339Utc"))
    return {
        "present": True,
        "success": bool(value.get("success")),
        "timestamp": ts,
        "age": max(0.0, now - ts) if ts is not None else None,
        "remote": value.get("remoteHost", ""),
        "message": value.get("message", ""),
    }

def build_metrics():
    now = unix_now()
    errors = []
    start = time.time()
    samples = [sample_ours(now, errors), sample_routinator(now, errors)]
    sync = sync_metrics(now)
    if not sync.get("success"):
        errors.append("sync: " + str(sync.get("message", "failed")))
    out = []
    out.append(metric_line("inter_rp_service_up", {"instance": INSTANCE}, 1))
    out.append(metric_line("inter_rp_service_last_scan_timestamp_seconds", {"instance": INSTANCE}, now))
    out.append(metric_line("inter_rp_service_last_scan_duration_seconds", {"instance": INSTANCE}, time.time() - start))
    out.append(metric_line("inter_rp_service_last_reload_success", {"instance": INSTANCE}, bool_num(len(errors) == 0)))
    out.append(metric_line("inter_rp_parse_errors", {"instance": INSTANCE}, len(errors)))
    out.append(metric_line("inter_rp_sync_present", {"instance": INSTANCE}, bool_num(sync.get("present"))))
    out.append(metric_line("inter_rp_sync_last_success", {"instance": INSTANCE}, bool_num(sync.get("success"))))
    if sync.get("age") is not None:
        out.append(metric_line("inter_rp_sync_age_seconds", {"instance": INSTANCE}, sync["age"]))
    if sync.get("timestamp") is not None:
        out.append(metric_line("inter_rp_sync_last_timestamp_seconds", {"instance": INSTANCE}, sync["timestamp"]))
    by_rp = {s["rp"]: s for s in samples}
    for s in samples:
        labels = {"instance": INSTANCE, "rp": s["rp"]}
        out.append(metric_line("inter_rp_run_present", labels, bool_num(s.get("present"))))
        out.append(metric_line("inter_rp_run_success", labels, bool_num(s.get("success"))))
        out.append(metric_line("inter_rp_sample_parse_errors", labels, s.get("errors", 0)))
        if s.get("run_seq") is not None:
            out.append(metric_line("inter_rp_run_seq", labels, s["run_seq"]))
        if s.get("wall") is not None:
            out.append(metric_line("inter_rp_run_wall_seconds", labels, s["wall"]))
        if s.get("age") is not None:
            out.append(metric_line("inter_rp_artifact_age_seconds", labels, s["age"]))
        if s.get("vrps") is not None:
            out.append(metric_line("inter_rp_vrps", labels, s["vrps"]))
        if s.get("vaps") is not None:
            out.append(metric_line("inter_rp_vaps", labels, s["vaps"]))
        for kind, value in s.get("max", {}).items():
            labels2 = dict(labels)
            labels2["kind"] = kind
            out.append(metric_line("inter_rp_run_max_rss_bytes", labels2, value))
    ours, rout = by_rp.get("ours-rp", {}), by_rp.get("routinator", {})
    if ours.get("vrp_set") is not None and rout.get("vrp_set") is not None:
        only_ours = ours["vrp_set"] - rout["vrp_set"]
        only_rout = rout["vrp_set"] - ours["vrp_set"]
        total_diff = len(only_ours) + len(only_rout)
        labels = {"instance": INSTANCE, "left": "ours-rp", "right": "routinator"}
        out.append(metric_line("inter_rp_vrps_diff", labels, total_diff))
        out.append(metric_line("inter_rp_vrps_diff_by_class", {**labels, "class": "total"}, total_diff))
        out.append(metric_line("inter_rp_vrps_diff_by_class", {**labels, "class": "only_ours"}, len(only_ours)))
        out.append(metric_line("inter_rp_vrps_diff_by_class", {**labels, "class": "only_routinator"}, len(only_rout)))
    elif ours.get("vrps") is not None and rout.get("vrps") is not None:
        labels = {"instance": INSTANCE, "left": "ours-rp", "right": "routinator"}
        total_diff = abs(int(ours["vrps"]) - int(rout["vrps"]))
        out.append(metric_line("inter_rp_vrps_diff", labels, total_diff))
        out.append(metric_line("inter_rp_vrps_diff_by_class", {**labels, "class": "total"}, total_diff))
    if ours.get("vap_set") is not None and rout.get("vap_set") is not None:
        only_ours = ours["vap_set"] - rout["vap_set"]
        only_rout = rout["vap_set"] - ours["vap_set"]
        total_diff = len(only_ours) + len(only_rout)
        labels = {"instance": INSTANCE, "left": "ours-rp", "right": "routinator"}
        out.append(metric_line("inter_rp_vaps_diff", labels, total_diff))
        out.append(metric_line("inter_rp_vaps_diff_by_class", {**labels, "class": "total"}, total_diff))
        out.append(metric_line("inter_rp_vaps_diff_by_class", {**labels, "class": "only_ours"}, len(only_ours)))
        out.append(metric_line("inter_rp_vaps_diff_by_class", {**labels, "class": "only_routinator"}, len(only_rout)))
    elif ours.get("vaps") is not None and rout.get("vaps") is not None:
        labels = {"instance": INSTANCE, "left": "ours-rp", "right": "routinator"}
        total_diff = abs(int(ours["vaps"]) - int(rout["vaps"]))
        out.append(metric_line("inter_rp_vaps_diff", labels, total_diff))
        out.append(metric_line("inter_rp_vaps_diff_by_class", {**labels, "class": "total"}, total_diff))
    repo_diff = emit_repo_diff_metrics(out, errors)
    return "".join(out), {"errors": errors, "samples": samples, "sync": sync, "repoDiff": repo_diff}

def get_metrics():
    now = time.time()
    with _cache_lock:
        if _cache["metrics"] and _cache["deadline"] > now:
            return _cache["metrics"]
        metrics, status = build_metrics()
        _cache["metrics"] = metrics
        _cache["status"] = status
        _cache["deadline"] = now + SCAN_TTL
        return metrics

class Handler(BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path == "/metrics":
            body = get_metrics().encode("utf-8")
            self.send_response(200)
            self.send_header("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return
        if self.path == "/status":
            with _cache_lock:
                if not _cache["status"] or _cache["deadline"] <= time.time():
                    build_metrics()
                body = json.dumps(_cache["status"], indent=2).encode("utf-8")
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return
        self.send_response(404)
        self.end_headers()
    def log_message(self, fmt, *args):
        return

def main():
    host, port = LISTEN.rsplit(":", 1)
    server = ThreadingHTTPServer((host, int(port)), Handler)
    print(f"inter-rp ours+routinator exporter listen={LISTEN} instance={INSTANCE}", flush=True)
    server.serve_forever()

if __name__ == "__main__":
    main()
