"""Opt-in sequential A/B through ARTEMIS's existing queue, never a second agent.

Usage: python tests/benchmark_plugins.py --traces <ARTEMIS traces> --output <new folder>
The user-started service must already run. Only this script's owned fixture is closed.
"""

import argparse
import json
import sqlite3
import subprocess
import time
import urllib.request
import uuid
from pathlib import Path

PROMPT = """只操作测试窗口 ARTEMIS Plugin Benchmark（进程 PID {pid}），不要观察或操作其他应用。先按此 PID 找到窗口。
1. 将 Network name 填为 ARTEMIS-LAB，勾选 Night mode，点击 Apply，确认显示 Saved: ARTEMIS-LAB; Night mode: True。
2. 只点击唯一的蓝色圆形一次，不要点击红色圆形。在下方 Scrollable settings 面板向下滚动，找到并点击 Finish，确认 Finished successfully。
3. 保存这个窗口的截图到任务 trace，文件名 proof.png，然后结束并如实说明实际结果。不要关闭窗口。
"""


def api(base, route, body=None):
    request = urllib.request.Request(
        base + route,
        data=json.dumps(body).encode() if body is not None else None,
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(request, timeout=30) as response:
        return json.load(response)


def collect(traces, sid):
    with sqlite3.connect(
        f"file:{(traces / 'data_engine.db').as_posix()}?mode=ro", uri=True
    ) as db:
        db.row_factory = sqlite3.Row
        session = dict(
            db.execute("SELECT * FROM sessions WHERE session_id=?", (sid,)).fetchone()
        )
        events = [
            dict(row)
            for row in db.execute(
                "SELECT name,type,status,duration,payload FROM traces WHERE session_id=? ORDER BY timestamp",
                (sid,),
            )
        ]
        steps = db.execute(
            "SELECT count(*) FROM steps WHERE session_id=?", (sid,)
        ).fetchone()[0]
    usage, calls, requests = [], [], []
    for event in events:
        payload = json.loads(event["payload"] or "{}")
        if event["name"] == "llm_usage":
            usage.append(payload)
        if "plugin_id" in payload and "tool_name" in payload:
            # Tool trace contains the real call receipt, not the separate log line.
            calls.append(
                {
                    "name": payload["tool_name"],
                    "status": event["status"],
                    "duration": event["duration"],
                    "args": payload.get("args"),
                    "result": payload.get("result"),
                }
            )
        if event["name"] == "llm_request":
            requests.append(payload)
    return {
        "session": session,
        "steps": steps,
        "llm_usage": usage,
        "plugin_calls": calls,
        "requests": requests,
    }


def main():
    parser = argparse.ArgumentParser(__doc__)
    parser.add_argument("--base", default="http://127.0.0.1:8000")
    parser.add_argument("--traces", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--model", default="mimo-v2.6-flash")
    parser.add_argument("--repetitions", type=int, default=3)
    parser.add_argument("--resume", action="store_true")
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=args.resume)
    config = {
        "model": args.model,
        "reasoning_effort": "high",
        "execution_target": "computer",
        "record_video": True,
        "recording": {"codec": "h264", "computer_scope": "window"},
        "analyze_video": False,
        "timeout_seconds": 480,
        "request_timeout_seconds": 120,
        "max_retries": 1,
    }
    protocol = {
        "prompt": PROMPT,
        "config": config,
        "repetitions": args.repetitions,
        "acceptance": "applied && finished && exactly one target click && no wrong clicks && verified proof.png",
        "order": "OCU,Cua / Cua,OCU / OCU,Cua",
        "fixture": "benchmark-fixture.ps1",
    }
    protocol_path = args.output / "protocol.json"
    if args.resume and protocol_path.exists():
        if json.loads(protocol_path.read_text("utf-8")) != protocol:
            raise ValueError(
                "Resume requires the original prompt, model and task settings"
            )
    else:
        protocol_path.write_text(
            json.dumps(protocol, ensure_ascii=False, indent=2), "utf-8"
        )
    for repetition in range(args.repetitions):
        order = (
            ["open-computer-use", "cua-driver"]
            if repetition % 2 == 0
            else ["cua-driver", "open-computer-use"]
        )
        for plugin in order:
            folder = args.output / f"{repetition + 1}-{plugin}"
            if args.resume and (folder / "result.json").is_file():
                continue
            queue = api(args.base, "/api/status")
            if queue.get("active_tasks") or queue.get("queue"):
                raise RuntimeError(
                    "Other tasks are active/queued; benchmark will not interfere"
                )
            folder.mkdir()
            state_file = folder / "fixture.json"
            stderr = (folder / "fixture-stderr.log").open("w")
            fixture = subprocess.Popen(
                [
                    "powershell.exe",
                    "-NoProfile",
                    "-File",
                    str(Path(__file__).with_name("benchmark-fixture.ps1").resolve()),
                    "-StateFile",
                    str(state_file.resolve()),
                ],
                creationflags=subprocess.CREATE_NO_WINDOW,
                stdout=subprocess.DEVNULL,
                stderr=stderr,
            )
            sid = str(uuid.uuid4())
            done = False
            try:
                for _ in range(100):
                    if state_file.exists():
                        break
                    time.sleep(0.1)
                if not state_file.exists():
                    raise RuntimeError("fixture did not start")
                body = {
                    "goal": PROMPT.format(pid=fixture.pid),
                    "session_id": sid,
                    "profile": "flash",
                    "ingress": "client",
                    "runtime_overrides": {
                        **config,
                        "plugins": [
                            {
                                "id": plugin,
                                "profile": "benchmark"
                                if plugin == "cua-driver"
                                else "default",
                                "version": "0.1.2"
                                if plugin == "cua-driver"
                                else "1.1.0",
                            }
                        ],
                    },
                }
                (folder / "request.json").write_text(
                    json.dumps(body, ensure_ascii=False, indent=2), "utf-8"
                )
                started = time.monotonic()
                response = api(args.base, "/api/run", body)
                (folder / "submission.json").write_text(json.dumps(response), "utf-8")
                print(
                    json.dumps(
                        {
                            "event": "submitted",
                            "plugin": plugin,
                            "round": repetition + 1,
                            "trace_id": sid,
                            "response": response,
                        }
                    ),
                    flush=True,
                )
                if response.get("status") in {"rejected", "error"}:
                    raise RuntimeError("task was rejected")
                while time.monotonic() - started < 600:
                    time.sleep(15)
                    status = api(args.base, f"/api/sessions/{sid}")
                    state = json.loads(state_file.read_text("utf-8-sig"))
                    print(
                        json.dumps(
                            {
                                "event": "progress",
                                "trace_id": sid,
                                "status": status["status"],
                                "elapsed": round(time.monotonic() - started),
                                "applied": state["applied"],
                                "finished": state["finished"],
                            }
                        ),
                        flush=True,
                    )
                    if status["status"] not in {
                        "running",
                        "pending",
                        "queued",
                        "stopping",
                        "starting",
                    }:
                        done = True
                        break
                if not done:
                    api(args.base, "/api/stop", {"session_id": sid})
                    time.sleep(3)
                elapsed = time.monotonic() - started
                result = collect(args.traces.resolve(), sid)
                state = json.loads(state_file.read_text("utf-8-sig"))
                attachment = (
                    args.traces / sid / "artifacts/plugins" / plugin / "proof.png"
                )
                result.update(
                    trace_id=sid,
                    plugin=plugin,
                    repetition=repetition + 1,
                    fixture=state,
                    observed_elapsed=elapsed,
                    watchdog_timeout=not done,
                    actual_pass=bool(
                        state["applied"]
                        and state["finished"]
                        and state["target_clicks"] == 1
                        and state["wrong_clicks"] == 0
                        and attachment.exists()
                    ),
                    attachment_exists=attachment.exists(),
                )
                (folder / "result.json").write_text(
                    json.dumps(result, ensure_ascii=False, indent=2), "utf-8"
                )
                print(
                    json.dumps(
                        {
                            "event": "done",
                            "trace_id": sid,
                            "actual_pass": result["actual_pass"],
                        }
                    ),
                    flush=True,
                )
                # Release/recording cleanup belongs to ARTEMIS, wait before next fixture.
                for _ in range(15):
                    if not api(args.base, "/api/status").get("active_tasks"):
                        break
                    time.sleep(2)
            finally:
                if not done:
                    try:
                        api(args.base, "/api/stop", {"session_id": sid})
                    except (OSError, ValueError) as error:
                        print(f"Cleanup request failed for {sid}: {error}", flush=True)
                if fixture.poll() is None:
                    fixture.terminate()
                fixture.wait(10)
                stderr.close()


if __name__ == "__main__":
    main()