#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Run one budget tier from configs/budgets.yaml (sequence of `av run-arm` runs, resumable, capped).

Usage: python3 scripts/33_run_tier.py <tier> [--live] [--dry-run]
Each run gets its own ledger (results/ledger/<run_id>.jsonl); the tier cap is enforced on the sum of the
ledgers of the tier's runs (estimated list-price spend).
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from agentic_valuation.config import RESULTS, load_config  # noqa: E402
from agentic_valuation.experiments.run_arm import RunSpec, run_arm  # noqa: E402


def tier_spent(run_ids: list[str]) -> float:
    tot = 0.0
    for rid in run_ids:
        f = RESULTS / "ledger" / f"{rid}.jsonl"
        if f.exists():
            for line in f.read_text().splitlines():
                try:
                    tot += float(json.loads(line).get("cost_usd", 0.0))
                except Exception:
                    pass
    return tot


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("tier")
    ap.add_argument("--live", action="store_true")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--n", type=int, default=None, help="override n for every run")
    ap.add_argument("--only", default=None, help="comma-separated substrings; keep runs whose run_id contains any of them")
    ap.add_argument("--concurrency", type=int, default=6, help="properties processed concurrently per run")
    a = ap.parse_args(argv)
    cfg = load_config("budgets")
    t = cfg["tiers"][a.tier]
    model_tier = t.get("tier", cfg.get("tier_default", "mid"))
    specs: list[RunSpec] = []
    for r in t["runs"]:
        for c in r["committees"]:
            n = a.n or r.get("n", t["n"])
            rounds = r.get("rounds")
            rid = "_".join([a.tier, r["arm"] + (f"_r{rounds}" if rounds is not None else ""), c, model_tier]) + ("" if a.live else "_mock")
            if not a.live:
                rid = "_" + rid
            specs.append(RunSpec(arm=r["arm"], rounds=rounds, committee=c, split=t["split"], n=n, live=a.live, tier=model_tier, run_id=rid, property_concurrency=a.concurrency))
    if a.only:
        keys = [k.strip() for k in a.only.split(",") if k.strip()]
        specs = [s for s in specs if any(k in s.run_id for k in keys)]
    all_ids = []
    for r in t["runs"]:
        for c in r["committees"]:
            rr = r.get("rounds")
            all_ids.append("_".join([a.tier, r["arm"] + (f"_r{rr}" if rr is not None else ""), c, model_tier]) + ("" if a.live else "_mock"))
    ids = [s.run_id for s in specs]
    print(f"[tier {a.tier}] {len(specs)} runs, split={t['split']}, model tier={model_tier}, cap={t['cap_usd']} USD, spent so far={tier_spent(ids):.2f}")
    for s in specs:
        print(f"  - {s.run_id}: n={s.n}")
    if a.dry_run:
        return
    for s in specs:
        spent = tier_spent(all_ids)
        if spent >= t["cap_usd"]:
            print(f"[tier {a.tier}] cap reached ({spent:.2f} >= {t['cap_usd']}); stopping before {s.run_id}")
            break
        import os

        os.environ["AV_MAX_SPEND_USD"] = str(max(t["cap_usd"] - spent, 1.0))
        run_arm(s)
    print(f"[tier {a.tier}] done; estimated spend of the whole tier {tier_spent(all_ids):.2f} USD")
    (RESULTS / f"tier_{a.tier}.json").write_text(json.dumps({"tier": a.tier, "runs": ids, "spent_usd": tier_spent(ids), "live": a.live}, indent=1))


if __name__ == "__main__":
    main()