import io
import re
import logging
import numpy as np
import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
from pydantic import BaseModel
import scipy.io.wavfile as wavfile
from pydub import AudioSegment

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="Silero TTS")

# ── Config ────────────────────────────────────────────────────────────────────
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
SAMPLE_RATE = 24000
MAX_CHUNK = 800  # chars per Silero call

# Model identifiers (passed as `speaker` to torch.hub.load — selects model file)
MODEL_ID = {"ru": "v3_1_ru", "en": "v3_en"}

# Silero speakers (passed to apply_tts)
RU_SPEAKERS = ["aidar", "baya", "kseniya", "xenia", "eugene"]
EN_SPEAKERS = [f"en_{i}" for i in range(10)]

# OpenAI voice → Silero speaker
VOICE_MAP = {
    "ru": {"alloy": "eugene", "echo": "aidar", "fable": "baya",
           "onyx": "eugene", "nova": "kseniya", "shimmer": "xenia"},
    "en": {"alloy": "en_3",  "echo": "en_1",  "fable": "en_2",
           "onyx": "en_3",   "nova": "en_4",   "shimmer": "en_5"},
}

# ── Model cache ───────────────────────────────────────────────────────────────
_models: dict[str, object] = {}


def _get_model(language: str):
    if language not in _models:
        logger.info(f"Loading Silero model {MODEL_ID[language]} lang={language} device={DEVICE}")
        model, _ = torch.hub.load(
            repo_or_dir="snakers4/silero-models",
            model="silero_tts",
            language=language,
            speaker=MODEL_ID[language],
            trust_repo=True,
        )
        model.to(DEVICE)
        _models[language] = model
        logger.info(f"Model ready: lang={language}")
    return _models[language]


@app.on_event("startup")
async def preload():
    """Preload both language models to avoid cold-start on first request."""
    for lang in ("ru", "en"):
        try:
            _get_model(lang)
        except Exception as e:
            logger.warning(f"Preload failed for lang={lang}: {e}")


# ── Helpers ───────────────────────────────────────────────────────────────────
def _is_russian(text: str) -> bool:
    return bool(re.search(r"[а-яёА-ЯЁ]", text))


def _split_sentences(text: str) -> list[str]:
    """Split on sentence boundaries, keeping chunks under MAX_CHUNK chars."""
    if len(text) <= MAX_CHUNK:
        return [text]
    parts = re.split(r"(?<=[.!?;])\s+", text.strip())
    chunks, cur = [], ""
    for part in parts:
        if len(cur) + len(part) + 1 <= MAX_CHUNK:
            cur = f"{cur} {part}" if cur else part
        else:
            if cur:
                chunks.append(cur)
            # If single part is too long, split mid-word as last resort
            cur = part[:MAX_CHUNK] if len(part) > MAX_CHUNK else part
    if cur:
        chunks.append(cur)
    return chunks or [text[:MAX_CHUNK]]


def _to_bytes(audio: torch.Tensor, fmt: str) -> bytes:
    pcm = (audio.cpu().numpy() * 32767).astype(np.int16)
    if fmt == "pcm":
        return pcm.tobytes()
    buf = io.BytesIO()
    wavfile.write(buf, SAMPLE_RATE, pcm)
    if fmt == "wav":
        return buf.getvalue()
    seg = AudioSegment.from_wav(io.BytesIO(buf.getvalue()))
    out = io.BytesIO()
    seg.export(out, format="mp3")
    return out.getvalue()


# ── API ───────────────────────────────────────────────────────────────────────
class SpeechRequest(BaseModel):
    model: str = "silero"
    input: str
    voice: str = "alloy"
    response_format: str = "mp3"
    speed: float = 1.0


@app.get("/health")
async def health():
    return {"status": "ok", "device": DEVICE}


@app.get("/v1/models")
async def list_models():
    return {
        "object": "list",
        "data": [{"id": "silero", "object": "model", "owned_by": "silero"}],
    }


@app.post("/v1/audio/speech")
async def speech(req: SpeechRequest):
    text = req.input.strip()
    if not text:
        raise HTTPException(status_code=400, detail="input is empty")

    language = "ru" if _is_russian(text) else "en"
    vm = VOICE_MAP[language]
    speaker = vm.get(req.voice, vm["alloy"])

    try:
        model = _get_model(language)
        chunks = _split_sentences(text)
        parts = [
            model.apply_tts(text=chunk, speaker=speaker, sample_rate=SAMPLE_RATE)
            for chunk in chunks
        ]
        audio = parts[0] if len(parts) == 1 else torch.cat(parts)
    except Exception as e:
        logger.error(f"TTS error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

    fmt = req.response_format.lower()
    audio_bytes = _to_bytes(audio, fmt)
    media_types = {"wav": "audio/wav", "pcm": "audio/pcm", "mp3": "audio/mpeg"}
    media_type = media_types.get(fmt, "audio/mpeg")
    return Response(content=audio_bytes, media_type=media_type)