serve.py

"""
Annator full stack for Hugging Face Docker Spaces.
- FastAPI backend (main_api_app) on one process
- AIMONEYFLOW + client HTML static frontend
- Listens on 0.0.0.0:7860
"""
from __future__ import annotations

import logging
import os
import sys
from pathlib import Path

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
log = logging.getLogger("annator-hf")

ROOT = Path(__file__).resolve().parent
BACKEND = ROOT / "backend"
FRONTEND = ROOT / "frontend"

# --- HF / Space environment defaults (before backend import) ---
os.environ.setdefault("ENVIRONMENT", "development")  # keep /docs available
os.environ.setdefault("ALLOWED_HOSTS", "*")
os.environ.setdefault(
    "ALLOWED_ORIGINS",
    "*,https://techprotrade-annator-atom.hf.space,http://localhost:7860",
)
if not os.environ.get('DATABASE_URL') and os.environ.get('NEON_DATABASE_URL'):
    os.environ['DATABASE_URL'] = os.environ['NEON_DATABASE_URL']
os.environ.setdefault("DATABASE_URL", f"sqlite:///{(ROOT / 'data' / 'atom.db').as_posix()}")
os.environ.setdefault("SKIP_USER_BOOTSTRAP", "true")
os.environ.setdefault("ATOM_MOCK_DATABASE", "false")
os.environ.setdefault("HF_SPACE", "1")
os.environ.setdefault("PORT", "7860")

(ROOT / "data").mkdir(parents=True, exist_ok=True)

if str(BACKEND) not in sys.path:
    sys.path.insert(0, str(BACKEND))

os.chdir(BACKEND)

app = None
mode = "unknown"

try:
    from main_api_app import app as _app  # type: ignore

    app = _app
    mode = "full"
    log.info("Loaded main_api_app (full backend)")
except Exception as exc:  # noqa: BLE001
    log.exception("Full backend failed to import: %s", exc)
    try:
        from main_api_app_safe import app as _app  # type: ignore

        app = _app
        mode = "safe"
        log.warning("Fell back to main_api_app_safe")
    except Exception as exc2:  # noqa: BLE001
        log.exception("Safe backend also failed: %s", exc2)
        from fastapi import FastAPI

        app = FastAPI(title="Annator Atom (degraded)")
        mode = "degraded"

        @app.get("/api/health")
        def _deg_health():
            return {"status": "degraded", "error": str(exc), "fallback_error": str(exc2)}


from fastapi import HTTPException
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles

# Drop TrustedHost so HF Space host is accepted
try:
    from starlette.middleware.trustedhost import TrustedHostMiddleware

    app.user_middleware = [  # type: ignore[attr-defined]
        m
        for m in getattr(app, "user_middleware", [])
        if getattr(m, "cls", None) is not TrustedHostMiddleware
    ]
    app.middleware_stack = None
except Exception as exc:  # noqa: BLE001
    log.warning("Could not adjust TrustedHostMiddleware: %s", exc)


# Final response headers for HF iframe + CSP (runs outermost when added last)
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request as StarletteRequest
from starlette.responses import Response as StarletteResponse


class HfBrowserCompatMiddleware(BaseHTTPMiddleware):
    CSP = (
        "default-src 'self'; "
        "script-src 'self' 'unsafe-inline' 'unsafe-eval' https://cdn.jsdelivr.net blob:; "
        "style-src 'self' 'unsafe-inline' https://fonts.googleapis.com; "
        "font-src 'self' https://fonts.gstatic.com data:; "
        "img-src 'self' data: https: blob:; "
        "connect-src 'self' https: wss: ws:; "
        "worker-src 'self' blob:; "
        "frame-ancestors 'self' https://huggingface.co https://*.huggingface.co https://*.hf.space; "
        "base-uri 'self'; "
        "object-src 'none'"
    )

    async def dispatch(self, request: StarletteRequest, call_next):
        response: StarletteResponse = await call_next(request)
        # Remove frame deny so Space can load in HF shell iframe
        if "x-frame-options" in response.headers:
            del response.headers["x-frame-options"]
        response.headers["Content-Security-Policy"] = self.CSP
        response.headers.setdefault("X-Content-Type-Options", "nosniff")
        return response


app.add_middleware(HfBrowserCompatMiddleware)

# Guarantee PDF Workflow Hub is mounted even if main_api_app import order skipped it
try:
    from api.pdf_workflow_routes import router as _pdf_hub  # type: ignore

    already = any(
        getattr(r, "path", "").startswith("/api/pdf")
        for r in app.router.routes
    )
    if not already:
        app.include_router(_pdf_hub)
        log.info("PDF Workflow Hub attached from serve.py")
    else:
        log.info("PDF Workflow Hub already present on app")
except Exception as exc:  # noqa: BLE001
    log.warning("Could not attach PDF Workflow Hub: %s", exc)


def _strip_root_routes() -> None:
    """main_api_app registers GET / — remove so frontend can own the landing page."""
    kept = []
    removed = 0
    for route in list(app.router.routes):
        path = getattr(route, "path", None)
        methods = getattr(route, "methods", None) or set()
        if path == "/" and (not methods or "GET" in methods or "HEAD" in methods):
            removed += 1
            continue
        kept.append(route)
    if removed:
        app.router.routes = kept
        log.info("Removed %s existing root route(s) for frontend landing", removed)


_strip_root_routes()


@app.get("/api/hf/status")
async def hf_status():
    return {
        "status": "ok",
        "mode": mode,
        "frontend": FRONTEND.exists(),
        "backend_path": str(BACKEND),
        "port": int(os.getenv("PORT", "7860")),
        "space": "techprotrade/annator-atom",
    }


@app.get("/health")
@app.get("/api/health")
async def health():
    return {"status": "ok", "mode": mode, "service": "annator-full-hf"}


@app.get("/api/platform")
async def platform_info():
    """Former root JSON payload still available under /api/platform."""
    return {
        "name": "ATOM Platform API",
        "version": "2.1.0",
        "status": "running",
        "mode": mode,
        "docs": "/docs",
        "frontend": "/",
        "clients": "/clients/",
    }


# Static frontend (API routes already registered — these come after)
if FRONTEND.exists():
    assets = FRONTEND / "assets"
    clients = FRONTEND / "clients"
    if assets.exists():
        app.mount("/assets", StaticFiles(directory=str(assets)), name="assets")
    if clients.exists():
        app.mount("/clients", StaticFiles(directory=str(clients), html=True), name="clients")

    index_file = FRONTEND / "index.html"

    @app.get("/")
    async def spa_root():
        if index_file.exists():
            return FileResponse(index_file)
        return JSONResponse(
            {"message": "Annator Atom", "mode": mode, "hint": "/docs", "clients": "/clients/"}
        )

    # Explicit HTML pages from AIMONEYFLOW
    for html_path in FRONTEND.glob("*.html"):
        name = html_path.name

        async def _serve_html(path: Path = html_path):  # noqa: B023
            return FileResponse(path)

        # Register /dashboard.html etc. (avoid double-register index)
        if name != "index.html":
            app.add_api_route(f"/{name}", _serve_html, methods=["GET"], name=f"html_{name}")

    # Fallback for other static files under frontend (css/js next to html)
    @app.get("/{filename:path}")
    async def frontend_fallback(filename: str):
        # Never shadow API / docs
        if (
            filename.startswith("api/")
            or filename.startswith("docs")
            or filename.startswith("redoc")
            or filename.startswith("openapi")
            or filename.startswith("health")
        ):
            raise HTTPException(status_code=404, detail="Not found")
        candidate = (FRONTEND / filename).resolve()
        try:
            candidate.relative_to(FRONTEND.resolve())
        except ValueError as exc:
            raise HTTPException(status_code=404, detail="Not found") from exc
        if candidate.is_file():
            return FileResponse(candidate)
        raise HTTPException(status_code=404, detail="Not found")

    log.info("Mounted frontend from %s", FRONTEND)
else:
    log.warning("Frontend directory missing: %s", FRONTEND)


if __name__ == "__main__":
    import uvicorn

    port = int(os.getenv("PORT", "7860"))
    uvicorn.run("serve:app", host="0.0.0.0", port=port, factory=False)