import spaces  # must be imported before torch

import os
import time

import gradio as gr
import torch
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration

from gamescaling import (
    MOUSE_AIM_LINE,
    OPENING,
    OUTPUT_FORMAT,
    TAIL,
    USER_TURN_TEXT,
    parse_action,
    split_reasoning,
)

ROOT = os.path.dirname(os.path.abspath(__file__))

# ---------------------------------------------------------------------------
# Models
# ---------------------------------------------------------------------------
MODELS = {
    "GameScaling-9B-S3 (mixed games)": "THUSI-Lab/GameScaling-9B-S3",
    "GameScaling-2B-S3 (mixed games, fast)": "THUSI-Lab/GameScaling-2B-S3",
}
DEFAULT_MODEL = "GameScaling-9B-S3 (mixed games)"

LOADED = {}
for label, repo in MODELS.items():
    proc = AutoProcessor.from_pretrained(repo)
    model = Qwen3_5ForConditionalGeneration.from_pretrained(repo, dtype=torch.bfloat16).to("cuda").eval()
    LOADED[label] = (model, proc)

THINK_PREFILL = "<think>"
NO_THINK_PREFILL = "<think>\n\n</think>"
MAX_SIDE = 1280  # reference frames are 1280x720

# ---------------------------------------------------------------------------
# Game presets (system-prompt slots: game name, rules, key bindings)
# ---------------------------------------------------------------------------
_doom = open(os.path.join(ROOT, "prompts", "example_doom_battle1.txt")).read().rstrip("\n")
_DOOM_RULES = _doom.split("**Game Rules**\n", 1)[1].split("\n**Common Keys in", 1)[0]
_DOOM_KEYS = _doom.split("**Common Keys in Doom Battle-1**\n", 1)[1].split("\n**Output Format**", 1)[0]

PRESETS = {
    "Doom Battle-1": dict(
        name="Doom Battle-1",
        rules=_DOOM_RULES,
        keys=_DOOM_KEYS,
    ),
    "Minecraft (voxel sandbox)": dict(
        name="Minecraft",
        rules=(
            "You are playing Minecraft, a first-person voxel sandbox game.\n\n"
            "## What you see on screen\n"
            "- First-person view; the crosshair at screen centre marks the block you are looking at.\n"
            "- The hotbar along the bottom shows the items you carry; the selected slot is highlighted.\n\n"
            "## Task\n"
            "- Chop down a tree to collect wood logs: walk up to a tree trunk, aim the crosshair at it "
            "and hold the left mouse button until the block breaks."
        ),
        keys="\n".join([
            "- W -> Move forward",
            "- A -> Strafe left",
            "- S -> Move backward",
            "- D -> Strafe right",
            "- Space -> Jump",
            "- Shift -> Sneak",
            "- Ctrl -> Sprint",
            "- LB -> Attack / break block",
            "- RB -> Use item / place block",
            "- E -> Open inventory",
            "- one -> Select hotbar slot 1",
            "- two -> Select hotbar slot 2",
            "- three -> Select hotbar slot 3",
            MOUSE_AIM_LINE,
        ]),
    ),
    "SuperTuxKart (racing)": dict(
        name="SuperTuxKart",
        rules=(
            "You are playing SuperTuxKart, a 3D kart racing game, seen from behind your kart.\n\n"
            "## Rules\n"
            "- Finish the laps as fast as possible and overtake the other karts.\n"
            "- Stay on the track: driving off-road slows you down.\n"
            "- Item boxes on the track give you an item that you can fire at other karts."
        ),
        keys="\n".join([
            "- Up -> Accelerate",
            "- Down -> Brake / reverse",
            "- Left -> Steer left",
            "- Right -> Steer right",
            "- Space -> Fire item",
            "- N -> Use nitro",
            "- V -> Drift",
        ]),
    ),
}
PRESET_NAMES = list(PRESETS) + ["Custom"]


def build_system_prompt(name: str, rules: str, keys: str) -> str:
    """Same skeleton as gamescaling.build_system_prompt (official prompt template)."""
    parts = [OPENING.format(name=name)]
    if rules.strip():
        parts.append("**Game Rules**\n" + rules.strip())
    if keys.strip():
        parts.append(f"**Common Keys in {name}**\n" + keys.strip())
    parts.append("**Output Format**\n" + OUTPUT_FORMAT)
    parts.append(TAIL)
    return "\n".join(parts)


# sanity check: preset reproduces the authors' Doom prompt byte-for-byte
assert build_system_prompt(**PRESETS["Doom Battle-1"]) == _doom


# ---------------------------------------------------------------------------
# Visualisation
# ---------------------------------------------------------------------------
def _font(size):
    for p in ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
              "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf"):
        if os.path.exists(p):
            return ImageFont.truetype(p, size)
    try:
        return ImageFont.load_default(size=size)
    except TypeError:
        return ImageFont.load_default()


def render_overlay(frame: Image.Image, action) -> Image.Image:
    """Draw the predicted mouse move on the frame and a 6-step key timeline below it."""
    img = frame.convert("RGB").copy()
    W, H = img.size
    s = max(W, H) / 1280
    d = ImageDraw.Draw(img, "RGBA")
    cx, cy = W / 2, H / 2
    x, y, z = action.mouse
    tx, ty = cx + x / 1000 * W, cy + y / 1000 * H
    r = 10 * s
    d.ellipse([cx - r, cy - r, cx + r, cy + r], outline=(255, 255, 255, 230), width=max(2, int(3 * s)))
    if abs(tx - cx) + abs(ty - cy) > 2:
        d.line([cx, cy, tx, ty], fill=(255, 60, 60, 255), width=max(3, int(5 * s)))
        import math
        ang = math.atan2(ty - cy, tx - cx)
        L = 18 * s
        for da in (2.6, -2.6):
            d.line([tx, ty, tx + L * math.cos(ang + da), ty + L * math.sin(ang + da)],
                   fill=(255, 60, 60, 255), width=max(3, int(5 * s)))
    f = _font(int(22 * s))
    label = f"mouse X={x:+.0f}  Y={y:+.0f}  scroll={z:+.0f}"
    d.rectangle([10 * s, 10 * s, 10 * s + d.textlength(label, font=f) + 16 * s, 46 * s], fill=(0, 0, 0, 170))
    d.text((18 * s, 14 * s), label, fill=(255, 255, 255), font=f)

    # timeline strip
    strip_h = int(150 * s)
    out = Image.new("RGB", (W, H + strip_h), (24, 24, 27))
    out.paste(img, (0, 0))
    d = ImageDraw.Draw(out)
    fh = _font(int(18 * s))
    fk = _font(int(20 * s))
    col_w = W / 6
    for i, keys in enumerate(action.steps):
        x0 = i * col_w
        d.rectangle([x0 + 4 * s, H + 6 * s, x0 + col_w - 4 * s, H + strip_h - 6 * s],
                    outline=(80, 80, 90), width=max(1, int(2 * s)))
        d.text((x0 + 12 * s, H + 10 * s), f"#{i + 1}  {i * 33}-{(i + 1) * 33} ms", fill=(170, 170, 180), font=fh)
        ky = H + 40 * s
        if not keys:
            d.text((x0 + 12 * s, ky), "(no keys)", fill=(110, 110, 120), font=fk)
        for k in keys:
            tw = d.textlength(k, font=fk)
            d.rounded_rectangle([x0 + 12 * s, ky, x0 + 28 * s + tw, ky + 26 * s], radius=6 * s,
                                fill=(250, 140, 40) if k in ("LB", "RB", "MB") else (70, 160, 90))
            d.text((x0 + 20 * s, ky + 2 * s), k, fill=(255, 255, 255), font=fk)
            ky += 30 * s
    return out


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------
def _gpu_duration(frame, preset="", game_name="", game_rules="", key_bindings="", model_name="",
                  reasoning=True, temperature=1.0, max_new_tokens=512, *args, **kwargs):
    # measured on the Space (9B): ~3 s prefill + ~0.036 s/token decode; +15% margin
    return min(120, int(4 + 0.04 * int(max_new_tokens)) + 1)


@spaces.GPU(duration=_gpu_duration)
def predict(
    frame: Image.Image,
    preset: str = "Doom Battle-1",
    game_name: str = "",
    game_rules: str = "",
    key_bindings: str = "",
    model_name: str = DEFAULT_MODEL,
    reasoning: bool = True,
    temperature: float = 1.0,
    max_new_tokens: int = 512,
    seed: int = 0,
):
    """Predict the next 200 ms of keyboard/mouse actions for a game frame.

    Args:
        frame: the current game screen.
        preset: game preset used for any empty prompt field ("Doom Battle-1", "Minecraft (voxel sandbox)", "SuperTuxKart (racing)", "Custom").
        game_name: game name for the system prompt (empty = from preset).
        game_rules: game rules / goal text (empty = from preset).
        key_bindings: one "- <key> -> <effect>" line per key (empty = from preset).
        model_name: which GameScaling checkpoint to use.
        reasoning: prefill "<think>" so the model may think before acting; otherwise act directly.
        temperature: sampling temperature (authors use 1.0).
        max_new_tokens: generation budget.
        seed: random seed.

    Returns:
        Annotated frame, parsed action summary (markdown), the model's reasoning, raw model reply.
    """
    if frame is None:
        raise gr.Error("Please upload a game frame.")
    p = PRESETS.get(preset, PRESETS["Doom Battle-1"])
    name = game_name.strip() or p["name"]
    rules = game_rules if game_rules.strip() else p["rules"]
    keys = key_bindings if key_bindings.strip() else p["keys"]
    system = build_system_prompt(name, rules, keys)

    img = frame.convert("RGB")
    if max(img.size) > MAX_SIDE:
        img.thumbnail((MAX_SIDE, MAX_SIDE), Image.BILINEAR)

    model, proc = LOADED.get(model_name, LOADED[DEFAULT_MODEL])
    messages = [
        {"role": "system", "content": system},
        {"role": "user", "content": [{"type": "text", "text": USER_TURN_TEXT}, {"type": "image"}]},
    ]
    prefill = THINK_PREFILL if reasoning else NO_THINK_PREFILL
    text = proc.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + prefill
    inputs = proc(text=[text], images=[img], return_tensors="pt").to("cuda")

    torch.manual_seed(int(seed))
    t0 = time.perf_counter()
    with torch.inference_mode():
        out = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=temperature > 0,
            temperature=max(float(temperature), 1e-5),
            top_p=1.0,
            top_k=0,
            stop_strings=["<|action_end|>"],
            tokenizer=proc.tokenizer,
        )
    dt = time.perf_counter() - t0
    gen = out[0, inputs["input_ids"].shape[1]:]
    reply = proc.tokenizer.decode(gen, skip_special_tokens=False)
    reply = reply.replace("<|im_end|>", "").replace("<|endoftext|>", "")
    full = prefill + reply

    action = parse_action(full)
    reasoning_text, _, closed = split_reasoning(full)
    vis = render_overlay(img, action)

    if action.parsed:
        x, y, z = action.mouse
        rows = "\n".join(
            f"| {i + 1} | {i * 33}-{(i + 1) * 33} ms | {' '.join(f'`{k}`' for k in g) or '—'} |"
            for i, g in enumerate(action.steps)
        )
        summary = (
            f"**Mouse:** X = {x:+.0f}‰ of width ({'right' if x > 0 else 'left' if x < 0 else 'none'}), "
            f"Y = {y:+.0f}‰ of height ({'down' if y > 0 else 'up' if y < 0 else 'none'}), scroll = {z:+.0f}\n\n"
            f"| step | time | keys held |\n|---|---|---|\n{rows}\n\n"
            f"*{model_name} · {len(gen)} tokens · {dt:.1f}s*"
        )
        if action.dropped:
            summary += f"\n\nDropped out-of-vocabulary keys: {action.dropped}"
    else:
        summary = ("**Could not parse an action** (" + ("generation was cut off while thinking — raise max tokens"
                   if action.truncated else "malformed reply") + "). Executed as a no-op.")
    return vis, summary, (reasoning_text if closed else reasoning_text + " …") or "(acted directly, no thinking)", full


def fill_preset(preset):
    if preset == "Custom":
        return gr.update(), gr.update(), gr.update()
    p = PRESETS[preset]
    return p["name"], p["rules"], p["keys"]


CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🎮 GameScaling: game frame → next 200 ms of keyboard & mouse\n"
            "GameScaling models are Qwen3.5 VLMs continued-pretrained on 1,600 h of keyboard-and-mouse gameplay. "
            "Give it a game screen and a description of the game; it returns one **200 ms action chunk**: "
            "a mouse move plus the keys held in each of six 33 ms steps. "
            "[Model](https://huggingface.co/THUSI-Lab/GameScaling) · "
            "[Project page](https://gamescaling.github.io) · [Paper](https://arxiv.org/abs/2610.28698)"
        )
        with gr.Row():
            with gr.Column():
                frame = gr.Image(type="pil", label="Game frame", height=360)
                with gr.Row():
                    preset = gr.Dropdown(PRESET_NAMES, value="Doom Battle-1", label="Game preset", scale=2)
                    model_name = gr.Dropdown(list(MODELS), value=DEFAULT_MODEL, label="Checkpoint", scale=2)
                reasoning = gr.Checkbox(value=True, label="Allow thinking (prefill <think>)")
                run = gr.Button("Predict action", variant="primary")
            with gr.Column():
                vis = gr.Image(type="pil", label="Predicted action (red arrow = mouse move, bottom = 6 key steps)")
                summary = gr.Markdown()
        with gr.Accordion("Game description (system prompt slots)", open=False):
            game_name = gr.Textbox(label="Game name", value=PRESETS["Doom Battle-1"]["name"])
            game_rules = gr.Textbox(label="Game rules", lines=8, value=PRESETS["Doom Battle-1"]["rules"])
            key_bindings = gr.Textbox(
                label="Key bindings (one '- <key> -> <effect>' per line; training key names, e.g. W, LB, Space, one, Up)",
                lines=8, value=PRESETS["Doom Battle-1"]["keys"],
            )
        with gr.Accordion("Model output", open=False):
            thinking = gr.Textbox(label="Reasoning", lines=4)
            raw = gr.Textbox(label="Raw reply", lines=4)
        with gr.Accordion("Advanced settings", open=False):
            temperature = gr.Slider(0.0, 1.5, value=1.0, step=0.05, label="Temperature (authors use 1.0)")
            max_new_tokens = gr.Slider(64, 2048, value=512, step=64, label="Max new tokens (only used up when the model thinks)")
            seed = gr.Number(value=0, precision=0, label="Seed")

        preset.change(fill_preset, preset, [game_name, game_rules, key_bindings], api_visibility="private")
        outputs = [vis, summary, thinking, raw]
        run.click(
            predict,
            [frame, preset, game_name, game_rules, key_bindings, model_name, reasoning, temperature, max_new_tokens, seed],
            outputs,
            api_name="predict",
        )
        gr.Examples(
            examples=[
                ["examples/doom_battle1.jpg", "Doom Battle-1"],
                ["examples/minetest.jpg", "Minecraft (voxel sandbox)"],
                ["examples/supertuxkart.jpg", "SuperTuxKart (racing)"],
            ],
            inputs=[frame, preset],
            outputs=outputs,
            fn=predict,
            cache_examples=True,
            cache_mode="lazy",
        )
        gr.Markdown(
            "<small>Single-frame demo: in a real game loop the model also sees the last 5 frames and its own past "
            "actions. Doom frame from the authors' repo; Minetest screenshot by Kotolegokot (CC BY-SA 3.0) and "
            "SuperTuxKart screenshot by the STK team (CC BY 3.0), via Wikimedia Commons.</small>"
        )

demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)