"""Load test a running design-canvas server: many concurrent episodes over the OpenEnv client, no model calls.

    DESIGN_CANVAS_CORPUS=<same as the server> python scripts/loadtest.py --url http://127.0.0.1:8007 \\
        --surface tools --episodes 64 --concurrency 32

Each episode is what training does to the server, minus the model:
- tools: replays the expert's tool calls, then submit (should score ~1.0);
- html: one write_html of the expert page, then submit;
- gui: `--steps` rounds of screenshot + click on the page + escape, then submit (built-in editor, or a real app with
  --skin penpot|impress|inkscape-desktop|gimp).

Prints throughput, reset/step/submit latency percentiles and every error, and writes them as JSON with --json.
"""

from __future__ import annotations

import argparse
import concurrent.futures as cf
import json
import random
import statistics as st
import sys
import threading
import time
from pathlib import Path

from openenv.core.env_server.mcp_types import CallToolAction
from openenv.core.mcp_client import MCPToolClient

try:
    from design_canvas.core import expert
    from design_canvas.core.pack import load_pack
except ImportError:  # run from the package directory
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from design_canvas.core import expert
    from design_canvas.core.pack import load_pack


def authorize(token: str) -> None:
    """Send an HF token with every request the OpenEnv client makes (a private Space needs it). The client has no
    option for headers, so its WebSocket connect and its HTTP client are wrapped here, in this script only."""
    import functools

    import httpx
    import openenv.core.env_client as ec

    auth = {"Authorization": f"Bearer {token}"}
    ec.ws_connect = functools.partial(ec.ws_connect, additional_headers=auth)
    base = httpx.AsyncClient.__init__

    def init(self, *a, **kw):
        kw["headers"] = {**auth, **(kw.get("headers") or {})}
        base(self, *a, **kw)

    httpx.AsyncClient.__init__ = init


def pct(xs: list[float], q: float) -> float | None:
    if not xs:
        return None
    xs = sorted(xs)
    return round(xs[min(len(xs) - 1, int(q * len(xs)))], 3)


def episode(url: str, pack, task: dict, surface: str, skin: str | None, steps: int, seed: int) -> dict:
    out = {"task_id": task["task_id"], "reset_s": None, "steps_s": [], "submit_s": None, "reward": None, "error": None,
           "stage": "connect"}
    rng = random.Random(seed)
    try:
        with MCPToolClient(base_url=url, message_timeout_s=600).sync() as env:
            out["stage"] = "reset"
            t = time.time()
            env.reset(task_id=task["task_id"], surface=surface, mode="reference",
                      **({"skin": skin} if skin else {}))
            out["reset_s"] = time.time() - t
            gt = pack.gt_doc(task)
            if surface == "tools":
                calls = [c for c in expert.tool_calls(gt) if c["name"] != "submit"]
            elif surface == "html":
                calls = [{"name": "write_html", "arguments": {"html": expert.html_page(gt)}}]
            else:
                calls = []
                for _ in range(steps):
                    calls += [{"name": "screenshot", "arguments": {}},
                              {"name": "left_click", "arguments": {"coordinate": [rng.randrange(300, 980),
                                                                                  rng.randrange(150, 650)]}},
                              {"name": "key", "arguments": {"keys": "escape"}}]
            for c in calls:
                out["stage"] = c["name"]
                t = time.time()
                r = env.step(CallToolAction(tool_name=c["name"], arguments=c["arguments"]))
                out["steps_s"].append(time.time() - t)
                err = getattr(r.observation, "error", None)
                if err is not None:
                    raise RuntimeError(f"{c['name']}: {getattr(err, 'message', err)}")
            out["stage"] = "submit"
            t = time.time()
            r = env.step(CallToolAction(tool_name="submit", arguments={}))
            out["submit_s"] = time.time() - t
            out["reward"] = r.reward
            out["stage"] = "close"
        out["stage"] = "done"
    except Exception as e:  # noqa: BLE001 - every failure is a result here
        out["error"] = f"{type(e).__name__} at {out['stage']}: {str(e)[:300]}"
    return out


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--url", required=True)
    ap.add_argument("--surface", choices=["tools", "html", "gui"], default="tools")
    ap.add_argument("--skin", default=None)
    ap.add_argument("--split", default="test")
    ap.add_argument("--episodes", type=int, default=32)
    ap.add_argument("--concurrency", type=int, default=16)
    ap.add_argument("--steps", type=int, default=5, help="gui: screenshot+click+key rounds per episode")
    ap.add_argument("--json", default=None, help="write the results here")
    ap.add_argument("--hf-token-env", default=None, help="env var holding an HF token, for a private Space")
    a = ap.parse_args(argv)
    if a.hf_token_env:
        import os

        authorize(os.environ[a.hf_token_env])

    pack = load_pack()
    tasks = [pack.at(a.split, i % pack.count(a.split)) for i in range(a.episodes)]
    lock, done = threading.Lock(), []
    t0 = time.time()
    with cf.ThreadPoolExecutor(a.concurrency) as pool:
        futs = [pool.submit(episode, a.url, pack, t, a.surface, a.skin, a.steps, i) for i, t in enumerate(tasks)]
        for f in cf.as_completed(futs):
            r = f.result()
            with lock:
                done.append(r)
                if r["error"]:
                    print(f"  error {r['task_id']}: {r['error']}", flush=True)
    wall = time.time() - t0
    ok = [r for r in done if not r["error"]]
    steps = [s for r in ok for s in r["steps_s"]]
    summary = {
        "surface": a.surface, "skin": a.skin, "episodes": a.episodes, "concurrency": a.concurrency,
        "ok": len(ok), "errors": len(done) - len(ok), "wall_s": round(wall, 1),
        "episodes_per_min": round(60 * len(ok) / wall, 1), "steps_per_s": round(len(steps) / wall, 1),
        "reset_p50": pct([r["reset_s"] for r in ok], .5), "reset_p95": pct([r["reset_s"] for r in ok], .95),
        "step_p50": pct(steps, .5), "step_p95": pct(steps, .95),
        "submit_p50": pct([r["submit_s"] for r in ok], .5), "submit_p95": pct([r["submit_s"] for r in ok], .95),
        "reward_mean": round(st.fmean([r["reward"] or 0 for r in ok]), 4) if ok else None,
        "error_kinds": sorted({r["error"].split(":")[0] for r in done if r["error"]}),
        "results_lost": sum(1 for r in done if r["error"] and r["reward"] is not None),
    }
    print(json.dumps(summary))
    if a.json:
        Path(a.json).write_text(json.dumps({"summary": summary, "episodes": done}, indent=1))
    return 0 if not summary["errors"] else 1


if __name__ == "__main__":
    sys.exit(main())