MediaManager/media_manager/main.py

import asyncio
import logging
import os
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager

import uvicorn
from asgi_correlation_id import CorrelationIdMiddleware
from fastapi import (
    APIRouter,
    Depends,
    FastAPI,
    HTTPException,
    Request,
    Response,
    status,
)
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from psycopg.errors import UniqueViolation
from sqlalchemy.exc import IntegrityError
from starlette.responses import FileResponse, RedirectResponse
from taskiq.receiver import Receiver
from taskiq_fastapi import populate_dependency_context
from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware

import media_manager.movies.router as movies_router
import media_manager.torrent.router as torrent_router
import media_manager.tv.router as tv_router
from media_manager.auth.router import (
    auth_metadata_router,
    get_openid_router,
)
from media_manager.auth.router import (
    users_router as custom_users_router,
)
from media_manager.auth.schemas import UserCreate, UserRead, UserUpdate
from media_manager.auth.users import (
    bearer_auth_backend,
    cookie_auth_backend,
    fastapi_users,
)
from media_manager.config import MediaManagerConfig
from media_manager.database import init_engine
from media_manager.exceptions import (
    ConflictError,
    InvalidConfigError,
    MediaAlreadyExistsError,
    NotFoundError,
    conflict_error_handler,
    invalid_config_error_exception_handler,
    media_already_exists_exception_handler,
    not_found_error_exception_handler,
    sqlalchemy_integrity_error_handler,
)
from media_manager.filesystem_checks import run_filesystem_checks
from media_manager.logging import LOGGING_CONFIG, setup_logging
from media_manager.notification.router import router as notification_router
from media_manager.scheduler import (
    broker,
    build_scheduler_loop,
    import_all_movie_torrents_task,
    import_all_show_torrents_task,
    update_all_movies_metadata_task,
    update_all_non_ended_shows_metadata_task,
)

setup_logging()

config = MediaManagerConfig()
log = logging.getLogger(__name__)


if config.misc.development:
    log.warning("Development Mode activated!")

run_filesystem_checks(config, log)

BASE_PATH = os.getenv("BASE_PATH", "")
FRONTEND_FILES_DIR = os.getenv("FRONTEND_FILES_DIR")
DISABLE_FRONTEND_MOUNT = os.getenv("DISABLE_FRONTEND_MOUNT", "").lower() == "true"
FRONTEND_FOLLOW_SYMLINKS = os.getenv("FRONTEND_FOLLOW_SYMLINKS", "").lower() == "true"

log.info("Hello World!")


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator:
    init_engine(config.database)
    broker_started = False
    started_sources: list = []
    finish_event: asyncio.Event | None = None
    receiver_task: asyncio.Task | None = None
    loop_task: asyncio.Task | None = None
    try:
        if not broker.is_worker_process:
            await broker.startup()
            broker_started = True
        populate_dependency_context(broker, app)
        scheduler_loop = build_scheduler_loop()
        for source in scheduler_loop.scheduler.sources:
            await source.startup()
            started_sources.append(source)
        finish_event = asyncio.Event()
        receiver = Receiver(broker, run_startup=False, max_async_tasks=10)
        receiver_task = asyncio.create_task(receiver.listen(finish_event))
        loop_task = asyncio.create_task(scheduler_loop.run(skip_first_run=True))
        try:
            await asyncio.gather(
                import_all_movie_torrents_task.kiq(),
                import_all_show_torrents_task.kiq(),
                update_all_movies_metadata_task.kiq(),
                update_all_non_ended_shows_metadata_task.kiq(),
            )
        except Exception:
            log.exception("Failed to submit initial background tasks during startup.")
            raise
        yield
    finally:
        if loop_task is not None:
            loop_task.cancel()
            try:
                await loop_task
            except asyncio.CancelledError:
                pass
        if finish_event is not None and receiver_task is not None:
            finish_event.set()
            await receiver_task
        for source in started_sources:
            await source.shutdown()
        if broker_started:
            await broker.shutdown()


app = FastAPI(root_path=BASE_PATH, lifespan=lifespan)
app.add_middleware(ProxyHeadersMiddleware, trusted_hosts="*")
origins = config.misc.cors_urls
log.info(f"CORS URLs activated for following origins: {origins}")
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["GET", "PUT", "POST", "DELETE", "PATCH", "HEAD", "OPTIONS"],
)
app.add_middleware(CorrelationIdMiddleware, header_name="X-Correlation-ID")
api_app = APIRouter(prefix="/api/v1")


@api_app.get("/health")
async def hello_world() -> dict:
    return {"message": "Hello World!", "version": os.getenv("PUBLIC_VERSION")}


api_app.include_router(
    fastapi_users.get_auth_router(bearer_auth_backend),
    prefix="/auth/jwt",
    tags=["auth"],
)
api_app.include_router(
    fastapi_users.get_auth_router(cookie_auth_backend),
    prefix="/auth/cookie",
    tags=["auth"],
)


def reject_if_registration_disabled() -> None:
    if not config.auth.registration_enabled:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="User registration is disabled.",
        )


api_app.include_router(
    fastapi_users.get_register_router(UserRead, UserCreate),
    prefix="/auth",
    tags=["auth"],
    dependencies=[Depends(reject_if_registration_disabled)],
)
api_app.include_router(
    fastapi_users.get_reset_password_router(), prefix="/auth", tags=["auth"]
)
api_app.include_router(
    fastapi_users.get_verify_router(UserRead), prefix="/auth", tags=["auth"]
)
api_app.include_router(custom_users_router, tags=["users"])
api_app.include_router(
    fastapi_users.get_users_router(UserRead, UserUpdate),
    prefix="/users",
    tags=["users"],
)
api_app.include_router(auth_metadata_router, tags=["openid"])

if get_openid_router():
    api_app.include_router(get_openid_router(), tags=["openid"], prefix="/auth/oauth")

api_app.include_router(tv_router.router, prefix="/tv", tags=["tv"])
api_app.include_router(torrent_router.router, prefix="/torrent", tags=["torrent"])
api_app.include_router(movies_router.router, prefix="/movies", tags=["movie"])
api_app.include_router(
    notification_router, prefix="/notification", tags=["notification"]
)

# serve static image files
app.mount(
    "/api/v1/static/image",
    StaticFiles(directory=config.misc.image_directory),
    name="static-images",
)
app.include_router(api_app)

# handle static frontend files
if not DISABLE_FRONTEND_MOUNT:
    app.mount(
        "/web",
        StaticFiles(
            directory=FRONTEND_FILES_DIR,
            html=True,
            follow_symlink=FRONTEND_FOLLOW_SYMLINKS,
        ),
        name="frontend",
    )
    log.debug(f"Mounted frontend at /web from {FRONTEND_FILES_DIR}")
else:
    log.info("Frontend mounting disabled (DISABLE_FRONTEND_MOUNT is set)")


@app.get("/")
async def root() -> RedirectResponse:
    return RedirectResponse(url="/web/")


@app.get("/dashboard")
async def dashboard() -> RedirectResponse:
    return RedirectResponse(url="/web/")


@app.get("/login")
async def login() -> RedirectResponse:
    return RedirectResponse(url="/web/")


# this will serve the custom 404 page for frontend routes, so SvelteKit can handle routing
@app.exception_handler(404)
async def not_found_handler(request: Request, _exc: Exception) -> Response:
    if not DISABLE_FRONTEND_MOUNT and any(
        base_path in ["/web", "/dashboard", "/login"] for base_path in request.url.path
    ):
        return FileResponse(f"{FRONTEND_FILES_DIR}/404.html")
    return Response(content="Not Found", status_code=404)


# Register exception handlers for custom exceptions
app.add_exception_handler(NotFoundError, not_found_error_exception_handler)
app.add_exception_handler(
    MediaAlreadyExistsError, media_already_exists_exception_handler
)
app.add_exception_handler(InvalidConfigError, invalid_config_error_exception_handler)
app.add_exception_handler(IntegrityError, sqlalchemy_integrity_error_handler)
app.add_exception_handler(UniqueViolation, sqlalchemy_integrity_error_handler)
app.add_exception_handler(ConflictError, conflict_error_handler)

if __name__ == "__main__":
    uvicorn.run(
        app,
        host="127.0.0.1",
        port=5049,
        log_config=LOGGING_CONFIG,
        proxy_headers=True,
        forwarded_allow_ips="*",
    )