"""Online-Mind2Web on live sites: the jev agent + laya on the tasks whose domain never appears in training.

    python apps/browser_suite_om2w.py [out.jsonl=finetune/out/om2w_run.jsonl]
    env: OM2W_LEVELS=easy,medium,hard  OM2W_LIMIT=0  OM2W_MAX_STEPS=30  SHARD=i/n
         (+ the usual jev switches: JEV_VERIFY_DONE / JEV_ESCALATE / JEV_PLANNER ...)

Resumable: tasks already in the output are skipped.  Every trajectory keeps what a judge needs (the action history with
URLs and typed text, the final URL / title / visible text); judge them with finetune/judge_om2w.py.
"""
import json, os, sys, time, urllib.parse
sys.path.insert(0, "/home/ckl/projects/S/jev-ultrafast")
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import browser_suite  # noqa: F401  (jev env: CDP, laya server, text model)
from jev_ultrafast import Agent

L = "/home/ckl/projects/S/laya/finetune/out"


def base(host):
    return ".".join((host or "").split(".")[-2:])


def held_out_tasks():
    tasks = json.load(open(f"{L}/online_mind2web.json"))
    trained = {base(d) for d in json.load(open(f"{L}/train_domains.json"))}
    # tasks whose completion would act on a real third party (message a seller, file a government request, send a gift
    # card) are not run on live sites; JEV_SAFE=1 additionally hides pay / order / send / submit-request controls
    side_effects = ("contact the cheapest", "Submit a request for vehicle registration", "send Christene")
    # OM2W_SITES (finetune/probe_sites.py output): only sites that load for our headless browser; the rest show an
    # anti-bot wall or "access denied" and would measure the wall, not the agent
    sites = json.load(open(os.environ["OM2W_SITES"])) if os.environ.get("OM2W_SITES") else None
    return [t for t in tasks if base(urllib.parse.urlparse(t["website"]).hostname) not in trained
            and not any(x in t["confirmed_task"] for x in side_effects) and (sites is None or sites.get(t["website"], {}).get("ok"))]


def run(task, max_steps):
    t0 = time.time(); status, hist, final, extra = "error", [], {}, {}
    try:
        with Agent(task["website"], task["confirmed_task"]) as agent:
            for state in agent.run():
                status = state["status"]
                if len(state["history"]) >= max_steps:
                    break
            s = agent.state
            hist = [{k: h.get(k) for k in ("action", "kind", "text", "url", "page_changed")} for h in s["history"]]
            p = s["page"]
            final = {"url": p["url"], "title": p["title"], "text": p["text"][:4000]}
            extra = {k: s.get(k, 0) for k in ("escalations", "done_rejections", "select_overrides")}
            status = s["status"]
    except Exception as e:
        status = f"error:{type(e).__name__}:{str(e)[:80]}"
    return {"task_id": task["task_id"], "level": task["level"], "website": task["website"], "task": task["confirmed_task"],
            "status": status, "steps": len(hist), "wall": round(time.time() - t0, 1), "history": hist, "final": final, **extra}


def main():
    out = sys.argv[1] if len(sys.argv) > 1 else f"{L}/om2w_run.jsonl"
    levels = os.environ.get("OM2W_LEVELS", "easy,medium,hard").split(",")
    tasks = [t for t in held_out_tasks() if t["level"] in levels]
    if os.environ.get("SHARD"):
        i, n = map(int, os.environ["SHARD"].split("/")); tasks = tasks[i::n]
    if int(os.environ.get("OM2W_LIMIT", "0")):
        tasks = tasks[:int(os.environ["OM2W_LIMIT"])]
    done = {json.loads(l)["task_id"] for l in open(out)} if os.path.exists(out) else set()
    max_steps = int(os.environ.get("OM2W_MAX_STEPS", "30"))
    print(f"== {len(tasks)} held-out tasks ({len(done)} already run)", flush=True)
    with open(out, "a") as f:
        for t in tasks:
            if t["task_id"] in done:
                continue
            r = run(t, max_steps)
            f.write(json.dumps(r, ensure_ascii=False) + "\n"); f.flush()
            print(f"  {r['level']:6s} {r['status'][:24]:24s} steps={r['steps']:2d} esc={r.get('escalations', 0)} {r['wall']:5.1f}s  {t['confirmed_task'][:80]}", flush=True)


if __name__ == "__main__":
    main()