"""
Nawah-Router — FastAPI backend for the interactive routing demo.

One endpoint. The model scores every supplied category in a single forward pass, so the frontend
can re-route on each keystroke without a per-category cost.
"""
import os
import time

import torch
from fastapi import FastAPI
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from transformers import AutoTokenizer

from routing_model import MAX_ROUTES, RouterModel, build_text, detect_lang, route

HF_TOKEN = os.environ.get("MODEL_HF_TOKEN") or os.environ.get("HF_TOKEN")

# Three backbones on the same head and the same task. The two 6M's are bidirectional BERT
# encoders (one Arabic-only, one pretrained jointly on Arabic+English); the 52M is a Llama
# decoder. Measured on the same eval sets in one session each - see the model cards.
REGISTRY = {
    "52M":         {"repo": "oddadmix/Nawah-Router-v3",
                     "label": "Nawah-Router-v3 (Llama decoder, Arabic)"},
    "6M-BERT":     {"repo": "oddadmix/Nawah-Router-BERT-6M-v2",
                     "label": "Nawah-Router-BERT-6M-v2 (BERT encoder, Arabic)"},
    "6M-BILINGUAL": {"repo": "oddadmix/Nawah-Router-BERT-6M-bilingual-pretrained",
                      "label": "Nawah-Router-BERT-6M-bilingual-pretrained (BERT encoder, English+Arabic)"},
}
DEFAULT = os.environ.get("MODEL_KEY", "6M-BERT")

LOADED = {}
for key, spec in REGISTRY.items():
    print(f"[*] loading {spec['repo']}", flush=True)
    tok = AutoTokenizer.from_pretrained(spec["repo"], token=HF_TOKEN)
    mdl = RouterModel.from_pretrained(spec["repo"], token=HF_TOKEN)
    LOADED[key] = {"tok": tok, "model": mdl, "label": spec["label"],
                   "repo": spec["repo"],
                   "params": sum(p.numel() for p in mdl.parameters())}
    print(f"[+] {key}: {LOADED[key]['params']/1e6:.2f}M params", flush=True)
torch.set_num_threads(int(os.environ.get("OMP_NUM_THREADS", 4)))

MODEL_ID = REGISTRY[DEFAULT]["repo"]
TOK, MODEL = LOADED[DEFAULT]["tok"], LOADED[DEFAULT]["model"]

app = FastAPI()


class RouteReq(BaseModel):
    text: str = ""
    cats: list[str] = []
    model: str = DEFAULT


@app.get("/")
def index():
    return FileResponse("static/index.html")


@app.get("/api/ready")
def ready():
    return {"ready": True, "model": MODEL_ID, "default": DEFAULT,
            "params": LOADED[DEFAULT]["params"], "max_routes": MAX_ROUTES,
            "models": [{"key": k, "label": v["label"], "repo": v["repo"],
                        "params": v["params"]} for k, v in LOADED.items()]}


@app.post("/api/route")
def api_route(req: RouteReq):
    cats = [c.strip() for c in req.cats if c and c.strip()][:MAX_ROUTES]
    text = (req.text or "").strip()
    if not text or not cats:
        return JSONResponse({"results": [], "ms": 0, "tokens": 0})
    sel = LOADED.get(req.model) or LOADED[DEFAULT]
    t0 = time.perf_counter()
    res = route(sel["model"], sel["tok"], text, cats)
    full, _ = build_text(text, cats)
    ntok = len(sel["tok"].encode(full, add_special_tokens=False))
    order = {r["route"]: r["score"] for r in res}
    return JSONResponse({
        # returned in the caller's order so the UI does not reshuffle rows under the cursor
        "results": [{"route": c, "score": order.get(c, 0.0)} for c in cats],
        "top": max(range(len(cats)), key=lambda i: order.get(cats[i], 0.0)),
        "ms": round((time.perf_counter() - t0) * 1000), "tokens": ntok, "lang": detect_lang(text),
        "model": sel["label"], "params": sel["params"]})


app.mount("/static", StaticFiles(directory="static"), name="static")