"""gradio.Server demo: 4-class GitHub issue classifier with a fully custom frontend.

Backend: gradio.Server (FastAPI + Gradio's API engine — queue, gradio_client
compatibility). One queued endpoint `classify` plus a custom `/` route serving
the hand-written Vercel-style index.html.
Models: pngwn/github-issue-classifier-setfit-minilm and
pngwn/github-issue-classifier-distilbert, trained on pngwn/github-issues-4class.
"""
from pathlib import Path

import numpy as np
import torch
from fastapi.responses import HTMLResponse
from gradio import Server
from setfit import SetFitModel
from transformers import AutoModelForSequenceClassification, AutoTokenizer

LABELS = ["bug", "feature", "question", "support"]
SETFIT_ID = "pngwn/github-issue-classifier-setfit-minilm"
DISTIL_ID = "pngwn/github-issue-classifier-distilbert"
INDEX = Path(__file__).parent / "index.html"

print("loading models (first start takes ~1 minute)...")
setfit_model = SetFitModel.from_pretrained(SETFIT_ID)
distil_tokenizer = AutoTokenizer.from_pretrained(DISTIL_ID)
distil_model = AutoModelForSequenceClassification.from_pretrained(DISTIL_ID).eval()
print("models loaded.")

app = Server()


def run_classify(text: str, model: str) -> dict:
    text = (text or "").strip()
    if not text:
        return {"labels": LABELS, "probs": [0.25, 0.25, 0.25, 0.25], "error": "empty_input"}
    if model == "distilbert":
        inputs = distil_tokenizer(text, truncation=True, max_length=256, return_tensors="pt")
        with torch.no_grad():
            logits = distil_model(**inputs).logits
        probs = torch.softmax(logits, dim=-1)[0].numpy()
    else:
        probs = np.asarray(setfit_model.predict_proba([text])).ravel()
    probs = np.asarray(probs, dtype=float).ravel()
    assert probs.shape == (4,)
    return {"labels": LABELS, "probs": [round(float(p), 4) for p in probs]}


@app.api(name="classify")
def classify(text: str, model: str) -> dict:
    """Classify a GitHub issue as bug / feature / question / support.

    Args:
        text: issue title and body.
        model: "setfit" (22M) or "distilbert" (66M); anything else falls back to setfit.
    """
    return run_classify(text, model if model in ("setfit", "distilbert") else "setfit")


@app.get("/", response_class=HTMLResponse)
async def index() -> str:
    return INDEX.read_text()


if __name__ == "__main__":
    app.launch()