import concurrent.futures
import datetime
import json
import pathlib
import shlex
import subprocess
import sys
import urllib.parse


SSH = [
    "ssh", "-i", str(pathlib.Path.home() / ".ssh/waow_ed25519"),
    "-o", "IdentitiesOnly=yes", "-o", "BatchMode=yes",
    "-o", "ConnectTimeout=5", "root@89.167.122.160",
]


def run(args):
    return subprocess.run(args, check=True, capture_output=True, text=True, timeout=30).stdout.strip()


def prometheus(path, query=None):
    url = "http://127.0.0.1:9090/api/v1/" + path
    if query is not None:
        url += "?" + urllib.parse.urlencode({"query": query})
    result = json.loads(run(SSH + [
        "docker exec stream-experiment-prometheus-1 wget -T 5 -qO- " + shlex.quote(url)
    ]))
    if result.get("status") != "success":
        raise RuntimeError(result)
    return result["data"]


def query(expression):
    result = prometheus("query", expression)
    return [{"labels": row["metric"], "value": row["value"][1]} for row in result["result"]]


def increases(names, window):
    return " or ".join(
        f'label_replace(increase({name}{{job="stream"}}[{window}]), "counter", "{name}", "", "")'
        for name in names.split()
    )


QUERIES = {
    "ingest_events_per_second_5m": 'rate(jetstream_livestream_events_received_total{job="stream"}[5m])',
    "ingest_5m_rate_min_max_24h": 'label_replace(min_over_time(rate(jetstream_livestream_events_received_total{job="stream"}[5m])[24h:5m]), "stat", "min", "", "") or label_replace(max_over_time(rate(jetstream_livestream_events_received_total{job="stream"}[5m])[24h:5m]), "stat", "max", "", "")',
    "target_scrape_success_fraction_24h": 'avg_over_time(up{job="stream"}[24h])',
    "target_scrape_samples_24h": 'count_over_time(up{job="stream"}[24h])',
    "process_start_changes_24h": 'changes(process_start_time_seconds{job="stream"}[24h])',
    "rss_bytes_now": 'process_resident_memory_bytes{job="stream"}',
    "rss_bytes_max_24h": 'max_over_time(process_resident_memory_bytes{job="stream"}[24h])',
    "upstream_event_age_seconds": 'time() - jetstream_livestream_last_seen_upstream_event_timestamp_seconds{job="stream"}',
    "counters_increase_1h": increases("jetstream_livestream_reconnects_total stream_upstream_slow_reconnects_total jetstream_livestream_decode_errors_total jetstream_livestream_verify_queue_dropped_events_total jetstream_subscribe_encode_errors_total stream_resync_pending_dropped_total stream_backfill_failed_repo_retry_terminal_failures_total", "1h"),
    "ingest_drops_increase_1h": 'increase(jetstream_ingest_dropped_events_total{job="stream"}[1h])',
    "maintenance_increase_24h": increases("jetstream_compaction_passes_total jetstream_backfill_failed_repo_retry_passes_total jetstream_backfill_failed_repo_retry_succeeded_total jetstream_backfill_failed_repo_retry_failed_total", "24h"),
}


def main():
    tasks = {
        "public_status": lambda: run(["curl", "--max-time", "5", "-fsS", "https://stream.waow.tech/status"]),
        "public_build": lambda: "\n".join(line for line in run(["curl", "--max-time", "5", "-fsS", "https://stream.waow.tech/metrics"]).splitlines() if line.startswith("stream_build_info")),
        "host": lambda: run(SSH + ["set -e; uptime; df -h / /opt/stream-experiment; df -i /; free -h; docker inspect --format '{{.Name}} image={{.Config.Image}} started={{.State.StartedAt}} status={{.State.Status}} restarts={{.RestartCount}} oom={{.State.OOMKilled}}' stream-experiment-stream-1 stream-experiment-caddy-1"]),
        "alerts": lambda: prometheus("alerts")["alerts"],
        "targets": lambda: [{"labels": row["labels"], "health": row["health"], "lastScrape": row["lastScrape"]} for row in prometheus("targets")["activeTargets"]],
    }
    tasks.update({name: lambda expression=expression: query(expression) for name, expression in QUERIES.items()})
    report = {"observed_at": datetime.datetime.now(datetime.UTC).isoformat(), "errors": {}}
    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as pool:
        futures = {pool.submit(task): name for name, task in tasks.items()}
        for future in concurrent.futures.as_completed(futures):
            name = futures[future]
            try:
                report[name] = future.result()
            except Exception as error:
                report["errors"][name] = str(error)
    print(json.dumps(report, indent=2, sort_keys=True))
    return bool(report["errors"])


if __name__ == "__main__":
    sys.exit(main())