"""Nawah-ASR-50M 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 25 per second. A 3-second clip costs 75
tokens rather than the 375 a fixed 30-second window would spend on silence.
"""
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 transformers import (AutoTokenizer, Qwen2AudioForConditionalGeneration,
                          WhisperFeatureExtractor)

REPO = os.environ.get("ASR_MODEL", "oddadmix/Nawah-ASR-50M-v2")
TOKEN = os.environ.get("MODEL_HF_TOKEN") or os.environ.get("HF_TOKEN")
SR, MAX_SEC = 16000, 30

print(f"[*] loading {REPO}", flush=True)
tok = AutoTokenizer.from_pretrained(REPO, token=TOKEN)
fe = WhisperFeatureExtractor.from_pretrained(REPO, token=TOKEN)
model = Qwen2AudioForConditionalGeneration.from_pretrained(REPO, token=TOKEN,
                                                           dtype=torch.float32).eval()
torch.set_num_threads(int(os.environ.get("OMP_NUM_THREADS", 4)))
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())
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])}


@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")