"""Command-line interface of the PGAR harness.
  pgar assign  --split pilot --tasks F1-dev-01,... --repeats 1 --seed 1
  pgar run     --split pilot [--runs R...] [--workers N]
  pgar status  [--split ...]
  pgar export
  pgar probe   (sentinel isolation probes on all nodes)
  pgar deploy  (install gate + env file on nodes)
"""
from __future__ import annotations
import argparse, json, os, sys, yaml
from .state import StateStore
from .scheduler import Controller
from .exports import export_runs
from . import isolation, launcher

REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))

def load_cfg(split):
    base = yaml.safe_load(open(os.path.join(REPO, "configs", "budgets.yaml")))
    cfg = yaml.safe_load(open(os.path.join(REPO, "configs", f"{split}.yaml")))
    out = dict(base); out.update(cfg); return out

def get_store():
    return StateStore(os.path.join(REPO, "runs", "state.sqlite"))

def main(argv=None):
    ap = argparse.ArgumentParser(prog="pgar"); sub = ap.add_subparsers(dest="cmd", required=True)
    a = sub.add_parser("assign"); a.add_argument("--split", required=True); a.add_argument("--tasks", required=True); a.add_argument("--repeats", type=int, default=1); a.add_argument("--seed", type=int, required=True); a.add_argument("--arms", default="A,B,C,D"); a.add_argument("--start-seq", type=int, default=None)
    r = sub.add_parser("run"); r.add_argument("--split", required=True); r.add_argument("--runs", default=None); r.add_argument("--workers", type=int, default=None); r.add_argument("--fault", default=None, help="JSON fault spec applied to all selected runs"); r.add_argument("--rerun-errors", action="store_true", help="also select runs in status 'error' (infrastructure rerun policy)")
    s = sub.add_parser("status"); s.add_argument("--split", default=None)
    sub.add_parser("export")
    sub.add_parser("probe")
    sub.add_parser("deploy")
    e = sub.add_parser("evaluate"); e.add_argument("--runs", required=True)
    args = ap.parse_args(argv)
    store = get_store()
    if args.cmd == "assign":
        cfg = load_cfg(args.split); ctl = Controller(REPO, store, cfg, cfg["nodes"])
        existing = store.query("SELECT COUNT(*) AS n FROM runs WHERE split=?", (args.split,))[0]["n"]
        runs = ctl.make_assignment(args.split, args.tasks.split(","), args.repeats, args.seed, arms=tuple(args.arms.split(",")), start_seq=(args.start_seq or existing + 1))
        ctl.register(runs)
        path = os.path.join(REPO, "runs", args.split, f"assignment_seed{args.seed}.json"); os.makedirs(os.path.dirname(path), exist_ok=True)
        json.dump([json.loads(m.to_json()) for m in runs], open(path, "w"), indent=1)
        print(f"registered {len(runs)} runs for split {args.split}: {runs[0].run_id}..{runs[-1].run_id} -> {path}")
    elif args.cmd == "run":
        cfg = load_cfg(args.split); ctl = Controller(REPO, store, cfg, cfg["nodes"])
        ids = args.runs.split(",") if args.runs else [m["run_id"] for m in store.list_runs(split=args.split) if m["status"] == "not_started" or (args.rerun_errors and m["status"] == "error")]
        faults = {rid: json.loads(args.fault) for rid in ids} if args.fault else None
        res = ctl.run_many(ids, max_workers=args.workers, faults=faults)
        print(json.dumps([{k: v for k, v in x.items() if k != "evaluation"} for x in res], indent=1, default=str))
        export_runs(store, REPO)
    elif args.cmd == "status":
        for m in store.list_runs(split=args.split):
            ev = store.query("SELECT scientific_success, n_violations FROM evaluations WHERE run_id=?", (m["run_id"],))
            u = m.get("usage") or {}
            print(f"{m['run_id']:8} {m['task_id']:12} {m['treatment']} r{m['repeat_index']} {m.get('node') or '-':8} {m['status']:16} cost={u.get('spent_usd', 0):.3f} success={ev[0]['scientific_success'] if ev else '-'} viol={ev[0]['n_violations'] if ev else '-'} {m.get('terminal_reason') or ''}")
    elif args.cmd == "export":
        print(export_runs(store, REPO))
    elif args.cmd == "probe":
        cfg = load_cfg("development")
        for n in cfg["nodes"]:
            print(json.dumps(isolation.sentinel_probe(n, "~/wp29")))
    elif args.cmd == "deploy":
        cfg = load_cfg("development")
        for n in cfg["nodes"]:
            launcher.ssh(n, "mkdir -p ~/wp29/gates ~/wp29/runs && chmod 755 ~/wp29/gates", check=True)
            import subprocess
            subprocess.run(["scp", "-q", "-o", "BatchMode=yes", os.path.join(REPO, "src", "pgar", "runtime_gates.py"), f"{n}:~/wp29/gates/pgar_gate.py"], check=True)
            launcher.ssh(n, "chmod 444 ~/wp29/gates/pgar_gate.py && shasum -a 256 ~/wp29/gates/pgar_gate.py", check=True)
            print(n, "gate deployed")
    elif args.cmd == "evaluate":
        from .evaluation import evaluate_run
        for rid in args.runs.split(","):
            print(rid, json.dumps({k: evaluate_run(REPO, store, rid).get(k) for k in ("scientific_success", "quality_score")}))
        export_runs(store, REPO)

if __name__ == "__main__":
    main()