openmusic_analysis/api.py

from __future__ import annotations

import asyncio
import json
import logging
import os
import tempfile
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Annotated, Any

from fastapi import FastAPI, File, Form, Request, UploadFile
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse

from openmusic_analysis.application import MusicAnalysisService, build_service
from openmusic_analysis.domain import AnalysisResponse, ModelsResponse
from openmusic_analysis.errors import AnalysisError
from openmusic_analysis.settings import Settings


log = logging.getLogger(__name__)
SUPPORTED_EXTENSIONS = {".mp3", ".m4a", ".flac", ".wav", ".aac", ".ogg", ".opus"}


def create_app(
    *,
    service: MusicAnalysisService | None = None,
    settings: Settings | None = None,
) -> FastAPI:
    settings = settings or Settings.from_env()
    service = service or build_service(settings)

    @asynccontextmanager
    async def lifespan(_: FastAPI):
        if settings.eager_model_loading:
            await service.load_models()
        yield

    app = FastAPI(
        title="OpenMusic Music Analysis API",
        version="1.0.0",
        lifespan=lifespan,
    )
    app.state.analysis_service = service
    app.state.settings = settings

    @app.middleware("http")
    async def request_limits(request: Request, call_next: Any):
        content_length = request.headers.get("content-length")
        request_limit = (
            settings.limits.max_upload_bytes
            + settings.limits.max_lyrics_characters * 4
            + 1024 * 1024
        )
        if content_length:
            try:
                if int(content_length) > request_limit:
                    return _error_response(
                        "REQUEST_TOO_LARGE", "Request body exceeds the configured limit.", 413
                    )
            except ValueError:
                return _error_response("INVALID_CONTENT_LENGTH", "Invalid Content-Length.", 400)
        try:
            async with asyncio.timeout(settings.limits.request_timeout_seconds):
                return await call_next(request)
        except TimeoutError:
            return _error_response("REQUEST_TIMEOUT", "Analysis timed out.", 504)

    @app.exception_handler(AnalysisError)
    async def analysis_error_handler(_: Request, exc: AnalysisError):
        return _error_response(exc.code, exc.message, exc.status_code, exc.details)

    @app.exception_handler(RequestValidationError)
    async def validation_error_handler(_: Request, exc: RequestValidationError):
        missing_audio = any(
            error.get("type") == "missing" and tuple(error.get("loc", ())) == ("body", "audio")
            for error in exc.errors()
        )
        if missing_audio:
            return _error_response("MISSING_AUDIO", "Multipart field 'audio' is required.", 422)
        return _error_response(
            "INVALID_REQUEST",
            "Request validation failed.",
            422,
            {"errors": _safe_validation_errors(exc.errors())},
        )

    @app.exception_handler(Exception)
    async def unhandled_error_handler(request: Request, exc: Exception):
        log.exception("Unhandled error for %s", request.url.path, exc_info=exc)
        return _error_response("INTERNAL_ERROR", "Internal server error.", 500)

    @app.get("/v1/models", response_model=ModelsResponse)
    async def models() -> ModelsResponse:
        return ModelsResponse(models=service.registry.models())

    @app.get("/v1/status")
    async def status() -> dict[str, Any]:
        return {
            "status": "ok",
            "device": service.device,
            "loaded_models": service.registry.loaded_representations(),
            "available_representations": list(service.registry.representations),
        }

    @app.get("/health")
    async def health() -> dict[str, str]:
        return {"status": "ok"}

    @app.post("/v1/tracks/analyze", response_model=AnalysisResponse)
    async def analyze_track(
        audio: Annotated[UploadFile, File(description="Audio file")],
        lyrics: Annotated[str | None, Form()] = None,
        track_id: Annotated[str | None, Form()] = None,
        content_identity: Annotated[str | None, Form()] = None,
        requested_representations: Annotated[list[str] | None, Form()] = None,
    ) -> AnalysisResponse:
        if lyrics is not None and len(lyrics) > settings.limits.max_lyrics_characters:
            raise AnalysisError(
                "LYRICS_TOO_LARGE",
                "Lyrics exceed the configured character limit.",
                status_code=413,
            )
        suffix = Path(audio.filename or "").suffix.lower()
        if suffix not in SUPPORTED_EXTENSIONS:
            raise AnalysisError(
                "UNSUPPORTED_AUDIO_FORMAT",
                f"Supported extensions: {', '.join(sorted(SUPPORTED_EXTENSIONS))}.",
                status_code=415,
            )
        representations = _parse_representations(requested_representations)
        temp_path = await _store_upload(audio, suffix, settings.limits.max_upload_bytes)
        try:
            return await service.analyze(
                temp_path,
                lyrics=lyrics,
                requested_representations=representations,
                track_id=track_id,
                content_identity=content_identity,
            )
        finally:
            try:
                os.unlink(temp_path)
            except FileNotFoundError:
                pass

    return app


async def _store_upload(upload: UploadFile, suffix: str, max_bytes: int) -> str:
    size = 0
    path: str | None = None
    try:
        with tempfile.NamedTemporaryFile(prefix="openmusic-", suffix=suffix, delete=False) as file:
            path = file.name
            while chunk := await upload.read(1024 * 1024):
                size += len(chunk)
                if size > max_bytes:
                    raise AnalysisError(
                        "AUDIO_TOO_LARGE",
                        "Audio upload exceeds the configured byte limit.",
                        status_code=413,
                    )
                file.write(chunk)
        if size == 0:
            raise AnalysisError("EMPTY_AUDIO", "Uploaded audio is empty.", status_code=422)
        return path
    except Exception:
        if path:
            try:
                os.unlink(path)
            except FileNotFoundError:
                pass
        raise
    finally:
        await upload.close()


def _parse_representations(values: list[str] | None) -> list[str] | None:
    if values is None:
        return None
    parsed: list[str] = []
    for value in values:
        stripped = value.strip()
        if stripped.startswith("["):
            try:
                decoded = json.loads(stripped)
            except json.JSONDecodeError as exc:
                raise AnalysisError(
                    "INVALID_REPRESENTATIONS",
                    "requested_representations contains invalid JSON.",
                    status_code=422,
                ) from exc
            if not isinstance(decoded, list) or not all(isinstance(item, str) for item in decoded):
                raise AnalysisError(
                    "INVALID_REPRESENTATIONS",
                    "requested_representations JSON must be an array of strings.",
                    status_code=422,
                )
            parsed.extend(decoded)
        else:
            parsed.extend(part.strip() for part in stripped.split(",") if part.strip())
    return parsed


def _error_response(
    code: str,
    message: str,
    status_code: int,
    details: dict[str, Any] | None = None,
) -> JSONResponse:
    body: dict[str, Any] = {"error": {"code": code, "message": message}}
    if details is not None:
        body["error"]["details"] = details
    return JSONResponse(status_code=status_code, content=body)


def _safe_validation_errors(errors: list[dict[str, Any]]) -> list[dict[str, Any]]:
    return [
        {
            "type": error.get("type"),
            "location": list(error.get("loc", ())),
            "message": error.get("msg"),
        }
        for error in errors
    ]