app/api/audiobook_routes.py

"""
Audiobook creation routes for the CSM-1B TTS API.
"""
import os
import uuid
import json
import shutil
import logging
from datetime import datetime
from typing import Optional, List
from fastapi import APIRouter, Request, HTTPException, BackgroundTasks, UploadFile, File
from fastapi.responses import FileResponse
from pydantic import BaseModel
from motor.motor_asyncio import AsyncIOMotorDatabase

from app.db import get_db, AUDIOBOOKS_COLLECTION
from app.config import AUDIO_DIR, TEXT_DIR, TEMP_DIR

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

router = APIRouter()

class AudiobookBase(BaseModel):
    title: str
    author: str
    voice_id: str
    status: str = "pending"
    created_at: datetime = datetime.utcnow()
    updated_at: datetime = datetime.utcnow()

class Audiobook(AudiobookBase):
    id: str
    file_path: Optional[str] = None
    text_path: Optional[str] = None
    error: Optional[str] = None

class TextChunk(BaseModel):
    text: str
    start_time: float
    end_time: float

async def process_audiobook(book_id: str, db: AsyncIOMotorDatabase):
    """Process the audiobook in the background."""
    try:
        # Update status to processing
        await db[AUDIOBOOKS_COLLECTION].update_one(
            {"id": book_id},
            {"$set": {"status": "processing", "updated_at": datetime.utcnow()}}
        )

        # Get the audiobook data
        audiobook = await db[AUDIOBOOKS_COLLECTION].find_one({"id": book_id})
        if not audiobook:
            raise HTTPException(status_code=404, detail="Audiobook not found")

        # TODO: Implement TTS processing logic here
        # For now, we'll just simulate processing
        logger.info(f"Processing audiobook {book_id}")

        # Update status to completed
        await db[AUDIOBOOKS_COLLECTION].update_one(
            {"id": book_id},
            {
                "$set": {
                    "status": "completed",
                    "file_path": f"{AUDIO_DIR}/{book_id}.mp3",
                    "updated_at": datetime.utcnow()
                }
            }
        )

    except Exception as e:
        logger.error(f"Error processing audiobook {book_id}: {str(e)}")
        await db[AUDIOBOOKS_COLLECTION].update_one(
            {"id": book_id},
            {
                "$set": {
                    "status": "failed",
                    "error": str(e),
                    "updated_at": datetime.utcnow()
                }
            }
        )

@router.post("/", response_model=Audiobook)
async def create_audiobook(
    background_tasks: BackgroundTasks,
    title: str,
    author: str,
    voice_id: str,
    text_file: Optional[UploadFile] = File(None),
    text_content: Optional[str] = None,
    request: Request = None
):
    """Create a new audiobook."""
    db = await get_db()
    book_id = str(uuid.uuid4())

    # Validate input
    if not text_file and not text_content:
        raise HTTPException(
            status_code=400,
            detail="Either text_file or text_content must be provided"
        )

    # Create audiobook document
    audiobook = {
        "id": book_id,
        "title": title,
        "author": author,
        "voice_id": voice_id,
        "status": "pending",
        "created_at": datetime.utcnow(),
        "updated_at": datetime.utcnow()
    }

    # Handle text input
    if text_file:
        text_path = f"{TEXT_DIR}/{book_id}.txt"
        with open(text_path, "wb") as f:
            shutil.copyfileobj(text_file.file, f)
        audiobook["text_path"] = text_path
    else:
        text_path = f"{TEXT_DIR}/{book_id}.txt"
        with open(text_path, "w") as f:
            f.write(text_content)
        audiobook["text_path"] = text_path

    # Insert audiobook into database
    await db[AUDIOBOOKS_COLLECTION].insert_one(audiobook)

    # Start background processing
    background_tasks.add_task(process_audiobook, book_id, db)

    return audiobook

@router.get("/{book_id}", response_model=Audiobook)
async def get_audiobook(book_id: str):
    """Get audiobook information."""
    db = await get_db()
    audiobook = await db[AUDIOBOOKS_COLLECTION].find_one({"id": book_id})
    if not audiobook:
        raise HTTPException(status_code=404, detail="Audiobook not found")
    return audiobook

@router.get("/{book_id}/audio")
async def get_audiobook_audio(book_id: str):
    """Get audiobook audio file."""
    db = await get_db()
    audiobook = await db[AUDIOBOOKS_COLLECTION].find_one({"id": book_id})
    
    if not audiobook:
        raise HTTPException(status_code=404, detail="Audiobook not found")
    
    if audiobook["status"] != "completed":
        raise HTTPException(
            status_code=400,
            detail=f"Audiobook is not ready (status: {audiobook['status']})"
        )
    
    file_path = audiobook.get("file_path")
    if not file_path or not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="Audio file not found")
    
    return FileResponse(
        file_path,
        media_type="audio/mpeg",
        filename=f"{audiobook['title']}.mp3"
    )

@router.get("/", response_model=List[Audiobook])
async def list_audiobooks():
    """List all audiobooks."""
    db = await get_db()
    audiobooks = await db[AUDIOBOOKS_COLLECTION].find().to_list(length=None)
    return audiobooks

@router.delete("/{book_id}")
async def delete_audiobook(book_id: str):
    """Delete an audiobook."""
    db = await get_db()
    audiobook = await db[AUDIOBOOKS_COLLECTION].find_one({"id": book_id})
    
    if not audiobook:
        raise HTTPException(status_code=404, detail="Audiobook not found")
    
    # Delete associated files
    if audiobook.get("file_path") and os.path.exists(audiobook["file_path"]):
        os.remove(audiobook["file_path"])
    if audiobook.get("text_path") and os.path.exists(audiobook["text_path"]):
        os.remove(audiobook["text_path"])
    
    # Delete from database
    await db[AUDIOBOOKS_COLLECTION].delete_one({"id": book_id})
    
    return {"message": "Audiobook deleted successfully"}