import gc
import os
import re
import shutil
import subprocess
import tempfile
import time
import uuid
from pathlib import Path

import torch
import uvicorn
from fastapi import FastAPI, File, Form, UploadFile
from fastapi.responses import JSONResponse
from lhotse import Recording
from lhotse.dataset import DynamicCutSampler
from optimum.quanto import quantize, freeze, qint8

from nemo.collections.speechlm2.models import SALM

MODEL_ID = "nvidia/canary-qwen-2.5b"
SAMPLE_RATE = 16000
CHUNK_SECONDS = 40.0
BATCH_SIZE = 1

app = FastAPI(title="Canary-Qwen Quanto STT", version="0.1")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = None


def clean_text(s: str) -> str:
    s = s.replace("<|im_start|>", "").replace("<|im_end|>", "")
    s = re.sub(r"<think>.*?</think>", "", s, flags=re.S)
    s = s.replace("<think>", "").replace("</think>", "")
    s = re.sub(r"\s+", " ", s).strip()
    return s


def load_model():
    global model
    if model is not None:
        return model

    print("[CANARY] loading SALM", flush=True)
    m = SALM.from_pretrained(MODEL_ID).eval()

    print("[CANARY] quantizing LLM with Quanto qint8", flush=True)
    quantize(m.llm, weights=qint8)
    freeze(m.llm)

    m.embed_tokens = m.embed_tokens.to(device)
    m.llm = m.llm.to(device)
    m.perception = m.perception.to(device).eval()

    model = m
    print("[CANARY] ready", flush=True)
    return model


@app.get("/health")
def health():
    free = total = None
    if torch.cuda.is_available():
        free_b, total_b = torch.cuda.mem_get_info()
        free = round(free_b / 1024**3, 3)
        total = round(total_b / 1024**3, 3)
    return {
        "ok": True,
        "backend": "canary-qwen-2.5b-quanto-int8",
        "device": str(device),
        "cuda": torch.cuda.is_available(),
        "vram_free_gib": free,
        "vram_total_gib": total,
    }


@app.get("/v1/models")
def models():
    return {
        "object": "list",
        "data": [
            {
                "id": "canary-qwen-2.5b-quanto-int8",
                "object": "model",
                "owned_by": "local",
            }
        ],
    }


@app.post("/v1/audio/transcriptions")
async def transcriptions(
    file: UploadFile = File(...),
    model_name: str = Form(default="canary-qwen-2.5b-quanto-int8", alias="model"),
    language: str = Form(default="en"),
    response_format: str = Form(default="json"),
):
    m = load_model()
    start = time.time()

    suffix = Path(file.filename or "audio.wav").suffix or ".wav"
    tmpdir = Path(tempfile.mkdtemp(prefix="canary_stt_"))
    src = tmpdir / f"input{suffix}"

    try:
        with src.open("wb") as f:
            shutil.copyfileobj(file.file, f)

        # Open WebUI browser mic uploads usually arrive as webm/opus.
        # Lhotse/Torchaudio can choke on that in this venv, so normalize first.
        wav = tmpdir / "input_16k_mono.wav"
        cmd = [
            "ffmpeg",
            "-y",
            "-hide_banner",
            "-loglevel",
            "error",
            "-i",
            str(src),
            "-vn",
            "-ar",
            str(SAMPLE_RATE),
            "-ac",
            "1",
            "-c:a",
            "pcm_s16le",
            str(wav),
        ]
        subprocess.run(cmd, check=True)

        rec = Recording.from_file(str(wav), recording_id=str(uuid.uuid4()))
        cut = rec.to_cut()
        if cut.num_channels > 1:
            cut = cut.to_mono(mono_downmix=True)

        sampler = DynamicCutSampler(cut.cut_into_windows(CHUNK_SECONDS), max_cuts=BATCH_SIZE)

        pred_text = []

        for batch in sampler:
            audio, audio_lens = batch.load_audio(collate=True)

            with torch.inference_mode():
                output_ids = m.generate(
                    prompts=[
                        [
                            {
                                "role": "user",
                                "content": f"Transcribe the following: {m.audio_locator_tag}",
                            }
                        ]
                    ] * len(batch),
                    audios=torch.as_tensor(audio).to(device, non_blocking=True),
                    audio_lens=torch.as_tensor(audio_lens).to(device, non_blocking=True),
                    max_new_tokens=256,
                )

            for oids in output_ids.cpu():
                pred_text.append(clean_text(m.tokenizer.ids_to_text(oids)))

        text = " ".join(x for x in pred_text if x).strip()
        elapsed = round(time.time() - start, 3)

        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        if response_format == "text":
            return JSONResponse(content=text)

        return {
            "text": text,
            "model": "canary-qwen-2.5b-quanto-int8",
            "backend": "canary-qwen-quanto",
            "language": language,
            "elapsed_sec": elapsed,
        }

    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)


if __name__ == "__main__":
    load_model()
    uvicorn.run(app, host="0.0.0.0", port=8022)