"""Scripted multi-step trajectories in the real browser, round 2: the skills the 16-task suite still fails on.

  search    : TYPE_TEXT query -> PRESS_ENTER (or CLICK a visible submit button) -> DONE on the results page
  open      : "Open the article/page about X": TYPE_TEXT -> PRESS_ENTER -> (CLICK the matching result if a results
              list appears) -> DONE
  next-page : "Go to page 2 / next page": SCROLL_DOWN until a next/2/older link is visible -> CLICK -> DONE
  scroll    : goal names an element only visible after scrolling -> SCROLL_DOWN -> CLICK -> DONE
  select    : choose an option of a <select> -> SELECT -> DONE

    python finetune/rollouts2.py sites.txt out/rollout2_cases.jsonl [max_sites=400] [seed=7]

sites.txt: one URL per line (a "#" starts a comment).  Suite URLs are skipped.  Goals are templated (varied phrasing),
no LLM needed.  Every case stores the observed page (page_obj) so build_items.py can rebuild the exact request.
Runs against the headless Chromium on BU_CDP_URL (default 127.0.0.1:9222) and needs the patched jev-ultrafast
(press_enter control, observe retry).
"""
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://search.marginalia.nu/", "https://arxiv.org/",
              "https://www.google.com/travel/flights?hl=en"}
SEARCH_T = ["Search for '{q}'.", "Search for '{q}' and show the results.", "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}.", "Run a search for \"{q}\".",
            "Show me search results for {q}.", "Search {site} for '{q}' and show the results list."]
OPEN_T = ["Open the page about {q}.", "Open the article about {q}.", "Go to the entry for '{q}'.", "Find and open '{q}'.",
          "Open '{q}' on this site.", "Search for {q} and open the matching page."]
NEXT_T = ["Go to page 2.", "Go to page 2 of the list.", "Show the next page of results.", "Open the next page.", "Go to the next page of the catalogue.",
          "Show page 2.", "Move on to the second page.", "Open the second page of the listing."]
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}'.", "Click '{x}'."]
INTENT_T = ["{o}.", "{o}, please.", "I want to {ol}.", "Change it to {o}.", "Show everything with {f} set to {o}.", "Use '{o}'.", "Switch to {o}."]
SELECT_T = ["Select '{o}' in the '{f}' dropdown.", "Choose {o} for {f}.", "Set '{f}' to '{o}'.", "Pick the option '{o}'.", "In the {f} dropdown, select {o}."]
QUERIES = ["python tutorial", "flash attention", "climate change", "rust async", "linear algebra", "tilelang", "modernbert", "black holes",
           "sourdough bread", "gpu kernels", "machine learning", "world war ii", "electric cars", "photosynthesis", "quantum computing",
           "jazz history", "docker compose", "typescript generics", "renewable energy", "albert einstein", "marie curie", "solar system",
           "coffee", "yoga", "kubernetes", "chess openings", "vitamin d", "roman empire", "neural networks", "recycling"]
OPEN_Q = ["Albert Einstein", "Python programming language", "Marie Curie", "Photosynthesis", "Black hole", "Isaac Newton", "Quantum computing",
          "Rust programming language", "Great Wall of China", "Ada Lovelace", "Solar System", "Machine learning", "Docker", "Jazz", "Coffee"]


def lab(a): return a["label"].split(" → ")[0].strip()
def site_of(url): return re.sub(r"^www\.", "", url.split("/")[2])
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 ctl(page, cid): return find(page, lambda a: a["id"] == cid)
def observe(b, tries=8):
    for i in range(tries):
        try: return b.observe(screenshot=False)
        except StalePage:
            if i == tries - 1: raise
            time.sleep(0.5)
def act(b, a, page, text=None, wait=0.5):
    b.act(a, page, text=text); time.sleep(wait); return observe(b)
def changed(p, q): return p["url"] != q["url"] or p["title"] != q["title"]
def search_field(page):
    return find(page, lambda a: a["kind"] == "fill" and re.search(r"search|query|find|keyword|suche|rechercher|buscar|cerca|zoek|検索|搜索", a["label"], re.I)) \
        or find(page, lambda a: a["kind"] == "fill" and a.get("role") in ("searchbox",))
def submit_button(page, q):
    return find(page, lambda a: a["kind"] == "click" and a.get("role") == "button" and re.search(r"^(search|go|submit|find|suchen|rechercher|buscar)\b", lab(a), re.I))


def do_search(b, page, rng, out, site):
    field = search_field(page)
    if not field: return 0
    q = rng.choice(QUERIES); goal = rng.choice(SEARCH_T).format(q=q, site=site)
    out.write(json.dumps(case(page, goal, "TYPE_TEXT", field["id"], "fill", field["label"], [], "search"), ensure_ascii=False) + "\n"); n = 1
    p2 = act(b, field, page, text=q)
    pe, sub = ctl(p2, "press_enter"), submit_button(p2, q)
    if not pe and not sub: return n
    # after typing the goal is not done: submit.  Enter is the default; a visible submit button is used 30% of the time
    use_click = bool(sub) and (not pe or rng.random() < 0.3)
    if use_click:
        out.write(json.dumps(case(p2, goal, "CLICK", sub["id"], "click", sub["label"], [hist_of(field, q, False)], "search"), ensure_ascii=False) + "\n")
        step = hist_of(sub); p3 = act(b, sub, p2, wait=1.0)
    else:
        out.write(json.dumps(case(p2, goal, "PRESS_ENTER", "press_enter", "key", pe["label"], [hist_of(field, q, False)], "search"), ensure_ascii=False) + "\n")
        step = hist_of(pe); p3 = act(b, pe, p2, wait=1.0)
    n += 1
    if changed(p3, p2) or q.split()[0].lower() in p3["text"].lower()[:3000]:
        out.write(json.dumps(case(p3, goal, "DONE", "DONE", "done", "", [hist_of(field, q, False), step], "search"), ensure_ascii=False) + "\n"); n += 1
    return n


def do_open(b, page, rng, out, site):
    field = search_field(page)
    if not field: return 0
    q = rng.choice(OPEN_Q); goal = rng.choice(OPEN_T).format(q=q)
    out.write(json.dumps(case(page, goal, "TYPE_TEXT", field["id"], "fill", field["label"], [], "open"), ensure_ascii=False) + "\n"); n = 1
    p2 = act(b, field, page, text=q)
    pe = ctl(p2, "press_enter")
    if not pe: return n
    out.write(json.dumps(case(p2, goal, "PRESS_ENTER", "press_enter", "key", pe["label"], [hist_of(field, q, False)], "open"), ensure_ascii=False) + "\n"); n += 1
    p3 = act(b, pe, p2, wait=1.2)
    if not changed(p3, p2): return n
    hist = [hist_of(field, q, False), hist_of(pe)]
    words = [w.lower() for w in q.split() if len(w) > 2]
    tl = p3["title"].lower()
    title_hit = all(w in tl for w in words[:2]) and not re.search(r"search|results|suche|recherche", tl)
    if title_hit:   # landed on the page itself (Wikipedia-style direct hit)
        out.write(json.dumps(case(p3, goal, "DONE", "DONE", "done", "", hist, "open"), ensure_ascii=False) + "\n"); return n + 1
    # a results list: click the result whose label matches the query best
    res = [a for a in p3["actions"] if a["kind"] == "click" and a.get("role") == "link" and sum(w in a["label"].lower() for w in words) >= max(1, len(words) - 1)]
    if not res: return n
    tgt = res[0]
    out.write(json.dumps(case(p3, goal, "CLICK", tgt["id"], "click", tgt["label"], hist, "open"), ensure_ascii=False) + "\n"); n += 1
    try:
        p4 = act(b, tgt, p3, wait=1.2)
        if changed(p4, p3):
            out.write(json.dumps(case(p4, goal, "DONE", "DONE", "done", "", hist + [hist_of(tgt)], "open"), ensure_ascii=False) + "\n"); n += 1
    except Exception: pass
    return n


NEXT_RE = re.compile(r"^(next\b.*|[»›>→]+\s*next.*|2|page 2|go to page 2|older( posts| entries| \d+)?|older\s*[»›>→]|more|more results|show more|load more|see more|[›»]|next »|next page.*|nächste.*|suivant.*|siguiente.*|次へ|下一页)$", re.I)
def next_link(page):
    return find(page, lambda a: a["kind"] == "click" and a.get("role") in ("link", "button") and NEXT_RE.match(lab(a)) is not None)

def do_next(b, page, rng, out, site):
    if not page.get("scroll") or page["scroll"]["height"] < 1000: return 0
    goal = rng.choice(NEXT_T); hist = []; p = page; n = 0; buf = []
    for _ in range(int(os.environ.get("MAX_SCROLLS", "10"))):
        nl = next_link(p)
        if nl: break
        sd = ctl(p, "scroll_down")
        if not sd: return 0
        buf.append(json.dumps(case(p, goal, "SCROLL_DOWN", "scroll_down", "scroll", "Scroll down", list(hist), "next"), ensure_ascii=False) + "\n")
        p = act(b, sd, p, wait=0.4); hist.append(hist_of(sd, changed=False))
    else:
        return 0              # no next-page link on this page: the scroll steps would teach scrolling for nothing
    out.writelines(buf); n = len(buf)     # (a visible next link without scrolling still teaches the 'next' semantics)
    out.write(json.dumps(case(p, goal, "CLICK", nl["id"], "click", nl["label"], list(hist), "next"), ensure_ascii=False) + "\n"); n += 1
    try:
        p2 = act(b, nl, p, wait=1.0)
        if changed(p2, p):
            out.write(json.dumps(case(p2, goal, "DONE", "DONE", "done", "", hist + [hist_of(nl)], "next"), ensure_ascii=False) + "\n"); n += 1
    except Exception: pass
    return n


def do_scroll(b, page, rng, out, site):
    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 = ctl(page, "scroll_down")
    if not sd: return 0
    p2 = act(b, sd, page, wait=0.4)
    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
    tgt = new[0]
    try:
        p3 = act(b, tgt, p2, wait=0.8)
        if changed(p3, p2):
            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_select(b, page, rng, out, site):
    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()
    # half the goals state the intent in the option's own words ("Sort by price: low to high."), the way users phrase it
    if rng.random() < 0.5 and 3 <= len(o) <= 60:
        goal = rng.choice(INTENT_T).format(o=o, ol=o[0].lower() + o[1:], f=f)
    else:
        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:
        p2 = act(b, opt, page, wait=0.3)
        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



CITIES = ["Zurich", "London", "Paris", "Berlin", "Tokyo", "New York", "Madrid", "Rome", "Amsterdam", "Vienna", "Munich", "Singapore",
          "Sydney", "Toronto", "Chicago", "Barcelona", "Lisbon", "Prague", "Dublin", "Seoul", "Bangkok", "Istanbul", "Oslo", "Helsinki"]
AUTO_T = ["Set the destination to {x}.", "Choose {x} as the departure city.", "Set where from to {x}.", "Set where to to {x}.",
          "Pick {x} from the suggestions.", "Enter {x} and select it from the list.", "Search for {x} and choose the matching suggestion.",
          "Select {x} as the location."]

def is_combo(a):
    return a["kind"] == "fill" and (a.get("role") in ("combobox", "searchbox") or re.search(r"from|to\b|where|destination|origin|city|location|airport|search", a["label"], re.I))

def do_autocomplete(b, page, rng, out, site):
    """TYPE_TEXT a partial value into a combobox -> CLICK the suggestion that appeared -> DONE."""
    field = find(page, is_combo)
    if not field: return 0
    travel = re.search(r"from|to\b|where|destination|origin|city|location|airport|hotel|stay", field["label"], re.I)
    full = rng.choice(CITIES) if travel else rng.choice(QUERIES)
    typed = full if rng.random() < 0.6 else full[:max(3, len(full) // 2)]
    goal = (rng.choice(AUTO_T) if travel else rng.choice(["Search for {x} and choose the matching suggestion.",
            "Type {x} and pick it from the suggestions.", "Find {x} using the search suggestions."])).format(x=full)
    before = {a["label"] for a in page["actions"] if a["kind"] == "click"}
    p2 = act(b, field, page, text=typed, wait=1.2)
    key = full.split()[0].lower()[:5]
    sugg = [a for a in p2["actions"] if a["kind"] == "click" and a["label"] not in before and key in a["label"].lower()
            and a.get("role") in ("option", "listitem", "link", "button", "menuitem", "gridcell", "treeitem")]
    if not sugg: return 0
    tgt = sugg[0]
    out.write(json.dumps(case(page, goal, "TYPE_TEXT", field["id"], "fill", field["label"], [], "autocomplete"), ensure_ascii=False) + "\n")
    out.write(json.dumps(case(p2, goal, "CLICK", tgt["id"], "click", tgt["label"], [hist_of(field, typed, True)], "autocomplete"), ensure_ascii=False) + "\n")
    n = 2
    try:
        p3 = act(b, tgt, p2, wait=1.0)
        out.write(json.dumps(case(p3, goal, "DONE", "DONE", "done", "", [hist_of(field, typed, True), hist_of(tgt)], "autocomplete"), ensure_ascii=False) + "\n"); n += 1
    except Exception: pass
    return n

SKILLS = (("autocomplete", do_autocomplete), ("search", do_search), ("open", do_open), ("next", do_next), ("scroll", do_scroll), ("select", do_select))
if os.environ.get("SKILLS"):                     # SKILLS=next,search : run only these skills
    SKILLS = tuple(x for x in SKILLS if x[0] in os.environ["SKILLS"].split(","))

def main():
    sites = [l.strip() for l in open(sys.argv[1]) if l.strip() and not l.startswith("#")]
    out_f = sys.argv[2]; mx = int(sys.argv[3]) if len(sys.argv) > 3 else 400; seed = int(sys.argv[4]) if len(sys.argv) > 4 else 7
    rng = random.Random(seed); rng.shuffle(sites)
    if len(sys.argv) > 5:                      # "i/n": this worker takes every n-th site starting at i (parallel workers)
        i, n = map(int, sys.argv[5].split("/")); sites = sites[i::n]
    counts = {k: 0 for k, _ in SKILLS}; fails = {}; done = 0; t0 = time.time()
    with open(out_f, "w") as out:
        for url in sites:
            if done >= mx: break
            if url in SUITE_URLS: continue
            site = site_of(url)
            for skill, fn in SKILLS:
                try:
                    b = Browser(url)
                    try:
                        page = observe(b); n = fn(b, page, rng, out, site)
                    finally:
                        b.close()
                    counts[skill] += n
                except Exception as e:
                    fails[type(e).__name__] = fails.get(type(e).__name__, 0) + 1
            done += 1
            if done % 20 == 0: print(f"  {done} sites {time.time()-t0:5.0f}s  {counts}  fails={fails}", flush=True)
    print("wrote", counts, "fails", fails, "->", out_f)

if __name__ == "__main__":
    main()