"""DAgger-style harvesting: run real tasks with the current laya agent, and at every step ask the local Qwen teacher which
action is right given the page's element table. Disagreements (and agreements) become training cases with the *agent's*
on-policy states, so the next model learns exactly where this one goes wrong.

    python finetune/dagger.py out/dagger_cases.jsonl [tasks.jsonl]   (services: chromium 9222, laya 8791, sglang 30000)

tasks.jsonl lines: {"url": ..., "goal": ...}; default = apps/browser_suite.TASKS plus extra goals below.
"""
import json, os, sys, time
import httpx
sys.path.insert(0, "/home/ckl/projects/S/jev-ultrafast"); sys.path.insert(0, "/home/ckl/projects/S/laya/apps")
os.environ.update(BU_CDP_URL="http://127.0.0.1:9222", TYPESAFE_BASE_URL="http://127.0.0.1:8791", TYPESAFE_API_KEY="local",
                  TEXT_MODEL_API_KEY="local", TEXT_MODEL_BASE_URL="http://127.0.0.1:30000/v1", TEXT_MODEL="Qwen/Qwen3-8B-AWQ",
                  TEXT_MODEL_EXTRA_JSON='{"chat_template_kwargs": {"enable_thinking": false}}')
from jev_ultrafast import Agent
from jev_ultrafast.model import action_space
from browser_suite import TASKS

EXTRA = [
    ("https://en.wikipedia.org/wiki/Main_Page", "Open the article about Albert Einstein."),
    ("https://news.ycombinator.com/", "Open the 'past' page."),
    ("https://news.ycombinator.com/", "Open the comments of the first story on the front page."),
    ("https://github.com/tile-ai/tilelang", "Open the Pull requests tab."),
    ("https://github.com/tile-ai/tilelang", "Open the README's 'examples' folder."),
    ("https://www.python.org/", "Open the documentation page."),
    ("https://docs.python.org/3/", "Open the tutorial."),
    ("https://books.toscrape.com/", "Open the 'Mystery' category and then open the first book in it."),
    ("https://books.toscrape.com/", "Go to page 2 of the catalogue."),
    ("https://the-internet.herokuapp.com/", "Open the 'Dropdown' example and select 'Option 2'."),
    ("https://the-internet.herokuapp.com/", "Open the 'Checkboxes' example and tick the first checkbox."),
    ("https://quotes.toscrape.com/", "Open the quotes tagged 'love'."),
    ("https://quotes.toscrape.com/", "Go to the next page of quotes."),
    ("https://arxiv.org/", "Open the listing of new submissions in cs.CL."),
    ("https://pypi.org/", "Search PyPI for 'tilelang' and open the project page."),
    ("https://duckduckgo.com/", "Search for 'ModernBERT paper'."),
    ("https://www.saucedemo.com/", "Log in with username 'standard_user' and password 'secret_sauce', then add the 'Sauce Labs Backpack' to the cart."),
    ("https://demo.opencart.com/", "Open the 'Desktops' category from the top menu."),
    ("https://www.demoblaze.com/", "Open the 'Laptops' category."),
    ("https://huggingface.co/models", "Search models for 'laya'."),
]
TEACHER = os.environ["TEXT_MODEL_BASE_URL"] + "/chat/completions"
client = httpx.Client(timeout=180)
SYS = """You are the teacher for a browser agent. You see the user's goal, the actions taken so far, the current page (title, url, text excerpt) and a
numbered table of the controls on it. Decide the single best NEXT step:
- {"operation": "CLICK", "index": n}  click control n
- {"operation": "TYPE_TEXT", "index": n}  type into text control n (a separate helper supplies the value)
- {"operation": "SELECT", "index": n, "value": "..."}  choose that option of select control n
- {"operation": "DONE"}  every requirement of the goal is already visibly satisfied on this page
- {"operation": "WAIT"} / {"operation": "SCROLL_DOWN"} / {"operation": "SCROLL_UP"}
Do not re-do satisfied steps; a field that already shows the requested value is done. Return JSON only."""

def teach(goal, history, page, elements):
    table = [{"index": e["index"], "label": e["label"][:70], "role": e.get("role"), "ops": e["operations"], **({"value": e["value"]} if e.get("value") else {})} for e in elements[:120]]
    user = {"goal": goal, "actions_so_far": [{k: h.get(k) for k in ("action", "kind", "text")} for h in history[-8:]],
            "page": {"title": page["title"], "url": page["url"], "text": page["text"][:2500]}, "controls": table}
    body = {"model": os.environ["TEXT_MODEL"], "max_tokens": 120, "temperature": 0.0, "response_format": {"type": "json_object"},
            "chat_template_kwargs": {"enable_thinking": False},
            "messages": [{"role": "system", "content": SYS}, {"role": "user", "content": json.dumps(user, ensure_ascii=False)}]}
    r = client.post(TEACHER, json=body).json()
    return json.loads(r["choices"][0]["message"]["content"])

def to_case(goal, history, page, verdict):
    elements, targets, controls = action_space(page["actions"])
    op = str(verdict.get("operation", "")).upper()
    if op in ("DONE",):
        return {"page": -1, "url": page["url"], "title": page["title"], "goal": goal, "gold_op": "DONE", "gold_id": "DONE", "kind": "done", "label": "",
                "history": history, "source": "dagger", "page_obj": {k: page[k] for k in ("url", "title", "text", "actions")}}
    if op in ("CLICK", "TYPE_TEXT", "SELECT"):
        idx = str(verdict.get("index"))
        cands = targets.get(op, {})
        if op == "SELECT":
            hit = [k for k, a in cands.items() if k.split(":")[0] == idx and (a["value"] == verdict.get("value") or a["label"].endswith(str(verdict.get("value"))))]
            key = hit[0] if hit else None
        else:
            key = idx if idx in cands else None
        if key is None: return None
        a = cands[key]
        return {"page": -1, "url": page["url"], "title": page["title"], "goal": goal, "gold_op": op, "gold_id": a["id"], "kind": a["kind"], "label": a["label"],
                "history": history, "source": "dagger", "page_obj": {k: page[k] for k in ("url", "title", "text", "actions")}}
    return None

def main():
    out = sys.argv[1]
    tasks = [(u, g) for _, u, g, _ in TASKS] + EXTRA
    if len(sys.argv) > 2:
        tasks += [(json.loads(l)["url"], json.loads(l)["goal"]) for l in open(sys.argv[2])]
    n_cases = n_dis = 0
    with open(out, "a") as f:
        for url, goal in tasks:
            t0 = time.time()
            try:
                with Agent(url, goal) as agent:
                    steps = 0
                    while agent.state["status"] not in ("done", "blocked") and steps < 12:
                        page = agent.state["page"]; history = list(agent.state["history"])
                        elements = action_space(page["actions"])[0]
                        try:
                            verdict = teach(goal, history, page, elements)
                        except Exception as e:
                            print("  teacher fail", str(e)[:60]); break
                        case = to_case(goal, [{k: h.get(k) for k in ("action", "kind", "text", "page_changed")} for h in history], page, verdict)
                        if case:
                            f.write(json.dumps(case, ensure_ascii=False) + "\n"); f.flush(); n_cases += 1
                        # step the agent with its own policy
                        try:
                            st = agent.command("tick")
                        except Exception as e:
                            print("  tick fail", type(e).__name__, str(e)[:50]); break
                        d = st["decisions"][-1] if st["decisions"] else None
                        agent_choice = (d["operation"], d.get("target")) if d else None
                        teacher_choice = (str(verdict.get("operation", "")).upper(), str(verdict.get("index")) if verdict.get("index") is not None else None)
                        if agent_choice and agent_choice[0] != teacher_choice[0] or (agent_choice and agent_choice[1] != teacher_choice[1] and teacher_choice[0] in ("CLICK", "TYPE_TEXT")):
                            n_dis += 1
                        steps += 1
            except Exception as e:
                print(f"  task fail {type(e).__name__}: {str(e)[:60]}")
            print(f"{goal[:60]:60s} cases={n_cases} disagreements={n_dis} {time.time()-t0:.0f}s", flush=True)
    print("wrote", n_cases, "cases ->", out)

if __name__ == "__main__":
    main()