# A/B harness copy of stacks/parakeet/app.py (2026-09-30): identical decode path; adds an
# x-ab-decode-ms header, an optional injected delay (AB_DELAY_MS) and a model-suffix knob.
"""Thin FastAPI wrapper around sherpa-onnx's OfflineRecognizer for Parakeet-TDT.

Load the encoder/decoder/joiner/tokens once at startup; serve:
  POST /transcribe                   — our native shape
  POST /v1/audio/transcriptions      — OpenAI-compatible alias (returns {"text": ...})
  GET  /healthz                      — used by the docker healthcheck

No VAD chunking, no Silero preprocessing — parakeet-tdt handles long-form natively
and the int8 ONNX model is a rounding error against this host's 96 GB cards.
"""

from __future__ import annotations

import io
import logging
import os
import time
from pathlib import Path

import numpy as np
import sherpa_onnx
import soundfile as sf
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.responses import JSONResponse

MODEL_DIR = Path(os.environ.get("MODEL_DIR", "/models"))
PROVIDER = os.environ.get("PROVIDER", "cuda")
NUM_THREADS = int(os.environ.get("NUM_THREADS", "1"))
# A/B harness only (not in the seat): an injected delay for the positive control, and the
# model filenames so a non-int8 export can be served by the same code.
AB_DELAY_S = float(os.environ.get("AB_DELAY_MS", "0")) / 1000.0
AB_SUFFIX = os.environ.get("AB_SUFFIX", "int8.onnx")

REQUIRED_FILES = (
    f"encoder.{AB_SUFFIX}",
    f"decoder.{AB_SUFFIX}",
    f"joiner.{AB_SUFFIX}",
    "tokens.txt",
)

logger = logging.getLogger("parakeet")
logging.basicConfig(level=os.environ.get("LOG_LEVEL", "INFO"))


def _ensure_model_present() -> None:
    missing = [f for f in REQUIRED_FILES if not (MODEL_DIR / f).exists()]
    if missing:
        raise RuntimeError(
            f"Missing model files in {MODEL_DIR}: {missing}. "
            "The entrypoint is responsible for downloading them before the server starts."
        )


def _load_recognizer() -> sherpa_onnx.OfflineRecognizer:
    _ensure_model_present()
    logger.info("loading OfflineRecognizer (provider=%s, threads=%d)", PROVIDER, NUM_THREADS)
    return sherpa_onnx.OfflineRecognizer.from_transducer(
        encoder=str(MODEL_DIR / f"encoder.{AB_SUFFIX}"),
        decoder=str(MODEL_DIR / f"decoder.{AB_SUFFIX}"),
        joiner=str(MODEL_DIR / f"joiner.{AB_SUFFIX}"),
        tokens=str(MODEL_DIR / "tokens.txt"),
        model_type="nemo_transducer",
        provider=PROVIDER,
        num_threads=NUM_THREADS,
    )


def _warm(rec: "sherpa_onnx.OfflineRecognizer") -> None:
    """Decode one throwaway buffer before the server accepts traffic.

    ⚠ NOT an optimisation — it moves a 45 s stall out of the first real request.
    ONNX Runtime's CUDA EP compiles and autotunes its kernels lazily, on the first
    decode, and on this host (RTX PRO 6000 Blackwell, sm_120) that measured **45.7 s**
    while every subsequent call was ~0.48 s. Without this, the first caller after any
    container restart sees a 45 s hang and most clients — LiteLLM's default request
    timeout included — give up long before it returns, which reads as "the service is
    broken" rather than "the service is warming".

    The healthcheck's `start_period` (300 s) is what makes paying it here safe.
    """
    try:
        t0 = time.monotonic()
        stream = rec.create_stream()
        # 1 s of silence at 16 kHz: enough to force the full encoder/decoder/joiner
        # path to compile, cheap enough not to matter.
        stream.accept_waveform(16000, np.zeros(16000, dtype=np.float32))
        rec.decode_stream(stream)
        logger.info("warmup decode complete in %.1fs — CUDA kernels compiled", time.monotonic() - t0)
    except Exception:
        # A failed warmup must not stop the server: the model is loaded and real
        # requests would still work, just with the stall back on the first caller.
        logger.exception("warmup decode failed; first real request will absorb the stall")


app = FastAPI(title="Parakeet ASR (sherpa-onnx)")
recognizer = _load_recognizer()
_warm(recognizer)


def _decode(raw: bytes) -> str:
    try:
        samples, sample_rate = sf.read(io.BytesIO(raw), dtype="float32")
    except Exception as exc:
        raise HTTPException(400, f"Could not decode audio: {exc}") from exc
    if samples.ndim > 1:
        samples = samples.mean(axis=1).astype(np.float32)

    stream = recognizer.create_stream()
    stream.accept_waveform(sample_rate, samples)
    recognizer.decode_stream(stream)
    return stream.result.text


@app.get("/healthz")
def healthz() -> dict[str, str]:
    return {"status": "ok"}


def _timed(raw: bytes) -> JSONResponse:
    t0 = time.perf_counter()
    text = _decode(raw)
    if AB_DELAY_S:
        time.sleep(AB_DELAY_S)  # positive control: blocks the loop exactly as a slower decode would
    ms = (time.perf_counter() - t0) * 1000.0
    return JSONResponse({"text": text}, headers={"x-ab-decode-ms": f"{ms:.3f}"})


@app.post("/transcribe")
async def transcribe(file: UploadFile = File(...)):
    return _timed(await file.read())


@app.post("/v1/audio/transcriptions")
async def openai_transcriptions(file: UploadFile = File(...)):
    return _timed(await file.read())