app.py

from __future__ import annotations

from datetime import datetime, timezone
from pathlib import Path
import socket
import time
from uuid import uuid4

from fastapi import FastAPI, HTTPException, Request as FastApiRequest
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

from tts_engine import (
    DEFAULT_KOKORO_LANG,
    DEFAULT_KOKORO_VOICE,
    DEFAULT_TEXT,
    DEFAULT_VOICE,
    get_kokoro_voices,
    get_supertonic_voices,
    get_voices,
    prewarm_kokoro,
    synthesize_kokoro_to_file,
    synthesize_supertonic_to_file,
    synthesize_to_file,
)


APP_DIR = Path(__file__).resolve().parent
GENERATED_DIR = APP_DIR / "generated"
STATIC_DIR = APP_DIR / "static"
TEMPLATE_DIR = APP_DIR / "templates"
INDEX_FILE = TEMPLATE_DIR / "index.html"

GENERATED_DIR.mkdir(parents=True, exist_ok=True)

MAX_TEXT_LENGTH = 3000
DEFAULT_LANG = "en"
DEFAULT_SUPERTONIC_VOICE = "M1"
VOICE_CACHE_TTL_SECONDS = 900

TRANSLATE_LANGUAGES = [
    {"code": "en", "label": "English"},
    {"code": "ar", "label": "Arabic"},
    {"code": "af", "label": "Afrikaans"},
    {"code": "sq", "label": "Albanian"},
    {"code": "am", "label": "Amharic"},
    {"code": "hy", "label": "Armenian"},
    {"code": "az", "label": "Azerbaijani"},
    {"code": "bn", "label": "Bengali"},
    {"code": "bs", "label": "Bosnian"},
    {"code": "bg", "label": "Bulgarian"},
    {"code": "ca", "label": "Catalan"},
    {"code": "zh-TW", "label": "Chinese (Traditional)"},
    {"code": "zh-CN", "label": "Chinese (Simplified)"},
    {"code": "hr", "label": "Croatian"},
    {"code": "cs", "label": "Czech"},
    {"code": "da", "label": "Danish"},
    {"code": "nl", "label": "Dutch"},
    {"code": "et", "label": "Estonian"},
    {"code": "fil", "label": "Filipino"},
    {"code": "fi", "label": "Finnish"},
    {"code": "fr", "label": "French"},
    {"code": "de", "label": "German"},
    {"code": "el", "label": "Greek"},
    {"code": "gu", "label": "Gujarati"},
    {"code": "he", "label": "Hebrew"},
    {"code": "hi", "label": "Hindi"},
    {"code": "hu", "label": "Hungarian"},
    {"code": "id", "label": "Indonesian"},
    {"code": "it", "label": "Italian"},
    {"code": "ja", "label": "Japanese"},
    {"code": "kn", "label": "Kannada"},
    {"code": "kk", "label": "Kazakh"},
    {"code": "ko", "label": "Korean"},
    {"code": "lv", "label": "Latvian"},
    {"code": "lt", "label": "Lithuanian"},
    {"code": "mk", "label": "Macedonian"},
    {"code": "ms", "label": "Malay"},
    {"code": "ml", "label": "Malayalam"},
    {"code": "mr", "label": "Marathi"},
    {"code": "ne", "label": "Nepali"},
    {"code": "no", "label": "Norwegian"},
    {"code": "fa", "label": "Persian"},
    {"code": "pl", "label": "Polish"},
    {"code": "pt", "label": "Portuguese"},
    {"code": "pa", "label": "Punjabi"},
    {"code": "ro", "label": "Romanian"},
    {"code": "ru", "label": "Russian"},
    {"code": "sr", "label": "Serbian"},
    {"code": "sk", "label": "Slovak"},
    {"code": "sl", "label": "Slovenian"},
    {"code": "es", "label": "Spanish"},
    {"code": "sw", "label": "Swahili"},
    {"code": "sv", "label": "Swedish"},
    {"code": "ta", "label": "Tamil"},
    {"code": "te", "label": "Telugu"},
    {"code": "th", "label": "Thai"},
    {"code": "tr", "label": "Turkish"},
    {"code": "uk", "label": "Ukrainian"},
    {"code": "ur", "label": "Urdu"},
    {"code": "uz", "label": "Uzbek"},
    {"code": "vi", "label": "Vietnamese"},
    {"code": "cy", "label": "Welsh"},
]

VOICE_CACHE: dict[str, dict[str, object]] = {}

app = FastAPI(title="EG Autonomous TTS Studio", version="3.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")


@app.on_event("startup")
async def prewarm_models_on_startup() -> None:
    # Preload all runtime-critical assets so the UI is fully ready on first open.
    try:
        prewarm_kokoro()
    except Exception as error:
        print(f"[startup] Kokoro prewarm failed: {error}")

    for provider in ("edge", "kokoro"):
        try:
            await get_cached_voices_(provider, None)
        except Exception as error:
            print(f"[startup] voices prewarm failed for {provider}: {error}")


class SynthesizeRequest(BaseModel):
    provider: str = Field(default="edge", pattern="^(edge|kokoro|supertonic)$")
    text: str = Field(..., min_length=1, max_length=MAX_TEXT_LENGTH)
    voice: str = Field(default=DEFAULT_VOICE, min_length=1)
    rate: str = Field(default="+0%")
    volume: str = Field(default="+0%")
    pitch: str = Field(default="+0Hz")
    lang: str = Field(default=DEFAULT_LANG, min_length=2, max_length=12)
    slow: bool = Field(default=False)
    speed: float = Field(default=1.0, ge=0.5, le=2.0)


@app.get("/", response_class=HTMLResponse)
async def index() -> HTMLResponse:
    return HTMLResponse(INDEX_FILE.read_text(encoding="utf-8"))


@app.get("/health")
async def health() -> dict[str, str]:
    return {"status": "ok"}


@app.get("/api/defaults")
async def get_defaults() -> dict[str, object]:
    return {
        "text": DEFAULT_TEXT,
        "provider": "edge",
        "voice": DEFAULT_VOICE,
        "kokoroVoice": DEFAULT_KOKORO_VOICE,
        "rate": "+0%",
        "volume": "+0%",
        "pitch": "+0Hz",
        "lang": DEFAULT_LANG,
        "kokoroLang": DEFAULT_KOKORO_LANG,
        "supertonicVoice": DEFAULT_SUPERTONIC_VOICE,
        "slow": False,
        "translateLanguages": TRANSLATE_LANGUAGES,
    }


@app.get("/api/bootstrap")
async def get_bootstrap_data(request: FastApiRequest) -> dict[str, object]:
    return {
        "defaults": await get_defaults(),
        "providers": (await get_providers())["items"],
        "network": (await get_network_info(request)),
    }


@app.get("/api/providers")
async def get_providers() -> dict[str, object]:
    return {
        "items": [
            {
                "id": "edge",
                "displayName": "Edge TTS Engine",
                "description": "Generate MP3 audio with edge-tts through the internal synthesis engine.",
            },
            {
                "id": "kokoro",
                "displayName": "Kokoro ONNX",
                "description": "Higher-quality ONNX speech model with multiple bundled voices.",
            },
            {
                "id": "supertonic",
                "displayName": "Supertonic 3",
                "description": "State-of-the-art on-device TTS with emotional and natural voice styles.",
            },
        ]
    }


@app.get("/api/network")
async def get_network_info(request: FastApiRequest) -> dict[str, object]:
    host_header = request.headers.get("host", "")
    current_host = host_header.split(":", 1)[0] if host_header else request.url.hostname or "127.0.0.1"
    current_port = request.url.port or 7860
    lan_urls = build_lan_urls_(current_port)

    return {
        "ok": True,
        "host": current_host,
        "port": current_port,
        "localUrl": f"http://127.0.0.1:{current_port}",
        "currentUrl": str(request.base_url).rstrip("/"),
        "lanUrls": lan_urls,
        "lanEnabledHint": f"python -m uvicorn app:app --host 0.0.0.0 --port {current_port}",
        "isLocalRequest": current_host in {"127.0.0.1", "localhost", "0.0.0.0"},
    }


@app.get("/api/voices")
async def list_voice_catalog(provider: str = "edge", filter: str | None = None) -> dict[str, object]:
    voices = await get_cached_voices_(provider, filter)
    return {
        "items": [
            {
                "name": voice["ShortName"],
                "locale": voice["Locale"],
                "gender": voice["Gender"],
                "friendlyName": voice.get("FriendlyName", voice["ShortName"]),
            }
            for voice in voices
        ]
    }


@app.post("/api/synthesize")
async def synthesize_audio(payload: SynthesizeRequest, request: FastApiRequest) -> dict[str, object]:
    text = normalize_text_(payload.text)
    if not text:
        raise HTTPException(status_code=400, detail="Text is required.")

    selected_voice = payload.voice.strip()

    extension = ".mp3" if payload.provider == "edge" else ".wav"
    file_name = datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S") + "-" + uuid4().hex + extension
    output_path = GENERATED_DIR / file_name

    if payload.provider == "edge":
        await synthesize_to_file(
            text=text,
            output=output_path,
            voice=selected_voice,
            rate=payload.rate.strip(),
            volume=payload.volume.strip(),
            pitch=payload.pitch.strip(),
        )
    elif payload.provider == "kokoro":
        # Guard against stale UI values (e.g., an Edge voice sent while Kokoro is selected).
        if "_" not in selected_voice:
            selected_voice = DEFAULT_KOKORO_VOICE
        kokoro_lang = normalize_kokoro_lang_(payload.lang)
        await synthesize_kokoro_to_file(
            text=text,
            output=output_path,
            voice=(selected_voice or DEFAULT_KOKORO_VOICE).strip(),
            speed=payload.speed,
            lang=kokoro_lang,
        )
    elif payload.provider == "supertonic":
        # Ensure we have a valid supertonic voice style
        if not selected_voice or selected_voice not in {"F1", "F2", "F3", "F4", "F5", "M1", "M2", "M3", "M4", "M5"}:
            selected_voice = DEFAULT_SUPERTONIC_VOICE
        await synthesize_supertonic_to_file(
            text=text,
            output=output_path,
            voice=selected_voice,
            lang=payload.lang[:2], # Supertonic expects 2-letter lang usually
        )
    else:
        raise HTTPException(status_code=400, detail="Unsupported provider.")

    audio_url = str(request.url_for("get_audio_file", file_name=file_name))
    return {
        "ok": True,
        "provider": payload.provider,
        "audioUrl": audio_url,
        "downloadUrl": audio_url,
        "fileName": file_name,
        "contentType": "audio/mpeg" if payload.provider == "edge" else "audio/wav",
        "voice": selected_voice,
        "rate": payload.rate.strip(),
        "volume": payload.volume.strip(),
        "pitch": payload.pitch.strip(),
        "speed": payload.speed,
        "lang": normalize_kokoro_lang_(payload.lang) if payload.provider == "kokoro" else payload.lang,
        "textLength": len(text),
    }


@app.get("/audio/{file_name}", name="get_audio_file")
async def get_audio_file(file_name: str) -> FileResponse:
    if "/" in file_name or "\\" in file_name:
        raise HTTPException(status_code=400, detail="Invalid file name.")

    file_path = GENERATED_DIR / file_name
    if not file_path.is_file():
        raise HTTPException(status_code=404, detail="Audio file not found.")

    media_type = "audio/mpeg" if file_path.suffix.lower() == ".mp3" else "audio/wav"
    return FileResponse(file_path, media_type=media_type, content_disposition_type="inline")


def normalize_text_(value: str) -> str:
    text = str(value or "").replace("\x00", "").replace("\r\n", "\n")
    text = "\n".join(line.strip() for line in text.split("\n"))
    text = text.strip()
    return text[:MAX_TEXT_LENGTH]


def normalize_lang_(value: str) -> str:
    lang = str(value or DEFAULT_LANG).strip()
    if not lang:
        return DEFAULT_LANG
    return lang[:12]


def normalize_kokoro_lang_(value: str) -> str:
    raw = (value or "").strip().lower()
    if not raw:
        return DEFAULT_KOKORO_LANG
    mapping = {
        "en": "en-us",
        "en-us": "en-us",
        "en_us": "en-us",
        "en-gb": "en-gb",
        "en_gb": "en-gb",
        "fr": "fr-fr",
        "fr-fr": "fr-fr",
        "fr_fr": "fr-fr",
        "ja": "ja",
        "zh": "zh",
    }
    return mapping.get(raw, DEFAULT_KOKORO_LANG)


async def get_cached_voices_(provider: str, keyword: str | None) -> list[dict[str, object]]:
    cache_key = provider + "::" + (keyword or "").strip().lower()
    cached = VOICE_CACHE.get(cache_key)
    now = time.time()

    if cached and now - float(cached["timestamp"]) < VOICE_CACHE_TTL_SECONDS:
        return cached["voices"]  # type: ignore[return-value]

    if provider == "edge":
        voices = await get_voices(keyword)
    elif provider == "kokoro":
        voices = get_kokoro_voices(keyword)
    elif provider == "supertonic":
        voices = get_supertonic_voices(keyword)
    else:
        raise HTTPException(status_code=400, detail="Unsupported provider for voices.")

    VOICE_CACHE[cache_key] = {
        "timestamp": now,
        "voices": voices,
    }
    return voices


def build_lan_urls_(port: int) -> list[str]:
    urls: list[str] = []
    candidates = {"127.0.0.1"}

    try:
        hostname = socket.gethostname()
        for result in socket.getaddrinfo(hostname, None, family=socket.AF_INET):
            ip = result[4][0]
            if ip:
                candidates.add(ip)
    except Exception:
        pass

    try:
        udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        udp_socket.connect(("8.8.8.8", 80))
        candidates.add(udp_socket.getsockname()[0])
        udp_socket.close()
    except Exception:
        pass

    for ip in sorted(candidates):
        if ip.startswith("127."):
            continue
        urls.append(f"http://{ip}:{port}")

    return urls