"""`ma` command line: run one config, run a design, analyze, estimate budget."""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path


def main(argv=None) -> None:
    ap = argparse.ArgumentParser(prog="ma")
    sub = ap.add_subparsers(dest="cmd", required=True)
    r = sub.add_parser("run", help="run one configuration")
    r.add_argument("config"); r.add_argument("--seed", type=int); r.add_argument("--out"); r.add_argument("--set", nargs="*", default=[])
    r.add_argument("--cache"); r.add_argument("--progress", action="store_true")
    d = sub.add_parser("design", help="run a design file (cells x seeds)")
    d.add_argument("design"); d.add_argument("--workers", type=int, default=4); d.add_argument("--cache"); d.add_argument("--cells", nargs="*")
    d.add_argument("--max-seeds", type=int); d.add_argument("--root", default="results/runs")
    b = sub.add_parser("budget", help="estimate API calls and USD for a design")
    b.add_argument("design"); b.add_argument("--tokens-in", type=int, default=800); b.add_argument("--tokens-out", type=int, default=250)
    args = ap.parse_args(argv)
    from dotenv import load_dotenv
    load_dotenv(Path(".env"), override=False)
    if args.cmd == "run":
        from .experiments.config import RunConfig
        from .experiments.runner import run_config
        from .metrics.outcomes import run_outcomes
        ov = {}
        for kv in args.set:
            k, v = kv.split("=", 1)
            ov[k] = json.loads(v) if v[:1] in "[{\"0123456789-tfn" and v not in ("true", "false") else v
            if v in ("true", "false"):
                ov[k] = v == "true"
        cfg = RunConfig.from_yaml(args.config, **ov)
        if args.seed is not None:
            cfg = cfg.with_overrides(seed=args.seed)
        res = run_config(cfg, out_dir=args.out, cache_path=args.cache, progress=args.progress)
        oc = run_outcomes(res)
        print(json.dumps({k: (round(v, 4) if isinstance(v, float) else v) for k, v in oc.items()}, indent=1, default=str))
    elif args.cmd == "design":
        from .experiments.design import run_design
        df = run_design(args.design, args.root, args.workers, args.cache, args.cells, args.max_seeds)
        print(df.groupby("name")[[c for c in ("msd_log", "spread_bp", "n_bubbles", "llm_usd") if c in df]].mean().round(4))
    elif args.cmd == "budget":
        from .experiments.design import cell_config, expand_cells, load_design, seeds_of
        design = load_design(args.design)
        total_calls, total_usd = 0, 0.0
        for c in expand_cells(design):
            for s in seeds_of(design):
                cfg = cell_config(design, c, s, Path(args.design).parent)
                p = cfg.population
                n_llm = int(round(p.llm_share * p.n_traders))
                n_dec = (cfg.n_ticks - cfg.warmup // 2) // p.llm.decision_interval
                if p.llm.decision_cap:
                    n_dec = min(n_dec, p.llm.decision_cap)
                calls = n_llm * n_dec
                usd = 0.0
                for prov, share in p.llm.provider_mix.items():
                    model = p.llm.models.get(prov, "")
                    pr = p.llm.price_assumptions.get(model, {"input": 0, "output": 0})
                    usd += calls * share * (args.tokens_in * pr["input"] + args.tokens_out * pr["output"]) / 1e6
                total_calls += calls
                total_usd += usd
        print(f"design {design['name']}: {total_calls} LLM calls, est. ${total_usd:.2f} at {args.tokens_in} in / {args.tokens_out} out tokens per call")


if __name__ == "__main__":
    main(sys.argv[1:])