docqube_backend/app/main.py

import os
import logging
import asyncio
import tempfile
import shutil
import uuid
from contextlib import asynccontextmanager
from datetime import datetime, timezone

from fastapi import FastAPI, UploadFile, File, HTTPException, Depends, Request, Response, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from sqlalchemy import text
from sqlalchemy.orm import Session
from app.middleware.csrf import CSRFMiddleware
from app.middleware.ip_middleware import IPMiddleware

from app.core.settings import settings
from app.db.database import init_db, init_db_pool, close_db_pool, get_db

import app.db.all_models  # noqa: F401
from app.middleware.auth import get_current_user
from app.modules.auth.models.user_model import User
from app.modules.auth.models.access_model import Access
from app.modules.auth.models.role_model import Role
from app.modules.auth.models.role_access_model import RoleAccess
from app.modules.documents.models.document_model import Project

from app.modules.auth.routes.auth_routes import router as auth_router
from app.modules.auth.routes.user_routes import router as user_router
from app.modules.auth.routes.me_routes import router as me_router
from app.modules.auth.routes.access_routes import router as access_router
from app.modules.auth.routes.role_routes import router as role_router
from app.modules.auth.routes.admin_dashboard_routes import router as admin_dashboard_router
from app.modules.auth.routes.admin_user_routes import router as admin_user_router
from app.modules.documents.routes.document_routes import router as projects_router
from app.modules.documents.routes.sharing_routes import router as sharing_router
from app.modules.documents.routes.version_routes import router as version_router
from app.modules.documents.routes.comment_routes import (
    router as project_comments_router,
)
from app.modules.documents.routes.details_routes import router as project_details_router
from app.modules.drive.routes.drive_routes import (
    router as drive_router,
    public_router as public_drive_router,
)
from app.modules.storage.routes.storage_routes import router as storage_router
from app.modules.chat.routes.chat_routes import router as chatbot_router
from app.tasks.routes import router as task_results_router
from app.modules.notifications.routes.notification_routes import (
    router as notifications_router,
)
from app.modules.documents.routes.export_routes import router as export_router

from app.modules.speech.routes.speech_routes import router as speech_router
from app.api.ws_router import router as ws_router
from app.core.pubsub_listener import start_listener
from app.core.redis import redis_pubsub
from app.modules.collab.routes.chat_routes import (
    router as collab_chat_router,
    rest_router as collab_rest_router,
)
from app.modules.tenant.routes.tenant_routes import router as tenant_router
from app.modules.tenant.routes.tenant_contact_routes import router as tenant_contact_router
from app.modules.signing.routes.signing_routes import router as signing_router
from app.modules.signing.routes.config_routes import router as signing_config_router
from app.modules.configuration.routes.system_configuration_routes import (
    router as system_configuration_router,
)
from app.modules.activity_logs.routes import router as activity_logs_router
from app.infrastructure.realtime.chat_subscriber import (
    start_chat_subscriber,
    stop_chat_subscriber,
)
from app.modules.drive.routes.scan_routes import router as scan_router
from app.modules.extraction.routes.extraction_routes import router as extraction_router
from app.modules.editor.pdf_router import router as pdf_editor_router
from app.modules.auth.sso_router import router as sso_router
from app.modules.security.routes.security_routes import router as security_router

TEMP_DIR = tempfile.mkdtemp(prefix="ai_uploads_")
processing_semaphore = asyncio.Semaphore(1)

logging.basicConfig(
    level=getattr(logging, settings.LOG_LEVEL),
    format="%(asctime)s - %(levelname)s - %(message)s",
)


@asynccontextmanager
async def lifespan(app: FastAPI):
    logging.info(f"Server is starting up in {settings.APP_ENV} mode...")
    try:
        init_db_pool()
        init_db()

        asyncio.create_task(start_listener())

        await redis_pubsub.connect()

        await start_chat_subscriber()

        from app.infrastructure.vector.vector_client import preload as preload_vector
        loop = asyncio.get_running_loop()
        loop.run_in_executor(None, preload_vector)

    except Exception as e:
        logging.error(f"Critical Startup failure: {e}")
        raise e

    yield

    logging.info("Server is shutting down...")
    try:
        await stop_chat_subscriber()
        await redis_pubsub.disconnect()
        logging.info("✅ Chat subscriber & Redis Pub/Sub disconnected")
    except Exception as e:
        logging.error(f"Shutdown error during subscriber/redis stop: {e}")
    try:
        close_db_pool()
        logging.info("✅ Database connection pool closed")
    except Exception as e:
        logging.error(f"Error closing database pool: {e}")

    try:
        shutil.rmtree(TEMP_DIR, ignore_errors=True)
        logging.info("🧹 Temp directory cleaned")
    except Exception as e:
        logging.error(f"Shutdown error during temp cleanup: {e}")


app = FastAPI(
    lifespan=lifespan,
    title="DocQube Document Processing Pipeline API",
    description="An interactive API and UI to process PDFs into JATS XML.",
    version="2.0.0",
)

app.add_middleware(CSRFMiddleware)
app.add_middleware(IPMiddleware)

from app.middleware.tenant_context_middleware import TenantContextMiddleware  # noqa: E402

app.add_middleware(TenantContextMiddleware)

app.add_middleware(
    CORSMiddleware,
    allow_origins=settings.CORS_ORIGINS_LIST,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    # Browsers only expose a small header safelist to JS on cross-origin
    # responses by default (Content-Range/Accept-Ranges/Content-Length are
    # NOT in it) — pdf.js reads these itself to decide whether it can use
    # HTTP Range requests for progressive PDF loading, so without this it
    # falls back to downloading the entire file before rendering anything.
    expose_headers=["Content-Range", "Accept-Ranges", "Content-Length"],
)


# File Serving Internal Helper
async def _get_s3_file_response(key: str, db: Session = None, tenant_id: str = None):
    from app.infrastructure.storage.local_storage_handler import (
        get_storage_client,
    )
    from fastapi.responses import StreamingResponse
    import mimetypes

    storage, quarantine_bucket, clean_bucket = get_storage_client(
        db=db, tenant_id=tenant_id)
    try:
        try:
            response = storage.get_object(Bucket=clean_bucket, Key=key)
        except Exception:
            response = storage.get_object(Bucket=quarantine_bucket, Key=key)

        body = response["Body"]

        content_type, _ = mimetypes.guess_type(key)
        if not content_type:
            content_type = "application/octet-stream"

        return StreamingResponse(body, media_type=content_type)
    except HTTPException:
        raise
    except Exception as e:
        logging.error(f"Storage File Proxy Error: {e}")
        raise HTTPException(status_code=404, detail="File not found")


@app.get("/api/storage_drive/{key:path}")
async def serve_storage_file(key: str, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    if ".." in key:
        raise HTTPException(status_code=400, detail="Invalid path")

    tenant_part = f"tenants/{current_user.tenant_id}/" if current_user.tenant_id else ""
    user_prefix1 = f"{tenant_part}user_{current_user.id}/"
    user_prefix2 = f"{tenant_part}users/{current_user.id}/"
    conv_prefix = f"{tenant_part}doc_conversion/{current_user.id}/"
    chatbot_prefix = f"{tenant_part}chatbot_documents/{current_user.id}/"
    images_prefix = f"extracted_images/{current_user.id}/"

    allowed_prefixes = [
        user_prefix1,
        user_prefix2,
        conv_prefix,
        chatbot_prefix,
        images_prefix,
        f"user_{current_user.id}/",
        f"users/{current_user.id}/",
        f"doc_conversion/{current_user.id}/",
        f"chatbot_documents/{current_user.id}/",
        f"extracted_images/{current_user.id}/",
    ]

    if not any(key.startswith(p) for p in allowed_prefixes):
        raise HTTPException(
            status_code=403, detail="Access denied to this storage key")

    return await _get_s3_file_response(key, db=db, tenant_id=current_user.tenant_id)


@app.get("/api/serve-image/{session_id}/{filename}")
async def serve_image(
    session_id: str,
    filename: str,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    """
    Alias for serving images extracted during document processing.
    """
    if ".." in session_id or ".." in filename:
        raise HTTPException(status_code=400, detail="Invalid parameters")

    project = (
        db.query(Project)
        .filter(Project.session_id == session_id, Project.user_id == current_user.id)
        .first()
    )

    if not project:
        raise HTTPException(
            status_code=403, detail="Access denied to this session's resources"
        )

    key = f"extracted_images/{current_user.id}/{session_id}/{filename}"
    return await _get_s3_file_response(key, db=db, tenant_id=current_user.tenant_id)


def _check_database_ready(db: Session) -> bool:
    """Lightweight DB reachability check shared by /api/health/ready and the
    legacy /api/health. On failure it rolls back before returning — without
    this, the session handed back to Depends(get_db)'s own cleanup (which
    unconditionally calls db.commit() after a successful return from this
    function) is left in an aborted-transaction state, and that commit()
    itself raises. That turns a DB failure this function already caught
    into an unhandled 500 from the dependency's own teardown — which is
    exactly what took the original /api/health down during the PgBouncer
    incident, even though it appeared to catch the error internally.
    """
    try:
        db.execute(text("SELECT 1"))
        return True
    except Exception:
        db.rollback()
        return False


@app.get("/api/health/live")
def health_live():
    """Liveness: is the API process able to serve HTTP at all. No database,
    no PgBouncer, no external dependency of any kind — this is what the
    Docker HEALTHCHECK now uses, so a temporarily unreachable database can
    no longer get the whole API container killed as unhealthy.
    """
    return {
        "status": "ok",
        "service": "api",
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }


@app.get("/api/health/ready")
def health_ready(response: Response, db: Session = Depends(get_db)):
    """Readiness: is the API ready to serve real traffic, i.e. can it reach
    the database through the application's actual configured pool/engine
    (app.db.database.engine via get_db) — the same one every request uses,
    not a separate health-only connection.
    """
    if _check_database_ready(db):
        return {
            "status": "ready",
            "database": "connected",
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }
    response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE
    return {
        "status": "not_ready",
        "database": "disconnected",
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }


@app.get("/api/health")
def health_check(db: Session = Depends(get_db)):
    """Backward-compatible general health endpoint — response shape kept
    identical to before this refactor. Reuses the same DB-reachability
    check as /api/health/ready instead of duplicating it; the Docker
    HEALTHCHECK no longer depends on this endpoint (see /api/health/live).
    """
    if not _check_database_ready(db):
        return {"status": "degraded", "database": "error", "error": "database unreachable"}

    health = {"status": "ok", "database": "connected"}
    try:
        version = db.execute(text("SELECT version_num FROM alembic_version")).scalar()
        health["migration_version"] = version
    except Exception:
        db.rollback()
        health["status"] = "degraded"
        health["error"] = "migration version lookup failed"

    return health


app.include_router(tenant_router, prefix="/api")
app.include_router(tenant_contact_router, prefix="/api/tenant")
app.include_router(auth_router, prefix="/api")
app.include_router(sso_router, prefix="/api")
app.include_router(me_router, prefix="/api")
app.include_router(access_router)
app.include_router(role_router)
app.include_router(admin_dashboard_router)
app.include_router(admin_user_router)
app.include_router(user_router, prefix="/api")
app.include_router(projects_router, prefix="/api")
app.include_router(sharing_router, prefix="/api")
app.include_router(version_router, prefix="/api")
app.include_router(project_comments_router, prefix="/api")
app.include_router(project_details_router, prefix="/api")
app.include_router(drive_router, prefix="/api")
app.include_router(public_drive_router, prefix="/api")
app.include_router(storage_router, prefix="/api")
app.include_router(task_results_router, prefix="/api")
app.include_router(chatbot_router)
app.include_router(notifications_router, prefix="/api")
app.include_router(system_configuration_router)
app.include_router(activity_logs_router)
app.include_router(export_router, prefix="/api")

app.include_router(speech_router)
app.include_router(ws_router, prefix="/api")
app.include_router(collab_chat_router, prefix="/api")
app.include_router(collab_rest_router, prefix="/api")
app.include_router(signing_router, prefix="/api")
app.include_router(signing_config_router, prefix="/api")
app.include_router(pdf_editor_router, prefix="/api")
app.include_router(scan_router, prefix="/api")
app.include_router(extraction_router, prefix="/api")

from app.modules.org.routes.org_routes import router as org_router  # noqa: E402

app.include_router(org_router, prefix="/api")

from app.modules.billing.routes.billing_routes import router as billing_router  # noqa: E402

app.include_router(billing_router, prefix="/api")
app.include_router(security_router, prefix="/api")


@app.post("/api/extract")
async def extract(
    request: Request,
    file: UploadFile = File(...),
    current_user: User = Depends(get_current_user),
):
    """
    On-demand document extraction endpoint.
    Loads models only when called and unloads after.
    """
    file_id = str(uuid.uuid4())
    ext = os.path.splitext(file.filename)[1].lower()
    temp_path = os.path.join(TEMP_DIR, f"{file_id}{ext}")

    MAX_FILE_SIZE = 50 * 1024 * 1024

    try:
        total_size = 0
        with open(temp_path, "wb") as buffer:
            while True:
                chunk = await file.read(1024 * 1024)
                if not chunk:
                    break
                total_size += len(chunk)
                if total_size > MAX_FILE_SIZE:
                    raise HTTPException(
                        status_code=413, detail="File too large (Max 50MB)"
                    )
                buffer.write(chunk)

        # Convert PDF to HTML locally, then derive markdown
        async with processing_semaphore:
            from app.modules.editor.converters import process_pdf_to_html

            loop = asyncio.get_running_loop()

            def _convert():
                html = process_pdf_to_html(temp_path)
                from bs4 import BeautifulSoup
                markdown = BeautifulSoup(
                    html, "html.parser").get_text("\n")
                return {"html": html, "markdown": markdown}

            result = await loop.run_in_executor(None, _convert)

        return {
            "status": "success",
            "file_id": file_id,
            "data": result,
        }

    except HTTPException:
        raise
    except Exception as e:
        error_id = str(uuid.uuid4())
        logging.error(
            f"Extraction failed [Ref: {error_id}]: {e}", exc_info=True)
        raise HTTPException(
            status_code=500,
            detail=f"An internal error occurred during processing. Please contact support and provide reference: {error_id}",
        )

    finally:
        if os.path.exists(temp_path):
            try:
                os.remove(temp_path)
            except Exception as e:
                logging.error(f"Cleanup error removing {temp_path}: {e}")


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "main:app", host=settings.HOST, port=settings.PORT, reload=settings.DEBUG
    )