import asyncio
import os
import secrets
from pathlib import Path
from threading import Lock
from typing import Literal

import httpx
from fastapi import FastAPI, Header, HTTPException
from fastapi.responses import Response
from gradio_client import Client
from pydantic import BaseModel, Field

INFERENCE_SPACE_ID = os.getenv("INFERENCE_SPACE_ID", "ramu23/ryza-zerogpu")
HF_TOKEN = os.getenv("HF_TOKEN", "")
API_KEY = os.getenv("CLOUD_TTS_API_KEY", "")

VOICES = {
    "ja_ryza": {"display_name": "莱莎 JP", "language": "ja"},
    "ja_ryza_asmr": {"display_name": "莱莎 JP（ASMR）", "language": "ja"},
    "zh_ryza": {"display_name": "莱莎 CN", "language": "zh"},
    "zh_ryza_asmr": {"display_name": "莱莎 CN（ASMR）", "language": "zh"},
}
LANGUAGES = {"zh": "中文", "ja": "日文", "en": "英文"}
CUT_METHODS = {
    "cut0": "不切",
    "cut1": "凑四句一切",
    "cut2": "凑50字一切",
    "cut3": "按中文句号。切",
    "cut4": "按英文句号.切",
    "cut5": "按标点符号切",
}

app = FastAPI(title="Ryza Cloud TTS API", version="1.0.0")
_inference_lock = asyncio.Lock()
_client_lock = Lock()
_client: Client | None = None


class TTSRequest(BaseModel):
    text: str = Field(min_length=1, max_length=2000)
    voice_id: Literal["ja_ryza", "ja_ryza_asmr", "zh_ryza", "zh_ryza_asmr"]
    language: Literal["auto", "zh", "ja", "en"] = "auto"
    response_format: Literal["wav"] = "wav"
    speed_factor: float = Field(default=1.0, ge=0.6, le=1.65)
    text_split_method: Literal["cut0", "cut1", "cut2", "cut3", "cut4", "cut5"] = "cut5"
    top_k: int = Field(default=15, ge=1, le=100)
    top_p: float = Field(default=1.0, ge=0.0, le=1.0)
    temperature: float = Field(default=1.0, ge=0.0, le=1.0)


def require_api_key(authorization: str | None) -> None:
    if not API_KEY:
        raise HTTPException(status_code=503, detail="API key is not configured")
    scheme, _, token = (authorization or "").partition(" ")
    if scheme.lower() != "bearer" or not secrets.compare_digest(token, API_KEY):
        raise HTTPException(status_code=401, detail="Invalid API key")


def get_client() -> Client:
    global _client
    if not HF_TOKEN:
        raise RuntimeError("HF_TOKEN is not configured")
    with _client_lock:
        if _client is None:
            _client = Client(INFERENCE_SPACE_ID, hf_token=HF_TOKEN, verbose=False)
        return _client


def read_audio(result: object) -> bytes:
    if isinstance(result, (list, tuple)) and len(result) == 1:
        result = result[0]
    if isinstance(result, dict):
        result = result.get("path") or result.get("url")
    if not isinstance(result, str):
        raise RuntimeError(f"Unexpected inference result: {type(result).__name__}")
    if result.startswith(("http://", "https://")):
        response = httpx.get(result, timeout=60, follow_redirects=True)
        response.raise_for_status()
        return response.content
    path = Path(result)
    audio = path.read_bytes()
    path.unlink(missing_ok=True)
    return audio


def synthesize(payload: TTSRequest) -> bytes:
    global _client
    language = payload.language
    if language == "auto":
        language = VOICES[payload.voice_id]["language"]
    try:
        result = get_client().predict(
            profile=payload.voice_id,
            text=payload.text.strip(),
            text_language=LANGUAGES[language],
            how_to_cut=CUT_METHODS[payload.text_split_method],
            top_k=payload.top_k,
            top_p=payload.top_p,
            temperature=payload.temperature,
            speed=payload.speed_factor,
            api_name="/synthesize_by_profile",
        )
        return read_audio(result)
    except Exception:
        with _client_lock:
            _client = None
        raise


@app.get("/health")
def health() -> dict:
    return {
        "status": "ok" if HF_TOKEN and API_KEY else "not_configured",
        "inference_space": INFERENCE_SPACE_ID,
    }


@app.get("/v1/voices")
def list_voices() -> dict:
    return {
        "voices": [
            {"voice_id": voice_id, **voice}
            for voice_id, voice in VOICES.items()
        ]
    }


@app.post("/v1/tts")
async def tts(payload: TTSRequest, authorization: str | None = Header(default=None)) -> Response:
    require_api_key(authorization)
    async with _inference_lock:
        try:
            audio = await asyncio.to_thread(synthesize, payload)
        except Exception as exc:
            raise HTTPException(status_code=502, detail=f"Cloud inference failed: {exc}") from exc
    return Response(
        content=audio,
        media_type="audio/wav",
        headers={"Content-Disposition": f'inline; filename="{payload.voice_id}.wav"'},
    )