import spaces  # must be imported before torch

import html
import json
import sys
import time
from pathlib import Path

import gradio as gr
import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration

MODEL_ID = "Cloudflare/clef-flash"
MODEL_PATH = Path(snapshot_download(MODEL_ID))
sys.path.insert(0, str(MODEL_PATH))
from joint_schema_model import (  # noqa: E402
    QUESTION_TYPES,
    ClefModel,
    JointSchemaHead,
    collate_records,
    encode_record,
    systemone_answer,
)

# Authors' load_release_model uses device_map, which initializes CUDA before the
# ZeroGPU patch can intercept it. Load eagerly with .to("cuda") instead.
backbone = Qwen3_5ForConditionalGeneration.from_pretrained(
    MODEL_PATH,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda")
backbone.config.use_cache = False
head = JointSchemaHead(**json.loads((MODEL_PATH / "joint_head_config.json").read_text()))
head.load_state_dict(load_file(MODEL_PATH / "joint_head.safetensors"), strict=True)
head = head.to(device="cuda", dtype=torch.bfloat16)
model = ClefModel(backbone, head).eval()
processor = AutoProcessor.from_pretrained(MODEL_PATH)

MAX_IMAGE_SIDE = 1024

EX_SUPPORT_STATE = "Our checkout started returning errors and orders are blocked."
EX_SUPPORT_Q = {
    "department": {
        "type": "choice",
        "instructions": "Which team should handle the message?",
        "criteria": {"billing": "Payments or invoices", "technical": "Bugs or outages"},
    },
    "urgency": {"type": "score", "criteria": ["Can wait", "This week", "Today"]},
    "outage": {"type": "noul", "instructions": "Is a service down?"},
}
EX_INVOICE_STATE = {
    "invoice": {"vendor": "Acme", "total": 1250.0, "currency": "USD", "status": "overdue"}
}
EX_INVOICE_Q = {
    "status": {
        "type": "choice",
        "instructions": "What is the invoice status?",
        "criteria": {
            "paid": "Invoice is paid.",
            "overdue": "Invoice is past due.",
            "draft": "Not sent.",
        },
    },
    "large": {"type": "noul", "instructions": "Is the total above 1000 USD?"},
}
EX_REVIEW_STATE = (
    "Ordered the noise-cancelling headphones two weeks ago. Sound is great but the left ear cup "
    "started crackling yesterday and the support chat keeps disconnecting. I want a replacement, not a refund."
)
EX_REVIEW_Q = {
    "sentiment": {
        "type": "score",
        "instructions": "Overall customer sentiment",
        "criteria": ["Very negative", "Negative", "Mixed", "Positive", "Very positive"],
    },
    "action": {
        "type": "choice",
        "instructions": "What should the agent do next?",
        "criteria": {
            "replace": "Ship a replacement unit",
            "refund": "Issue a refund",
            "troubleshoot": "Send troubleshooting steps",
            "escalate": "Escalate to a supervisor",
        },
    },
    "defect": {"type": "noul", "instructions": "Does the customer report a hardware defect?"},
}
EX_FOOD_STATE = {
    "task": "Moderate a photo uploaded to a restaurant listing.",
    "listing": "Vegan Garden Bistro",
}
EX_FOOD_Q = {
    "category": {
        "type": "choice",
        "instructions": "What does the photo mainly show?",
        "criteria": {
            "food": "A dish or meal",
            "interior": "Restaurant interior",
            "menu": "A menu or text",
            "people": "People",
        },
    },
    "matches_listing": {
        "type": "noul",
        "instructions": "Is the dish plausibly vegan, consistent with the listing?",
    },
    "appetizing": {
        "type": "score",
        "instructions": "How appetizing is the photo?",
        "criteria": ["Not at all", "Somewhat", "Very"],
    },
}
EX_ROOM_STATE = "A guest submitted this photo with a listing for a short-term rental apartment."
EX_ROOM_Q = {
    "room": {
        "type": "choice",
        "instructions": "Which room is shown?",
        "criteria": {
            "bedroom": "Bedroom",
            "kitchen": "Kitchen",
            "living_room": "Living room",
            "bathroom": "Bathroom",
        },
    },
    "has_tv": {"type": "noul", "instructions": "Is there a television in the room?"},
    "tidiness": {
        "type": "score",
        "instructions": "How tidy is the room?",
        "criteria": ["Messy", "Average", "Spotless"],
    },
}


def _j(obj) -> str:
    return json.dumps(obj, indent=2, ensure_ascii=False)


def render_html(questions: dict, answers: dict, probs: dict, latency_ms: float, n_tokens: int) -> str:
    parts = ['<div class="clef-out">']
    for qid, q in questions.items():
        qtype = q["type"]
        p = probs[qid]
        ans = answers[qid]
        if qtype == "noul":
            headline = f"{'true' if p['true'] >= 0.5 else 'false'} · P(true) = {p['true']:.3f}"
            labels = {"true": "true", "false": "false"}
        elif qtype == "choice":
            headline = f"{html.escape(str(ans['choice']))} · {ans['confidence']:.3f}"
            labels = {k: html.escape(str(k)) for k in p}
        else:
            headline = f"expected score {ans['score']:.2f} / {len(q['criteria']) - 1}"
            labels = {str(i): f"{i} · {html.escape(str(c))}" for i, c in enumerate(q["criteria"])}
        best = max(p, key=p.get)
        instr = html.escape(str(q.get("instructions") or qid))
        parts.append(
            f'<div class="clef-q"><div class="clef-qh"><span class="clef-id">{html.escape(str(qid))}</span>'
            f'<span class="clef-type">{html.escape(qtype)}</span></div>'
            f'<div class="clef-instr">{instr}</div><div class="clef-head">{headline}</div>'
        )
        order = list(p.keys()) if qtype in ("score", "noul") else sorted(p, key=p.get, reverse=True)
        for opt in order:
            pct = 100 * p[opt]
            cls = "clef-bar best" if opt == best else "clef-bar"
            parts.append(
                f'<div class="clef-row"><div class="clef-lbl">{labels.get(opt, html.escape(str(opt)))}</div>'
                f'<div class="clef-track"><div class="{cls}" style="width:{pct:.1f}%"></div></div>'
                f'<div class="clef-pct">{pct:.1f}%</div></div>'
            )
        parts.append("</div>")
    parts.append(
        f'<div class="clef-meta">{n_tokens} input tokens · forward pass {latency_ms:.0f} ms</div></div>'
    )
    return "".join(parts)


def _parse_state(state: str):
    text = (state or "").strip()
    if text[:1] in "{[":
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            pass
    return text


def _validate(questions):
    if not isinstance(questions, dict) or not questions:
        raise gr.Error("Questions must be a non-empty JSON object mapping question IDs to questions.")
    for qid, question in questions.items():
        if not isinstance(question, dict) or question.get("type") not in QUESTION_TYPES:
            raise gr.Error(f"{qid}: type must be noul, choice, or score.")
        if question["type"] == "choice" and not (
            isinstance(question.get("criteria"), dict) and question["criteria"]
        ):
            raise gr.Error(f"{qid}: a choice question needs criteria as option id → description.")
        if question["type"] == "score" and not (
            isinstance(question.get("criteria"), list) and question["criteria"]
        ):
            raise gr.Error(f"{qid}: a score question needs criteria as a list of level descriptions.")


@spaces.GPU(duration=45)
def decide(state: str, questions_json: str, image=None):
    """Score every allowed option of every typed question about a state with Clef-Flash.

    Args:
        state: The situation to decide on, as free text or a JSON value.
        questions_json: JSON object mapping question IDs to questions. Each question has a
            type (`noul` for true/false, `choice` for named options, `score` for ordered
            levels), optional instructions, and criteria.
        image: Optional image that is part of the state.

    Returns:
        An HTML view of per-option probabilities and a SystemOne-style JSON response.
    """
    try:
        questions = json.loads(questions_json)
    except json.JSONDecodeError as exc:
        raise gr.Error(f"Questions JSON is invalid: {exc}") from exc
    _validate(questions)
    if not (state or "").strip() and image is None:
        raise gr.Error("Provide a state, an image, or both.")

    record = {"model": "clef-flash", "state": _parse_state(state), "questions": questions}
    if image is not None:
        image = image.convert("RGB")
        image.thumbnail((MAX_IMAGE_SIDE, MAX_IMAGE_SIDE))
        record["images"] = [image]

    try:
        encoded = encode_record(processor.tokenizer, record, max_length=16384, processor=processor)
    except Exception as exc:
        raise gr.Error(f"Could not encode the request: {exc}") from exc

    batch = collate_records([encoded], processor.tokenizer.pad_token_id, torch.device("cuda"))
    torch.cuda.synchronize()
    started = time.perf_counter()
    with torch.inference_mode():
        logits = model(batch)[0]
    torch.cuda.synchronize()
    latency_ms = (time.perf_counter() - started) * 1000

    probs, answers = {}, {}
    for encoded_question, question_logits in zip(encoded.questions, logits):
        probabilities = dict(
            zip(encoded_question.option_ids, question_logits.float().softmax(-1).tolist())
        )
        probs[encoded_question.question_id] = probabilities
        answers[encoded_question.question_id] = systemone_answer(
            questions[encoded_question.question_id], probabilities
        )
    response = {
        "model": "clef-flash",
        "answers": answers,
        "usage": {"input_tokens": len(encoded.input_ids), "output_tokens": 0},
    }
    print(f"clef-flash forward {latency_ms:.1f} ms tokens={len(encoded.input_ids)}", flush=True)
    return render_html(questions, answers, probs, latency_ms, len(encoded.input_ids)), response


CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.clef-out { display: flex; flex-direction: column; gap: 12px; }
.clef-q { border: 1px solid var(--border-color-primary); border-radius: 10px; padding: 12px 14px; background: var(--block-background-fill); }
.clef-qh { display: flex; gap: 8px; align-items: center; }
.clef-id { font-weight: 700; font-family: var(--font-mono); }
.clef-type { font-size: 0.75em; padding: 1px 8px; border-radius: 99px; background: var(--color-accent-soft); }
.clef-instr { font-size: 0.9em; opacity: 0.8; margin: 4px 0 6px; }
.clef-head { font-weight: 600; margin-bottom: 8px; }
.clef-row { display: grid; grid-template-columns: minmax(90px, 38%) 1fr 56px; gap: 8px; align-items: center; margin: 3px 0; }
.clef-lbl { font-size: 0.85em; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
.clef-track { height: 12px; border-radius: 6px; background: var(--background-fill-secondary); overflow: hidden; }
.clef-bar { height: 100%; background: var(--neutral-400); border-radius: 6px; }
.clef-bar.best { background: var(--color-accent); }
.clef-pct { font-size: 0.85em; text-align: right; font-family: var(--font-mono); }
.clef-meta { font-size: 0.8em; opacity: 0.7; }
"""

with gr.Blocks(title="Clef-Flash") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# Clef-Flash\n"
            "Give [Cloudflare/clef-flash](https://huggingface.co/Cloudflare/clef-flash) a **state** "
            "(text, JSON, or an image) and a **schema of typed questions**. It returns a probability "
            "for every allowed option of every question in one forward pass — no text generation and "
            "no output parsing.\n\n"
            "`noul` is true/false, `choice` is named options, `score` is ordered levels. "
            "[Announcement](https://blog.cloudflare.com/clef-decision-models)"
        )
        with gr.Row():
            with gr.Column(scale=1):
                state = gr.Textbox(
                    label="State (free text or JSON)",
                    lines=5,
                    value=EX_SUPPORT_STATE,
                    placeholder="Describe the situation, or paste a JSON object.",
                )
                questions = gr.Code(
                    label="Questions schema (JSON)",
                    language="json",
                    value=_j(EX_SUPPORT_Q),
                    lines=14,
                )
                image = gr.Image(label="Image in the state (optional)", type="pil", height=240)
                run = gr.Button("Decide", variant="primary")
            with gr.Column(scale=1):
                out_html = gr.HTML(label="Decisions")
                with gr.Accordion("SystemOne JSON response", open=False):
                    out_json = gr.JSON(label="Response")

        gr.Markdown("### Examples — text or JSON")
        text_examples = gr.Examples(
            examples=[
                [EX_SUPPORT_STATE, _j(EX_SUPPORT_Q)],
                [_j(EX_INVOICE_STATE), _j(EX_INVOICE_Q)],
                [EX_REVIEW_STATE, _j(EX_REVIEW_Q)],
            ],
            inputs=[state, questions],
            outputs=[out_html, out_json],
            fn=decide,
            cache_examples=True,
            cache_mode="lazy",
        )
        gr.Markdown("### Examples — with an image")
        gr.Examples(
            examples=[
                [_j(EX_FOOD_STATE), _j(EX_FOOD_Q), "examples/burger.jpg"],
                [EX_ROOM_STATE, _j(EX_ROOM_Q), "examples/living_room.jpg"],
            ],
            inputs=[state, questions, image],
            outputs=[out_html, out_json],
            fn=decide,
            cache_examples=True,
            cache_mode="lazy",
        )
        text_examples.dataset.click(lambda: None, outputs=image, api_visibility="private")

    run.click(decide, inputs=[state, questions, image], outputs=[out_html, out_json], api_name="decide")

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)