app.py

from __future__ import annotations

import asyncio
import io
import logging
import re
import threading
import wave
from contextlib import asynccontextmanager
from functools import lru_cache
from typing import Any, Dict, List, Optional, Union

import numpy as np
import torch
from fastapi import BackgroundTasks, FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field, field_validator, model_validator
from transformers import AutoTokenizer, VitsModel
from langdetect import detect, DetectorFactory

# ------------------------- Logging -------------------------
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
DetectorFactory.seed = 0

# ------------------------- Config -------------------------
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
SAMPLE_RATE = 16000

MAX_CHARS_PER_CHUNK = 900
SILENCE_BETWEEN_SENTENCES_SEC = 0.04
SILENCE_BETWEEN_PARAGRAPHS_SEC = 0.12

# MMS currently exposes one voice per language checkpoint.
SUPPORTED_VOICES = {"default"}

LANGDETECT_TO_MMS = {
    "en": "eng", "hi": "hin", "bn": "ben", "as": "asm", "ta": "tam",
    "te": "tel", "mr": "mar", "gu": "guj", "kn": "kan", "ml": "mal",
    "pa": "pan", "ur": "urd", "or": "ori", "ne": "nep", "si": "sin",
    "ru": "rus", "fr": "fra", "de": "deu", "es": "spa", "it": "ita",
    "pt": "por", "ja": "jpn", "ko": "kor", "zh-cn": "zho", "ar": "ara",
    "fa": "fas", "tr": "tur", "pl": "pol", "uk": "ukr", "vi": "vie",
    "th": "tha", "id": "ind", "ms": "msa", "sw": "swa", "am": "amh",
    "yo": "yor", "ig": "ibo", "ha": "hau", "sn": "sna", "zu": "zul",
    "so": "som", "km": "khm", "my": "mya", "lo": "lao", "mn": "mon",
    "kk": "kaz", "uz": "uzb", "tl": "tgl", "ceb": "ceb", "hmn": "hmn",
    "nv": "nav", "sm": "smo", "ty": "tah", "haw": "haw", "mt": "mlt",
    "et": "est", "lv": "lav", "lt": "lit", "sl": "slv", "hr": "hrv",
    "sr": "srp", "bs": "bos", "mk": "mkd", "bg": "bul", "ro": "ron",
    "hu": "hun", "fi": "fin", "sv": "swe", "no": "nor", "da": "dan",
    "is": "isl", "ga": "gle", "cy": "cym", "eu": "eus", "ca": "cat",
    "gl": "glg", "he": "heb", "hy": "hye", "ka": "kat", "sa": "san",
}

LANGUAGE_MAP: Dict[str, str] = {
    "asm": "Assamese", "ben": "Bengali", "eng": "English", "hin": "Hindi",
    "tam": "Tamil", "tel": "Telugu", "mar": "Marathi", "guj": "Gujarati",
    "kan": "Kannada", "mal": "Malayalam", "pan": "Punjabi", "urd": "Urdu",
    "ori": "Odia", "nep": "Nepali", "sin": "Sinhala", "rus": "Russian",
    "fra": "French", "deu": "German", "spa": "Spanish", "ita": "Italian",
    "por": "Portuguese", "jpn": "Japanese", "kor": "Korean", "zho": "Chinese",
    "ara": "Arabic", "fas": "Persian", "tur": "Turkish", "pol": "Polish",
    "ukr": "Ukrainian", "vie": "Vietnamese", "tha": "Thai", "ind": "Indonesian",
    "msa": "Malay", "swa": "Swahili", "amh": "Amharic", "yor": "Yoruba",
    "ibo": "Igbo", "hau": "Hausa", "sna": "Shona", "zul": "Zulu",
    "som": "Somali", "khm": "Khmer", "mya": "Burmese", "lao": "Lao",
    "mon": "Mongolian", "kaz": "Kazakh", "uzb": "Uzbek", "tgl": "Tagalog",
    "ceb": "Cebuano", "hmn": "Hmong", "nav": "Navajo", "smo": "Samoan",
    "tah": "Tahitian", "haw": "Hawaiian", "mlt": "Maltese", "est": "Estonian",
    "lav": "Latvian", "lit": "Lithuanian", "slv": "Slovenian", "hrv": "Croatian",
    "srp": "Serbian", "bos": "Bosnian", "mkd": "Macedonian", "bul": "Bulgarian",
    "ron": "Romanian", "hun": "Hungarian", "fin": "Finnish", "swe": "Swedish",
    "nor": "Norwegian", "dan": "Danish", "isl": "Icelandic", "gle": "Irish",
    "cym": "Welsh", "eus": "Basque", "cat": "Catalan", "glg": "Galician",
    "heb": "Hebrew", "hye": "Armenian", "kat": "Georgian", "san": "Sanskrit",
}

# ------------------------- Thread-safe caches -------------------------
_model_cache: Dict[str, VitsModel] = {}
_tokenizer_cache: Dict[str, Any] = {}
_cache_lock = threading.Lock()


@lru_cache(maxsize=128)
def _build_model_id(language_code: str) -> str:
    return f"facebook/mms-tts-{language_code}"


def get_model_and_tokenizer(language_code: str):
    with _cache_lock:
        if language_code in _model_cache:
            return _model_cache[language_code], _tokenizer_cache[language_code]

        model_id = _build_model_id(language_code)
        logger.info("Loading model: %s", model_id)

        try:
            model = VitsModel.from_pretrained(model_id)
            tokenizer = AutoTokenizer.from_pretrained(model_id)
            model.to(DEVICE)
            model.eval()
        except Exception as e:
            logger.exception("Failed to load %s", model_id)
            raise HTTPException(
                status_code=404,
                detail=f"Language code '{language_code}' is not available in MMS TTS."
            ) from e

        _model_cache[language_code] = model
        _tokenizer_cache[language_code] = tokenizer
        return model, tokenizer


def detect_language(text: str) -> str:
    try:
        lang_code = detect(text)
        return LANGDETECT_TO_MMS.get(lang_code, lang_code)
    except Exception:
        return "eng"


def normalize_text(text: str) -> str:
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def split_paragraphs(text: str) -> List[str]:
    text = normalize_text(text)
    parts = [p.strip() for p in re.split(r"\n\s*\n+", text) if p.strip()]
    return parts if parts else [text]


def split_sentences(paragraph: str) -> List[str]:
    # English punctuation + Hindi danda.
    parts = re.split(r"(?<=[.!?。！？।॥])\s+", paragraph.strip())
    return [p.strip() for p in parts if p.strip()]


def chunk_text(text: str, max_chars: int = MAX_CHARS_PER_CHUNK) -> List[str]:
    paragraphs = split_paragraphs(text)
    chunks: List[str] = []
    current = ""

    def flush():
        nonlocal current
        if current.strip():
            chunks.append(current.strip())
        current = ""

    for paragraph in paragraphs:
        sentences = split_sentences(paragraph) or [paragraph]

        if current:
            flush()

        for sent in sentences:
            if len(sent) <= max_chars:
                if not current:
                    current = sent
                elif len(current) + 1 + len(sent) <= max_chars:
                    current += " " + sent
                else:
                    flush()
                    current = sent
            else:
                words = sent.split()
                temp = ""
                for word in words:
                    if not temp:
                        temp = word
                    elif len(temp) + 1 + len(word) <= max_chars:
                        temp += " " + word
                    else:
                        if temp.strip():
                            chunks.append(temp.strip())
                        temp = word
                if temp.strip():
                    if current and len(current) + 1 + len(temp) <= max_chars:
                        current += " " + temp
                    else:
                        if current:
                            flush()
                        current = temp

    flush()
    return chunks


def sanitize_chunks(chunks: List[str]) -> List[str]:
    return [normalize_text(c) for c in chunks if normalize_text(c)]


def resolve_language(
    request_language: Optional[str],
    auto_detect: bool,
    text: Union[str, List[str]],
) -> tuple[str, Optional[str]]:
    if not auto_detect:
        if not request_language:
            raise HTTPException(status_code=400, detail="Provide 'language' or set 'auto_detect=true'.")
        return request_language, None

    sample = ""
    if isinstance(text, str):
        sample = next((p for p in split_paragraphs(text) if p.strip()), "")
    else:
        sample = next((p for p in text if isinstance(p, str) and p.strip()), "")

    if not sample:
        raise HTTPException(status_code=400, detail="Cannot auto-detect language from empty text.")

    detected = detect_language(sample)
    return detected, detected


def synthesize_single_chunk(text: str, language_code: str) -> np.ndarray:
    model, tokenizer = get_model_and_tokenizer(language_code)
    inputs = tokenizer(text, return_tensors="pt")
    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}

    with torch.inference_mode():
        outputs = model(**inputs)
        waveform = outputs.waveform[0].detach().cpu().numpy()

    return waveform.astype(np.float32)


def merge_audio_chunks(chunk_texts: List[str], language_code: str) -> np.ndarray:
    segments: List[np.ndarray] = []

    for idx, txt in enumerate(chunk_texts):
        audio = synthesize_single_chunk(txt, language_code)
        segments.append(audio)

        if idx < len(chunk_texts) - 1:
            silence_sec = SILENCE_BETWEEN_SENTENCES_SEC
            segments.append(np.zeros(int(SAMPLE_RATE * silence_sec), dtype=np.float32))

    return np.concatenate(segments) if segments else np.zeros(0, dtype=np.float32)


def wav_bytes_from_float32(audio: np.ndarray) -> bytes:
    audio_int16 = (audio * 32767.0).clip(-32768, 32767).astype(np.int16)
    buffer = io.BytesIO()
    with wave.open(buffer, "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(2)
        wav_file.setframerate(SAMPLE_RATE)
        wav_file.writeframes(audio_int16.tobytes())
    return buffer.getvalue()


# ------------------------- Pydantic models -------------------------
class TTSRequest(BaseModel):
    text: Union[str, List[str]] = Field(
        ...,
        description="Text to speak. A single string is auto-chunked; list items are treated as separate blocks."
    )
    language: Optional[str] = Field(
        None,
        description="ISO 639-3 language code, for example 'hin', 'eng', 'ben'."
    )
    auto_detect: bool = Field(
        False,
        description="Detect language from the first non-empty chunk."
    )
    voice: str = Field(
        "default",
        description="Compatibility field. MMS currently exposes only the 'default' voice."
    )
    speed: float = Field(1.0, ge=0.5, le=2.0)

    @field_validator("text")
    @classmethod
    def validate_text(cls, v):
        if isinstance(v, str) and not v.strip():
            raise ValueError("Text cannot be empty.")
        if isinstance(v, list):
            cleaned = [x for x in v if isinstance(x, str) and x.strip()]
            if not cleaned:
                raise ValueError("Text array cannot be empty.")
        return v

    @model_validator(mode="after")
    def validate_voice(self):
        if self.voice not in SUPPORTED_VOICES:
            raise ValueError(
                f"Unsupported voice '{self.voice}'. Available voices: {sorted(SUPPORTED_VOICES)}"
            )
        return self


class TTSResponse(BaseModel):
    audio_format: str = "audio/wav"
    language_detected: Optional[str] = None
    language_used: str
    num_chunks: int
    duration_seconds: float
    chunks: List[str] = Field(default_factory=list)


# ------------------------- FastAPI -------------------------
@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Starting MMS-TTS API on %s", DEVICE)
    yield
    if DEVICE == "cuda":
        torch.cuda.empty_cache()


app = FastAPI(
    title="Enhanced MMS-TTS API",
    description="Production-grade Text-to-Speech for 1,100+ languages using Facebook's MMS models.",
    version="2.1.0",
    lifespan=lifespan,
)


@app.get("/health", tags=["System"])
async def health_check() -> Dict[str, str]:
    return {"status": "ok", "device": DEVICE, "sample_rate": str(SAMPLE_RATE)}


@app.get("/languages", tags=["Metadata"])
async def get_languages() -> Dict[str, Any]:
    languages = [{"code": code, "name": name} for code, name in sorted(LANGUAGE_MAP.items())]
    return {
        "total_languages": len(languages),
        "languages": languages,
        "voices": sorted(SUPPORTED_VOICES),
    }


@app.post("/tts", tags=["Synthesis"])
async def text_to_speech(request: TTSRequest, background_tasks: BackgroundTasks) -> StreamingResponse:
    try:
        language_code, detected_lang = resolve_language(request.language, request.auto_detect, request.text)

        if isinstance(request.text, str):
            raw_chunks = chunk_text(request.text)
        else:
            raw_chunks = []
            for part in request.text:
                raw_chunks.extend(chunk_text(part))

        chunks = sanitize_chunks(raw_chunks)
        if not chunks:
            raise HTTPException(status_code=400, detail="No valid chunks after processing.")

        audio_np = await asyncio.to_thread(merge_audio_chunks, chunks, language_code)

        if request.speed != 1.0 and len(audio_np) > 1:
            new_len = max(1, int(len(audio_np) / request.speed))
            x_old = np.linspace(0, 1, num=len(audio_np), endpoint=True)
            x_new = np.linspace(0, 1, num=new_len, endpoint=True)
            audio_np = np.interp(x_new, x_old, audio_np).astype(np.float32)

        wav_data = wav_bytes_from_float32(audio_np)
        duration = float(len(audio_np) / SAMPLE_RATE)

        headers = {
            "Content-Disposition": "attachment; filename=speech.wav",
            "X-Language-Used": language_code,
            "X-Detected-Language": detected_lang or "",
            "X-Num-Chunks": str(len(chunks)),
            "X-Duration-Seconds": f"{duration:.3f}",
            "X-Voice-Used": request.voice,
        }

        return StreamingResponse(
            io.BytesIO(wav_data),
            media_type="audio/wav",
            headers=headers,
            background=background_tasks,
        )

    except HTTPException:
        raise
    except Exception as e:
        logger.exception("Synthesis error")
        raise HTTPException(status_code=500, detail=f"Speech synthesis failed: {e}")


@app.post("/tts/info", tags=["Synthesis"], response_model=TTSResponse)
async def tts_info(request: TTSRequest) -> TTSResponse:
    language_code, detected_lang = resolve_language(request.language, request.auto_detect, request.text)

    if isinstance(request.text, str):
        raw_chunks = chunk_text(request.text)
    else:
        raw_chunks = []
        for part in request.text:
            raw_chunks.extend(chunk_text(part))

    chunks = sanitize_chunks(raw_chunks)

    return TTSResponse(
        language_detected=detected_lang,
        language_used=language_code,
        num_chunks=len(chunks),
        duration_seconds=0.0,
        chunks=chunks,
    )