"""Scripted multi-step trajectories executed in the real browser, for the three skills the model lacks:
  scroll   : goal targets an element that is only visible after scrolling  -> [SCROLL_DOWN, CLICK, DONE]
  search   : goal asks to search for a phrase in a text field               -> [TYPE_TEXT, CLICK submit/suggestion, DONE]
  select   : goal asks to choose an option of a <select>                    -> [SELECT, DONE]

    python finetune/rollouts.py out/pages.jsonl out/rollout_cases.jsonl [max_pages=300]

Goals are templated (varied phrasing); no LLM needed.  Pages are the crawled ones; the suite's exact URLs are skipped.
"""
import json, os, random, re, sys, time
sys.path.insert(0, "/home/ckl/projects/S/jev-ultrafast")
os.environ.setdefault("BU_CDP_URL", "http://127.0.0.1:9222")
from jev_ultrafast.browser import Browser, StalePage

SUITE_URLS = {"https://en.wikipedia.org/wiki/Main_Page", "https://news.ycombinator.com/", "https://github.com/tile-ai/tilelang", "https://www.python.org/",
              "https://books.toscrape.com/", "https://the-internet.herokuapp.com/dropdown", "https://the-internet.herokuapp.com/checkboxes",
              "https://quotes.toscrape.com/", "https://duckduckgo.com/", "https://arxiv.org/", "https://www.google.com/travel/flights?hl=en"}
SCROLL_T = ["Open '{x}'.", "I want to see '{x}', take me there.", "Go to {x}.", "Find and click '{x}' on this page.", "Open the '{x}' link further down the page.", "Navigate to '{x}'."]
SEARCH_T = ["Search for '{q}'.", "Look up '{q}' using the search box and show the results.", "Find results for '{q}'.", "Use the site search to find '{q}'.", "Search this site for {q}."]
SELECT_T = ["Select '{o}' in the '{f}' dropdown.", "Choose {o} for {f}.", "Set '{f}' to '{o}'.", "Pick the option '{o}'."]
QUERIES = ["python tutorial", "flash attention", "climate change", "rust async", "linear algebra", "tilelang", "modernbert", "black holes", "sourdough bread", "gpu kernels"]

def lab(a): return a["label"].split(" → ")[0].strip()
def case(page, goal, gold_op, gold_id, kind, label, hist, tag):
    return {"page": -1, "url": page["url"], "title": page["title"], "goal": goal, "gold_op": gold_op, "gold_id": gold_id, "kind": kind, "label": label,
            "history": hist, "source": "rollout", "skill": tag, "page_obj": {k: page[k] for k in ("url", "title", "text", "actions")}}
def hist_of(a, text=None, changed=True): return {"action": a["label"], "kind": a["kind"], "text": text, "page_changed": changed}
def find(page, pred):
    return next((a for a in page["actions"] if pred(a)), None)

def do_scroll(b, page, rng, out):
    if not page.get("scroll") or page["scroll"]["height"] < 1300: return 0
    before = {a["label"] for a in page["actions"] if a["kind"] == "click"}
    sd = find(page, lambda a: a["id"] == "scroll_down")
    if not sd: return 0
    b.act(sd, page); time.sleep(0.3); p2 = b.observe(screenshot=False)
    new = [a for a in p2["actions"] if a["kind"] == "click" and a["label"] not in before and 3 <= len(lab(a)) <= 60 and a.get("role") in ("link", "button")]
    if not new: return 0
    n = 0
    for tgt in rng.sample(new, min(2, len(new))):
        goal = rng.choice(SCROLL_T).format(x=lab(tgt))
        out.write(json.dumps(case(page, goal, "SCROLL_DOWN", "scroll_down", "scroll", "Scroll down", [], "scroll"), ensure_ascii=False) + "\n")
        out.write(json.dumps(case(p2, goal, "CLICK", tgt["id"], "click", tgt["label"], [hist_of(sd, changed=False)], "scroll"), ensure_ascii=False) + "\n")
        n += 2
    # execute one click for a DONE state
    tgt = new[0]
    try:
        b.act(tgt, p2); time.sleep(0.4); p3 = b.observe(screenshot=False)
        if p3["url"] != p2["url"]:
            goal = rng.choice(SCROLL_T).format(x=lab(tgt))
            out.write(json.dumps(case(p3, goal, "DONE", "DONE", "done", "", [hist_of(sd, changed=False), hist_of(tgt)], "scroll"), ensure_ascii=False) + "\n"); n += 1
    except Exception: pass
    return n

def do_search(b, page, rng, out):
    field = find(page, lambda a: a["kind"] == "fill" and re.search(r"search|query|find|keyword", a["label"], re.I))
    if not field: return 0
    q = rng.choice(QUERIES); goal = rng.choice(SEARCH_T).format(q=q)
    out.write(json.dumps(case(page, goal, "TYPE_TEXT", field["id"], "fill", field["label"], [], "search"), ensure_ascii=False) + "\n")
    b.act(field, page, text=q); time.sleep(0.4); p2 = b.observe(screenshot=False)
    # submit: a search/go button, or a suggestion containing the query
    sub = find(p2, lambda a: a["kind"] == "click" and (re.search(r"^(search|go|submit|find)\b", lab(a), re.I) or (a.get("role") in ("option", "listitem", "link") and q.split()[0].lower() in a["label"].lower())))
    if not sub: return 1
    out.write(json.dumps(case(p2, goal, "CLICK", sub["id"], "click", sub["label"], [hist_of(field, q, False)], "search"), ensure_ascii=False) + "\n")
    try:
        b.act(sub, p2); time.sleep(0.6); p3 = b.observe(screenshot=False)
        if p3["url"] != p2["url"] or p3["title"] != p2["title"]:
            out.write(json.dumps(case(p3, goal, "DONE", "DONE", "done", "", [hist_of(field, q, False), hist_of(sub)], "search"), ensure_ascii=False) + "\n"); return 3
    except Exception: pass
    return 2

def do_select(b, page, rng, out):
    sels = [a for a in page["actions"] if a["kind"] == "select" and a.get("value") and a["value"] != a.get("current_value")]
    if not sels: return 0
    by_field = {}
    for a in sels: by_field.setdefault(a["node"], []).append(a)
    node, opts = rng.choice(list(by_field.items()))
    opt = rng.choice(opts); f = lab(opt); o = opt["label"].split(" → ")[-1].strip()
    goal = rng.choice(SELECT_T).format(o=o, f=f)
    out.write(json.dumps(case(page, goal, "SELECT", opt["id"], "select", opt["label"], [], "select"), ensure_ascii=False) + "\n")
    try:
        b.act(opt, page); time.sleep(0.3); p2 = b.observe(screenshot=False)
        out.write(json.dumps(case(p2, goal, "DONE", "DONE", "done", "", [hist_of(opt, changed=False)], "select"), ensure_ascii=False) + "\n"); return 2
    except Exception: return 1

def main():
    pages = [json.loads(l) for l in open(sys.argv[1])]; out_f = sys.argv[2]; mx = int(sys.argv[3]) if len(sys.argv) > 3 else 300
    rng = random.Random(11); rng.shuffle(pages)
    counts = {"scroll": 0, "search": 0, "select": 0}; done = 0
    with open(out_f, "w") as out:
        for pg in pages:
            if done >= mx: break
            if pg["url"] in SUITE_URLS: continue
            for skill, fn in (("scroll", do_scroll), ("search", do_search), ("select", do_select)):
                try:
                    b = Browser(pg["url"])
                    try:
                        page = b.observe(screenshot=False); n = fn(b, page, rng, out)
                    finally:
                        b.close()
                    counts[skill] += n
                except Exception as e:
                    pass
            done += 1
            if done % 25 == 0: print(f"  {done} pages  {counts}", flush=True)
    print("wrote", counts, "->", out_f)

if __name__ == "__main__":
    main()