"""Run the plan's jobs against every actor through apify-cache, in parallel; save everything raw.

  uv run --project ~/prj/hyperbach/apify-cache python run_tests.py --plan plan.json [task ...]

Writes <out>/raw/<user__name>/<task>.{input,run,items}.json and <out>/raw/manifest.jsonl.
"""
import argparse, json, sys, time, datetime as dt
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))
from common import load_plan, slug_dir, apify_token
from apify_cache import run_cached, core

CALLER = "apifystats-comparison-post"


def fetch_run(run_id):
    return core._get(f"{core.BASE}/actor-runs/{run_id}?token={apify_token()}")["data"]


def one(plan, slug, task):
    cfg = plan["actors"][slug]
    actor, run_input = cfg["id"], cfg["inputs"][task]
    d = slug_dir(plan, slug)
    (d / f"{task}.input.json").write_text(json.dumps(run_input, indent=1))
    t0 = time.time()
    rec = {"slug": slug, "actor_id": actor, "task": task, "started_local": dt.datetime.now(dt.timezone.utc).isoformat()}
    try:
        # per-actor cap/memory (plan actors[slug].cap_usd / .memory_mb) override the plan-wide cap
        cap = (cfg.get("task_caps") or {}).get(task, cfg.get("cap_usd", plan.get("cap_usd", 2.0)))
        mem = (cfg.get("task_memory") or {}).get(task, cfg.get("memory_mb"))
        res = run_cached(actor, run_input, refresh=True, caller=CALLER, max_wait=plan.get("max_wait", 900), cap_usd=cap, memory_mb=mem)
        rec.update(run_id=res.run_id, cached=res.cached, cost_usd_first=res.cost_usd, n_items=len(res.items), status="SUCCEEDED", call_id=res.call_id)
        (d / f"{task}.items.json").write_text(json.dumps(res.items, ensure_ascii=False, indent=1))
    except core.ApifyError as e:
        rec.update(status="FAILED", error=str(e))
        parts = str(e).split(" run ")
        if len(parts) > 1:
            rec["run_id"] = parts[1].split(" ")[0]
    except Exception as e:  # noqa: BLE001
        rec.update(status="ERROR", error=repr(e))
    rec["wall_s"] = round(time.time() - t0, 1)
    if rec.get("run_id"):
        try:
            run = fetch_run(rec["run_id"])
            (d / f"{task}.run.json").write_text(json.dumps(run, indent=1))
            rec.update(run_status=run.get("status"), duration_s=round((run.get("stats") or {}).get("durationMillis", 0) / 1000, 1),
                       usage_total_usd=run.get("usageTotalUsd"), build_number=run.get("buildNumber"), memory_mb=(run.get("options") or {}).get("memoryMbytes"))
            if rec["status"] != "SUCCEEDED":
                try:
                    items = core.fetch_all_items(run["defaultDatasetId"], apify_token())
                    (d / f"{task}.items.json").write_text(json.dumps(items, ensure_ascii=False, indent=1)); rec["n_items"] = len(items)
                except Exception:
                    pass
        except Exception as e:  # noqa: BLE001
            rec["run_fetch_error"] = repr(e)
    with (plan["_raw"] / "manifest.jsonl").open("a") as f:
        f.write(json.dumps(rec) + "\n")
    print(json.dumps(rec), flush=True)
    return rec


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--plan", required=True)
    ap.add_argument("tasks", nargs="*")
    ap.add_argument("--wave", type=int, default=None, help="only actors whose plan entry has this wave number")
    a = ap.parse_args()
    plan = load_plan(a.plan)
    tasks = a.tasks or plan["tasks"]
    plan["_raw"].mkdir(parents=True, exist_ok=True)
    actors = [s for s in plan["actors"] if a.wave is None or plan["actors"][s].get("wave") == a.wave]
    jobs = [(s, t) for s in actors for t in tasks if t in plan["actors"][s]["inputs"]]
    print(f"{len(jobs)} runs, tasks={tasks}, out={plan['_raw']}", flush=True)
    with ThreadPoolExecutor(max_workers=len(jobs)) as ex:
        list(ex.map(lambda j: one(plan, *j), jobs))


if __name__ == "__main__":
    main()