src/api.py

"""
REST API Server for Multi-lingual TTS
FastAPI-based server with OpenAPI documentation

Hackathon API Specification:
- GET /Get_Inference with text, lang, speaker_wav parameters
"""

import os
import io
import time
import logging
import tempfile
import uuid
from typing import Optional, List, Dict
from pathlib import Path
import numpy as np

from fastapi import (
    FastAPI,
    HTTPException,
    Query,
    Response,
    BackgroundTasks,
    UploadFile,
    File,
    Form,
)
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, FileResponse, JSONResponse
from pydantic import BaseModel, Field
import soundfile as sf

from .engine import TTSEngine, TTSOutput
from .config import (
    LANGUAGE_CONFIGS,
    get_available_languages,
    get_available_voices,
    STYLE_PRESETS,
)
from .elevenlabs_service import ElevenLabsService

# Language mapping for XTTS voice cloning
XTTS_LANG_MAP = {
    "english": "en",
    "hindi": "hi",
    "bengali": "bn",
    "gujarati": "gu",
    "marathi": "mr",
    "telugu": "te",
    "kannada": "kn",
}

# Language name to voice key mapping (for hackathon API)
LANG_TO_VOICE = {
    "hindi": "hi_female",
    "bengali": "bn_female",
    "marathi": "mr_female",
    "telugu": "te_female",
    "kannada": "kn_female",
    "bhojpuri": "bho_female",
    "chhattisgarhi": "hne_female",
    "maithili": "mai_female",
    "magahi": "mag_female",
    "english": "en_female",
    "gujarati": "gu_mms",
}

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

# Initialize FastAPI app
app = FastAPI(
    title="Voice Tech for All - Multi-lingual TTS API",
    description="""
    A multi-lingual Text-to-Speech API supporting 10+ Indian languages.
    
    ## Features
    - 10 Indian languages with male/female voices
    - Real-time speech synthesis
    - Text normalization for Indian languages
    - Speed control
    - Multiple audio formats (WAV, MP3)
    
    ## Supported Languages
    Hindi, Bengali, Marathi, Telugu, Kannada, Bhojpuri, 
    Chhattisgarhi, Maithili, Magahi, English
    
    ## Use Case
    Built for an LLM-based healthcare assistant for pregnant mothers
    in low-income communities.
    """,
    version="1.0.0",
    contact={
        "name": "Harshil PAtel",
        "url": "https://harshilpatel.me/#contact",
    },
    license_info={
        "name": "CC BY 4.0",
        "url": "https://creativecommons.org/licenses/by/4.0/",
    },
)

# CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize TTS Engine (lazy loading)
_engine: Optional[TTSEngine] = None
_elevenlabs = ElevenLabsService()

# In-memory session voice cache for temporary cloned voice IDs
_voice_session_cache: Dict[str, str] = {}

ALLOWED_AUDIO_TYPES = {
    "audio/wav",
    "audio/x-wav",
    "audio/mpeg",
    "audio/mp3",
}
MAX_UPLOAD_BYTES = int(os.getenv("MAX_UPLOAD_BYTES", str(10 * 1024 * 1024)))


def get_engine() -> TTSEngine:
    """Get or create TTS engine instance"""
    global _engine
    if _engine is None:
        _engine = TTSEngine(device="auto")
    return _engine


# Request/Response Models
class SynthesizeRequest(BaseModel):
    """Request body for text synthesis"""

    text: str = Field(
        ..., description="Text to synthesize", min_length=1, max_length=5000
    )
    voice: str = Field(
        "hi_male", description="Voice key (e.g., hi_male, bn_female, gu_mms)"
    )
    speed: float = Field(1.0, description="Speech speed (0.5-2.0)", ge=0.5, le=2.0)
    pitch: float = Field(1.0, description="Pitch multiplier (0.5-2.0)", ge=0.5, le=2.0)
    energy: float = Field(1.0, description="Energy/volume (0.5-2.0)", ge=0.5, le=2.0)
    style: Optional[str] = Field(
        None, description="Style preset (happy, sad, calm, excited, etc.)"
    )
    normalize: bool = Field(True, description="Apply text normalization")

    class Config:
        schema_extra = {
            "example": {
                "text": "નમસ્તે, હું તમારી કેવી રીતે મદદ કરી શકું?",
                "voice": "gu_mms",
                "speed": 1.0,
                "pitch": 1.0,
                "energy": 1.0,
                "style": "calm",
                "normalize": True,
            }
        }


class SynthesizeResponse(BaseModel):
    """Response metadata for synthesis"""

    success: bool
    duration: float
    sample_rate: int
    voice: str
    text: str
    inference_time: float


class CloneResponse(BaseModel):
    """Response metadata for voice cloning"""

    success: bool
    duration: float
    sample_rate: int
    inference_time: float
    language: str


def _validate_audio_upload(upload: UploadFile, raw_bytes: bytes) -> None:
    if upload is None:
        raise HTTPException(status_code=400, detail="speaker_wav is required")

    filename = (upload.filename or "").lower()
    if not filename.endswith((".wav", ".mp3")):
        raise HTTPException(
            status_code=400, detail="Only .wav or .mp3 files are supported"
        )

    content_type = upload.content_type or ""
    if content_type and content_type not in ALLOWED_AUDIO_TYPES:
        raise HTTPException(
            status_code=400, detail=f"Unsupported content type: {content_type}"
        )

    if len(raw_bytes) == 0:
        raise HTTPException(status_code=400, detail="Uploaded audio file is empty")

    if len(raw_bytes) > MAX_UPLOAD_BYTES:
        raise HTTPException(
            status_code=400,
            detail=f"Audio file too large. Max allowed is {MAX_UPLOAD_BYTES // (1024 * 1024)} MB",
        )


class VoiceInfo(BaseModel):
    """Information about a voice"""

    key: str
    name: str
    language_code: str
    gender: str
    loaded: bool
    downloaded: bool
    model_type: str = "vits"


class HealthResponse(BaseModel):
    """Health check response"""

    status: str
    device: str
    loaded_voices: List[str]
    available_voices: int
    style_presets: List[str]


# API Endpoints
@app.get("/", response_class=JSONResponse)
async def root():
    """API root - welcome message"""
    return {
        "message": "Voice Tech for All - Multi-lingual TTS API",
        "docs": "/docs",
        "health": "/health",
        "synthesize": "/synthesize",
    }


@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check endpoint"""
    engine = get_engine()
    return HealthResponse(
        status="healthy",
        device=str(engine.device),
        loaded_voices=engine.get_loaded_voices(),
        available_voices=len(LANGUAGE_CONFIGS),
        style_presets=list(STYLE_PRESETS.keys()),
    )


@app.get("/voices", response_model=List[VoiceInfo])
async def list_voices():
    """List all available voices"""
    engine = get_engine()
    voices = engine.get_available_voices()

    return [
        VoiceInfo(
            key=key,
            name=info["name"],
            language_code=info["code"],
            gender=info["gender"],
            loaded=info["loaded"],
            downloaded=info["downloaded"],
            model_type=info.get("type", "vits"),
        )
        for key, info in voices.items()
    ]


@app.get("/styles")
async def list_styles():
    """List available style presets for prosody control"""
    return {
        "presets": STYLE_PRESETS,
        "description": {
            "speed": "Speech rate multiplier (0.5-2.0)",
            "pitch": "Pitch multiplier (0.5-2.0), >1 = higher",
            "energy": "Volume/energy multiplier (0.5-2.0)",
        },
    }


@app.get("/languages")
async def list_languages():
    """List supported languages"""
    return get_available_languages()


@app.post("/synthesize", response_class=Response)
async def synthesize_audio(request: SynthesizeRequest):
    """
    Synthesize speech from text

    Returns WAV audio file directly
    """
    engine = get_engine()

    # Validate voice
    if request.voice not in LANGUAGE_CONFIGS:
        raise HTTPException(
            status_code=400,
            detail=f"Unknown voice: {request.voice}. Use /voices to see available options.",
        )

    try:
        start_time = time.time()

        # Synthesize
        output = engine.synthesize(
            text=request.text,
            voice=request.voice,
            speed=request.speed,
            pitch=request.pitch,
            energy=request.energy,
            style=request.style,
            normalize_text=request.normalize,
        )

        inference_time = time.time() - start_time

        # Convert to WAV bytes
        buffer = io.BytesIO()
        sf.write(buffer, output.audio, output.sample_rate, format="WAV")
        buffer.seek(0)

        # Return audio with metadata headers
        return Response(
            content=buffer.read(),
            media_type="audio/wav",
            headers={
                "X-Duration": str(output.duration),
                "X-Sample-Rate": str(output.sample_rate),
                "X-Voice": output.voice,
                "X-Style": output.style or "default",
                "X-Inference-Time": str(inference_time),
            },
        )

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


@app.post("/synthesize/stream")
async def synthesize_stream(request: SynthesizeRequest):
    """
    Synthesize speech and stream the audio

    Returns streaming WAV audio
    """
    engine = get_engine()

    if request.voice not in LANGUAGE_CONFIGS:
        raise HTTPException(status_code=400, detail=f"Unknown voice: {request.voice}")

    try:
        output = engine.synthesize(
            text=request.text,
            voice=request.voice,
            speed=request.speed,
            pitch=request.pitch,
            energy=request.energy,
            style=request.style,
            normalize_text=request.normalize,
        )

        # Create streaming response
        buffer = io.BytesIO()
        sf.write(buffer, output.audio, output.sample_rate, format="WAV")
        buffer.seek(0)

        return StreamingResponse(
            buffer,
            media_type="audio/wav",
            headers={"Content-Disposition": "attachment; filename=speech.wav"},
        )

    except Exception as e:
        logger.error(f"Streaming error: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/generate", response_class=StreamingResponse)
async def generate_with_cloned_voice(
    text: str = Form(...),
    lang: str = Form("english"),
    session_id: Optional[str] = Form(None),
    speaker_wav: Optional[UploadFile] = File(None),
):
    """
    Production-ready temporary voice cloning + speech generation using ElevenLabs.
    """
    lang_lower = (lang or "english").lower().strip()
    local_session_id = session_id or uuid.uuid4().hex
    fallback_voice_id = os.getenv("ELEVENLABS_FALLBACK_VOICE_ID")

    voice_id = _voice_session_cache.get(local_session_id)

    if voice_id is None:
        if speaker_wav is None:
            raise HTTPException(
                status_code=400,
                detail="speaker_wav is required for first request in a session",
            )

        audio_bytes = await speaker_wav.read()
        _validate_audio_upload(speaker_wav, audio_bytes)

        # Re-wrap file for downstream service after validation read
        clone_file = UploadFile(
            filename=speaker_wav.filename,
            file=io.BytesIO(audio_bytes),
            headers=speaker_wav.headers,
        )

        try:
            voice_id = await _elevenlabs.clone_voice(clone_file)
            _voice_session_cache[local_session_id] = voice_id
            logger.info(
                "Cloned ElevenLabs voice for session=%s voice_id=%s",
                local_session_id,
                voice_id,
            )
        except HTTPException as exc:
            if fallback_voice_id:
                logger.warning(
                    "Clone failed for session=%s, using fallback voice",
                    local_session_id,
                )
                voice_id = fallback_voice_id
            else:
                raise exc

    try:
        audio_bytes = _elevenlabs.generate_speech(
            text=text, voice_id=voice_id, language=lang_lower
        )
    except HTTPException as exc:
        if fallback_voice_id and voice_id != fallback_voice_id:
            logger.warning(
                "TTS failed for voice=%s, retrying with fallback voice", voice_id
            )
            audio_bytes = _elevenlabs.generate_speech(
                text=text,
                voice_id=fallback_voice_id,
                language=lang_lower,
            )
        else:
            raise exc

    headers = {
        "Content-Disposition": "attachment; filename=generated.mp3",
        "X-Session-Id": local_session_id,
        "X-Voice-Id": voice_id,
        "X-Provider": "elevenlabs",
    }
    return StreamingResponse(
        io.BytesIO(audio_bytes), media_type="audio/mpeg", headers=headers
    )


@app.post("/clone", response_class=Response)
async def clone_voice(
    text: str = Query(..., description="Text to synthesize with cloned voice"),
    lang: str = Query(
        "english",
        description="Language name (english, hindi, bengali, gujarati, marathi, telugu, kannada)",
    ),
    speed: float = Query(1.0, description="Speech speed", ge=0.5, le=2.0),
    pitch: float = Query(1.0, description="Pitch", ge=0.5, le=2.0),
    energy: float = Query(1.0, description="Energy", ge=0.5, le=2.0),
    style: Optional[str] = Query(None, description="Style preset"),
    session_id: Optional[str] = Query(
        None, description="Session key to reuse cloned voice"
    ),
    speaker_wav: UploadFile = File(
        ..., description="Reference speaker WAV (3-15 seconds recommended)"
    ),
):
    """
    Backward-compatible clone endpoint using ElevenLabs voice cloning.
    """
    lang_lower = lang.lower().strip()
    local_session_id = session_id or uuid.uuid4().hex
    fallback_voice_id = os.getenv("ELEVENLABS_FALLBACK_VOICE_ID")

    if lang_lower not in XTTS_LANG_MAP:
        supported = ", ".join(sorted(XTTS_LANG_MAP.keys()))
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported clone language: {lang}. Supported: {supported}",
        )

    voice_id = _voice_session_cache.get(local_session_id)
    if voice_id is None:
        audio_bytes = await speaker_wav.read()
        _validate_audio_upload(speaker_wav, audio_bytes)
        clone_file = UploadFile(
            filename=speaker_wav.filename,
            file=io.BytesIO(audio_bytes),
            headers=speaker_wav.headers,
        )
        try:
            voice_id = await _elevenlabs.clone_voice(clone_file)
            _voice_session_cache[local_session_id] = voice_id
        except HTTPException as exc:
            if fallback_voice_id:
                voice_id = fallback_voice_id
            else:
                raise exc

    try:
        audio_bytes = _elevenlabs.generate_speech(
            text=text,
            voice_id=voice_id,
            language=lang_lower,
        )
    except HTTPException as exc:
        if fallback_voice_id and voice_id != fallback_voice_id:
            audio_bytes = _elevenlabs.generate_speech(
                text=text,
                voice_id=fallback_voice_id,
                language=lang_lower,
            )
        else:
            raise exc

    return Response(
        content=audio_bytes,
        media_type="audio/mpeg",
        headers={
            "Content-Disposition": "attachment; filename=cloned_output.mp3",
            "X-Language": lang_lower,
            "X-Voice": voice_id,
            "X-Session-Id": local_session_id,
            "X-Provider": "elevenlabs",
        },
    )


@app.get("/synthesize/get")
async def synthesize_get(
    text: str = Query(
        ..., description="Text to synthesize", min_length=1, max_length=1000
    ),
    voice: str = Query("hi_male", description="Voice key"),
    speed: float = Query(1.0, description="Speech speed", ge=0.5, le=2.0),
    pitch: float = Query(1.0, description="Pitch", ge=0.5, le=2.0),
    energy: float = Query(1.0, description="Energy", ge=0.5, le=2.0),
    style: Optional[str] = Query(None, description="Style preset"),
):
    """
    GET endpoint for simple synthesis

    Useful for testing and simple integrations
    """
    request = SynthesizeRequest(
        text=text, voice=voice, speed=speed, pitch=pitch, energy=energy, style=style
    )
    return await synthesize_audio(request)


@app.api_route("/Get_Inference", methods=["GET", "POST"])
async def get_inference(
    text: str = Query(
        ...,
        description="The input text to be converted into speech. For English, text must be lowercase.",
    ),
    lang: str = Query(
        ...,
        description="Language of input text. Supported: bhojpuri, bengali, english, gujarati, hindi, chhattisgarhi, kannada, magahi, maithili, marathi, telugu",
    ),
    speaker_wav: UploadFile = File(
        ...,
        description="A reference WAV file representing the speaker's voice (mandatory per hackathon spec).",
    ),
):
    """
    Hackathon API - Generate speech audio from text

    This endpoint follows the Voice Tech for All hackathon specification.

    Supports both GET and POST methods with multipart form data.

    Parameters:
    - text: Input text to synthesize (query param)
    - lang: Language (query param) - bhojpuri, bengali, english, gujarati, hindi, chhattisgarhi, kannada, magahi, maithili, marathi, telugu
    - speaker_wav: Reference WAV file (multipart file upload, mandatory)

    Returns:
    - 200 OK: WAV audio file as streaming response
    """
    engine = get_engine()

    # Normalize language name
    lang_lower = lang.lower().strip()

    # Enforce lowercase for English text (per spec)
    if lang_lower == "english":
        text = text.lower()

    # Map language to voice
    if lang_lower not in LANG_TO_VOICE:
        supported = list(LANG_TO_VOICE.keys())
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported language: {lang}. Supported languages: {', '.join(supported)}",
        )

    voice = LANG_TO_VOICE[lang_lower]

    # Read speaker_wav (mandatory per spec)
    # Note: Current VITS models don't support voice cloning, but we accept the file
    # for API compatibility and validation. In future, this could be used for voice adaptation.
    try:
        speaker_audio_bytes = await speaker_wav.read()
        logger.info(
            f"Received speaker reference WAV: {len(speaker_audio_bytes)} bytes, filename: {speaker_wav.filename}"
        )
        # Validate it's a valid audio file (basic check)
        if len(speaker_audio_bytes) < 44:  # Minimum WAV header size
            raise HTTPException(
                status_code=400,
                detail="Invalid speaker_wav: file too small to be a valid WAV",
            )
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Could not read speaker_wav: {e}")
        raise HTTPException(
            status_code=400, detail=f"Failed to read speaker_wav file: {str(e)}"
        )

    try:
        # Synthesize audio
        output = engine.synthesize(
            text=text,
            voice=voice,
            speed=1.0,
            normalize_text=True,
        )

        # Convert to WAV bytes
        buffer = io.BytesIO()
        sf.write(buffer, output.audio, output.sample_rate, format="WAV")
        buffer.seek(0)

        # Return as streaming response (per spec)
        return StreamingResponse(
            buffer,
            media_type="audio/wav",
            headers={
                "Content-Disposition": "attachment; filename=output.wav",
                "X-Duration": str(output.duration),
                "X-Sample-Rate": str(output.sample_rate),
                "X-Language": lang,
                "X-Voice": voice,
            },
        )

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


@app.post("/preload")
async def preload_voice(voice: str):
    """Preload a voice model into memory"""
    engine = get_engine()

    if voice not in LANGUAGE_CONFIGS:
        raise HTTPException(status_code=400, detail=f"Unknown voice: {voice}")

    try:
        engine.load_voice(voice)
        return {"message": f"Voice {voice} loaded successfully"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/unload")
async def unload_voice(voice: str):
    """Unload a voice model from memory"""
    engine = get_engine()
    engine.unload_voice(voice)
    return {"message": f"Voice {voice} unloaded"}


@app.post("/batch")
async def batch_synthesize(
    texts: List[str], voice: str = "hi_male", speed: float = 1.0
):
    """
    Synthesize multiple texts

    Returns a list of base64-encoded audio
    """
    import base64

    engine = get_engine()

    if voice not in LANGUAGE_CONFIGS:
        raise HTTPException(status_code=400, detail=f"Unknown voice: {voice}")

    results = []
    for text in texts:
        output = engine.synthesize(text, voice, speed)

        buffer = io.BytesIO()
        sf.write(buffer, output.audio, output.sample_rate, format="WAV")
        buffer.seek(0)

        results.append(
            {
                "text": text,
                "audio_base64": base64.b64encode(buffer.read()).decode(),
                "duration": output.duration,
            }
        )

    return results


# Startup/Shutdown events
@app.on_event("startup")
async def startup_event():
    """Initialize on startup"""
    logger.info("Starting TTS API server...")
    # Optionally preload default voice
    # get_engine().load_voice("hi_male")


@app.on_event("shutdown")
async def shutdown_event():
    """Cleanup on shutdown"""
    logger.info("Shutting down TTS API server...")


def start_server(host: str = "0.0.0.0", port: int = 8000, reload: bool = False):
    """Start the API server"""
    import uvicorn

    uvicorn.run("src.api:app", host=host, port=port, reload=reload, log_level="info")


if __name__ == "__main__":
    start_server()