"""Run the real agent (jev loop + the laya systemone server) on HELD-OUT webgym seeds (seed % 10 == 0, never generated
for training) and judge success from the page's own state.

    python finetune/webgym/eval_gym.py [n_per_kind=10] [kinds=flight,hotel,shop,car,restaurant,signup,filter] [max_steps=40]
    env: JEV_PLANNER=1 to plan sub-goals first;  JEV_MACROS=1 date/count widgets as one action;  S1_URL=http://...;  GYM_OUT=results.json

Success: flight/hotel/car = on the results page with every target value (+ filter/sort/opened car), shop/filter = the
listing/product state the task asks for, restaurant/signup = the confirmation page carries every target value.  The agent's DONE is not trusted; a run that ends without reaching the state fails.
"""
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__)))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "apps"))
import browser_suite  # noqa: F401  (sets the jev env: CDP, systemone URL, text model)
if os.environ.get("S1_URL"):   # a second laya server (e.g. another checkpoint while the main one is busy)
    os.environ["TYPESAFE_BASE_URL"] = os.environ["S1_URL"]
from jev_ultrafast import Agent
import spec as S

BASE = os.environ.get("GYM_URL", "http://127.0.0.1:8811")


def same(sp, key, got, want):
    """Value equality with the widget's semantics (autocomplete keys / free text, lists joined by '|', booleans as strings)."""
    f = next((x for x in sp.get("fields", []) if x["key"] == key), None)
    if f and f["w"] == "combo":
        if not isinstance(got, str): return False
        if got == want or got == want + "-ap": return True
        opt = next((o for o in f["options"] if o["key"] == want), None)
        name = (opt or {}).get("name") or (opt or {}).get("short") or want
        return got.startswith("text:") and name.lower() in got.lower()
    if isinstance(want, list):
        return sorted(got if isinstance(got, list) else [x for x in str(got or "").split("|") if x]) == sorted(want)
    if isinstance(want, bool):
        return str(got).lower() == str(want).lower()
    return str(got) == str(want)


def success(sp, url, meta):
    meta = meta or {}
    if sp["kind"] == "car":
        res = meta.get("results")
        if res is None or not all(same(sp, k, res.get(k), v) for k, v in sp["target"].items()):
            return False
        post = sp["post"]
        if "ctype" in post and post["ctype"] not in (res.get("ctype") or "").split("|"): return False
        if "trans" in post and res.get("trans") != post["trans"]: return False
        if "csort" in post and res.get("sort") != post["csort"]: return False
        if "open" in post and meta.get("car") != post["open"]: return False
        return True
    if sp["kind"] in ("restaurant", "signup"):
        car = meta.get("carried") or {}
        return bool(meta.get("done")) and all(same(sp, k, car.get(k), v) for k, v in sp["target"].items())
    if sp["kind"] == "filter":
        tg = sp["target"]
        if "open" in tg:
            return meta.get("product") == tg["open"]
        lst = meta.get("listing") or {}
        return (all(sorted(lst.get(k) or []) == sorted(tg[k]) for k in S.FACETS) and (lst.get("pmin") or None) == tg["pmin"] and (lst.get("pmax") or None) == tg["pmax"]
                and (lst.get("rating") or None) == tg["rating"] and (lst.get("sort") or None) == tg["sort"] and lst.get("page") == tg["page"])
    if sp["kind"] in ("flight", "hotel"):
        res = (meta or {}).get("results")
        if res is None:
            return False
        for k, v in sp["target"].items():
            got = (res.get(k) or "")
            if isinstance(v, str) and len(v) == 3 and v.isalpha():
                name = next(c[1] for c in S.CITIES if c[0] == v)
                if got.replace("-ap", "") != v and not (got.startswith("text:") and name.lower() in got.lower()): return False
            elif str(got) != str(v):
                return False
        post = sp.get("post") or {}
        if "filter" in post and post["filter"] not in (res.get("filter") or "").split("|"): return False
        if "sort" in post and res.get("sort") != post["sort"]: return False
        return True
    post = sp["post"]
    if "open" in post:
        return (meta or {}).get("product") == post["open"]
    lst = (meta or {}).get("listing") or {}
    return all((lst.get(k) or (1 if k == "page" else None)) == v if k != "q" else (lst.get("q") or "").lower() == v for k, v in post.items())


def run(kind, seed, max_steps):
    sp = S.make(kind, seed)
    t0 = time.time(); steps = 0; status = "error"; meta = None; url = ""; plan = None
    try:
        with Agent(f"{BASE}/t/{kind}/{seed}", sp["goal"]) as agent:
            plan = agent.state.get("plan")
            for state in agent.run():
                steps = len(state["history"]); status = state["status"]
                if steps >= max_steps: break
            time.sleep(0.8)
            meta = json.loads(agent.browser.evaluate("JSON.stringify(window.__gym ? window.__gym.meta : null)") or "null")
            url = agent.browser.evaluate("location.href")
    except Exception as e:
        status = f"error:{type(e).__name__}:{str(e)[:60]}"
    ok = success(sp, url, meta)
    return {"kind": kind, "seed": seed, "ok": ok, "steps": steps, "status": status, "wall": round(time.time() - t0, 1), "url": url,
            "goal": sp["goal"], "plan": plan}


def main():
    n = int(sys.argv[1]) if len(sys.argv) > 1 else 10
    kinds = (sys.argv[2] if len(sys.argv) > 2 else "flight,hotel,shop").split(",")
    max_steps = int(sys.argv[3]) if len(sys.argv) > 3 else 40
    rows = []
    for kind in kinds:
        for i in range(n):
            seed = 900000 + 10 * i          # held out: never generated for training (and disjoint from training ranges)
            r = run(kind, seed, max_steps); rows.append(r)
            print(f"{'PASS' if r['ok'] else 'FAIL'}  {kind:6s} seed={seed} steps={r['steps']:2d} {r['status'][:28]:28s} {r['wall']:5.1f}s  {r['goal'][:70]}", flush=True)
    for kind in kinds:
        k = [r for r in rows if r["kind"] == kind]
        print(f"== {kind}: {sum(r['ok'] for r in k)}/{len(k)}  median steps {sorted(r['steps'] for r in k)[len(k)//2]}")
    print(f"== total {sum(r['ok'] for r in rows)}/{len(rows)} ({100*sum(r['ok'] for r in rows)/max(1,len(rows)):.0f}%)")
    json.dump(rows, open(os.environ.get("GYM_OUT", "/tmp/gym_eval.json"), "w"), indent=1)


if __name__ == "__main__":
    main()