"""Step-2 negatives: on landing pages from done_cases (history = one executed click), reverse-generate NEW goals whose
next step is another element on that page. Breaks the 'any history => DONE' shortcut.

    python finetune/gen_step2.py out/done_cases.jsonl out/step2_cases.jsonl [per_page=3]
"""
import json, random, sys, threading
from concurrent.futures import ThreadPoolExecutor
sys.path.insert(0, "/home/ckl/projects/S/laya/finetune")
from gen_goals import ask

def main():
    src, out, per = sys.argv[1], sys.argv[2], int(sys.argv[3]) if len(sys.argv) > 3 else 3
    dones = [json.loads(l) for l in open(src)]
    rng = random.Random(7); jobs = []
    for d in dones:
        page = d["page_obj"]
        cands = [a for a in page["actions"] if a["kind"] in ("click", "fill", "select") and len(a["label"].split(" → ")[0].strip()) > 1
                 and a["label"] != d["label"]]
        seen, uniq = set(), []
        for a in cands:
            k = a["label"].lower()
            if k in seen: continue
            seen.add(k); uniq.append(a)
        rng.shuffle(uniq)
        fills = [a for a in uniq if a["kind"] == "fill"][:1]
        for a in fills + [a for a in uniq if a["kind"] != "fill"][: per - len(fills)]:
            jobs.append((d, a, rng.sample([o for o in uniq if o is not a], min(10, len(uniq) - 1))))
    print(f"{len(dones)} landing pages -> {len(jobs)} jobs", flush=True)
    lock = threading.Lock(); n = [0]
    def work(job):
        d, a, others = job
        try:
            goal = ask(d["page_obj"], a, others)
        except Exception as e:
            print("fail", str(e)[:60], flush=True); return None
        with lock:
            n[0] += 1
            if n[0] % 100 == 0: print(f"  {n[0]}/{len(jobs)}", flush=True)
        op = {"click": "CLICK", "fill": "TYPE_TEXT", "select": "SELECT"}[a["kind"]]
        # keep the history: the previous click is unrelated to the new goal, which is what happens mid-task all the time
        return {**{k: d[k] for k in ("page", "url", "title", "history", "page_obj")}, "goal": goal, "gold_op": op, "gold_id": a["id"],
                "gold_node": a.get("node"), "kind": a["kind"], "label": a["label"], "source": "live"}
    with ThreadPoolExecutor(16) as ex:
        cases = [c for c in ex.map(work, jobs) if c]
    with open(out, "w") as f:
        for c in cases: f.write(json.dumps(c, ensure_ascii=False) + "\n")
    print("wrote", len(cases))

if __name__ == "__main__":
    main()