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, GenerateArtworkRequest, 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, delete_all_clips, set_analysis
from lyria_client import generate_clip, save_clip_file
from waveform import generate_waveform_svg, get_waveform_path
from analysis import analyze_clip
from image_gen import generate_artwork, get_image_gen_config, get_available_models
from analysis import analyze_clip
from image_gen import generate_artwork, get_image_gen_config, get_available_models, SAFETY_LEVELS

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.get("/api/clips/{clip_id}/download")
async def download_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 not file_path.exists():
        raise HTTPException(status_code=404, detail="Audio file not found")
    
    # Build download filename from title or prompt
    base_name = clip.get("title") or clip["prompt"]
    # Sanitize for filesystem
    safe_name = "".join(c for c in base_name if c.isalnum() or c in " -_").strip()
    safe_name = safe_name[:60]  # Limit length
    if not safe_name:
        safe_name = clip_id
    download_filename = f"{safe_name}.mp3"
    
    return FileResponse(
        file_path,
        media_type="audio/mpeg",
        filename=download_filename,
        content_disposition_type="attachment"
    )

@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.post("/api/clips/{clip_id}/generate-artwork")
async def generate_clip_artwork(clip_id: str, request: GenerateArtworkRequest):
    """Generate album artwork using Gemini Image model."""
    clip = await get_clip(clip_id)
    if not clip:
        raise HTTPException(status_code=404, detail="Clip not found")

    try:
        # Generate image
        config = get_image_gen_config()
        image_bytes = await generate_artwork(request.prompt, config)

        # Save to disk
        import uuid
        artwork_filename = f"{clip_id}_{uuid.uuid4().hex[:8]}.png"
        artwork_path = ARTWORK_DIR / artwork_filename
        with open(artwork_path, "wb") as f:
            f.write(image_bytes)

        # Remove old artwork if exists
        if clip.get("artwork_filename"):
            old_path = ARTWORK_DIR / clip["artwork_filename"]
            if old_path.exists():
                old_path.unlink()

        # Update DB
        await set_artwork(clip_id, artwork_filename)
        updated = await get_clip(clip_id)
        return updated

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Image generation failed: {str(e)}")

@app.get("/api/image-gen/models")
async def list_image_gen_models():
    """List available image generation models."""
    return {"models": get_available_models()}

@app.get("/api/image-gen/config")
async def get_image_gen_configuration():
    """Get current image generation configuration (safe values only)."""
    config = get_image_gen_config()
    return {
        "model": config["model"],
        "safety_harassment": config["safety_harassment"],
        "safety_hate_speech": config["safety_hate_speech"],
        "safety_sexually_explicit": config["safety_sexually_explicit"],
        "safety_dangerous": config["safety_dangerous"],
        # API key is masked
        "has_api_key": bool(config.get("api_key")),
    }

@app.post("/api/image-gen/config")
async def update_image_gen_configuration(request: dict):
    """Update image generation configuration (in-memory only, not persisted to .env)."""
    # Update environment variables in-memory for this process
    if "model" in request:
        os.environ["IMAGE_GEN_MODEL"] = request["model"]
    if "safety_harassment" in request:
        os.environ["IMAGE_GEN_SAFETY_HARASSMENT"] = request["safety_harassment"]
    if "safety_hate_speech" in request:
        os.environ["IMAGE_GEN_SAFETY_HATE_SPEECH"] = request["safety_hate_speech"]
    if "safety_sexually_explicit" in request:
        os.environ["IMAGE_GEN_SAFETY_SEXUALLY_EXPLICIT"] = request["safety_sexually_explicit"]
    if "safety_dangerous" in request:
        os.environ["IMAGE_GEN_SAFETY_DANGEROUS"] = request["safety_dangerous"]
    
    return {"status": "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

# ─── Admin Endpoints ──────────────────────────────────────────────

@app.get("/api/clips/{clip_id}/analysis")
async def get_clip_analysis(clip_id: str):
    """Get BPM and key analysis for a clip. Computes on-demand if not cached."""
    clip = await get_clip(clip_id)
    if not clip:
        raise HTTPException(status_code=404, detail="Clip not found")
    if clip.get("status") != "ready":
        raise HTTPException(status_code=400, detail="Clip is not ready for analysis")
    
    # Return cached analysis if available
    if clip.get("bpm") and clip.get("key"):
        return {"bpm": clip["bpm"], "key": clip["key"]}
    
    # Compute analysis
    mp3_path = CLIPS_DIR / clip["filename"]
    if not mp3_path.exists():
        raise HTTPException(status_code=404, detail="Audio file not found")
    
    try:
        result = analyze_clip(mp3_path)
        # Cache in DB
        await set_analysis(clip_id, result["bpm"], result["key"])
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Analysis failed: {str(e)}")

@app.post("/api/admin/clear-waveforms")
async def clear_waveforms():
    """Delete all cached waveform SVGs. They regenerate on next view."""
    from waveform import WAVEFORM_DIR
    count = 0
    if WAVEFORM_DIR.exists():
        for f in WAVEFORM_DIR.glob("*.svg"):
            f.unlink()
            count += 1
    return {"cleared": count}

@app.post("/api/admin/delete-all-clips")
async def delete_all_clips_endpoint():
    """Delete all clips, MP3s, artwork, and waveforms. Nuclear option."""
    # Get all clips to clean up files
    all_clips = await get_clips()
    for clip in all_clips:
        # Delete MP3
        mp3_path = CLIPS_DIR / clip["filename"]
        if mp3_path.exists():
            mp3_path.unlink()
        # Delete artwork
        if clip.get("artwork_filename"):
            art_path = ARTWORK_DIR / clip["artwork_filename"]
            if art_path.exists():
                art_path.unlink()
        # Delete waveform
        wf_path = get_waveform_path(clip["id"])
        if wf_path and wf_path.exists():
            wf_path.unlink()
    # Clear DB
    await delete_all_clips()
    return {"deleted": len(all_clips)}

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