web-page-backend/music-lab/app/main.py

import os
import uuid
import requests
from typing import List, Optional
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel

from .db import (
    init_db,
    create_task, update_task, get_task,
    get_all_tracks, add_track, delete_track, get_track_file_path,
)

app = FastAPI()

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

MUSIC_AI_SERVER_URL = os.getenv("MUSIC_AI_SERVER_URL", "")
MUSIC_DATA_DIR = "/app/data/music"
MUSIC_MEDIA_BASE = os.getenv("MUSIC_MEDIA_BASE", "/media/music")


@app.on_event("startup")
def on_startup():
    init_db()
    os.makedirs(MUSIC_DATA_DIR, exist_ok=True)


@app.get("/health")
def health():
    return {"ok": True}


# ── 음악 생성 워커 ────────────────────────────────────────────────────────────

def _run_generation(task_id: str, params: dict) -> None:
    """BackgroundTask: AI 서버에 생성 요청 → 파일 저장 → 라이브러리 등록"""
    try:
        update_task(task_id, "processing", 10, "AI 서버에 연결 중...")

        if not MUSIC_AI_SERVER_URL:
            update_task(task_id, "failed", 0, "", error="MUSIC_AI_SERVER_URL이 설정되지 않았습니다")
            return

        update_task(task_id, "processing", 30, "음악 생성 중... (수 분 소요될 수 있습니다)")

        resp = requests.post(
            f"{MUSIC_AI_SERVER_URL}/generate",
            json=params,
            timeout=600,  # 10분
            stream=True,
        )

        if resp.status_code != 200:
            update_task(task_id, "failed", 0, "", error=f"AI 서버 오류: {resp.status_code} {resp.text[:200]}")
            return

        update_task(task_id, "processing", 80, "파일 저장 중...")

        # AI 서버 응답: binary audio 또는 JSON {"audio_url": "..."}
        content_type = resp.headers.get("content-type", "")
        filename = f"{task_id}.mp3"
        file_path = os.path.join(MUSIC_DATA_DIR, filename)

        if "application/json" in content_type:
            result = resp.json()
            remote_url = result.get("audio_url") or result.get("url")
            if not remote_url:
                update_task(task_id, "failed", 0, "", error="AI 서버 응답에 audio_url이 없습니다")
                return
            # 원격 URL에서 파일 다운로드
            dl = requests.get(remote_url, timeout=120, stream=True)
            with open(file_path, "wb") as f:
                for chunk in dl.iter_content(chunk_size=8192):
                    f.write(chunk)
        else:
            # binary audio 직접 저장
            with open(file_path, "wb") as f:
                for chunk in resp.iter_content(chunk_size=8192):
                    f.write(chunk)

        audio_url = f"{MUSIC_MEDIA_BASE}/{filename}"

        # 라이브러리 자동 등록
        genre = params.get("genre", "")
        title = f"{genre} {task_id[:8]}" if genre else task_id[:8]
        add_track({
            "title": title,
            "genre": genre,
            "moods": params.get("moods", []),
            "instruments": params.get("instruments", []),
            "duration_sec": params.get("duration_sec"),
            "bpm": params.get("bpm"),
            "key": params.get("key", ""),
            "scale": params.get("scale", ""),
            "prompt": params.get("prompt", ""),
            "audio_url": audio_url,
            "file_path": file_path,
            "task_id": task_id,
        })

        update_task(task_id, "succeeded", 100, "생성 완료", audio_url=audio_url)

    except requests.Timeout:
        update_task(task_id, "failed", 0, "", error="AI 서버 타임아웃 (10분 초과)")
    except Exception as e:
        update_task(task_id, "failed", 0, "", error=str(e))


# ── 음악 생성 API ─────────────────────────────────────────────────────────────

class GenerateRequest(BaseModel):
    genre: str = ""
    moods: List[str] = []
    instruments: List[str] = []
    duration_sec: Optional[int] = None
    bpm: Optional[int] = None
    key: str = ""
    scale: str = ""
    prompt: str = ""


@app.post("/api/music/generate")
def generate_music(req: GenerateRequest, background_tasks: BackgroundTasks):
    """
    음악 생성 작업 시작. task_id 즉시 반환 후 백그라운드에서 AI 서버 호출.
    생성 완료 시 music_library에 자동 등록됨.
    """
    task_id = str(uuid.uuid4())
    params = req.model_dump()
    create_task(task_id, params)
    background_tasks.add_task(_run_generation, task_id, params)
    return {"task_id": task_id}


@app.get("/api/music/status/{task_id}")
def get_status(task_id: str):
    """
    생성 작업 상태 조회. 프론트는 succeeded 또는 failed가 될 때까지 폴링.
    status: queued | processing | succeeded | failed
    """
    task = get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    return {
        "status":    task["status"],
        "progress":  task["progress"],
        "message":   task["message"],
        "audio_url": task["audio_url"],
        "error":     task["error"],
    }


# ── 라이브러리 API ────────────────────────────────────────────────────────────

class TrackCreate(BaseModel):
    title: str = ""
    genre: str = ""
    moods: List[str] = []
    instruments: List[str] = []
    duration_sec: Optional[int] = None
    bpm: Optional[int] = None
    key: str = ""
    scale: str = ""
    prompt: str = ""
    audio_url: str = ""
    file_path: str = ""
    task_id: Optional[str] = None
    tags: List[str] = []


@app.get("/api/music/library")
def list_library():
    """저장된 트랙 목록 전체 조회 (생성일 내림차순)"""
    return {"tracks": get_all_tracks()}


@app.post("/api/music/library", status_code=201)
def save_to_library(req: TrackCreate):
    """트랙 수동 추가 (외부 파일 등록 또는 프론트 직접 저장용)"""
    track = add_track(req.model_dump())
    return track


@app.delete("/api/music/library/{track_id}")
def remove_from_library(track_id: int):
    """
    라이브러리에서 트랙 삭제. 로컬 파일도 함께 삭제.
    """
    file_path = get_track_file_path(track_id)
    ok = delete_track(track_id)
    if not ok:
        raise HTTPException(status_code=404, detail="Track not found")

    # 생성된 파일이 있으면 함께 삭제
    if file_path and os.path.isfile(file_path):
        try:
            os.remove(file_path)
        except OSError:
            pass  # 파일 삭제 실패해도 DB에서는 이미 삭제됨

    return {"ok": True}