#!/usr/bin/env python3

import argparse
import json
import time
import urllib.error
import urllib.request


def http_json(method: str, url: str, body=None, timeout: float = 5.0):
    data = None
    headers = {
        "Accept": "application/json",
    }
    if body is not None:
        data = json.dumps(body).encode("utf-8")
        headers["Content-Type"] = "application/json"

    req = urllib.request.Request(url, data=data, headers=headers, method=method)
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        raw = resp.read().decode("utf-8", errors="replace")
        return json.loads(raw)


def main():
    ap = argparse.ArgumentParser(description="Poll /api/graphs and /api/graphs/{name} and write JSONL.")
    ap.add_argument("--api", default="http://127.0.0.1:9000", help="Base API URL")
    ap.add_argument("--interval", type=float, default=2.0, help="Polling interval seconds")
    ap.add_argument("--duration", type=float, default=60.0, help="Total duration seconds")
    ap.add_argument("--out", default="metrics.jsonl", help="Output jsonl path")
    ap.add_argument("--reload", action="store_true", help="POST /api/config/reload before polling")
    args = ap.parse_args()

    base = args.api.rstrip("/")

    if args.reload:
        try:
            http_json("POST", f"{base}/api/config/reload")
        except Exception as e:
            print(f"reload failed: {e}")

    end_ts = time.time() + max(0.0, args.duration)
    with open(args.out, "w", encoding="utf-8") as f:
        while time.time() <= end_ts:
            ts = int(time.time() * 1000)
            record: dict = {"ts_ms": ts}
            try:
                graphs = http_json("GET", f"{base}/api/graphs")
                record["graphs"] = graphs
                details = {}
                for g in graphs:
                    name = g.get("name")
                    if not name:
                        continue
                    try:
                        details[name] = http_json("GET", f"{base}/api/graphs/{name}")
                    except Exception as e:
                        details[name] = {"error": str(e)}
                record["graph_details"] = details
            except urllib.error.HTTPError as e:
                record["error"] = f"http {e.code}: {e.read().decode('utf-8', errors='replace')}"
            except Exception as e:
                record["error"] = str(e)

            f.write(json.dumps(record, ensure_ascii=False) + "\n")
            f.flush()
            time.sleep(max(0.1, args.interval))


if __name__ == "__main__":
    main()