"""The web agent's models on ZeroGPU: Qwen3.5-9B plans, DeskForge-Qwen3.5-4B clicks.

The planner is loaded here, next to the duel's four models from deskforge.py (about 65 GiB of
the 96 GB xlarge GPU), and DeskForge-Qwen3.5-4B is reused from deskforge.py. Each agent step
is ONE ZeroGPU call: the planner reads the screenshot and names an action, then DeskForge
locates every target the planner describes (web_plan.decide_with).

The paper's long-horizon study planned with Qwen3.6-27B. On our demo tasks Qwen3.5-9B did as well and
plans a step about four times faster, so it is the default; WEB_PLANNER=Qwen/Qwen3.6-27B-FP8 restores
the paper's planner (the FP8 loading below handles it).
"""

import os
import time
from pathlib import Path

import spaces  # MUST come before torch / transformers (ZeroGPU)

import torch
from PIL import Image
from transformers import AutoConfig, AutoModelForImageTextToText, AutoProcessor

import deskforge
from web_plan import PLANNER_TOKENS, decide_with, planner_user_text

PLANNER_ID = os.environ.get("WEB_PLANNER", "Qwen/Qwen3.5-9B")
GROUNDER_FAMILY = "Qwen3.5-4B"
GPU_SECONDS = int(os.environ.get("WEB_GPU_SECONDS", "20"))  # per step (~3 s with 9B); ZeroGPU checks quota against this


def _load_planner():
    config = AutoConfig.from_pretrained(PLANNER_ID)
    q = getattr(config, "quantization_config", None)
    if not q:  # the default Qwen3.5-9B: plain bf16 weights
        model = AutoModelForImageTextToText.from_pretrained(PLANNER_ID, dtype=torch.bfloat16)
        print(f"[web] planner {PLANNER_ID} loaded", flush=True)
        return model.eval().to("cuda")
    # An FP8 checkpoint (e.g. Qwen3.6-27B-FP8). Keep the FP8 weights: transformers dequantizes to bf16 when the
    # device reports a compute capability below 8.9, which ZeroGPU's CUDA emulation may do outside a GPU call.
    # The checkpoint also lists MoE-router modules (`mlp.gate`) as unquantized, and transformers matches those
    # names as prefixes, so every dense `mlp.gate_proj` would lose its FP8 scales (garbage output).
    q["modules_to_not_convert"] = [m for m in q["modules_to_not_convert"]
                                   if not m.endswith((".mlp.gate", ".mlp.shared_expert_gate"))]
    try:
        fp8_ok = tuple(torch.cuda.get_device_capability()) >= (8, 9)
    except Exception:
        fp8_ok = False
    real = torch.cuda.get_device_capability
    if not fp8_ok:
        torch.cuda.get_device_capability = lambda *a, **k: (12, 0)  # ZeroGPU: RTX Pro 6000 Blackwell
    try:
        model = AutoModelForImageTextToText.from_pretrained(PLANNER_ID, config=config, dtype=torch.bfloat16)
    finally:
        torch.cuda.get_device_capability = real
    print(f"[web] planner {PLANNER_ID} loaded as FP8 (capability patched: {not fp8_ok})", flush=True)
    return model.eval().to("cuda")


PLANNER_PROCESSOR = AutoProcessor.from_pretrained(PLANNER_ID)
PLANNER = _load_planner()


def _plan(system, user_text, image):
    messages = [{"role": "system", "content": system},
                {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": user_text}]}]
    prompt = PLANNER_PROCESSOR.apply_chat_template(messages, tokenize=False, add_generation_prompt=True,
                                                   enable_thinking=False)
    inputs = PLANNER_PROCESSOR(text=[prompt], images=[image], return_tensors="pt").to(PLANNER.device)
    t0 = time.perf_counter()
    with torch.inference_mode():
        out = PLANNER.generate(**inputs, max_new_tokens=PLANNER_TOKENS, do_sample=False)  # greedy, see web_plan
    n_in = inputs["input_ids"].shape[1]
    print(f"[web] plan {n_in} -> {out.shape[1] - n_in} tokens in {time.perf_counter() - t0:.2f}s", flush=True)
    return PLANNER_PROCESSOR.decode(out[0, n_in:], skip_special_tokens=True).strip()


def _ground(instruction, image):
    code, _ = deskforge._generate(GROUNDER_FAMILY, "deskforge", instruction, image, 128)
    return code


@spaces.GPU(duration=GPU_SECONDS, size="xlarge")
def decide(user_text: str, image):
    """One agent step: plan with the planner (Qwen3.5-9B), then ground each target with DeskForge."""
    t0 = time.perf_counter()
    out = decide_with(_plan, _ground, user_text, image.convert("RGB"))
    print(f"[web] step {time.perf_counter() - t0:.2f}s (plan {out['t_plan']:.2f}s, ground {out['t_ground']:.2f}s)",
          flush=True)
    return out


@spaces.GPU(duration=240, size="xlarge")
def _warm_up():
    """One step at startup: compiles the Triton kernels and fetches the FP8 kernel before the first visitor."""
    image = Image.open(Path(__file__).with_name("examples") / "web_warmup.jpg").convert("RGB")
    text = planner_user_text("Open the HTML version of the DeskForge paper.", 4, 30, "", [],
                             "https://arxiv.org/abs/2610.02320", "DeskForge on arXiv")
    for attempt in ("cold", "warm"):
        t0 = time.perf_counter()
        out = decide_with(_plan, _ground, text, image)
        print(f"[web] warm-up ({attempt}) {time.perf_counter() - t0:.2f}s: {out.get('plan') or out.get('error')} "
              f"{out.get('codes')}", flush=True)
    print(f"[web] GPU memory in use {torch.cuda.memory_allocated() / 2**30:.1f} GiB", flush=True)


if os.environ.get("WEB_WARMUP", "1") == "1":
    try:
        _warm_up()
    except Exception as exc:  # the app still starts; the first visitor then pays the compile time
        print(f"[web] warm-up failed: {exc!r}", flush=True)