from __future__ import annotations

import argparse
import gc
import subprocess
import tempfile
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]] = {}

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 _load_model(model_id: str) -> Tuple[Any, Any]:
    if model_id in _MODEL_CACHE:
        return _MODEL_CACHE[model_id]

    import torch
    from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor

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

    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()

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


@app.on_event("startup")
def startup() -> None:
    ensure_runtime_dirs()


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

    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)


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)