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

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

import paramiko
import psycopg2
import requests


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


def _safe_json_response(url: str) -> dict:
    try:
        response = requests.get(url, timeout=8)
        payload = response.json()
        if isinstance(payload, dict):
            return payload.get("data") or payload
    except Exception as exc:
        return {"error": str(exc), "url": url}
    return {}


def _query_local_db() -> dict:
    try:
        conn = psycopg2.connect(
            host="127.0.0.1",
            port=5432,
            dbname="domain",
            user="postgres",
            password="Qazwe123,./",
        )
        with conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    select node_code, status, current_load, last_heartbeat_at
                    from detect_worker_nodes
                    where node_code in ('mainland-controller-01', 'mainland-worker-01')
                    order by node_code
                    """
                )
                worker_nodes = [
                    {
                        "node_code": row[0],
                        "status": row[1],
                        "current_load": int(row[2] or 0),
                        "last_heartbeat_at": row[3].isoformat(sep=" ", timespec="seconds") if row[3] else "",
                    }
                    for row in cur.fetchall()
                ]
        conn.close()
        return {"worker_nodes": worker_nodes}
    except Exception as exc:
        return {"error": str(exc)}


def _run_remote_sample(host: str, user: str, password: str) -> dict:
    if not str(host or "").strip() or not str(user or "").strip() or not str(password or "").strip():
        return {"skipped": True, "reason": "remote ssh credentials not provided"}
    cli = paramiko.SSHClient()
    cli.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        cli.connect(hostname=host, username=user, password=password, timeout=10)
        command = r"""printf 'proc='; pgrep -fc '[d]etect_worker.py'
printf 'load='; ps -eo pcpu,cmd | grep '[d]etect_worker.py' | awk '{s+=$1} END {print s+0}'
printf 'db_err_recent='; journalctl -u domaincheck-worker -u 'domaincheck-worker@*' --since '5 min ago' --no-pager | egrep -c '初始化数据库连接池失败|连接池耗尽|等待超时|数据库连接失败' || true
cd /www/wwwroot/getDomain/domain-api && /opt/domaincheck/domainCheck/.venv/bin/python - <<'INNER'
import json
from app.services.detect_service import get_detect_status
result = get_detect_status() or {}
print('detect_status=' + json.dumps({
    'job_code': result.get('job_code'),
    'active_thread_count': result.get('active_thread_count'),
    'max_thread_count': result.get('max_thread_count'),
    'aggregate_process_count': result.get('aggregate_process_count'),
    'display_items_running': result.get('display_items_running'),
    'display_active_threads': result.get('display_active_threads'),
    'display_max_threads': result.get('display_max_threads'),
}, ensure_ascii=False))
INNER"""
        _, stdout, stderr = cli.exec_command(command, timeout=45)
        output = stdout.read().decode("utf-8", errors="replace")
        error = stderr.read().decode("utf-8", errors="replace").strip()
        sample: dict[str, object] = {"raw": output.strip()}
        if error:
            sample["stderr"] = error
        for line in output.splitlines():
            if "=" not in line:
                continue
            key, value = line.split("=", 1)
            if key == "detect_status":
                try:
                    sample[key] = json.loads(value)
                except Exception:
                    sample[key] = value
                continue
            try:
                sample[key] = int(value)
            except Exception:
                try:
                    sample[key] = float(value)
                except Exception:
                    sample[key] = value
        return sample
    except Exception as exc:
        return {"error": str(exc)}
    finally:
        try:
            cli.close()
        except Exception:
            pass


def main() -> int:
    parser = argparse.ArgumentParser(description="Collect overnight worker metrics into JSONL.")
    parser.add_argument("--worker-host", default="")
    parser.add_argument("--worker-user", default="")
    parser.add_argument("--worker-password", default="")
    parser.add_argument("--interval-seconds", type=int, default=60)
    parser.add_argument("--duration-hours", type=float, default=6.0)
    parser.add_argument("--label", default="mainland-worker-01-night")
    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"night_run_{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),
                "worker_host": args.worker_host,
                "remote_sampling_enabled": bool(
                    str(args.worker_host or "").strip()
                    and str(args.worker_user or "").strip()
                    and str(args.worker_password or "").strip()
                ),
            },
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
    )

    base_url = "http://127.0.0.1:8100/api/v1"
    sample_index = 0
    while _now() < end_time:
        sample_index += 1
        timestamp = _now().isoformat(sep=" ", timespec="seconds")
        row = {
            "timestamp": timestamp,
            "sample_index": sample_index,
            "local": {
                "detect_status": _safe_json_response(f"{base_url}/detect/status"),
                "runtime_status": _safe_json_response(f"{base_url}/runtime/status"),
                "dashboard_overview": _safe_json_response(f"{base_url}/dashboard/overview"),
                "db": _query_local_db(),
            },
            "remote_worker": _run_remote_sample(
                host=args.worker_host,
                user=args.worker_user,
                password=args.worker_password,
            ),
        }
        with metrics_path.open("a", encoding="utf-8") as fp:
            fp.write(json.dumps(row, ensure_ascii=False) + "\n")
        time.sleep(max(10, 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())