#!/usr/bin/env python3
"""Dispatch imba.yml workflow on selected fleet nodes (with SSL retries).
SAFEGATE: requires explicit --node, --limit, or --all.  No args = dry-run + exit.

Usage examples:
  python fleet_dispatch.py              # dry-run: print node count, exit
  python fleet_dispatch.py --node imba-node-0002
  python fleet_dispatch.py --limit 10
  python fleet_dispatch.py --all
"""
import os
KF_BASE = os.path.dirname(os.path.abspath(__file__))
os.makedirs(os.path.join(KF_BASE, "data"), exist_ok=True)
import argparse
import concurrent.futures as cf
import json
import time
import sys

import requests
from requests.adapters import HTTPAdapter, Retry

DEPLOY = os.path.join(KF_BASE, "data/fleet_deploy.json")
H = {"Accept": "application/vnd.github+json"}

s = requests.Session()
s.trust_env = False
s.mount("https://", HTTPAdapter(max_retries=Retry(total=5, backoff_factor=1)))


def dispatch(rec):
    full = f"{rec['login']}/{rec['repo']}"
    h = {**H, "Authorization": f"Bearer {rec['pat']}"}
    for attempt in range(3):
        try:
            r = s.post(
                f"https://api.github.com/repos/{full}/actions/workflows/imba.yml/dispatches",
                headers=h, json={"ref": "main"}, timeout=30)
            if r.status_code in (201, 204):
                return full, True, r.status_code
            return full, False, f"{r.status_code} {r.text[:120]}"
        except Exception as e:
            if attempt == 2:
                return full, False, str(e)[:120]
            time.sleep(3)


def main():
    ap = argparse.ArgumentParser(
        description="Dispatch imba.yml workflow on fleet nodes.",
        epilog="Without --node/--limit/--all: dry-run + exit (no dispatch).")
    target = ap.add_mutually_exclusive_group()
    target.add_argument("--node", metavar="REPO",
                        help="exact repo name or login/repo")
    target.add_argument("--limit", type=int, metavar="N",
                        help="first N ok nodes")
    target.add_argument("--all", action="store_true",
                        help="ALL ok nodes (1391)")
    ap.add_argument("--offset", type=int, default=0,
                    help="starting index")
    ap.add_argument("--dry-run", action="store_true",
                    help="print target count and exit")
    args, _ = ap.parse_known_args()

    recs = [r for r in json.load(open(DEPLOY, encoding="utf-8")) if r.get("ok")]
    if args.node:
        recs = [r for r in recs
                if r.get("repo") == args.node
                or f"{r.get('login')}/{r.get('repo')}" == args.node]
    elif args.limit:
        recs = recs[args.offset:args.offset + args.limit]
    elif args.all:
        recs = recs[args.offset:]  # support offset even with --all
    else:
        # dry-run default when no explicit target
        print(f"fleet_dispatch.py SAFEGATE: no target specified.\n"
              f"  --node REPO | --limit N | --all  required for live dispatch.\n"
              f"  ok nodes available: {len(recs)}\n"
              f"  (dry-run; use --dry-run to see target count)")
        sys.exit(0)

    if args.dry_run or not (args.node or args.limit or args.all):
        print(f"DRY-RUN: would dispatch {len(recs)} nodes")
        return

    if not recs:
        sys.exit("no nodes selected (--node/--limit matched nothing)")
    print(f"dispatching {len(recs)} nodes", flush=True)
    ok = 0
    t0 = time.time()
    with cf.ThreadPoolExecutor(max_workers=12) as ex:
        for i, (full, good, info) in enumerate(ex.map(dispatch, recs), 1):
            ok += good
            if not good:
                print(f"  FAIL {full}: {info}", flush=True)
            if i % 50 == 0:
                # ":.0s" is not a format code: every dispatch died here on the
                # final summary (and on the per-50 progress line), after the
                # ex.map loop had already run, so runs were dispatched but the
                # ok/total was never printed and the exit code was a traceback.
                print(f"  [{i}/{len(recs)}] ok={ok} {time.time()-t0:.0f}s", flush=True)
    print(f"\nDONE {time.time()-t0:.0f}s: {ok}/{len(recs)} dispatched", flush=True)


if __name__ == "__main__":
    main()