from __future__ import annotations

import argparse
import os
import tempfile
from pathlib import Path
from typing import Any, Dict

import sys
sys.path.insert(0, "/app")
from tts.config import TTS_MODEL_CACHE, ensure_runtime_dirs

# TTS.api를 import하기 전에 설정해야 적용된다.
os.environ.setdefault("TTS_HOME", str(TTS_MODEL_CACHE))
os.environ.setdefault("COQUI_TOS_AGREED", "1")  # XTTS(CPML 라이선스) 동의 프롬프트를 비대화식으로 통과
# 컨테이너는 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")

import soundfile as sf
import uvicorn
from fastapi import FastAPI, Form, HTTPException
from fastapi.responses import Response

app = FastAPI(title="TTS XTTS Worker")

MODEL_NAME = "tts_models/multilingual/multi-dataset/xtts_v2"
_MODEL_CACHE: Dict[str, Any] = {}


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


def _device() -> str:
    try:
        import torch
        return "cuda" if torch.cuda.is_available() else "cpu"
    except Exception:
        return "cpu"


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

    from TTS.api import TTS

    _log(f"loading model {MODEL_NAME}")
    tts = TTS(MODEL_NAME).to(_device())
    _MODEL_CACHE[MODEL_NAME] = tts
    _log("model loaded")
    return tts


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


@app.get("/health")
def health() -> Dict[str, Any]:
    return {"status": "ok", "device": _device(), "loaded": MODEL_NAME in _MODEL_CACHE}


@app.post("/synthesize")
def synthesize(
    text: str = Form(...),
    language: str = Form("ko"),
    speaker_wav: str = Form(...),
) -> Response:
    speaker_path = Path(speaker_wav)
    if not speaker_path.exists():
        raise HTTPException(status_code=400, detail=f"speaker_wav not found: {speaker_wav}")

    try:
        tts = _load_model()
    except Exception as e:
        _log(f"model load error: {type(e).__name__}: {e}")
        raise HTTPException(status_code=500, detail=f"모델 로드 실패: {type(e).__name__}: {e}")

    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
        out_path = Path(tmp.name)

    try:
        _log(f"synthesizing language={language} chars={len(text)}")
        tts.tts_to_file(
            text=text,
            speaker_wav=str(speaker_path),
            language=language,
            file_path=str(out_path),
        )
        audio_bytes = out_path.read_bytes()
        duration = sf.info(str(out_path)).duration
        return Response(
            content=audio_bytes,
            media_type="audio/wav",
            headers={"x-audio-duration": str(round(duration, 3))},
        )
    except HTTPException:
        raise
    except Exception as e:
        _log(f"synthesize error: {type(e).__name__}: {e}")
        raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")
    finally:
        out_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=8005)
    args = parser.parse_args()
    ensure_runtime_dirs()
    _log(f"starting host={args.host} port={args.port} device={_device()}")
    uvicorn.run(app, host=args.host, port=args.port)