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

import os
import time
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, get_track_by_task_id,
)

app = FastAPI()

_cors_origins = os.getenv("CORS_ALLOW_ORIGINS", "http://localhost:3007,http://localhost:8080").split(",")
app.add_middleware(
    CORSMiddleware,
    allow_origins=[o.strip() for o in _cors_origins],
    allow_credentials=False,
    allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
    allow_headers=["Content-Type"],
)

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, "음악 생성 중... (수 분 소요될 수 있습니다)")

        # 1단계: 생성 요청 → ai_task_id 반환
        resp = requests.post(
            f"{MUSIC_AI_SERVER_URL}/generate",
            json=params,
            timeout=30,
        )

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

        ai_task_id = resp.json().get("task_id")
        if not ai_task_id:
            update_task(task_id, "failed", 0, "", error="AI 서버 응답에 task_id가 없습니다")
            return

        # 2단계: 상태 폴링 (최대 10분, 5초 간격) — AI 서버 progress/message 그대로 반영
        remote_url = None
        for _ in range(120):
            time.sleep(5)
            status_resp = requests.get(f"{MUSIC_AI_SERVER_URL}/status/{ai_task_id}", timeout=10)
            status_data = status_resp.json()
            ai_status = status_data.get("status")

            # AI 서버의 progress/message를 로컬 task에 전달 (30~79 범위로 스케일)
            ai_progress = status_data.get("progress", 0)
            ai_message = status_data.get("message", "음악 생성 중...")
            scaled = 30 + int(ai_progress * 0.49)  # 30% ~ 79%
            update_task(task_id, "processing", scaled, ai_message)

            if ai_status == "succeeded":
                remote_url = status_data.get("audio_url")
                break
            elif ai_status == "failed":
                update_task(task_id, "failed", 0, "", error=status_data.get("error", "AI 서버 생성 실패"))
                return

        if not remote_url:
            update_task(task_id, "failed", 0, "", error="AI 서버 타임아웃 (10분 초과)")
            return

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

        filename = f"{task_id}.mp3"
        file_path = os.path.join(MUSIC_DATA_DIR, filename)

        # 3단계: 오디오 파일 다운로드
        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)

        # audio_url은 항상 Nginx 상대경로 (Mixed Content 방지)
        audio_url = f"/media/music/{filename}"

        # 라이브러리 자동 등록 — payload title 우선, 없으면 자동 생성
        genre = params.get("genre", "")
        moods = params.get("moods", [])
        mood_str = moods[0] if moods else "Original"
        title = params.get("title") or (f"{genre} — {mood_str} Mix" if genre else f"{mood_str} Mix")

        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):
    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 = ""


@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
    succeeded 시 track 메타데이터 포함 (라이브러리 별도 저장 불필요).
    """
    task = get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")

    resp = {
        "status":    task["status"],
        "progress":  task["progress"],
        "message":   task["message"],
        "audio_url": task["audio_url"],
        "error":     task["error"],
    }

    # succeeded 시 라이브러리에 저장된 트랙 메타데이터 포함
    # 프론트는 이 track 객체로 UI를 바로 업데이트하면 됨 (Save 버튼 불필요)
    if task["status"] == "succeeded":
        track = get_track_by_task_id(task_id)
        resp["track"] = track

    return resp


# ── 라이브러리 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}