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

import argparse
import sys
from datetime import datetime, timedelta
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


def main() -> None:
    parser = argparse.ArgumentParser(description="Prune stale/offline cluster node records from detect_worker_nodes.")
    parser.add_argument("--minutes", type=int, default=30, help="delete nodes whose heartbeat is older than this many minutes")
    parser.add_argument("--node-code", action="append", default=[], help="delete a specific node code; can be passed multiple times")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    cutoff = datetime.now() - timedelta(minutes=max(1, int(args.minutes or 30)))
    removed: list[dict] = []
    with get_db() as conn:
        with conn.cursor() as cur:
            if args.node_code:
                cur.execute(
                    """
                    SELECT node_code, region, role, status, last_heartbeat_at
                    FROM detect_worker_nodes
                    WHERE node_code = ANY(%s)
                    ORDER BY node_code ASC
                    """,
                    (list(args.node_code),),
                )
            else:
                cur.execute(
                    """
                    SELECT node_code, region, role, status, last_heartbeat_at
                    FROM detect_worker_nodes
                    WHERE last_heartbeat_at < %s
                    ORDER BY last_heartbeat_at ASC, node_code ASC
                    """,
                    (cutoff,),
                )
            rows = cur.fetchall()
            for row in rows:
                removed.append(
                    {
                        "node_code": row[0],
                        "region": row[1],
                        "role": row[2],
                        "status": row[3],
                        "last_heartbeat_at": row[4].isoformat(sep=" ", timespec="seconds") if row[4] else "",
                    }
                )
            if removed and not args.dry_run:
                cur.execute(
                    "DELETE FROM detect_worker_nodes WHERE node_code = ANY(%s)",
                    ([item["node_code"] for item in removed],),
                )
        if removed and not args.dry_run:
            conn.commit()

    if not removed:
        print("no cluster nodes matched prune conditions")
        return

    print("matched cluster nodes:")
    for item in removed:
        print(
            f"- {item['node_code']} | {item['region']} | {item['role']} | {item['status']} | {item['last_heartbeat_at']}"
        )
    if args.dry_run:
        print("dry-run only, nothing deleted")
    else:
        print(f"deleted {len(removed)} cluster node record(s)")


if __name__ == "__main__":
    main()