"""CLI: python -m wp42.cli <command> [...]"""
from __future__ import annotations
import argparse, json, os, signal, sys
from .config import DB_PATH, ALL_ARMS, PRIMARY_MODEL, NODES
from .store import Store
from . import scheduler, watchdog, export as export_mod, bundle

def main(argv=None):
    ap = argparse.ArgumentParser(prog="wp42"); sub = ap.add_subparsers(dest="cmd", required=True)
    a = sub.add_parser("assign"); a.add_argument("--tasks", nargs="+", required=True); a.add_argument("--arms", nargs="+", default=ALL_ARMS); a.add_argument("--reps", type=int, default=2)
    a.add_argument("--seed", type=int, default=4242); a.add_argument("--model", default=PRIMARY_MODEL); a.add_argument("--stage", default="pilot"); a.add_argument("--model-tag", default=""); a.add_argument("--interrupt-s", type=int, default=None)
    r = sub.add_parser("run"); r.add_argument("--where", default="status='assigned'"); r.add_argument("--parallel", type=int, default=12); r.add_argument("--no-judge", action="store_true"); r.add_argument("--envelope", type=float, default=None)
    s = sub.add_parser("status"); s.add_argument("--json", action="store_true")
    e = sub.add_parser("evaluate"); e.add_argument("--run-ids", nargs="*"); e.add_argument("--where", default="status='collected'"); e.add_argument("--no-judge", action="store_true")
    sub.add_parser("export").add_argument("--tag", default="pilot")
    sub.add_parser("kill-all"); sub.add_parser("sync-tools"); sub.add_parser("ledger")
    i = sub.add_parser("intervention"); i.add_argument("--run-id", required=True); i.add_argument("--minutes", type=float, required=True); i.add_argument("--desc", required=True)
    x = sub.add_parser("reset-run"); x.add_argument("--run-ids", nargs="+", required=True)
    args = ap.parse_args(argv); store = Store(DB_PATH)
    if args.cmd == "assign":
        n = scheduler.make_assignment(store, args.tasks, args.arms, args.reps, args.seed, args.model, args.stage, args.model_tag, args.interrupt_s); print(f"assigned {n} new runs")
    elif args.cmd == "run":
        signal.signal(signal.SIGINT, lambda *_: scheduler.STOP.set())
        kw = {"envelope": args.envelope} if args.envelope else {}
        scheduler.run_queue(store, args.where, args.parallel, do_judge=not args.no_judge, **kw)
    elif args.cmd == "status":
        rows = store.q("SELECT status, arm, COUNT(*) AS n, ROUND(SUM(COALESCE(cost_usd,0)+COALESCE(recovery_cost_usd,0)),2) AS cost FROM runs GROUP BY status, arm ORDER BY status, arm")
        if args.json: print(json.dumps(rows, indent=1))
        else:
            for x in rows: print(f"{x['status']:<12} {x['arm']:<3} n={x['n']:<4} cost={x['cost']}")
            print(f"ledger total: {store.ledger_total():.2f} USD")
    elif args.cmd == "evaluate":
        ids = args.run_ids or [x["run_id"] for x in store.runs(args.where)]
        for rid in ids: scheduler.evaluate_and_judge(store, rid, do_judge=not args.no_judge)
    elif args.cmd == "export":
        print(export_mod.export(store, args.tag))
    elif args.cmd == "kill-all":
        scheduler.STOP.set(); watchdog.kill_everything()
    elif args.cmd == "sync-tools":
        for n in NODES: bundle.sync_tools(n); print(f"{n}: tools synced")
    elif args.cmd == "ledger":
        for x in store.q("SELECT kind, COUNT(*) n, ROUND(SUM(cost_usd),3) usd FROM ledger GROUP BY kind"): print(x)
        print("total", round(store.ledger_total(), 3))
    elif args.cmd == "intervention":
        store.add_intervention(args.run_id, args.minutes, args.desc); print("recorded")
    elif args.cmd == "reset-run":
        for rid in args.run_ids:
            store.exec("DELETE FROM events WHERE run_id=?", (rid,)); store.exec("DELETE FROM receipts WHERE run_id=?", (rid,)); store.exec("DELETE FROM judgments WHERE run_id=?", (rid,))
            store.exec("DELETE FROM metrics WHERE run_id=?", (rid,)); store.exec("DELETE FROM artifacts WHERE run_id=?", (rid,))
            store.update_run(rid, status="assigned", node=None, session_id=None, cost_usd=None, recovery_cost_usd=None, recovery_attempted=0, recovery_success=None, selfcheck_pass_first=None, selfcheck_pass_final=None, notes="reset")
            print("reset", rid)

if __name__ == "__main__":
    main()