#!/usr/bin/env python3
import argparse
import json
import os
import re
import subprocess
import time
from datetime import datetime, timezone
from typing import Dict, List, Optional


def run_ssh(remote_cmd: str, timeout: int) -> str:
    proc = subprocess.run(
        [
            "python3",
            "-c",
            (
                "import paramiko,sys;"
                "host,user,password,cmd,timeout=sys.argv[1:6];"
                "client=paramiko.SSHClient();"
                "client.set_missing_host_key_policy(paramiko.AutoAddPolicy());"
                "client.connect(host, username=user, password=password, timeout=int(timeout));"
                "stdin,stdout,stderr=client.exec_command(cmd, timeout=int(timeout));"
                "sys.stdout.write(stdout.read().decode('utf-8','ignore'));"
                "err=stderr.read().decode('utf-8','ignore');"
                "client.close();"
                "sys.stderr.write(err)"
            ),
            os.environ["REMOTE_HOST"],
            os.environ["REMOTE_USER"],
            os.environ["REMOTE_PASS"],
            remote_cmd,
            str(timeout),
        ],
        capture_output=True,
        text=True,
        timeout=timeout + 10,
    )
    if proc.returncode != 0 and proc.stderr.strip():
        raise RuntimeError(proc.stderr.strip())
    return proc.stdout


def parse_cpu_busy(text: str) -> Optional[float]:
    match = re.search(r"([0-9]+(?:\.[0-9]+)?)\s*id", text)
    if not match:
        return None
    idle = float(match.group(1))
    return round(max(0.0, 100.0 - idle), 2)


def collect_snapshot(gray_workers: List[str], since: str) -> dict:
    units = " ".join(f"-u domaincheck-worker@{worker}.service" for worker in gray_workers)
    remote_cmd = (
        "echo '###CPU'; "
        "top -bn1 | grep '%Cpu' || true; "
        "echo '###GRAY'; "
        f"journalctl {units} --since '{since}' --no-pager "
        "| grep -E 'proxy_direct_fallback|站长之家调用边界|single_step_finalized|重试预算耗尽' "
        "| tail -n 1200 || true; "
        "echo '###GLOBAL'; "
        "journalctl -u 'domaincheck-worker@*.service' "
        f"--since '{since}' --no-pager "
        "| grep 'detect_order=detect_chinaz' | tail -n 1200 || true"
    )
    out = run_ssh(remote_cmd, timeout=40)
    cpu_part = ""
    gray_part = ""
    global_part = ""
    if "###GRAY" in out:
        cpu_part, rest = out.split("###GRAY", 1)
        if "###GLOBAL" in rest:
            gray_part, global_part = rest.split("###GLOBAL", 1)
        else:
            gray_part = rest
    else:
        cpu_part = out
    cpu_busy = parse_cpu_busy(cpu_part)

    counts = {
        "proxy_direct_fallback": 0,
        "boundary_proxy_yes": 0,
        "boundary_proxy_no": 0,
        "single_step_finalized": 0,
        "retry_budget_exhausted": 0,
    }
    per_worker: Dict[str, Dict[str, int]] = {}
    for line in gray_part.splitlines():
        worker_match = re.search(r"domaincheck-worker@([a-z0-9]+)\[", line)
        worker = worker_match.group(1) if worker_match else "-"
        bucket = per_worker.setdefault(
            worker,
            {
                "proxy_direct_fallback": 0,
                "boundary_proxy_yes": 0,
                "boundary_proxy_no": 0,
                "single_step_finalized": 0,
                "retry_budget_exhausted": 0,
            },
        )
        if "proxy_direct_fallback" in line:
            counts["proxy_direct_fallback"] += 1
            bucket["proxy_direct_fallback"] += 1
        if "站长之家调用边界" in line and "proxy=yes" in line:
            counts["boundary_proxy_yes"] += 1
            bucket["boundary_proxy_yes"] += 1
        if "站长之家调用边界" in line and "proxy=no" in line:
            counts["boundary_proxy_no"] += 1
            bucket["boundary_proxy_no"] += 1
        if "single_step_finalized" in line and "detect_key=detect_chinaz" in line:
            counts["single_step_finalized"] += 1
            bucket["single_step_finalized"] += 1
        if "重试预算耗尽" in line and "站长之家检测" in line:
            counts["retry_budget_exhausted"] += 1
            bucket["retry_budget_exhausted"] += 1

    global_worker_counts: Dict[str, int] = {}
    if global_part.strip():
        top_counts: Dict[str, int] = {}
        for line in global_part.splitlines():
            worker_match = re.search(r"domaincheck-worker@([a-z0-9]+)\[", line)
            if not worker_match:
                continue
            worker = worker_match.group(1)
            top_counts[worker] = top_counts.get(worker, 0) + 1
        global_worker_counts = dict(
            sorted(top_counts.items(), key=lambda item: item[1], reverse=True)[:12]
        )

    return {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "cpu_busy_pct": cpu_busy,
        "counts": counts,
        "per_worker": per_worker,
        "global_detect_chinaz_orders": sum(global_worker_counts.values()),
        "global_detect_chinaz_top_workers": global_worker_counts,
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", required=True)
    parser.add_argument("--workers", required=True, help="comma-separated worker suffixes")
    parser.add_argument("--duration-seconds", type=int, default=3600)
    parser.add_argument("--interval-seconds", type=int, default=30)
    parser.add_argument("--since-window-seconds", type=int, default=45)
    args = parser.parse_args()

    gray_workers = [item.strip() for item in args.workers.split(",") if item.strip()]
    os.makedirs(os.path.dirname(args.output), exist_ok=True)
    deadline = time.time() + max(0, args.duration_seconds)
    while time.time() <= deadline:
        since_ts = datetime.fromtimestamp(
            time.time() - max(1, args.since_window_seconds),
            tz=timezone.utc,
        ).astimezone().strftime("%Y-%m-%d %H:%M:%S")
        snapshot = collect_snapshot(gray_workers, since_ts)
        with open(args.output, "a", encoding="utf-8") as fh:
            fh.write(json.dumps(snapshot, ensure_ascii=False) + "\n")
        time.sleep(max(1, args.interval_seconds))
    return 0


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