from __future__ import annotations

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


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


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

    import torch
    from vibevoice.modular.modeling_vibevoice_asr import VibeVoiceASRForConditionalGeneration
    from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor

    _log(f"loading model {model_id}")
    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()

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

    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)


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)