"""Nawah-ASR-118M-v5 demo.

The browser does all audio handling: decodeAudioData accepts wav/mp3/m4a/webm, resamples to
16 kHz mono and posts raw float32 samples. So this server needs no ffmpeg, no soundfile, and no
format guessing -- it only ever sees the array the model wants.

The audio token count is per-clip, not fixed: `<audio>` is repeated exactly
`_get_feat_extract_output_lengths(...)` times, which is 50 per second in v5 (v1/v2 were 25). A 3-second clip costs 150
tokens rather than the 375 a fixed 30-second window would spend on silence.

Loading goes through load_nawah, NOT from_pretrained: this model's projector is an MLP and its
LM nests differently, so from_pretrained silently random-initialises ~a third of the tensors and
serves fluent nonsense. See load_model.py.
"""
import os, time

import numpy as np
import torch
from fastapi import FastAPI
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

from load_v5 import load_v5

REPO = os.environ.get("ASR_MODEL", "oddadmix/Nawah-ASR-118M-v5")
TOKEN = os.environ.get("MODEL_HF_TOKEN") or os.environ.get("HF_TOKEN")
SR, MAX_SEC = 16000, 30
torch.set_num_threads(int(os.environ.get("OMP_NUM_THREADS", 4)))

_env_seen = {k: (len(v) if v else 0) for k, v in os.environ.items()
             if "TOKEN" in k.upper() or k.upper().startswith("HF")}
print(f"[*] token env vars visible: {_env_seen}", flush=True)
print(f"[*] loading {REPO} (token {'present' if TOKEN else 'MISSING'})", flush=True)
model, tok, fe = load_v5(REPO, token=TOKEN)   # sets the 50 Hz rate and asserts it
CV = tok.convert_tokens_to_ids
AUDIO, A_START, A_END = CV("<audio>"), CV("<|audio_start|>"), CV("<|audio_end|>")
PARAMS = sum(p.numel() for p in model.parameters())
# Greedy always -- this IS temperature 0. transformers rejects temperature=0.0 (must be > 0) and
# ignores it when do_sample is False, so determinism is pinned on the config, not the call site.
gc = model.generation_config
gc.do_sample, gc.num_beams = False, 1
gc.temperature = gc.top_p = gc.top_k = None
print(f"[+] {PARAMS/1e6:.1f}M params ready", flush=True)

app = FastAPI()


class Req(BaseModel):
    samples: list[float] = []
    max_new: int = 96


@torch.no_grad()
def run(wav: np.ndarray, max_new: int):
    wav = wav[: SR * MAX_SEC]
    feats = fe([wav], sampling_rate=SR, return_attention_mask=True, return_tensors="pt")
    _, n = model.model.audio_tower._get_feat_extract_output_lengths(feats.attention_mask.sum(-1))
    n = int(n[0])
    ids = torch.tensor([[tok.bos_token_id, A_START] + [AUDIO] * n + [A_END]])
    t0 = time.perf_counter()
    out = model.generate(input_ids=ids, attention_mask=torch.ones_like(ids),
                         input_features=feats.input_features,
                         feature_attention_mask=feats.attention_mask,
                         max_new_tokens=max_new, do_sample=False, num_beams=1,
                         pad_token_id=tok.pad_token_id, eos_token_id=tok.eos_token_id)
    ms = (time.perf_counter() - t0) * 1000
    text = tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True).strip()
    dur = len(wav) / SR
    return {"text": text, "audio_tokens": n, "duration": round(dur, 2), "ms": round(ms),
            "rtf": round(dur * 1000 / max(ms, 1), 1),
            "new_tokens": int(out.shape[1] - ids.shape[1]), "model": REPO, "params": PARAMS}


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


@app.get("/api/ready")
def ready():
    return {"ready": True, "model": REPO, "params": PARAMS, "max_sec": MAX_SEC}


@app.post("/api/asr")
def asr(req: Req):
    if not req.samples:
        return JSONResponse({"error": "no audio"}, status_code=400)
    wav = np.asarray(req.samples, dtype=np.float32)
    if wav.size < SR // 4:
        return JSONResponse({"error": "clip too short"}, status_code=400)
    return JSONResponse(run(wav, max(8, min(req.max_new, 200))))


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