"""Minimal reference harness for KORMo-Web-10B (see ../HARNESS.md).

    pip install openai playwright && playwright install chromium
    python ../vllm/convert_to_vllm.py .. ../../KORMo-Web-10B-vllm
    PYTHONPATH=../vllm vllm serve ../../KORMo-Web-10B-vllm --served-model-name KORMo-VL/KORMo-Web-10B --max-model-len 65536 --limit-mm-per-prompt '{"image":1,"video":0}'
    python agent_playwright.py --task "Search for 'KAIST' on Google and open the official homepage."

It reproduces the message layout the model was trained on:
  system  = system_prompt.md          tools = tools.json (22 tools, rendered by chat_template.jinja)
  user    = <user_request>…</user_request> + <observation> (+ screenshot on the first step)
  then per step: assistant (raw <think>…</think> + <tool_call>) → tool (feedback + <observation>)
  Full action/observation history is kept; only the *latest* observation carries a screenshot.
"""
from __future__ import annotations

import argparse
import asyncio
import base64
import json
import re
from pathlib import Path

from openai import OpenAI
from playwright.async_api import async_playwright

ROOT = Path(__file__).resolve().parent.parent
SYSTEM = (ROOT / "system_prompt.md").read_text(encoding="utf-8").strip()
TOOLS = json.loads((ROOT / "tools.json").read_text(encoding="utf-8"))
SPEC = {t["function"]["name"]: t["function"]["parameters"] for t in TOOLS}
TERMINAL = {"finished", "fail", "EXIT"}
SAMPLING = dict(temperature=0.6, top_p=0.95, max_tokens=16384,
                extra_body={"top_k": 20, "repetition_penalty": 1.0,
                            "chat_template_kwargs": {"enable_thinking": True}})


# ── parsing: exactly one <tool_call> in Qwen XML form ───────────────────────────
def parse_call(text: str) -> tuple[str, dict]:
    body = re.sub(r"^.*</think>", "", text, flags=re.S)
    calls = re.findall(r"<tool_call>\s*(.*?)\s*</tool_call>", body, re.S)
    if len(calls) != 1:
        raise ValueError("exactly one tool call required")
    m = re.fullmatch(r"<function=([^>\s]+)>\s*(.*?)\s*</function>", calls[0], re.S)
    if not m or m[1] not in SPEC:
        raise ValueError("malformed or unknown tool call")
    name, props, args = m[1], SPEC[m[1]]["properties"], {}
    for key, value in re.findall(r"<parameter=([^>\s]+)>\n?(.*?)\n?</parameter>", m[2], re.S):
        args[key] = value if props.get(key, {}).get("type") == "string" else json.loads(value)
    missing = [k for k in SPEC[name].get("required", []) if k not in args]
    if missing:
        raise ValueError(f"missing args {missing}")
    return name, args


# ── observation text (same layout as training) ─────────────────────────────────
async def observation(ctx, page, with_shot: bool) -> list[dict]:
    w, h = page.viewport_size["width"], page.viewport_size["height"]
    pages = ctx.pages
    lines = ["<observation>", f"screen size: {w} x {h}", f"Current tab: {pages.index(page)}", "Available tabs:"]
    for i, p in enumerate(pages):
        lines.append(f"- Tab {i}{' (active)' if p is page else ''}: {p.url[:300]}")
    lines += ["", f"url: {page.url[:300]}"]
    if not with_shot:
        return [{"type": "text", "text": "\n".join(lines) + "\n</observation>"}]
    png = base64.b64encode(await page.screenshot()).decode()
    return [{"type": "text", "text": "\n".join(lines) + "\nscreenshot:\n"},
            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{png}"}},
            {"type": "text", "text": "\n</observation>"}]


def drop_images(content: list[dict]) -> list[dict]:
    """Past turns keep the observation text but lose the screenshot."""
    out = [c for c in content if c["type"] == "text"]
    text = "".join(c["text"] for c in out).replace("\nscreenshot:\n\n</observation>", "\n</observation>")
    return [{"type": "text", "text": text}]


# ── execution: coordinates are 0-1000 normalized, scroll is in pixels ───────────
async def execute(ctx, page, name: str, a: dict):
    w, h = page.viewport_size["width"], page.viewport_size["height"]
    px = lambda x, y: (x / 1000 * w, y / 1000 * h)
    before, url0 = len(ctx.pages), page.url
    try:
        if name in ("click", "double_click", "right_click"):
            x, y = px(a["x"], a["y"])
            await page.mouse.click(x, y, button="right" if name == "right_click" else "left",
                                   click_count=2 if name == "double_click" else 1)
            fb = f"Succeed: `{name}` at ({a['x']}, {a['y']}) executed."
        elif name in ("hover", "focus"):
            x, y = px(a["x"], a["y"])
            await page.mouse.move(x, y)
            if name == "focus":
                await page.evaluate("([x,y]) => document.elementFromPoint(x,y)?.focus()", [x, y])
            fb = f"Succeed: `{name}` at ({a['x']}, {a['y']}) executed."
        elif name == "drag":
            await page.mouse.move(*px(a["x"], a["y"])); await page.mouse.down()
            await page.mouse.move(*px(a["x2"], a["y2"]), steps=10); await page.mouse.up()
            fb = "Succeed: `drag` executed."
        elif name == "type":
            await page.keyboard.type(a["text"], delay=30)
            fb = f'Succeed: `type` typed "{a["text"][:200]}".'
        elif name == "clear":
            await page.keyboard.press("Control+A"); await page.keyboard.press("Backspace")
            fb = "Succeed: `clear` executed."
        elif name == "hotkey":
            await page.keyboard.press("+".join(a["keys"]))
            fb = f"Succeed: `hotkey` {a['keys']} executed."
        elif name == "scroll":
            dx, dy = a.get("scroll_dx", 0), a["scroll_dy"]
            await page.mouse.wheel(dx, dy)
            fb = f"Succeed: `scroll` by (dx={dx}, dy={dy}) executed."
        elif name == "wait":
            await asyncio.sleep(float(a["seconds"]))
            fb = f"Succeed: waited {float(a['seconds'])} seconds."
        elif name == "goto":
            await page.goto(a["url"], timeout=60000)
            return page, f"Succeed: `goto` navigated to {page.url}."
        elif name in ("go_back", "go_forward"):
            await getattr(page, name)()
            fb = f"Succeed: `{name}` executed."
        elif name == "new_tab":
            page = await ctx.new_page()
            return page, f"Succeed: `new_tab` opened blank tab (tab {ctx.pages.index(page)})."
        elif name == "switch_tab":
            page = ctx.pages[int(a["index"])]; await page.bring_to_front()
            return page, f"Succeed: `switch_tab` switched to tab {a['index']} ({page.url})."
        elif name == "close_tab":
            await page.close(); page = ctx.pages[-1]
            return page, "Succeed: `close_tab` executed."
        elif name == "upload":
            return page, "Failed: `upload` execution failed: not supported by this harness"
        elif name == "unknown":
            return page, "No action executed."
        else:
            raise ValueError(name)
    except Exception as e:  # noqa: BLE001
        return page, f"Failed: `{name}` execution failed: {e}"
    await page.wait_for_timeout(1000)
    if len(ctx.pages) > before:                       # follow a newly opened tab
        page = ctx.pages[-1]; await page.bring_to_front()
        return page, fb + f" New tab opened (tab {len(ctx.pages) - 1})."
    if page.url != url0:
        return page, fb + f" Page navigated to {page.url}."
    return page, fb + (" Note: no visible navigation or new tab detected." if name.endswith("click") else "")


async def run(task: str, start_url: str, endpoint: str, model: str, max_steps: int):
    client = OpenAI(base_url=endpoint, api_key="EMPTY")
    async with async_playwright() as pw:
        browser = await pw.chromium.launch()
        ctx = await browser.new_context(viewport={"width": 1280, "height": 720})
        page = await ctx.new_page()
        await page.goto(start_url)
        first = await observation(ctx, page, with_shot=True)
        first[0]["text"] = f"<user_request>\n{task}\n</user_request>\n" + first[0]["text"]
        msgs = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": first}]
        for step in range(max_steps):
            out = client.chat.completions.create(model=model, messages=msgs, tools=TOOLS, **SAMPLING)
            text = out.choices[0].message.content or ""
            if "<think>" not in text and "</think>" not in text and getattr(out.choices[0].message, "reasoning_content", None):
                text = f"<think>\n{out.choices[0].message.reasoning_content}\n</think>\n\n{text}"
            msgs.append({"role": "assistant", "content": text})   # keep the full <think> + <tool_call>
            try:
                name, args = parse_call(text)
            except ValueError as e:
                fb = f"Failed: invalid tool call ({e})."
            else:
                print(f"[{step}] {name} {json.dumps(args, ensure_ascii=False)}")
                if name in TERMINAL:
                    print("DONE:", args.get("answer", name)); break
                page, fb = await execute(ctx, page, name, args)
            for m in msgs:                                       # only the newest observation keeps its screenshot
                if m["role"] in ("user", "tool") and isinstance(m["content"], list):
                    m["content"] = drop_images(m["content"])
            obs = await observation(ctx, page, with_shot=True)
            obs[0]["text"] = fb + "\n" + obs[0]["text"]
            msgs.append({"role": "tool", "content": obs})
        await browser.close()


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--task", required=True)
    ap.add_argument("--start-url", default="https://www.google.com/")
    ap.add_argument("--endpoint", default="http://localhost:8000/v1")
    ap.add_argument("--model", default="KORMo-VL/KORMo-Web-10B")
    ap.add_argument("--max-steps", type=int, default=50)
    a = ap.parse_args()
    asyncio.run(run(a.task, a.start_url, a.endpoint, a.model, a.max_steps))