from __future__ import annotations

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

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 VibeVoice Worker")

_MODEL_CACHE: Dict[str, 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


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


def _free_gpu() -> None:
    gc.collect()
    try:
        import torch
        if DEVICE == "cuda" and torch.cuda.is_available():
            torch.cuda.empty_cache()
    except Exception:
        pass


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) -> 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 vibevoice.modular.modeling_vibevoice_asr import VibeVoiceASRForConditionalGeneration
    from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor

    _log(f"loading model {model_id}")
    try:
        processor = VibeVoiceASRProcessor.from_pretrained(
            model_id,
            language_model_pretrained_name="Qwen/Qwen2.5-1.5B",
            cache_dir=str(MODEL_CACHE),
        )
        model = VibeVoiceASRForConditionalGeneration.from_pretrained(
            model_id,
            dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32,
            attn_implementation="sdpa",
            trust_remote_code=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("microsoft/VibeVoice-ASR"),
    max_new_tokens: int = Form(4096),
) -> 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
    try:
        import torch

        vv_model, processor = _load_model(model)

        inputs = processor(
            audio=[str(tmp_path)],
            sampling_rate=None,
            return_tensors="pt",
            padding=True,
            add_generation_prompt=True,
        )
        inputs = {k: v.to(DEVICE) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}

        gen_config = {
            "max_new_tokens": max_new_tokens,
            "pad_token_id": processor.pad_id,
            "eos_token_id": processor.tokenizer.eos_token_id,
            "do_sample": False,
        }

        _log(f"transcribing model={model}")
        with torch.no_grad():
            output_ids = vv_model.generate(**inputs, **gen_config)

        input_length = inputs["input_ids"].shape[1]
        generated_ids = output_ids[0, input_length:]
        raw_text = processor.decode(generated_ids, skip_special_tokens=True)

        try:
            raw_segments = processor.post_process_transcription(raw_text)
        except Exception as e:
            _log(f"post_process_transcription failed: {e}")
            raw_segments = []

        segments: List[Dict[str, Any]] = []
        full_text_parts: List[str] = []
        for i, seg in enumerate(raw_segments):
            text = (seg.get("text") or "").strip()
            speaker_id = seg.get("speaker_id")
            start = seg.get("start_time")
            end = seg.get("end_time")
            segments.append({
                "id": i,
                "start": round(float(start), 3) if start is not None else None,
                "end": round(float(end), 3) if end is not None else None,
                "text": text,
                "speaker": f"SPEAKER_{speaker_id:02d}" if speaker_id is not None else "UNKNOWN",
            })
            if text:
                full_text_parts.append(text)

        duration = segments[-1]["end"] if segments and segments[-1]["end"] is not None else None
        full_text = " ".join(full_text_parts)

        return JSONResponse({
            "backend": "vibevoice",
            "model": model,
            "language": None,
            "duration": duration,
            "text": full_text,
            "segments": segments,
            "diarized": True,
        })

    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)
        # 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 output_ids
        except NameError:
            pass
        try:
            del generated_ids
        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=8006)
    args = parser.parse_args()
    _log(f"starting host={args.host} port={args.port} device={DEVICE}")
    uvicorn.run(app, host=args.host, port=args.port)