"""NVIDIA Parakeet transcription service."""

import io
import logging
import os
import tempfile
import time

import modal
from fastapi import File, HTTPException, UploadFile

MODEL_NAME = "nvidia/parakeet-tdt-0.6b-v3"
CACHE_DIR = "/cache"
MAX_UPLOAD_BYTES = 25 * 1024 * 1024
MAX_AUDIO_SECONDS = 5 * 60

app = modal.App("gtrox-transcribe")
model_cache = modal.Volume.from_name("gtrox-model-cache", create_if_missing=True)

image = (
    modal.Image.from_registry(
        "nvidia/cuda:13.0.1-cudnn-devel-ubuntu24.04",
        add_python="3.12",
    )
    .env(
        {
            "HF_HOME": CACHE_DIR,
            "HF_XET_HIGH_PERFORMANCE": "1",
            "CC": "g++",
            "CXX": "g++",
        }
    )
    .apt_install("ffmpeg")
    .uv_pip_install(
        "huggingface_hub[hf-xet]==0.31.2",
        "nemo_toolkit[asr]==2.3.2",
        "cuda-python==13.0.1",
        "soundfile==0.13.1",
        "fastapi[standard]==0.115.4",
        "librosa==0.11.0",
    )
    .entrypoint([])
)


@app.cls(
    gpu="L40S",
    image=image,
    timeout=600,
    max_containers=2,
    scaledown_window=300,
    volumes={CACHE_DIR: model_cache},
)
class Transcriber:
    @modal.enter()
    def setup(self):
        import nemo.collections.asr as nemo_asr

        logging.getLogger("nemo_logger").setLevel(logging.CRITICAL)
        self.model = nemo_asr.models.ASRModel.from_pretrained(MODEL_NAME)
        self.model.eval()

    def _transcribe(self, audio, sample_rate):
        import soundfile as sf

        temp_path = None
        try:
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
                temp_path = temp_file.name
            sf.write(temp_path, audio, sample_rate)
            result = self.model.transcribe([temp_path])[0]
            return result.text if hasattr(result, "text") else str(result)
        finally:
            if temp_path and os.path.exists(temp_path):
                os.remove(temp_path)

    @modal.fastapi_endpoint(method="POST")
    async def analyze(self, file: UploadFile = File(...)):
        import librosa
        import soundfile as sf

        contents = await file.read()
        if len(contents) > MAX_UPLOAD_BYTES:
            raise HTTPException(status_code=413, detail="Audio upload exceeds 25 MB.")

        audio, sample_rate = sf.read(io.BytesIO(contents), dtype="float32")

        if len(audio.shape) > 1:
            audio = audio.mean(axis=1)
        if sample_rate != 16000:
            audio = librosa.resample(
                audio,
                orig_sr=sample_rate,
                target_sr=16000,
            )
            sample_rate = 16000

        if len(audio) / sample_rate > MAX_AUDIO_SECONDS:
            raise HTTPException(status_code=413, detail="Audio exceeds five minutes.")

        started_at = time.time()
        transcript = self._transcribe(audio, sample_rate)
        return {
            "text": transcript,
            "duration_seconds": round(len(audio) / sample_rate, 3),
            "model": MODEL_NAME,
            "inference_seconds": round(time.time() - started_at, 3),
        }