audiobook/backend/main.py

"""
FastAPI backend for Text to Audiobook application.
"""
import logging
from typing import List, Optional
from fastapi import FastAPI, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.exceptions import RequestValidationError
from pydantic import BaseModel, Field

# Import services
from services.role_detector import RoleDetector
from services.tts_service import TTSService

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Pydantic Models
class RoleDetectionRequest(BaseModel):
    """Request model for role detection endpoint."""
    text: str = Field(..., description="Text to analyze for role detection")

class RoleInfo(BaseModel):
    """Model representing a detected role."""
    id: str = Field(..., description="Unique identifier for the role")
    label: str = Field(..., description="Display name for the role")

class RolesResponse(BaseModel):
    """Response model for role detection endpoint."""
    roles: List[RoleInfo] = Field(..., description="List of detected roles")

class ScriptSegment(BaseModel):
    """Model representing a script segment with speaker assignment."""
    id: str = Field(..., description="Unique segment identifier")
    speaker: str = Field(..., description="Assigned speaker/role")
    text: str = Field(..., description="Text content for this segment")

class ScriptRequest(BaseModel):
    """Request model for script generation endpoint."""
    text: str = Field(..., description="Text to segment and analyze")
    roles: List[str] = Field(..., description="List of detected role IDs")

class ScriptResponse(BaseModel):
    """Response model for script generation endpoint."""
    segments: List[ScriptSegment] = Field(..., description="List of script segments with speaker assignments")

class RoleVoiceMapping(BaseModel):
    """Model for role-to-voice mapping."""
    id: str = Field(..., description="Role identifier")
    voice: str = Field(..., description="Kokoro voice ID")

class TTSRequest(BaseModel):
    """Request model for TTS generation endpoint."""
    segments: Optional[List[ScriptSegment]] = Field(None, description="Optional edited script segments")
    text: Optional[str] = Field(None, description="Text to convert to speech (used if segments not provided)")
    lang: str = Field(..., description="Language code (a, b, f, j)")
    roles: List[RoleVoiceMapping] = Field(..., description="Role to voice mappings")
    speed: Optional[float] = Field(1.0, description="Playback speed (0.5-1.5)")

class VoicePreviewRequest(BaseModel):
    """Request model for voice preview endpoint."""
    voice: str = Field(..., description="Kokoro voice ID to preview")
    lang: str = Field(..., description="Language code (a, b, f, j)")
    speed: Optional[float] = Field(1.0, description="Playback speed (0.5-1.5)")
    text: Optional[str] = Field(None, description="Custom text to preview (uses default if not provided)")

class ErrorResponse(BaseModel):
    """Standard error response model."""
    error: str = Field(..., description="Error type")
    message: str = Field(..., description="Human-readable error message")
    details: Optional[dict] = Field(None, description="Additional error details")

# Initialize FastAPI app
app = FastAPI(
    title="Text to Audiobook API",
    description="API for converting text to multi-voice audiobooks using Kokoro TTS",
    version="1.0.0"
)

# Configure CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:5173", "http://localhost:5174", "http://localhost:3000"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Error handling middleware
@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    """Handle validation errors with custom error response."""
    logger.error(f"Validation error: {exc.errors()}")
    return JSONResponse(
        status_code=status.HTTP_400_BAD_REQUEST,
        content={
            "error": "ValidationError",
            "message": "Invalid request data",
            "details": {"errors": exc.errors()}
        }
    )

@app.exception_handler(Exception)
async def general_exception_handler(request: Request, exc: Exception):
    """Handle general exceptions with custom error response."""
    logger.error(f"Unhandled exception: {str(exc)}", exc_info=True)
    return JSONResponse(
        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        content={
            "error": "InternalServerError",
            "message": "An unexpected error occurred",
            "details": {"error": str(exc)}
        }
    )

# Initialize services
role_detector = RoleDetector()
tts_service = TTSService()

# API Endpoints
@app.get("/api/health")
async def health_check():
    """Health check endpoint."""
    logger.info("Health check requested")
    return {"status": "ok"}

@app.post("/api/tts/roles", response_model=RolesResponse)
async def detect_roles(request: RoleDetectionRequest):
    """
    Analyzes text and returns detected speaker roles.
    
    Uses AI-based detection (Ollama) first, then falls back to pattern matching.
    
    Supports two pattern types:
    - Colon pattern: "RoleName:" at line start
    - Bracket pattern: "[RoleName]" preceding dialogue
    
    Returns default "Narrator" role if no patterns found.
    """
    logger.info(f"Role detection requested for text of length {len(request.text)}")
    
    try:
        roles = await role_detector.detect_roles(request.text)
        logger.info(f"Detected {len(roles)} roles: {[r['label'] for r in roles]}")
        return RolesResponse(roles=roles)
    except Exception as e:
        logger.error(f"Error during role detection: {str(e)}", exc_info=True)
        raise

@app.post("/api/tts/script", response_model=ScriptResponse)
async def generate_script(request: ScriptRequest):
    """
    Generates script preview with AI-segmented text and speaker assignments.
    
    Uses AI (Ollama) to intelligently segment the text and assign speakers to each segment.
    Falls back to pattern-based segmentation if AI is unavailable.
    
    Args:
        request: ScriptRequest containing text and list of detected roles
        
    Returns:
        ScriptResponse with list of segments, each containing id, speaker, and text
        
    Raises:
        400: If validation fails (empty text or roles)
        500: If script generation fails
    """
    logger.info(
        f"Script generation requested for text of length {len(request.text)} "
        f"with {len(request.roles)} roles"
    )
    
    # Validate text is not empty
    if not request.text or not request.text.strip():
        logger.error("Empty text provided for script generation")
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={
                "error": "ValidationError",
                "message": "Text field cannot be empty",
                "details": {"field": "text"}
            }
        )
    
    # Validate roles list is not empty
    if not request.roles:
        logger.error("No roles provided for script generation")
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={
                "error": "ValidationError",
                "message": "At least one role is required",
                "details": {"field": "roles"}
            }
        )
    
    try:
        # Generate script segments using TTS service
        segments = await tts_service.generate_script_segments(
            text=request.text,
            roles=request.roles
        )
        
        if not segments:
            logger.warning("Script generation produced no segments")
            return JSONResponse(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                content={
                    "error": "ScriptGenerationError",
                    "message": "Failed to generate script segments",
                    "details": {"reason": "No segments produced"}
                }
            )
        
        # Convert to response format with unique IDs
        script_segments = [
            ScriptSegment(
                id=f"seg_{i}",
                speaker=speaker,
                text=text
            )
            for i, (speaker, text) in enumerate(segments)
        ]
        
        logger.info(f"Generated script with {len(script_segments)} segments")
        return ScriptResponse(segments=script_segments)
        
    except Exception as e:
        logger.error(f"Error during script generation: {str(e)}", exc_info=True)
        return JSONResponse(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            content={
                "error": "ScriptGenerationError",
                "message": "Failed to generate script",
                "details": {"error": str(e)}
            }
        )

@app.post("/api/tts/preview")
async def preview_voice(request: VoicePreviewRequest):
    """
    Generates a short audio preview of a voice.
    
    Args:
        request: VoicePreviewRequest containing voice ID, language, speed, and optional text
        
    Returns:
        StreamingResponse with audio/wav content
        
    Raises:
        400: If validation fails
        500: If audio generation fails
    """
    logger.info(f"Voice preview requested: voice={request.voice}, lang={request.lang}, speed={request.speed}")
    
    # Validate language code
    supported_langs = ['a', 'b', 'e', 'f', 'h', 'i', 'j', 'p', 'z']
    if request.lang not in supported_langs:
        logger.error(f"Invalid language code: {request.lang}")
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={
                "error": "ValidationError",
                "message": f"Invalid language code: {request.lang}",
                "details": {"field": "lang", "supported": supported_langs}
            }
        )
    
    # Validate speed range
    if request.speed < 0.5 or request.speed > 1.5:
        logger.error(f"Invalid speed value: {request.speed}")
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={
                "error": "ValidationError",
                "message": "Speed must be between 0.5 and 1.5",
                "details": {"field": "speed", "value": request.speed, "range": [0.5, 1.5]}
            }
        )
    
    # Use provided text or default preview text
    preview_text = request.text or "Hello! This is a preview of this voice. How does it sound?"
    
    try:
        # Get pipeline for language
        pipeline = tts_service.get_pipeline(request.lang)
        
        # Generate preview audio
        audio = tts_service.synthesize_segment(
            pipeline=pipeline,
            text=preview_text,
            voice=request.voice,
            speed=request.speed
        )
        
        # Convert to WAV bytes
        import io
        import soundfile as sf
        buffer = io.BytesIO()
        sf.write(buffer, audio, tts_service.SAMPLE_RATE, format='WAV')
        buffer.seek(0)
        audio_bytes = buffer.read()
        
        logger.info(f"Generated voice preview: {len(audio_bytes)} bytes")
        
        return StreamingResponse(
            iter([audio_bytes]),
            media_type="audio/wav",
            headers={"Content-Disposition": "inline; filename=preview.wav"}
        )
        
    except Exception as e:
        logger.error(f"Error during voice preview: {str(e)}", exc_info=True)
        return JSONResponse(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            content={
                "error": "PreviewGenerationError",
                "message": "Failed to generate voice preview",
                "details": {"error": str(e)}
            }
        )

@app.post("/api/tts")
async def generate_audiobook(request: TTSRequest):
    """
    Generates multi-voice audiobook from text with role-based voice assignments.
    
    Validates request fields, integrates with TTSService to generate audio,
    and streams the response with appropriate headers.
    
    Args:
        request: TTSRequest containing text, language, role mappings, and speed
        
    Returns:
        StreamingResponse with audio/wav content
        
    Raises:
        400: If validation fails (missing fields, invalid language/voice)
        500: If audio generation fails
    """
    logger.info(
        f"Audio generation requested: lang={request.lang}, "
        f"roles={len(request.roles)}, speed={request.speed}, "
        f"segments={len(request.segments) if request.segments else 0}"
    )
    
    # Validate that either segments or text is provided
    if not request.segments and not request.text:
        logger.error("Neither segments nor text provided")
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={
                "error": "ValidationError",
                "message": "Either segments or text must be provided",
                "details": {"fields": ["segments", "text"]}
            }
        )
    
    # Validate text if provided and no segments
    if not request.segments and request.text and not request.text.strip():
        logger.error("Empty text provided")
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={
                "error": "ValidationError",
                "message": "Text field cannot be empty",
                "details": {"field": "text"}
            }
        )
    
    # Validate language code
    supported_langs = ['a', 'b', 'e', 'f', 'h', 'i', 'j', 'p', 'z']
    if request.lang not in supported_langs:
        logger.error(f"Invalid language code: {request.lang}")
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={
                "error": "ValidationError",
                "message": f"Invalid language code: {request.lang}",
                "details": {
                    "field": "lang",
                    "supported": supported_langs
                }
            }
        )
    
    # Validate speed range
    if request.speed < 0.5 or request.speed > 1.5:
        logger.error(f"Invalid speed value: {request.speed}")
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={
                "error": "ValidationError",
                "message": "Speed must be between 0.5 and 1.5",
                "details": {
                    "field": "speed",
                    "value": request.speed,
                    "range": [0.5, 1.5]
                }
            }
        )
    
    # Validate roles list is not empty
    if not request.roles:
        logger.error("No role mappings provided")
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={
                "error": "ValidationError",
                "message": "At least one role mapping is required",
                "details": {"field": "roles"}
            }
        )
    
    # Convert roles list to dictionary for TTS service
    role_voice_map = {role.id: role.voice for role in request.roles}
    
    try:
        # Generate audiobook using segments if provided, otherwise use text
        if request.segments:
            logger.info(f"Generating audiobook from {len(request.segments)} edited segments")
            audio_bytes = await tts_service.generate_audiobook_from_segments(
                segments=[(seg.speaker, seg.text) for seg in request.segments],
                lang=request.lang,
                role_voice_map=role_voice_map,
                speed=request.speed
            )
        else:
            logger.info("Generating audiobook from text with automatic segmentation")
            audio_bytes = await tts_service.generate_audiobook(
                text=request.text,
                lang=request.lang,
                role_voice_map=role_voice_map,
                speed=request.speed
            )
        
        # Create streaming response with proper headers
        return StreamingResponse(
            iter([audio_bytes]),
            media_type="audio/wav",
            headers={
                "Content-Disposition": "attachment; filename=audiobook.wav"
            }
        )
        
    except ValueError as e:
        logger.error(f"Validation error during generation: {str(e)}")
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={
                "error": "ValidationError",
                "message": str(e),
                "details": {}
            }
        )
    except Exception as e:
        logger.error(f"Error during audio generation: {str(e)}", exc_info=True)
        return JSONResponse(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            content={
                "error": "AudioGenerationError",
                "message": "Failed to generate audiobook",
                "details": {"error": str(e)}
            }
        )