"""
Inference for the Nawah-Router family.

The text and every category share one sequence. Each category's character span is mapped to token
indices through the tokenizer's offset mapping and mean-pooled into its own vector; a shared
scorer turns each into one logit, and the softmax runs over the categories actually supplied.
Because the scorer is shared across positions it reads category *content*, not slot index — which
is what makes the label set free text chosen at inference.

Layout is text-first, categories-second on purpose. The backbone is causal for the 52M decoder, so
this ordering is what lets every category token attend to the whole text; reversed, the categories
would be encoded blind to it.

`lang` picks the wrapper the text is embedded in: "النص:/الفئات:" for Arabic, "Text:/Categories:"
for English. The two Arabic-only backbones (`Nawah-Router-v3`, `Nawah-Router-BERT-6M-v2`) were
only ever trained on the Arabic wrapper, so English input routed through them uses a template they
never saw — expect degraded results there. `Nawah-Router-BERT-6M-bilingual-pretrained` was trained
on both. Left on "auto" (default), the wrapper is picked from the text's script.
"""
import re

import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer

MAX_ROUTES = 9
MAX_LENGTH = 320

WRAPPERS = {
    "en": ("Text:\n{text}\n\nCategories:\n", "- "),
    "ar": ("النص:\n{text}\n\nالفئات:\n", "- "),
}
_ARABIC_RE = re.compile(r"[؀-ۿ]")


class RouterModel(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.backbone = AutoModel.from_pretrained(base_model, dtype=torch.float32)
        h = self.backbone.config.hidden_size
        self.score = nn.Sequential(nn.Linear(h, h), nn.GELU(), nn.Linear(h, 1))
        self.config = self.backbone.config

    def forward(self, input_ids, attention_mask, cat_pool, n_routes):
        hs = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
        logits = self.score(torch.bmm(cat_pool.to(hs.dtype), hs)).squeeze(-1)
        ar = torch.arange(logits.size(1), device=logits.device)[None, :]
        return logits.masked_fill(ar >= n_routes[:, None], torch.finfo(logits.dtype).min)

    @classmethod
    def from_pretrained(cls, path, token=None):
        import os
        from huggingface_hub import hf_hub_download
        m = cls(path)
        w = (os.path.join(path, "router_model.pt") if os.path.isdir(path)
             else hf_hub_download(path, "router_model.pt", token=token))
        m.load_state_dict(torch.load(w, map_location="cpu", weights_only=True))
        return m.eval()


def detect_lang(text):
    return "ar" if _ARABIC_RE.search(text) else "en"


def build_text(text, routes, lang="auto"):
    if lang == "auto":
        lang = detect_lang(text)
    head_fmt, bullet = WRAPPERS[lang]
    s, spans = head_fmt.format(text=text), []
    for c in routes:
        s += bullet
        spans.append((len(s), len(s) + len(c)))
        s += c + "\n"
    return s, spans


@torch.no_grad()
def route(model, tok, text, routes, lang="auto"):
    """-> [{'route': str, 'score': float}] sorted high to low. `lang`: 'en', 'ar', or 'auto'."""
    routes = [r for r in routes if r and r.strip()][:MAX_ROUTES]
    if not text.strip() or not routes:
        return []
    full, spans = build_text(text, routes, lang)
    enc = tok(full, return_offsets_mapping=True, add_special_tokens=False,
              truncation=True, max_length=MAX_LENGTH)
    ids, offs = enc["input_ids"], enc["offset_mapping"]
    pool = torch.zeros(1, MAX_ROUTES, len(ids))
    for ci, (s, e) in enumerate(spans):
        idx = [t for t, (a, b) in enumerate(offs) if a < e and b > s and a != b]
        if idx:
            pool[0, ci, idx] = 1.0 / len(idx)
    logits = model(torch.tensor([ids]), torch.ones(1, len(ids), dtype=torch.long),
                   pool, torch.tensor([len(routes)]))
    probs = logits.softmax(-1)[0][: len(routes)].tolist()
    out = [{"route": r, "score": p} for r, p in zip(routes, probs)]
    return sorted(out, key=lambda x: -x["score"])


if __name__ == "__main__":
    M = "oddadmix/Nawah-Router-BERT-6M-bilingual-pretrained"
    tok = AutoTokenizer.from_pretrained(M)
    model = RouterModel.from_pretrained(M)
    for r in route(model, tok, "The order is an hour late and the driver isn't answering",
                   ["delivery enquiry", "delayed-order complaint", "payment issue"]):
        print(f"{r['score']:.3f}  {r['route']}")
    for r in route(model, tok, "الطلب تأخر ساعة والسائق ما رد على الاتصال",
                   ["استفسار عن التوصيل", "شكوى تأخير", "مشكلة في الدفع"]):
        print(f"{r['score']:.3f}  {r['route']}")