app/main.py

"""
MiningNiti Enterprise Backend
FastAPI Application Entry Point

AI-Powered Document Intelligence for the Coal Mining Industry
"""

import logging
import re
from contextlib import asynccontextmanager
from datetime import datetime, timezone

from fastapi import FastAPI, HTTPException, Request, status
from fastapi.exceptions import RequestValidationError
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.errors import RateLimitExceeded
from slowapi.middleware import SlowAPIMiddleware
from slowapi.util import get_remote_address

from app.api.v1 import api_router
from app.config import settings
from app.core.exceptions import MiningNitiException
from app.core.url_guard import UnsafeURLError
from app.db.session import check_db_connection, init_db

# Configure logging
logging.basicConfig(
    level=logging.DEBUG if settings.DEBUG else logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)

# Suppress SQLAlchemy's extremely verbose SQL echo in debug mode —
# it drowns out real application logs. Set to WARNING to only see errors.
logging.getLogger("sqlalchemy.engine").setLevel(logging.WARNING)
logging.getLogger("sqlalchemy.pool").setLevel(logging.WARNING)
logging.getLogger("sqlalchemy.dialects").setLevel(logging.WARNING)
# Also suppress httpcore connection-level debug spam
logging.getLogger("httpcore").setLevel(logging.WARNING)
logging.getLogger("httpx").setLevel(logging.INFO)


# Arbitrary but fixed key so every replica contends for the same lock.
_RECOVERY_LOCK_KEY = 8712_2024


async def _recover_interrupted_work() -> None:
    """
    Requeue work that was in flight when the process last died.

    Two things were previously wrong here:

    1. Rows were reset to PENDING but never re-enqueued. Nothing scans for
       PENDING documents — enqueue_document_task() is only called from the
       upload/create/reanalyze endpoints — so "recovered" documents sat at
       pending forever and needed a manual re-analyze.

    2. With more than one worker process, every replica ran recovery
       concurrently and would requeue the same rows N times. A Postgres
       advisory lock now elects a single recoverer; the others skip.
    """
    from sqlalchemy import text

    from app.db.session import get_db_context
    from app.models.compliance import AuditStatus, ComplianceAudit
    from app.models.document import Document, DocumentStatus
    from app.services.queue import enqueue_compliance_task, enqueue_document_task

    # Advisory locks are a PostgreSQL feature. On SQLite (unit tests) there is
    # only ever one process, so the election is unnecessary.
    try:
        with get_db_context() as db:
            # Ask the session we actually got, not the module-level engine —
            # they can differ (tests inject a SQLite session while DATABASE_URL
            # points at PostgreSQL), and issuing pg_try_advisory_lock against
            # the wrong backend fails the whole recovery.
            use_lock = db.get_bind().dialect.name == "postgresql"

            if use_lock:
                acquired = db.execute(
                    text("SELECT pg_try_advisory_lock(:key)"),
                    {"key": _RECOVERY_LOCK_KEY},
                ).scalar()

                if not acquired:
                    logger.info("Recovery: another worker holds the lock, skipping")
                    return

            try:
                stuck_docs = (
                    db.query(Document)
                    .filter(
                        Document.status.in_(
                            [DocumentStatus.PROCESSING, DocumentStatus.ANALYZING]
                        )
                    )
                    .all()
                )
                doc_ids = [str(doc.id) for doc in stuck_docs]
                for doc in stuck_docs:
                    doc.status = DocumentStatus.PENDING
                    doc.processing_error = "Reset after server restart"

                stuck_audits = (
                    db.query(ComplianceAudit)
                    .filter(ComplianceAudit.status == AuditStatus.RUNNING)
                    .all()
                )
                audit_ids = [str(a.id) for a in stuck_audits]
                for audit in stuck_audits:
                    audit.status = AuditStatus.PENDING
                    audit.processing_error = "Reset after server restart"

                db.commit()
            finally:
                if use_lock:
                    db.execute(
                        text("SELECT pg_advisory_unlock(:key)"),
                        {"key": _RECOVERY_LOCK_KEY},
                    )

        # Requeue only after the reset is durably committed, so a crash here
        # leaves rows in PENDING for the next startup rather than in a state
        # no one will pick up.
        for document_id in doc_ids:
            enqueue_document_task(document_id)
        for audit_id in audit_ids:
            await enqueue_compliance_task(audit_id)

        if doc_ids or audit_ids:
            logger.info(
                f"Recovery: requeued {len(doc_ids)} document(s) and "
                f"{len(audit_ids)} audit(s) interrupted by the last restart"
            )
        else:
            logger.info("Recovery: nothing was interrupted")

    except Exception as e:
        # Recovery is best-effort; a failure here must not stop the app booting.
        logger.warning(f"Recovery failed: {e}", exc_info=True)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan handler for startup and shutdown events"""
    # Startup
    logger.info(f"Starting {settings.APP_NAME} v{settings.APP_VERSION}")
    logger.info(f"Environment: {settings.ENVIRONMENT}")

    import asyncio

    from app.services.queue import compliance_worker, document_worker

    worker_task = asyncio.create_task(document_worker())
    compliance_worker_task = asyncio.create_task(compliance_worker())

    # Check database connection
    if check_db_connection():
        logger.info("Database connection verified")
        # Auto-create tables on startup (idempotent)
        try:
            init_db()
            logger.info("Database tables initialized")
        except Exception as e:
            logger.warning(f"Database table creation warning: {e}")

        await _recover_interrupted_work()
    else:
        logger.warning("Database connection failed - some features may not work")

    yield

    # Shutdown
    logger.info("Shutting down application")
    worker_task.cancel()
    compliance_worker_task.cancel()
    try:
        await worker_task
    except asyncio.CancelledError:
        pass
    try:
        await compliance_worker_task
    except asyncio.CancelledError:
        pass


# Create FastAPI application
app = FastAPI(
    title=settings.APP_NAME,
    description="""
## MiningNiti - AI Document Intelligence for Mining

Enterprise-grade document processing and AI chat platform 
specifically designed for the coal mining industry.

### Features
- 📄 **Smart Document Processing** - Upload PDF, DOCX, TXT with AI analysis
- 🤖 **Multi-Agent AI** - Classification, Safety Analysis, Entity Extraction
- 💬 **RAG Chat** - Context-aware conversations with document citations
- 📊 **Analytics Dashboard** - Safety metrics, compliance tracking
- 🔒 **Enterprise Security** - JWT auth, audit logging

### AI Agents
1. **Classifier Agent** - Categorizes mining documents
2. **Safety Analyzer** - Detects hazards and compliance issues
3. **Entity Extractor** - Extracts equipment, chemicals, regulations
4. **Summarizer** - Creates executive summaries
    """,
    version=settings.APP_VERSION,
    docs_url="/docs",
    redoc_url="/redoc",
    openapi_url="/openapi.json",
    lifespan=lifespan,
)

# Rate Limiter
limiter = Limiter(key_func=get_remote_address, default_limits=["120/minute"])

app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
app.add_middleware(SlowAPIMiddleware)

# CORS. Exact origins go in allow_origins; patterns (Vercel previews) go in
# allow_origin_regex — Starlette does NOT expand '*' inside allow_origins.
app.add_middleware(
    CORSMiddleware,
    allow_origins=settings.CORS_ORIGINS,
    allow_origin_regex=settings.CORS_ORIGIN_REGEX or None,
    allow_credentials=True,
    allow_methods=["GET", "POST", "PATCH", "DELETE", "OPTIONS"],
    allow_headers=["Authorization", "Content-Type", "X-Requested-With"],
)


# Exception Handlers


def _get_cors_headers(request: Request) -> dict:
    """
    Build CORS headers for the *unhandled* exception path only.

    Responses from FastAPI's normal exception handlers pass back out through
    CORSMiddleware, which attaches the headers itself — duplicating that logic
    here is how the two implementations drifted apart in the first place. But
    the catch-all Exception handler is installed in ServerErrorMiddleware,
    which sits *outside* CORSMiddleware, so a 500 would otherwise reach the
    browser with no CORS headers and surface as a misleading CORS error.
    """
    origin = request.headers.get("origin", "")
    if not origin:
        return {}

    allowed = origin in settings.CORS_ORIGINS
    if not allowed and settings.CORS_ORIGIN_REGEX:
        allowed = bool(re.fullmatch(settings.CORS_ORIGIN_REGEX, origin))

    if allowed:
        return {
            "Access-Control-Allow-Origin": origin,
            "Access-Control-Allow-Credentials": "true",
            "Vary": "Origin",
        }
    return {}


@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, exc: HTTPException):
    """Handle HTTP exceptions. CORSMiddleware attaches the CORS headers."""
    return JSONResponse(
        status_code=exc.status_code,
        content={"detail": exc.detail},
        headers=exc.headers or None,
    )


@app.exception_handler(MiningNitiException)
async def miningniti_exception_handler(request: Request, exc: MiningNitiException):
    """Handle custom application exceptions"""
    return JSONResponse(
        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        content={
            "error": exc.message,
            "code": exc.code,
            "details": exc.details,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        },
    )


@app.exception_handler(UnsafeURLError)
async def unsafe_url_exception_handler(request: Request, exc: UnsafeURLError):
    """A URL was rejected by the SSRF guard — that is a client error, not a bug."""
    logger.warning(f"Blocked unsafe URL request: {exc}")
    return JSONResponse(
        status_code=status.HTTP_400_BAD_REQUEST,
        content={
            "error": str(exc),
            "code": "UNSAFE_URL",
            "timestamp": datetime.now(timezone.utc).isoformat(),
        },
    )


@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    """Handle Pydantic validation errors"""
    return JSONResponse(
        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
        content={
            "error": "Validation failed",
            "code": "VALIDATION_ERROR",
            "details": exc.errors(),
            "timestamp": datetime.now(timezone.utc).isoformat(),
        },
    )


@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    """Global exception handler for unhandled errors"""
    logger.error(f"Unhandled exception: {exc}", exc_info=True)
    # This handler runs outside CORSMiddleware, so the headers are added by hand.
    return JSONResponse(
        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        content={
            "error": "Internal server error",
            "code": "INTERNAL_SERVER_ERROR",
            "timestamp": datetime.now(timezone.utc).isoformat(),
        },
        headers=_get_cors_headers(request),
    )


# Include API router
app.include_router(api_router, prefix=settings.API_V1_PREFIX)


# Root endpoint (without /api/v1 prefix for health checks)
@app.get("/", tags=["Root"])
async def root():
    """Root endpoint - application info"""
    return {
        "name": settings.APP_NAME,
        "version": settings.APP_VERSION,
        "description": "AI Document Intelligence for Mining Industry",
        "docs": "/docs",
        "health": f"{settings.API_V1_PREFIX}/health",
    }


@app.get("/health", tags=["Root"])
async def health():
    """Quick health check for load balancers"""
    return {"status": "healthy", "timestamp": datetime.now(timezone.utc).isoformat()}


# Run with: uvicorn app.main:app --reload
if __name__ == "__main__":
    import uvicorn

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