import spaces  # before torch, so ZeroGPU can manage CUDA

import gradio as gr
import torch
from peft import PeftModel
from PIL import Image
from transformers import AutoProcessor, AutoTokenizer, Qwen3_5ForConditionalGeneration

from jemm import DecisionModel
from jemm.model import _decision_config

BASE = "Qwen/Qwen3.8-27B"
ADAPTER = "MaestroYan/JEMM"
MAX_PIXELS = 1280 * 800

# Attach the adapter on CPU, then .to("cuda") once, so ZeroGPU can pack all weights at startup
base = Qwen3_5ForConditionalGeneration.from_pretrained(BASE, dtype=torch.bfloat16, attn_implementation="sdpa")
peft_model = PeftModel.from_pretrained(base, ADAPTER, is_trainable=False, torch_device="cpu").to("cuda")
model = DecisionModel(peft_model, AutoTokenizer.from_pretrained(BASE), name="JEMM",
                      processor=AutoProcessor.from_pretrained(BASE), **_decision_config(ADAPTER))


def fit(image: Image.Image) -> Image.Image:
    image = image.convert("RGB")
    scale = (MAX_PIXELS / (image.width * image.height)) ** 0.5
    if scale >= 1:
        return image
    return image.resize((round(image.width * scale), round(image.height * scale)), Image.Resampling.LANCZOS)


@spaces.GPU(size="xlarge", duration=20)
def decide(state: str, question: str, candidates: str, image: Image.Image | None = None):
    """Pick one candidate for the question and return a probability for every candidate.

    Args:
        state: Context the decision is based on, such as a conversation or a task description.
        question: The question to answer.
        candidates: 2 to 32 distinct candidates, one per line.
        image: Optional screenshot. Larger images are scaled to about 1280x800, the training size.
    """
    options = [line.strip() for line in candidates.splitlines() if line.strip()]
    if not 2 <= len(options) <= 32 or len(set(options)) != len(options):
        raise gr.Error("Enter 2 to 32 distinct candidates, one per line.")
    questions = {"q": {"type": "choice", "instructions": question, "criteria": {o: o for o in options}}}
    try:
        result = model.decide(state, questions, None if image is None else [fit(image)])["q"]
    except ValueError as exc:
        raise gr.Error(str(exc)) from exc
    verdict = "" if result["confidence"] >= model.threshold else " · below threshold, treat as undecided"
    return result["probabilities"], f"**{result['choice']}**  \nConfidence {result['confidence']:.3f}{verdict}"


TEXT_EXAMPLES = [
    ["User: Will it rain in Paris tomorrow?", "Which tool should handle the request?",
     "get_weather: weather forecast for a city\nsearch_flights: flight search\nNo tool applies"],
    ["Customer: I was charged twice for my March invoice. Please refund one of the charges.",
     "Which team should handle this ticket?", "Billing\nTechnical support\nSales\nAccount security"],
]
SCREENSHOT_EXAMPLES = [
    ["Task: Search for wireless headphones.\nPrevious actions: [textbox] Search -> TYPE: wireless headphones",
     "Which element should be clicked next? Candidates refer to the elements marked with the same letters on the screenshot.",
     "\n".join(f"CLICK the element marked {x}" for x in "ABCDE"), "examples/search.png"],
]

HEADER = """# JEMM
**Like [Jev](https://typesafe.ai/blog/introducing-system-one-models-and-jev), but multimodal and open-weight.** Picks one candidate per question and returns a probability for every candidate, from text or text plus a screenshot.
[Model card](https://huggingface.co/MaestroYan/JEMM) · [GitHub](https://github.com/ypcypc/JEMM) · Not affiliated with TypeSafe AI."""


def panel(name: str, examples: list, image: bool):
    with gr.Tab(name):
        with gr.Row(equal_height=False):
            with gr.Column():
                state = gr.Textbox(label="State", lines=4)
                question = gr.Textbox(label="Question")
                candidates = gr.Textbox(label="Candidates, one per line", lines=5)
                inputs = [state, question, candidates]
                if image:
                    inputs.append(gr.Image(label="Screenshot", type="pil"))
                button = gr.Button("Decide", variant="primary")
            with gr.Column():
                result = gr.Markdown()
                probabilities = gr.Label(label="Probabilities", show_label=True)
        button.click(decide, inputs, [probabilities, result], api_name="decide" if not image else "decide_screenshot")
        gr.Examples(examples, inputs, cache_examples=False)


with gr.Blocks(title="JEMM") as demo:
    gr.Markdown(HEADER)
    panel("Text", TEXT_EXAMPLES, image=False)
    panel("Screenshot", SCREENSHOT_EXAMPLES, image=True)

demo.launch(theme=gr.themes.Default(primary_hue="blue", font=[gr.themes.GoogleFont("Inter"), "sans-serif"]))