"""Streaming-capable ASR worker (runs in ~/venv_omni): Meta Omnilingual ASR CTC (Apache-2.0), language-agnostic decoding,
so code-switched speech needs no language hint. POST /asr  body = raw float32 mono 16 kHz  -> {"text", "ms"}.
Partial hypotheses = re-decoding the growing buffer (CTC is cheap: ~10-40 ms per call on H100)."""
import os, time, numpy as np, torch
import soundfile  # noqa: F401  (loads the bundled libsndfile before fairseq2n looks for it)
from fastapi import FastAPI, Request
from omnilingual_asr.models.inference.pipeline import ASRInferencePipeline
# ASR model selection (env):
#   ASR_MODEL=dynamic-asr (default) -> Dynamic ASR real-speech model  (Modularcomputing/dynamic-asr, real/model.pt)
#   ASR_MODEL=dynamic-asr-cs        -> Dynamic ASR experimental code-switch model (codeswitch-experimental/model.pt)
#   ASR_MODEL=omni                  -> stock Meta omniASR (ASR_CARD, default omniASR_CTC_1B)
#   ASR_CKPT=<local .pt>            -> any local fine-tuned state_dict on top of ASR_CARD
# Dynamic ASR = fine-tuned omniASR_CTC_1B_v2 (same CTC architecture; needs omnilingual-asr>=0.2.0, or patch_v2.sh on 0.1.0).
# If the Dynamic ASR weights cannot be loaded, the worker falls back to stock omniASR_CTC_1B.
LC = {"ibo": "ibo_Latn", "yor": "yor_Latn", "hau": "hau_Latn", "pcm": "pcm_Latn", "eng": "eng_Latn"}
import unicodedata, types
import omnilingual_asr.models.inference.pipeline as PL
DYN = {"dynamic-asr": "real/model.pt", "dynamic-asr-cs": "codeswitch-experimental/model.pt"}
MODEL = os.environ.get("ASR_MODEL", "dynamic-asr")
CKPT = os.environ.get("ASR_CKPT", "")
CARD = os.environ.get("ASR_CARD", "omniASR_CTC_1B_v2" if (MODEL in DYN or CKPT) else "omniASR_CTC_1B")

def _load():
    global CARD, CKPT, MODEL
    if MODEL in DYN:
        CARD = "omniASR_CTC_1B_v2"          # Dynamic ASR is always on the v2 base
    if MODEL in DYN and not CKPT:
        try:
            from huggingface_hub import hf_hub_download
            CKPT = hf_hub_download(os.environ.get("ASR_REPO", "Modularcomputing/dynamic-asr"), DYN[MODEL])
        except Exception as e:
            print("Dynamic ASR download failed, falling back to omniASR_CTC_1B:", repr(e)[:200], flush=True)
            CARD, MODEL, CKPT = "omniASR_CTC_1B", "omni", ""
    try:
        pipe = ASRInferencePipeline(model_card=CARD)
        if CKPT:
            ck = torch.load(CKPT, map_location="cpu", weights_only=False)
            print("ASR_CKPT", CKPT, pipe.model.load_state_dict({k: v.to(pipe.dtype) for k, v in ck["model"].items()}), flush=True)
            del ck
        return pipe
    except Exception as e:
        if CARD == "omniASR_CTC_1B":
            raise
        print("Dynamic ASR load failed (", repr(e)[:200], ") -> falling back to omniASR_CTC_1B", flush=True)
        CARD, MODEL, CKPT = "omniASR_CTC_1B", "omni", ""
        return ASRInferencePipeline(model_card=CARD)

p = _load()
LATIN_ONLY = os.environ.get("ASR_LATIN_ONLY", "1") == "1"

def _ok(sx):
    for c in sx:
        cat = unicodedata.category(c)
        if cat[0] == "M":                                   # combining marks: only the generic combining diacritics block
            if not ("\u0300" <= c <= "\u036f"):
                return False
        elif c.isalpha() and "LATIN" not in unicodedata.name(c, ""):
            return False
    return True

_BAD = {}
def _apply_latin(self, batch):
    """CTC greedy decoding restricted to Latin-script tokens: all five target languages are written in Latin script, and the
    language-agnostic CTC otherwise drifts into Arabic/other scripts on accented or synthetic speech."""
    bl = PL.BatchLayout(batch.source_seqs.shape, seq_lens=batch.source_seq_lens, device=batch.source_seqs.device)
    logits, bl_out = self.model(batch.source_seqs, bl)
    V = logits.shape[-1]
    if V not in _BAD:
        bad = [i for i in range(V) if not _ok(self.token_decoder(torch.tensor([i])))]
        _BAD[V] = torch.tensor(bad, dtype=torch.long, device=logits.device); print("latin mask: blocked", len(bad), "of", V, flush=True)
    logits[..., _BAD[V]] = -float("inf")
    pred = torch.argmax(logits, dim=-1); out = []
    for i in range(pred.shape[0]):
        seq = pred[i][: bl_out.seq_lens[i]]; m = torch.ones(seq.shape[0], dtype=torch.bool, device=seq.device); m[1:] = seq[1:] != seq[:-1]
        out.append(self.token_decoder(seq[m]))
    return out

if LATIN_ONLY and "CTC" in CARD:
    p._apply_model_wav2vec2asr = types.MethodType(_apply_latin, p)
app = FastAPI()
w = np.zeros(16000, np.float32); p.transcribe([{"waveform": w, "sample_rate": 16000}], batch_size=1)

@app.post("/asr")
async def asr(req: Request, lang: str = ""):
    y = np.frombuffer(await req.body(), dtype=np.float32).copy()
    if len(y) < 1600:
        return {"text": "", "ms": 0.0}
    t = time.perf_counter()
    kw = {} if "CTC" in CARD or not lang else dict(lang=[LC.get(lang, "eng_Latn")])
    h = p.transcribe([{"waveform": y, "sample_rate": 16000}], batch_size=1, **kw)[0]
    return {"text": h, "ms": round(1000 * (time.perf_counter() - t), 1), "card": CARD}

@app.get("/health")
def health():
    return {"ok": True, "model": MODEL, "card": CARD, "ckpt": CKPT}