"""Persian support-triage demo for tarfandoon/laya-fa-support (a Laya decision model).

A Persian support message goes in; one forward pass returns the department that should
handle it (billing / technical / cancel / other) with a calibrated probability for every
option. No text is generated, so there is nothing to parse and nothing to hallucinate.

The `laya` package is a small non-autoregressive encoder (mmBERT-base, 322M) with a
decision head: ~150-200 ms on CPU, ~35 ms on a GPU.
"""
# --------------------------------------------------------------------------------------
# `spaces` MUST be imported before anything that touches torch: it patches torch's CUDA
# initialisation so a GPU can be attached per call on ZeroGPU. laya imports torch.
# --------------------------------------------------------------------------------------
import spaces  # noqa: F401  (import order is the whole point of this line)

import json
import time

import gradio as gr

import laya

MODEL_ID = "tarfandoon/laya-fa-support"

# The question the model card ships as its recommended prompt: Persian instructions with a
# full description for every option. The card warns that one-word option descriptions make
# the model confidently wrong, so the descriptions below are kept verbatim.
QUESTION_FA = {
    "department": {
        "type": "choice",
        "instructions": "این پیام پشتیبانی باید به کدام بخش ارجاع داده شود؟",
        "criteria": {
            "billing": "پرداخت، کسر وجه، استرداد پول، فاکتور، قیمت و هزینه‌ی اشتراک",
            "technical": "خطا، باگ، کرش، مشکل ورود یا رمز، کندی و قطعی برنامه",
            "cancel": "کاربر می‌خواهد اشتراک را لغو کند، تمدید را متوقف کند یا حساب را ببندد یا حذف کند",
            "other": "هر چیز دیگر: تشکر، پیشنهاد، همکاری، استخدام، سؤال عمومی",
        },
    }
}

DEPARTMENTS = ("billing", "technical", "cancel", "other")

# Emoji + Persian label for the headline result.
DEPT_EMOJI = {
    "billing": "💳",
    "technical": "🔧",
    "cancel": "🚪",
    "other": "💬",
}
DEPT_FA = {
    "billing": "مالی / صورتحساب",
    "technical": "فنی",
    "cancel": "لغو اشتراک",
    "other": "سایر",
}

# Real messages from the held-out persian_fa benchmark (MIT licensed,
# github.com/alipyth/laya-persian-benchmark): each is one the model never saw in training.
# Every example below is one the fine-tuned model answered correctly.
EXAMPLES = [
    "با سلام و احترام، مبلغ اشتراک ماه جاری دو مرتبه از حساب اینجانب کسر گردیده است. خواهشمند است نسبت به استرداد وجه اضافی اقدام فرمایید.",
    "با سلام، از صبح امروز امکان ورود به حساب کاربری وجود ندارد و پیام «خطای داخلی سرور» نمایش داده می‌شود. لطفاً بررسی فرمایید.",
    "میخوام اشتراک ماهیانه م رو کنسل کنم",
    "خدا قوت، واقعاً پشتیبانی‌تون عالی بود، تشکر ویژه از همکارتون",
    "ramz ro avaz kardam vali bazam login nemishe",
    "کارت بانکیم دیکلاین شد ولی invoice صادر شده، refund می‌کنید؟",
    "قصد لغو ندارم، ولی اگه باگ ذخیره نشدن فایل‌ها درست نشه کارم لنگ می‌مونه",
    "دستتون درد نکنه، سه ماهه دارم پول دو تا اشتراک رو می‌دم و کسی جواب نمی‌ده!",
]


# Download and build the weights while the app boots, not on someone's first click.
# CPU placement at module scope: ZeroGPU attaches the real GPU per call, and the agent is
# moved onto it the first time a GPU is actually available (see _use_cuda below).
AGENT = laya.load(MODEL_ID, device="cpu")
print(f"[laya] {MODEL_ID} loaded on {AGENT.device}", flush=True)


def _use_cuda():
    """Move the model to the GPU the first time a GPU is actually attached."""
    import torch

    agent = AGENT
    if torch.cuda.is_available() and agent.device.type != "cuda":
        agent.device = torch.device("cuda")
        agent.model.to(agent.device)
        if torch.cuda.get_device_capability(0)[0] >= 8:
            agent.dtype = torch.bfloat16
        print(f"[laya] moved to {agent.device} (dtype {agent.dtype})", flush=True)
    return agent


# Measured on ZeroGPU: ~2.1 s on a cold worker (GPU attach + weight streaming), ~20 ms warm.
# 15 s covers the cold case with generous margin without wasting a visitor's quota.
@spaces.GPU(duration=15)
def triage(message: str, threshold: float = 0.6):
    """Route a Persian support message to a department with calibrated probabilities.

    Args:
        message: the customer's support message (formal or colloquial Persian, Finglish,
            or Persian mixed with English words).
        threshold: confidence needed to route automatically; below it the message is
            sent to a human instead.

    Returns:
        A headline (department + latency), a probability bar chart over the four
        departments, the confidence-gated action, and the raw JSON response.
    """
    if not message or not message.strip():
        raise gr.Error("لطفاً یک پیام پشتیبانی بنویسید — please enter a support message.")

    agent = _use_cuda()
    state = {"message": message.strip()}

    t = time.perf_counter()
    result = agent.predict(state, QUESTION_FA)
    latency_ms = round((time.perf_counter() - t) * 1000, 1)

    answer = result["answers"]["department"]
    choice = answer["choice"]
    probs = answer["probabilities"]
    confidence = answer["answer_confidence"]

    headline = "%s **%s** (`%s`) · %.0f ms" % (
        DEPT_EMOJI.get(choice, ""), DEPT_FA.get(choice, choice), choice, latency_ms
    )

    chart = {DEPT_FA.get(k, k): probs.get(k, 0.0) for k in DEPARTMENTS}

    if confidence >= threshold:
        action = "✅ **مسیر خودکار** — route automatically to **%s** (confidence %.2f ≥ %.2f)" % (
            choice, confidence, threshold
        )
    else:
        action = "🙋 **ارسال به کاربر انسانی** — send to a human (confidence %.2f < %.2f)" % (
            confidence, threshold
        )

    raw = json.dumps(result, indent=2, ensure_ascii=False)
    return headline, chart, action, raw


with gr.Blocks(title="laya-fa-support") as demo:
    gr.HTML(
        """
        <div style="display:flex;align-items:center;gap:12px;flex-wrap:wrap;">
          <h1 style="margin:0;font-size:1.6rem;">laya-fa-support</h1>
          <span style="font-size:1.5rem;">🇮🇷</span>
        </div>
        """
    )
    gr.Markdown(
        """
Persian customer-support triage, in a single forward pass. Type a support message and the
model decides which department should handle it — **billing**, **technical**, **cancel** or
**other** — with a calibrated probability for every option. It generates no text: nothing to
parse, nothing to hallucinate.

Fine-tuned from Laya `multilingual` (322M) · held-out Persian benchmark: **79.7%** accuracy,
ECE **0.103** · [model card](https://huggingface.co/tarfandoon/laya-fa-support) ·
[benchmark](https://github.com/alipyth/laya-persian-benchmark)
        """
    )

    with gr.Row():
        with gr.Column(scale=5):
            message = gr.Textbox(
                label="پیام پشتیبانی · support message",
                placeholder="مثلاً: دو بار پول از حسابم کم شده ولی اشتراکم فعال نشده…",
                lines=6,
                value="سلام، دو بار پول از حسابم کم شده ولی اشتراکم فعال نشده",
                rtl=True,
            )
            go = gr.Button("ارجاع · Route message", variant="primary")
            gr.Examples(
                examples=[[m] for m in EXAMPLES],
                inputs=[message],
                label="نمونه‌ها · real benchmark messages the model never saw in training",
            )
        with gr.Column(scale=4):
            headline = gr.Markdown("###")
            chart = gr.Label(label="probabilities per department", num_top_classes=4)
            action = gr.Markdown()
            with gr.Accordion("Raw response", open=False):
                raw = gr.Code(language="json", label="raw response")

    with gr.Accordion("Settings", open=False):
        threshold = gr.Slider(
            0.3, 0.95, value=0.6, step=0.05,
            label="Confidence needed to route automatically (below this, a human handles it)",
        )
        gr.Markdown(
            "The threshold gates the *action*, not the model — the model always returns every "
            "probability. Pick the threshold on your own labelled data. A `cancel` decision "
            "closes an account, so false `cancel`s are the costly ones."
        )

    go.click(triage, [message, threshold], [headline, chart, action, raw], api_name="triage")

    gr.Markdown(
        "Model: `%s` · runs on [ZeroGPU](https://huggingface.co/docs/hub/spaces-zerogpu) · "
        "example messages and option descriptions from the "
        "[persian_fa benchmark](https://github.com/alipyth/laya-persian-benchmark) (MIT)"
        % MODEL_ID
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch(mcp_server=True, theme=gr.themes.Citrus())