from __future__ import annotations

import argparse
import asyncio
import gc
import subprocess
import tempfile
import time
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

import uvicorn
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import JSONResponse

import sys
sys.path.insert(0, "/app")
from asr.config import DEVICE, MODEL_CACHE, ensure_runtime_dirs

import os
# 컨테이너는 983 유저로 실행되는데 HOME(/app)이 root 소유라 numba/matplotlib/torch가
# 각자 ~/.cache, ~/.config 아래에 쓰려다 실패한다. HOME을 쓰기 가능한 곳으로 돌린다.
os.environ["HOME"] = "/tmp"  # 컨테이너 기본 HOME=/app은 983 유저가 쓰기 불가 (setdefault로는 덮어쓰기 안 됨)
os.environ.setdefault("NUMBA_CACHE_DIR", "/tmp/numba_cache")

app = FastAPI(title="ASR Qwen3 Worker")

# model_id -> (model, processor). Qwen3-ASR is a chat-style audio LM, not a
# Whisper/CTC model, so it cannot go through the generic HF ASR `pipeline()`
# (no return_timestamps, no task/language generate_kwargs). Load it directly
# and drive it through Qwen3ASRProcessor.apply_transcription_request().
_MODEL_CACHE: Dict[str, Tuple[Any, Any]] = {}

# Idle-unload / model-switch eviction: only one model stays resident at a
# time, and the whole cache is dropped after IDLE_UNLOAD_SECONDS of no
# requests, so this backend doesn't permanently hog GPU memory shared with
# the other ASR workers.
IDLE_UNLOAD_SECONDS = 120
_last_used: float = 0.0
_active_requests: int = 0

LANG_MAP = {
    "ko": "korean", "en": "english", "ja": "japanese", "zh": "chinese",
    "fr": "french", "de": "german", "es": "spanish", "ru": "russian",
    "vi": "vietnamese", "th": "thai", "ar": "arabic", "pt": "portuguese",
    "it": "italian", "nl": "dutch", "pl": "polish", "tr": "turkish",
}


def _log(msg: str) -> None:
    print(f"[qwen3] {msg}", flush=True)


def _free_gpu(*objs: Any) -> None:
    for obj in objs:
        try:
            del obj
        except Exception:
            pass
    gc.collect()
    try:
        import torch
        if DEVICE == "cuda" and torch.cuda.is_available():
            torch.cuda.empty_cache()
    except Exception:
        pass


def _to_wav(src_path: Path) -> Path:
    """librosa/soundfile (used by Qwen3ASRFeatureExtractor) can't decode m4a/mp3/etc,
    so normalize everything to 16kHz mono WAV via ffmpeg before handing it to the
    processor, the same way faster-whisper does internally."""
    wav_path = src_path.with_suffix(".conv.wav")
    result = subprocess.run(
        ["ffmpeg", "-y", "-i", str(src_path), "-ar", "16000", "-ac", "1", "-f", "wav", str(wav_path)],
        capture_output=True,
    )
    if result.returncode != 0:
        raise RuntimeError(f"ffmpeg conversion failed: {result.stderr.decode(errors='replace')[-500:]}")
    return wav_path


def _unload_models() -> None:
    if not _MODEL_CACHE:
        return
    _log(f"unloading {list(_MODEL_CACHE.keys())}")
    _MODEL_CACHE.clear()
    _free_gpu()


async def _idle_unload_loop() -> None:
    while True:
        await asyncio.sleep(30)
        if _active_requests == 0 and _last_used and (time.monotonic() - _last_used) >= IDLE_UNLOAD_SECONDS:
            _unload_models()


def _load_model(model_id: str) -> Tuple[Any, Any]:
    if model_id in _MODEL_CACHE:
        return _MODEL_CACHE[model_id]

    if _MODEL_CACHE:
        # Only one model resident at a time — switching models frees the old one.
        _unload_models()

    import torch
    from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor

    _log(f"loading model {model_id}")
    dtype = torch.float16 if DEVICE == "cuda" else torch.float32

    try:
        processor = AutoProcessor.from_pretrained(model_id, cache_dir=str(MODEL_CACHE))
        model = AutoModelForSpeechSeq2Seq.from_pretrained(
            model_id,
            dtype=dtype,
            low_cpu_mem_usage=True,
            cache_dir=str(MODEL_CACHE),
        )
        model.to(DEVICE)
        model.eval()
    except Exception as e:
        # A load interrupted partway (e.g. OOM) can leave CUDA memory
        # fragmented/leaked in ways gc.collect()+empty_cache() don't reliably
        # reclaim, and since _MODEL_CACHE never got populated the idle-unload
        # loop has nothing to clean up either. Restarting the whole process
        # is the only guaranteed way to get that memory back — supervisord's
        # autorestart=true respawns it immediately.
        _log(f"model load failed, restarting process to reclaim GPU memory: {type(e).__name__}: {e}")
        os._exit(1)

    _MODEL_CACHE[model_id] = (model, processor)
    _log(f"model {model_id} loaded")
    return _MODEL_CACHE[model_id]


@app.on_event("startup")
async def startup() -> None:
    ensure_runtime_dirs()
    asyncio.create_task(_idle_unload_loop())


@app.get("/health")
def health() -> Dict[str, Any]:
    return {"status": "ok", "device": DEVICE, "loaded_models": list(_MODEL_CACHE.keys())}


@app.post("/transcribe")
async def transcribe(
    file: UploadFile = File(...),
    model: str = Form("Qwen/Qwen3-ASR-1.7B-hf"),
    language: Optional[str] = Form("ko"),
) -> JSONResponse:
    ensure_runtime_dirs()
    suffix = Path(file.filename or "audio.bin").suffix or ".wav"

    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        tmp.write(await file.read())
        tmp_path = Path(tmp.name)

    global _active_requests, _last_used
    _active_requests += 1
    wav_path: Optional[Path] = None
    try:
        import torch

        model_obj, processor = _load_model(model)

        lang_code = (language or "").strip().lower()
        lang_name = LANG_MAP.get(lang_code, lang_code) if lang_code else None

        _log(f"transcribing language={lang_code or 'auto'} model={model}")

        wav_path = _to_wav(tmp_path)
        inputs = processor.apply_transcription_request(audio=str(wav_path), language=lang_name)
        inputs = inputs.to(model_obj.device, dtype=model_obj.dtype)
        prompt_len = inputs["input_ids"].shape[1]

        with torch.no_grad():
            generated = model_obj.generate(**inputs)

        parsed = processor.decode(generated[0][prompt_len:], return_format="parsed")
        full_text = (parsed.get("transcription") or "").strip()
        detected_language = parsed.get("language") or lang_code or None

        segments = [{"id": 0, "start": None, "end": None, "text": full_text}] if full_text else []

        return JSONResponse({
            "backend": "qwen3",
            "model": model,
            "language": detected_language,
            "duration": None,
            "text": full_text,
            "segments": segments,
            "diarized": False,
        })

    except HTTPException:
        raise
    except Exception as e:
        _log(f"error: {type(e).__name__}: {e}")
        raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")
    finally:
        tmp_path.unlink(missing_ok=True)
        if wav_path is not None:
            wav_path.unlink(missing_ok=True)
        # Long files blow up KV-cache/activation memory; without this, that
        # memory stays reserved by this process and starves the other backends
        # sharing the same GPU until the container restarts. (del locals()[...]
        # does NOT work in CPython, hence the explicit names.)
        try:
            del inputs
        except NameError:
            pass
        try:
            del generated
        except NameError:
            pass
        _free_gpu()
        _active_requests -= 1
        _last_used = time.monotonic()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8004)
    args = parser.parse_args()
    _log(f"starting host={args.host} port={args.port} device={DEVICE}")
    uvicorn.run(app, host=args.host, port=args.port)