import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # must be imported before torch

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

import gradio as gr
import torch
from huggingface_hub import snapshot_download

MODEL_ID = "nace-ai/drex-v1.5"

# The checkpoint ships its runtime with the weights: code/kev (tokenizer-to-answer stack, Apache-2.0),
# a Qwen3.5 hybrid backbone (bf16 safetensors) and head.pt (pointer head).
CKPT = snapshot_download(
    MODEL_ID,
    allow_patterns=[
        "*.json", "*.jinja", "inference.py", "head.pt",
        "*.safetensors",
        "code/*", "code/**/*",
    ],
)
sys.path.insert(0, str(Path(CKPT) / "code"))

from kev.api import SystemOneRequest, output_tokens, to_answers, to_record  # noqa: E402
from kev.checkpoint import load  # noqa: E402
from kev.model import ContextOverflow, SERVE_MAX_BRANCH, SERVE_MAX_STATE  # noqa: E402

# checkpoint.load() defaults to CUDA when available and bf16 there (fp32 otherwise);
# .to("cuda") inside DecisionModel is intercepted by the ZeroGPU hijack and packed at startup.
tok, model = load(CKPT)
print(f"Loaded {MODEL_ID} on {model.device} ({model.dtype}); hybrid={model.hybrid}", flush=True)

MAX_STATE_CHARS = 40000
EX_REQUEST = {
    "state": {"ticket": "I was charged twice for the same order. Please refund the extra payment."},
    "questions": {
        "team": {
            "type": "choice",
            "instructions": "Which team should handle this ticket?",
            "criteria": {"billing": "Payments, charges, and refunds", "technical": "Bugs and outages", "other": "Anything else"},
        },
        "refund": {"type": "noul", "instructions": "Does the customer explicitly ask for a refund?"},
        "urgency": {"type": "score", "instructions": "How urgent is this ticket?", "criteria": ["Routine", "Soon", "Urgent"]},
    },
}
EX_SERVICE = {
    "state": {"alert": "checkout-api error_rate jumped from 0.2% to 34% five minutes ago; latency_p99 is 4.8 s and orders are failing."},
    "questions": {
        "sev": {
            "type": "score",
            "instructions": "How severe is this incident?",
            "criteria": ["Minor", "Elevated", "Major", "Critical"],
        },
        "page": {"type": "noul", "instructions": "Should the on-call engineer be paged right now?"},
        "area": {
            "type": "choice",
            "instructions": "Which area is most likely involved?",
            "criteria": {
                "payments": "Charging, refunds, invoices",
                "frontend": "Checkout UI and browser errors",
                "backend": "API failures and latency",
                "infra": "Deployments, scaling, cloud resources",
            },
        },
    },
}
EX_REVIEW = {
    "state": {"review": "The room was spotless and the staff were wonderful, but the room heater rattled all night and I barely slept.", "stars": 2},
    "questions": {
        "sentiment": {
            "type": "score",
            "instructions": "How positive is this review overall?",
            "criteria": ["Very negative", "Negative", "Neutral", "Positive", "Very positive"],
        },
        "would_return": {"type": "noul", "instructions": "Does the reviewer sound likely to return?"},
        "followup": {
            "type": "choice",
            "instructions": "What should the hotel do first?",
            "criteria": {
                "maintenance": "Fix the heater",
                "compensation": "Offer a refund or discount",
                "thanks": "Thank the reviewer",
                "nothing": "No action needed",
            },
        },
    },
}


def _dump(value) -> str:
    return value if isinstance(value, str) else json.dumps(value, indent=2, ensure_ascii=False)


def _first_error(error) -> str:
    """Pydantic wraps every failing question in one long message; keep the head of it."""
    text = str(error).replace("\n", " ")
    return text[:300] + ("…" if len(text) > 300 else "")


def _parse_state(state: str):
    """Plain text passes through; JSON-looking text is parsed into a structured state."""
    text = (state or "").strip()
    if not text:
        raise gr.Error("Please provide a state (plain text or JSON).")
    if text[0] in "{[":
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            pass
    return text


def _gpu_duration(state: str, questions_json: str, *args, **kwargs) -> int:
    # Measured on ZeroGPU: 87-token sample ~3.6 s cold / ~0.1 s warm; 8k-token state with
    # 3 questions ~4.7 s model time. Scale with input size, capped well under the tier limit.
    chars = len(state or "") + len(questions_json or "")
    return int(min(60, 10 + chars / 3000))


@spaces.GPU(duration=_gpu_duration)
def decide(state: str, questions_json: str):
    """Answer typed questions about a document with Drex v1.5.

    Args:
        state: The document to decide about. Plain text, or JSON (object/list) which is
            rendered as indented text for the model.
        questions_json: JSON object mapping question names to questions. Each question has
            a `type` — `choice` (named options), `noul` (yes/no) or `score` (ordered levels) —
            optional `instructions`, and `criteria` (object of option → description for
            choice, list of level descriptions for score, optional true/false descriptions
            for noul).

    Returns:
        Per-question probability distributions rendered as bars, plus the SystemOne JSON
        response (answer, confidence and probability per option for every question).
    """
    parsed_state = _parse_state(state)
    try:
        parsed_questions = json.loads(questions_json or "")
    except json.JSONDecodeError as error:
        raise gr.Error(f"Questions must be valid JSON: {error}")
    if not isinstance(parsed_questions, dict) or not parsed_questions:
        raise gr.Error("Questions must be a non-empty JSON object keyed by question name.")
    request = {"state": parsed_state, "questions": parsed_questions}
    if len(json.dumps(parsed_state, ensure_ascii=False)) > MAX_STATE_CHARS:
        raise gr.Error(f"State is too large (>{MAX_STATE_CHARS} characters). Use a shorter document.")
    try:
        req = SystemOneRequest.model_validate(request)
    except Exception as error:
        raise gr.Error(f"Invalid request: {_first_error(error)}")
    try:
        rec, meta = to_record(req)
        enc = model.encode(tok, rec, max_state=SERVE_MAX_STATE, max_branch=SERVE_MAX_BRANCH, strict=True)
    except ContextOverflow as error:
        raise gr.Error(f"Too long for the context window: {error}")
    torch.cuda.synchronize()
    started = time.perf_counter()
    with torch.no_grad():
        probs = [p.float().tolist() for p in model.probs(enc)]
    torch.cuda.synchronize()
    latency_ms = (time.perf_counter() - started) * 1000
    if len(probs) != len(meta) or any(
        len(row) != len(question["keys"]) or any(not math.isfinite(value) for value in row)
        for row, question in zip(probs, meta)
    ):
        raise gr.Error("Model readout is incomplete or non-finite; try a shorter request.")
    answers = to_answers(probs, meta)
    response = {
        "model": "drex-v1.5",
        "answers": answers,
        "usage": {
            "input_tokens": len(enc["ids"]),
            "output_tokens": output_tokens(tok, answers),
        },
        "latency_ms": round(latency_ms, 1),
    }
    return _render(parsed_questions, answers, latency_ms, len(enc["ids"])), response


def _render(questions: dict, answers: dict, latency_ms: float, n_tokens: int) -> str:
    cards = []
    for qid, question in questions.items():
        ans = answers[qid]
        if ans["type"] == "noul":
            options = [("yes", "yes", ans["noul"]), ("no", "no", 1.0 - ans["noul"])]
        elif ans["type"] == "choice":
            criteria = question.get("criteria") or {}
            options = [
                (k, f"{k} — {v}" if isinstance(v, str) and v else k, ans["probabilities"][k])
                for k, v in criteria.items()
            ]
        else:  # score
            legend = ans["legend"]
            options = [(f"{i}: {legend[i]}", legend[i], ans["probabilities"][str(i)]) for i in range(len(legend))]
        best = max(options, key=lambda o: o[2])[0]
        rows = []
        for key, label, value in options:
            highlight = "drex-best" if key == best else ""
            rows.append(
                f'<div class="drex-row {highlight}"><div class="drex-label">{html.escape(str(label))}</div>'
                f'<div class="drex-bar"><div class="drex-fill" style="width:{value * 100:.1f}%"></div></div>'
                f'<div class="drex-val">{value * 100:.1f}%</div></div>'
            )
        extra = ""
        if ans["type"] == "noul":
            extra = f' · P(yes) <b>{ans["noul"] * 100:.1f}%</b>'
        elif ans["type"] == "score":
            extra = f' · expected score <b>{ans["score"]:.2f}</b> of {len(ans["legend"]) - 1}'
        if "confidence" in ans:
            extra += f' · confidence <b>{ans["confidence"] * 100:.1f}%</b>'
        instr = question.get("instructions") or qid
        cards.append(
            f'<div class="drex-card"><div class="drex-head"><span class="drex-id">{html.escape(str(qid))}</span>'
            f'<span class="drex-type">{ans["type"]}</span></div>'
            f'<div class="drex-instr">{html.escape(str(instr))}</div>'
            f'<div class="drex-answer">→ <b>{html.escape(str(best))}</b>{extra}</div>{"".join(rows)}</div>'
        )
    footer = f'<div class="drex-foot">{n_tokens} input tokens · model time {latency_ms:.0f} ms · no text generated</div>'
    return '<div class="drex-wrap">' + "".join(cards) + footer + "</div>"


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
.drex-wrap { display: flex; flex-direction: column; gap: 12px; }
.drex-card { border: 1px solid var(--border-color-primary); border-radius: 10px; padding: 12px 14px; }
.drex-head { display: flex; gap: 8px; align-items: center; margin-bottom: 2px; }
.drex-id { font-weight: 700; font-family: var(--font-mono); }
.drex-type { font-size: 0.75em; padding: 1px 8px; border-radius: 999px; background: var(--color-accent-soft); }
.drex-instr { opacity: 0.8; font-size: 0.92em; }
.drex-answer { margin: 6px 0 8px; }
.drex-row { display: grid; grid-template-columns: minmax(120px, 40%) 1fr 56px; gap: 8px; align-items: center; margin: 3px 0; font-size: 0.9em; }
.drex-label { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
.drex-bar { height: 12px; border-radius: 6px; background: var(--block-background-fill); border: 1px solid var(--border-color-primary); overflow: hidden; }
.drex-fill { height: 100%; background: var(--neutral-400); }
.drex-best .drex-fill { background: var(--color-accent); }
.drex-best .drex-label { font-weight: 600; }
.drex-val { text-align: right; font-family: var(--font-mono); }
.drex-foot { font-size: 0.8em; opacity: 0.7; }
"""

with gr.Blocks(title="Drex v1.5") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# Drex v1.5 — typed decisions with probabilities\n"
            "[nace-ai/drex-v1.5](https://huggingface.co/nace-ai/drex-v1.5) is a 9B **decision model** from "
            "[Nace.AI](https://www.nace.ai/): give it a **state** (a document: text or JSON) and named typed questions — "
            "`choice` (pick an option), `noul` (yes/no) or `score` (ordered levels) — and it returns a "
            "probability for every option, one forward pass per question. No text is generated. "
            "Built on [MiMo-V2.6-Distill-Qwen-9B](https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Distill-Qwen-9B) "
            "(Qwen3.5 hybrid backbone) with a pointer-head readout."
        )
        with gr.Row():
            with gr.Column(scale=1):
                state = gr.Textbox(
                    label="State (the document: plain text or JSON)",
                    lines=5,
                    placeholder="Paste a support ticket, a log line, a JSON record, a contract clause…",
                )
                questions = gr.Code(
                    label="Questions (JSON: name → {type, instructions, criteria})",
                    language="json",
                    value=_dump(EX_REQUEST["questions"]),
                    lines=16,
                )
                run = gr.Button("Decide", variant="primary")
            with gr.Column(scale=1):
                result_html = gr.HTML(label="Decisions")
                with gr.Accordion("SystemOne JSON response", open=False):
                    result_json = gr.JSON(label="Response")

        run.click(
            decide,
            inputs=[state, questions],
            outputs=[result_html, result_json],
            api_name="decide",
        )

        gr.Examples(
            label="Examples",
            examples=[
                [_dump(EX_REQUEST["state"]), _dump(EX_REQUEST["questions"])],
                [_dump(EX_SERVICE["state"]), _dump(EX_SERVICE["questions"])],
                [_dump(EX_REVIEW["state"]), _dump(EX_REVIEW["questions"])],
            ],
            inputs=[state, questions],
            outputs=[result_html, result_json],
            fn=decide,
            cache_examples=True,
            cache_mode="lazy",
        )

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