"""
MeloTTS-Korean 테스트 서버.

POST /tts          { "text": "...", "speed": 1.0 }  -> audio/wav
GET  /             -> index.html (테스트 UI)

실행: python server.py    (기본 0.0.0.0:7860)
"""
import os
import io
import re
import time
import tempfile
import logging
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("melo-test")

# MeloTTS 모델 핸들 (lifespan 에서 로드)
state = {"model": None, "speaker_id": None}


@asynccontextmanager
async def lifespan(_: FastAPI):
    # import 가 무거우므로 lifespan 안에서 한 번만
    from melo.api import TTS
    log.info("MeloTTS-Korean 모델 로딩중... (첫 실행은 가중치 다운로드로 수 분 걸릴 수 있음)")
    t0 = time.time()
    device = os.environ.get("MELO_DEVICE", "auto")
    model = TTS(language="KR", device=device)
    spk = model.hps.data.spk2id
    # 한국어 모델은 보통 {"KR": 0}.
    # spk2id 는 dict 가 아니라 MeloTTS 의 HParams 인스턴스라
    # `.get()` 이 없고 `__contains__` / `__getitem__` 만 지원함.
    speaker_id = spk["KR"] if "KR" in spk else next(iter(spk.values()))
    state["model"] = model
    state["speaker_id"] = speaker_id
    log.info("모델 로딩 완료 (%.1fs). speakers=%s, 사용=%s",
             time.time() - t0, list(spk.keys()), speaker_id)
    yield
    state["model"] = None


app = FastAPI(lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


class TTSReq(BaseModel):
    text: str
    speed: float = 1.0


# MeloTTS-Korean 의 symbol vocab 은 완성형 음절(가, 나, ...)만 포함한다.
# 'ㅋㅋㅋ', 'ㅠㅠ', 'ㅎㅎ' 같이 한글 자모(Hangul Compatibility Jamo /
# Hangul Jamo) 가 단독으로 들어오면 cleaned_text_to_sequence 에서
# KeyError 로 죽는다. 합성 전에 미리 제거한다.
_STANDALONE_JAMO_RE = re.compile(r"[\u1100-\u11FF\u3130-\u318F\uA960-\uA97F\uD7B0-\uD7FF]+")


def _sanitize_for_melo(text: str) -> str:
    # 단독 자모를 공백으로 치환한 뒤 연속 공백을 하나로 정리.
    cleaned = _STANDALONE_JAMO_RE.sub(" ", text)
    return re.sub(r"\s+", " ", cleaned).strip()


@app.post("/tts")
def tts(req: TTSReq):
    model = state.get("model")
    if model is None:
        raise HTTPException(503, "model not ready")
    raw = (req.text or "").strip()
    if not raw:
        raise HTTPException(400, "empty text")
    if len(raw) > 500:
        raise HTTPException(400, "too long (<=500)")
    text = _sanitize_for_melo(raw)
    if not text:
        raise HTTPException(
            400,
            "text contains only standalone Hangul jamo (e.g. ㅋㅋㅋ); "
            "MeloTTS-Korean cannot synthesize these.",
        )
    if text != raw:
        log.info("sanitized standalone jamo: %r -> %r", raw[:60], text[:60])

    speed = max(0.5, min(2.0, float(req.speed or 1.0)))
    fd, out_path = tempfile.mkstemp(suffix=".wav")
    os.close(fd)
    try:
        t0 = time.time()
        model.tts_to_file(text, state["speaker_id"], out_path, speed=speed)
        dt_ms = int((time.time() - t0) * 1000)
        with open(out_path, "rb") as f:
            buf = f.read()
        log.info("tts ok: %dms, %dB, text=%r", dt_ms, len(buf), text[:30])
        return Response(
            content=buf,
            media_type="audio/wav",
            headers={"X-Synth-Ms": str(dt_ms), "Cache-Control": "no-store"},
        )
    finally:
        try:
            os.unlink(out_path)
        except OSError:
            pass


# 정적 파일 (index.html) 서빙: 반드시 라우트 등록 뒤에 mount
HERE = os.path.dirname(os.path.abspath(__file__))
app.mount("/", StaticFiles(directory=HERE, html=True), name="static")


if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", "7860"))
    uvicorn.run(app, host="0.0.0.0", port=port)