#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import re
import time
from datetime import datetime, timedelta
from pathlib import Path

import paramiko


def _now() -> datetime:
    return datetime.now()


def _parse_key_value_lines(raw: str) -> dict:
    result: dict[str, object] = {"raw": raw.strip()}
    for line in raw.splitlines():
        if "=" not in line:
            continue
        key, value = line.split("=", 1)
        key = key.strip()
        value = value.strip()
        if not key:
            continue
        if key in {"running_jobs", "step_counts"}:
            try:
                result[key] = json.loads(value)
            except Exception:
                result[key] = value
            continue
        for caster in (int, float):
            try:
                result[key] = caster(value)
                break
            except Exception:
                continue
        else:
            result[key] = value
    return result


def _remote_sample(host: str, user: str, password: str, recent_minutes: int) -> dict:
    cli = paramiko.SSHClient()
    cli.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    command = """
python3 - <<'PY'
import json, re, subprocess

def sh(cmd: str) -> str:
    return subprocess.check_output(cmd, shell=True, text=True, stderr=subprocess.DEVNULL)

cpu_line = sh("top -bn1 | sed -n '3p'").strip()
idle_match = re.search(r'([0-9]+(?:\.[0-9]+)?)\s+id', cpu_line)
busy = round(100.0 - float(idle_match.group(1)), 2) if idle_match else None
print(f"cpu_busy_pct={busy if busy is not None else ''}")

worker_stats = sh(r'''ps -eo pcpu,cmd | awk '/[d]etect_worker.py/ {cpu=$1+0; total++; sum+=cpu; if(cpu>1) gt1++; if(cpu>2) gt2++; if(cpu>3) gt3++;} END {printf("workers_total=%d\\nworkers_gt_1pct=%d\\nworkers_gt_2pct=%d\\nworkers_gt_3pct=%d\\nworkers_cpu_sum=%.1f\\n", total, gt1, gt2, gt3, sum)}' ''')
print(worker_stats.strip())

running_jobs_raw = sh(r'''/www/server/pgsql/bin/psql -U postgres -d domain -At -c "select id, job_code, status from detect_jobs where status='running' order by id desc limit 16;" ''')
running_jobs = []
for line in running_jobs_raw.splitlines():
    parts = line.split("|")
    if len(parts) >= 3:
        running_jobs.append({"job_id": int(parts[0]), "job_code": parts[1], "status": parts[2]})
print("running_jobs=" + json.dumps(running_jobs, ensure_ascii=False))

recent = __RECENT_MINUTES__
journal_cmd = f"journalctl -u domaincheck-worker -u 'domaincheck-worker@*' --since '{recent} min ago' --no-pager -o cat"
logs = sh(journal_cmd)
step_counts = {}
for match in re.finditer(r'detect_order=(detect_[a-zA-Z0-9_]+)', logs):
    step = match.group(1)
    step_counts[step] = step_counts.get(step, 0) + 1
print("step_counts=" + json.dumps(step_counts, ensure_ascii=False, sort_keys=True))
PY
"""
    command = command.replace("__RECENT_MINUTES__", str(int(recent_minutes)))
    try:
        cli.connect(hostname=host, username=user, password=password, timeout=15, banner_timeout=15, auth_timeout=15)
        _, stdout, stderr = cli.exec_command(command, timeout=90)
        output = stdout.read().decode("utf-8", errors="replace")
        error = stderr.read().decode("utf-8", errors="replace").strip()
        row = _parse_key_value_lines(output)
        if error:
            row["stderr"] = error
        return row
    except Exception as exc:
        return {"error": str(exc)}
    finally:
        try:
            cli.close()
        except Exception:
            pass


def main() -> int:
    parser = argparse.ArgumentParser(description="Collect remote CPU/worker/step-mix metrics into JSONL.")
    parser.add_argument("--worker-host", required=True)
    parser.add_argument("--worker-user", required=True)
    parser.add_argument("--worker-password", required=True)
    parser.add_argument("--interval-seconds", type=int, default=120)
    parser.add_argument("--duration-hours", type=float, default=8.0)
    parser.add_argument("--recent-minutes", type=int, default=3)
    parser.add_argument("--label", default="remote-step-mix")
    args = parser.parse_args()

    start_time = _now()
    end_time = start_time + timedelta(hours=max(0.25, float(args.duration_hours)))
    run_dir = Path("/www/wwwroot/getDomain/docs/ops_center_runtime/night_runs") / (
        f"step_mix_{start_time.strftime('%Y%m%d_%H%M%S')}_{args.label}"
    )
    run_dir.mkdir(parents=True, exist_ok=True)
    metrics_path = run_dir / "metrics.jsonl"
    meta_path = run_dir / "meta.json"
    meta_path.write_text(
        json.dumps(
            {
                "label": args.label,
                "started_at": start_time.isoformat(sep=" ", timespec="seconds"),
                "planned_end_at": end_time.isoformat(sep=" ", timespec="seconds"),
                "interval_seconds": int(args.interval_seconds),
                "recent_minutes": int(args.recent_minutes),
                "worker_host": args.worker_host,
            },
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
    )

    sample_index = 0
    while _now() < end_time:
        sample_index += 1
        row = {
            "timestamp": _now().isoformat(sep=" ", timespec="seconds"),
            "sample_index": sample_index,
            "remote_worker": _remote_sample(
                host=args.worker_host,
                user=args.worker_user,
                password=args.worker_password,
                recent_minutes=max(1, int(args.recent_minutes)),
            ),
        }
        with metrics_path.open("a", encoding="utf-8") as fp:
            fp.write(json.dumps(row, ensure_ascii=False) + "\n")
        time.sleep(max(30, int(args.interval_seconds)))

    (run_dir / "done.flag").write_text(_now().isoformat(sep=" ", timespec="seconds"), encoding="utf-8")
    return 0


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