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

import argparse
import json
import sys
import time
from pathlib import Path

API_ROOT = Path(__file__).resolve().parents[2]
if str(API_ROOT) not in sys.path:
    sys.path.insert(0, str(API_ROOT))

from app.core.db import get_db
from app.services.cluster_runtime_service import ensure_runtime_schema, register_node_heartbeat


def _metadata_from_args(args: argparse.Namespace) -> dict:
    metadata = {
        "service": "cluster-simulator",
        "simulated": True,
        "phase": args.phase or "",
        "detail": args.detail or "",
    }
    if args.job_id:
        metadata["job_id"] = args.job_id
    if args.job_code:
        metadata["job_code"] = args.job_code
    if args.cycle_token:
        metadata["cycle_token"] = args.cycle_token
    if args.active_threads is not None:
        metadata["active_threads"] = args.active_threads
    if args.available_proxy_count is not None:
        metadata["available_proxy_count"] = args.available_proxy_count
    if args.metadata_json:
        try:
            custom = json.loads(args.metadata_json)
            if isinstance(custom, dict):
                metadata.update(custom)
        except Exception as exc:
            raise SystemExit(f"invalid --metadata-json: {exc}") from exc
    return metadata


def _delete_node(node_code: str) -> None:
    with get_db() as conn:
        with conn.cursor() as cur:
            cur.execute("DELETE FROM detect_worker_nodes WHERE node_code = %s", (str(node_code),))
        conn.commit()


def main() -> None:
    parser = argparse.ArgumentParser(description="Simulate a cluster node heartbeat for multi-region rehearsal.")
    parser.add_argument("--node-code", required=True)
    parser.add_argument("--region", required=True)
    parser.add_argument("--role", required=True, choices=("control", "worker"))
    parser.add_argument("--status", default="online")
    parser.add_argument("--current-load", type=int, default=0)
    parser.add_argument("--phase", default="")
    parser.add_argument("--detail", default="")
    parser.add_argument("--job-id", type=int, default=0)
    parser.add_argument("--job-code", default="")
    parser.add_argument("--cycle-token", default="")
    parser.add_argument("--active-threads", type=int, default=None)
    parser.add_argument("--available-proxy-count", type=int, default=None)
    parser.add_argument("--metadata-json", default="")
    parser.add_argument("--interval", type=int, default=20)
    parser.add_argument("--iterations", type=int, default=0, help="0 means run forever")
    parser.add_argument("--clear", action="store_true", help="delete the node record and exit")
    args = parser.parse_args()

    ensure_runtime_schema()
    if args.clear:
        _delete_node(args.node_code)
        print(f"cleared simulated node: {args.node_code}")
        return

    metadata = _metadata_from_args(args)
    interval = max(5, int(args.interval or 20))
    iterations = max(0, int(args.iterations or 0))
    current = 0
    while True:
        current += 1
        register_node_heartbeat(
            node_code=args.node_code,
            region=args.region,
            role=args.role,
            status=args.status,
            current_load=max(0, int(args.current_load or 0)),
            metadata=metadata,
        )
        print(
            json.dumps(
                {
                    "node_code": args.node_code,
                    "region": args.region,
                    "role": args.role,
                    "status": args.status,
                    "current_load": max(0, int(args.current_load or 0)),
                    "iteration": current,
                    "metadata": metadata,
                },
                ensure_ascii=False,
            ),
            flush=True,
        )
        if iterations and current >= iterations:
            break
        time.sleep(interval)


if __name__ == "__main__":
    main()