"""Failure attribution on HELD-OUT webgym seeds: the real agent drives (as in eval_gym), a shadow expert labels every
state the agent visits, and each failed episode is attributed to its FIRST divergence from the expert.

    python finetune/webgym/attrib_gym.py [n_per_kind=10] [kinds=all] [max_steps=60]
    env: ATTRIB_OUT=attrib.json (per-episode steps);  S1_URL / JEV_* as in eval_gym

Category of a step = what the expert is doing there: the widget of the sub-goal's field (date:picker, count:stepper,
combo:..., toggle, choice, ...), or submit / sort / filter / page / open / done.  A divergence is
  wrong_target   the model acts on another element than the expert
  wrong_text     same field, different text
  premature_done the model says DONE / BLOCKED before the task is done
  missed_done    the task is done and the model keeps acting
The expert's action is never executed: this measures the policy as deployed, only labelled.
"""
import json, os, random, sys, time
from collections import Counter, defaultdict
sys.path.insert(0, "/home/ckl/projects/S/jev-ultrafast")
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import eval_gym as G                     # sets the jev env (and S1_URL override)
import expert as E
import spec as S
from dagger_gym import expert_gold
from jev_ultrafast import Agent


def category(sp, key, gold_op):
    if gold_op == "DONE":
        return "done"
    if key is None:
        return "recover"
    f = next((x for x in sp.get("fields", []) if x["key"] == key), None)
    if f:
        return f"{f['w']}:{f.get('impl', '-')}"
    for k in ("submit", "sort", "filter", "page", "open", "q", "go"):
        if key.startswith(k):
            return k
    return key


def run(kind, seed, max_steps):
    sp = S.make(kind, seed)
    steps, first, status, meta = [], None, "error", None
    try:
        with Agent(f"{G.BASE}/t/{kind}/{seed}", sp["goal"]) as agent:
            b = agent.browser
            ex = E.Expert(sp, b, random.Random(seed))
            while agent.state["status"] not in ("done", "blocked") and len(agent.state["history"]) < max_steps:
                page = agent.state["page"]
                info = E.status(b); st = info.get("st") or {}; m = info.get("meta")
                if st and "open" not in st: st["open"] = []
                try:
                    gold_op, ga, gtext, cur = expert_gold(ex, b, sp, page, st, m)
                    key = sp["subgoals"][cur][0] if cur is not None else None
                except E.Stop as e:
                    gold_op, ga, gtext, key = "STOP", None, None, None
                n_hist = len(agent.state["history"])
                agent.command("tick")
                d = agent.state["decisions"][-1]["choice"] if agent.state["decisions"] else None
                if len(agent.state["history"]) == n_hist and agent.state["status"] not in ("done", "blocked"):
                    continue                          # stale / rejected decision: nothing was executed, decide again
                h = agent.state["history"][-1] if len(agent.state["history"]) > n_hist else None
                cat = category(sp, key, gold_op)
                if gold_op == "STOP":
                    kind_ = None
                elif gold_op == "DONE":
                    kind_ = None if d == "DONE" else "missed_done"
                elif d in ("DONE", "BLOCKED"):
                    kind_ = "premature_done"
                elif d != ga["id"]:
                    kind_ = "wrong_target"
                elif gtext and h and (h.get("text") or "").strip().lower() != str(gtext).strip().lower() and cat.split(":")[0] not in ("combo",):
                    kind_ = "wrong_text"
                else:
                    kind_ = None
                row = {"i": len(steps), "cat": cat, "gold": gold_op, "gold_label": (ga or {}).get("label", "")[:60], "gold_text": gtext,
                       "model": d, "model_label": (h or {}).get("action", d)[:60] if h else d, "model_text": (h or {}).get("text"), "div": kind_}
                steps.append(row)
                if kind_ and first is None:
                    first = row
            status = agent.state["status"]
            time.sleep(0.8)
            meta = json.loads(b.evaluate("JSON.stringify(window.__gym ? window.__gym.meta : null)") or "null")
    except Exception as e:
        status = f"error:{type(e).__name__}:{str(e)[:60]}"
    ok = G.success(sp, "", meta)
    return {"kind": kind, "seed": seed, "ok": ok, "status": status, "steps": steps, "first": first, "goal": sp["goal"]}


def main():
    n = int(sys.argv[1]) if len(sys.argv) > 1 else 10
    kinds = (sys.argv[2] if len(sys.argv) > 2 and sys.argv[2] != "all" else ",".join(S.KINDS)).split(",")
    max_steps = int(sys.argv[3]) if len(sys.argv) > 3 else 60
    rows = []
    for kind in kinds:
        for i in range(n):
            r = run(kind, 900000 + 10 * i, max_steps); rows.append(r)
            f = r["first"] or {}
            print(f"{'PASS' if r['ok'] else 'FAIL'}  {kind:10s} {r['seed']} steps={len(r['steps']):2d} {r['status'][:20]:20s} "
                  f"first={f.get('div')}@{f.get('cat')} gold='{f.get('gold_label', '')[:30]}' model='{str(f.get('model_label'))[:30]}'", flush=True)
    json.dump(rows, open(os.environ.get("ATTRIB_OUT", "/tmp/attrib.json"), "w"), indent=1)
    fails = [r for r in rows if not r["ok"]]
    print(f"== success {len(rows) - len(fails)}/{len(rows)}")
    print("== failed episodes by FIRST divergence (category / kind):")
    for (c, k), v in Counter(((r["first"] or {}).get("cat", "none"), (r["first"] or {}).get("div", "none")) for r in fails).most_common():
        print(f"   {v:3d}  {c:22s} {k}")
    agg = defaultdict(lambda: [0, 0])
    for r in rows:
        for s in r["steps"]:
            if s["gold"] != "STOP":
                agg[s["cat"]][0] += s["div"] is None; agg[s["cat"]][1] += 1
    print("== per-step agreement with the expert, by category:")
    for c, (a, t) in sorted(agg.items(), key=lambda x: x[1][0] / max(1, x[1][1])):
        print(f"   {c:22s} {a:4d}/{t:<4d} {a / max(1, t):.2f}")


if __name__ == "__main__":
    main()