kithara-studio/backend/src/main.py

from fastapi import FastAPI, HTTPException, BackgroundTasks, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pathlib import Path
import os
import shutil
from dotenv import load_dotenv

# Load environment variables BEFORE importing lyria_client (it reads env at module level)
load_dotenv(Path(__file__).parent.parent / ".env")

from models import GenerateRequest, GenerateResponse, ClipList, RenameRequest, LYRIA_MODELS, DEFAULT_MODEL
from database import init_db, create_clip, get_clips, get_clip, delete_clip, update_clip_status, rename_clip, set_artwork, remove_artwork
from lyria_client import generate_clip, save_clip_file
from waveform import generate_waveform_svg, get_waveform_path

app = FastAPI(title="Lyria Studio", version="0.1.0")

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

CLIPS_DIR = Path(__file__).parent.parent / "data" / "clips"
CLIPS_DIR.mkdir(parents=True, exist_ok=True)

ARTWORK_DIR = Path(__file__).parent.parent / "data" / "artwork"
ARTWORK_DIR.mkdir(parents=True, exist_ok=True)

ALLOWED_ARTWORK_TYPES = {"image/jpeg", "image/png", "image/gif", "image/webp"}
MAX_ARTWORK_SIZE = 5 * 1024 * 1024  # 5 MB

@app.on_event("startup")
async def startup():
    await init_db()

@app.get("/api/health")
async def health():
    return {"status": "ok", "version": "0.1.0"}

@app.get("/api/models")
async def list_models():
    return {"models": LYRIA_MODELS, "default": DEFAULT_MODEL}

async def _generate_and_save(clip_id: str, prompt: str, model: str, duration: int):
    """Background task: call Lyria API and save the result."""
    try:
        result = await generate_clip(prompt, model, duration)
        
        if result["success"]:
            file_info = await save_clip_file(clip_id, result["audio_bytes"])
            await update_clip_status(
                clip_id, 
                status="ready", 
                file_size=file_info["file_size"]
            )
        else:
            await update_clip_status(
                clip_id, 
                status="failed", 
                error=result.get("error", "Unknown error")
            )
    except Exception as e:
        await update_clip_status(
            clip_id, 
            status="failed", 
            error=str(e)
        )

@app.post("/api/generate", response_model=GenerateResponse)
async def generate(request: GenerateRequest, background_tasks: BackgroundTasks):
    duration = request.get_effective_duration()
    clip_id = await create_clip(request.prompt, request.model, duration)
    
    # Kick off generation in background
    background_tasks.add_task(
        _generate_and_save,
        clip_id,
        request.prompt,
        request.model,
        duration
    )
    
    return GenerateResponse(
        id=clip_id,
        status="generating",
        message="Clip generation started. Check back for status."
    )

@app.get("/api/clips", response_model=ClipList)
async def list_clips(limit: int = 50, offset: int = 0):
    clips = await get_clips(limit, offset)
    return ClipList(clips=clips, total=len(clips))

@app.get("/api/clips/{clip_id}")
async def get_clip_by_id(clip_id: str):
    clip = await get_clip(clip_id)
    if not clip:
        raise HTTPException(status_code=404, detail="Clip not found")
    return clip

@app.get("/api/clips/{clip_id}/audio")
async def stream_audio(clip_id: str):
    clip = await get_clip(clip_id)
    if not clip:
        raise HTTPException(status_code=404, detail="Clip not found")
    file_path = CLIPS_DIR / clip["filename"]
    if not file_path.exists():
        raise HTTPException(status_code=404, detail="Audio file not found")
    return FileResponse(file_path, media_type="audio/mpeg", filename=clip["filename"])

@app.delete("/api/clips/{clip_id}")
async def remove_clip(clip_id: str):
    clip = await get_clip(clip_id)
    if not clip:
        raise HTTPException(status_code=404, detail="Clip not found")
    file_path = CLIPS_DIR / clip["filename"]
    if file_path.exists():
        file_path.unlink()
    # Also delete artwork if it exists
    if clip.get("artwork_filename"):
        artwork_path = ARTWORK_DIR / clip["artwork_filename"]
        if artwork_path.exists():
            artwork_path.unlink()
    # Also delete waveform if it exists
    waveform_path = get_waveform_path(clip_id)
    if waveform_path and waveform_path.exists():
        waveform_path.unlink()
    await delete_clip(clip_id)
    return {"deleted": True}


@app.get("/api/clips/{clip_id}/waveform")
async def get_waveform(clip_id: str):
    clip = await get_clip(clip_id)
    if not clip:
        raise HTTPException(status_code=404, detail="Clip not found")
    # Check if waveform already exists
    svg_path = get_waveform_path(clip_id)
    if svg_path and svg_path.exists():
        return FileResponse(svg_path, media_type="image/svg+xml")
    # Generate on demand
    try:
        svg_path_str = generate_waveform_svg(clip_id, clip["filename"])
        return FileResponse(svg_path_str, media_type="image/svg+xml")
    except FileNotFoundError:
        raise HTTPException(status_code=404, detail="Audio file not found")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Waveform generation failed: {str(e)}")


@app.get("/api/clips/{clip_id}/artwork")
async def get_artwork(clip_id: str):
    clip = await get_clip(clip_id)
    if not clip:
        raise HTTPException(status_code=404, detail="Clip not found")
    if not clip.get("artwork_filename"):
        raise HTTPException(status_code=404, detail="No artwork for this clip")
    artwork_path = ARTWORK_DIR / clip["artwork_filename"]
    if not artwork_path.exists():
        raise HTTPException(status_code=404, detail="Artwork file not found")
    # Determine media type from extension
    ext = artwork_path.suffix.lower()
    media_types = {".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".png": "image/png", ".gif": "image/gif", ".webp": "image/webp"}
    media_type = media_types.get(ext, "application/octet-stream")
    return FileResponse(artwork_path, media_type=media_type)


@app.patch("/api/clips/{clip_id}/artwork")
async def upload_artwork(clip_id: str, file: UploadFile = File(...)):
    clip = await get_clip(clip_id)
    if not clip:
        raise HTTPException(status_code=404, detail="Clip not found")
    # Validate content type
    if file.content_type not in ALLOWED_ARTWORK_TYPES:
        raise HTTPException(
            status_code=400,
            detail=f"Invalid file type. Allowed: {', '.join(ALLOWED_ARTWORK_TYPES)}"
        )
    # Read and validate size
    contents = await file.read()
    if len(contents) > MAX_ARTWORK_SIZE:
        raise HTTPException(status_code=400, detail="File too large. Maximum size is 5 MB.")
    # Determine extension
    ext_map = {"image/jpeg": ".jpg", "image/png": ".png", "image/gif": ".gif", "image/webp": ".webp"}
    ext = ext_map.get(file.content_type or "", ".jpg")
    artwork_filename = f"{clip_id}{ext}"
    artwork_path = ARTWORK_DIR / artwork_filename
    # Remove old artwork if it exists with a different extension
    if clip.get("artwork_filename") and clip["artwork_filename"] != artwork_filename:
        old_path = ARTWORK_DIR / clip["artwork_filename"]
        if old_path.exists():
            old_path.unlink()
    # Save the file
    with open(artwork_path, "wb") as f:
        f.write(contents)
    # Update DB
    await set_artwork(clip_id, artwork_filename)
    updated = await get_clip(clip_id)
    return updated


@app.delete("/api/clips/{clip_id}/artwork")
async def remove_artwork_endpoint(clip_id: str):
    clip = await get_clip(clip_id)
    if not clip:
        raise HTTPException(status_code=404, detail="Clip not found")
    if not clip.get("artwork_filename"):
        raise HTTPException(status_code=404, detail="No artwork for this clip")
    artwork_path = ARTWORK_DIR / clip["artwork_filename"]
    if artwork_path.exists():
        artwork_path.unlink()
    await remove_artwork(clip_id)
    updated = await get_clip(clip_id)
    return updated

@app.patch("/api/clips/{clip_id}")
async def rename_clip_endpoint(clip_id: str, request: RenameRequest):
    clip = await get_clip(clip_id)
    if not clip:
        raise HTTPException(status_code=404, detail="Clip not found")
    await rename_clip(clip_id, request.title)
    updated = await get_clip(clip_id)
    return updated

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)