backend/space_app.py

"""MuSProt-only FastAPI application for Hugging Face Docker Spaces."""
from __future__ import annotations

import logging
import os
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse, PlainTextResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles

from app.protein.api.routes.protein import router as protein_router
from app.protein.api.routes.protein import set_data_manager
from app.protein.config import get_database_path, get_docs_path, get_plots_dir
from app.protein.data_loader import DataManager
from app.protein import tsv_loader


logging.basicConfig(
    level=os.getenv("LOG_LEVEL", "INFO"),
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)

BACKEND_DIR = Path(__file__).resolve().parent
FRONTEND_DIR = Path(os.getenv("MUSPROT_FRONTEND_DIR", BACKEND_DIR / "frontend"))


@asynccontextmanager
async def lifespan(_: FastAPI):
    """Initialize the shared read-only database manager."""
    db_path = get_database_path()
    logger.info("Using read-only MuSProt database: %s", db_path)
    manager = DataManager(db_path)
    manager.load_data()
    set_data_manager(manager)

    if os.getenv("MUSPROT_PRELOAD_NODE_INDEX", "0") == "1":
        tsv_loader._get_index()

    yield
    set_data_manager(None)


app = FastAPI(
    title="MuSProt",
    description="Multistate protein structure dataset explorer",
    version="1.0.0",
    lifespan=lifespan,
    docs_url="/api/docs",
    redoc_url="/api/redoc",
    openapi_url="/api/openapi.json",
)
app.include_router(protein_router, prefix="/api/protein", tags=["MuSProt"])


@app.get("/health")
async def health():
    return {"status": "healthy", "database": str(get_database_path())}


@app.get("/api/protein/docs", response_class=PlainTextResponse)
async def protein_docs():
    path = get_docs_path()
    if not path:
        raise HTTPException(status_code=404, detail="MuSProt documentation is not available")
    return path.read_text(encoding="utf-8")


@app.get("/api/protein/download/db")
async def download_database():
    repo_id = os.getenv("MUSPROT_DATASET_REPO")
    if repo_id:
        revision = os.getenv("MUSPROT_DATASET_REVISION", "main")
        filename = os.getenv("MUSPROT_DB_FILENAME", "MuSProt.db")
        return RedirectResponse(
            f"https://huggingface.co/datasets/{repo_id}/resolve/{revision}/{filename}?download=true"
        )

    return FileResponse(
        get_database_path(),
        media_type="application/x-sqlite3",
        filename="MuSProt.db",
    )


plots_dir = get_plots_dir()
if plots_dir:
    app.mount("/api/protein/plots", StaticFiles(directory=plots_dir), name="protein-plots")

assets_dir = FRONTEND_DIR / "assets"
if assets_dir.is_dir():
    app.mount("/assets", StaticFiles(directory=assets_dir), name="frontend-assets")


@app.get("/{path:path}", include_in_schema=False)
async def frontend(path: str):
    """Serve the React SPA and its pre-rendered routes."""
    requested = FRONTEND_DIR / path
    if path and requested.is_file():
        return FileResponse(requested)

    prerendered = requested / "index.html"
    if path and prerendered.is_file():
        return FileResponse(prerendered)

    index = FRONTEND_DIR / "index.html"
    if not index.is_file():
        raise HTTPException(status_code=503, detail="Frontend build is not available")
    return FileResponse(index)