app/main.py

"""FastAPI application for the img2threejs Hugging Face Docker Space.

Routes
  GET  /                                  SPA (static/index.html)
  GET  /static/*                          static assets
  GET  /health                            liveness + readiness signal
  GET  /api/config                        public, credential-free config view
  POST /api/jobs                          start a conversion job (multipart image)
  GET  /api/jobs/{job_id}                 job status snapshot
  GET  /api/jobs/{job_id}/events          SSE stream of pipeline events
  GET  /api/jobs/{job_id}/artifacts/{n}   whitelisted per-job artifacts
  GET  /api/gallery                       persistent community gallery
  GET  /api/gallery/{item_id}             gallery item detail
  GET  /api/gallery/{item_id}/artifacts/* immutable gallery artifacts

Binding: 0.0.0.0 on $PORT (default 7860) per the HF Docker Space contract.
One uvicorn worker: the job registry is in-memory by design.
"""

from __future__ import annotations

import os

# Keep every third-party cache in writable Space storage. These must be set
# before importing the web stack or any optional model/runtime dependency.
os.environ.setdefault("HF_HOME", "/tmp/.cache/huggingface")
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

import asyncio
import json
import logging
import threading
import time
from pathlib import Path

from fastapi import FastAPI, Query, Request
from fastapi.responses import (FileResponse, JSONResponse, Response,
                               StreamingResponse)
from fastapi.staticfiles import StaticFiles
from starlette.datastructures import UploadFile
from starlette.types import ASGIApp, Receive, Scope, Send

from .config import load_settings
from .gallery import (
    GALLERY_ARTIFACT_NAMES,
    GalleryError,
    GalleryStore,
    valid_item_id,
)
from .image_guard import ImageRejected
from .llm import LLMClient
from .pipeline import ARTIFACT_NAMES, JobRegistry, fail_job_timeout, run_job
from .ratelimit import RateLimiter, client_key

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s %(message)s",
)
logger = logging.getLogger("img2threejs")

REPO_ROOT = Path(__file__).resolve().parents[1]
STATIC_DIR = Path(__file__).resolve().parent / "static"

settings = load_settings()
registry = JobRegistry(Path(settings.runs_dir))
gallery_store = GalleryStore(Path(settings.gallery_dir))
limiter = RateLimiter(settings.rate_limit_jobs_per_hour)
job_semaphore = asyncio.Semaphore(settings.max_concurrent_jobs)

# Multipart boundary + the small ``hint`` field. The image itself still has
# the tighter ``max_upload_bytes`` limit below.
MULTIPART_OVERHEAD_BYTES = 256 * 1024
QUEUE_PROGRESS_INTERVAL_S = 25.0
SUCCESS_ONLY_JOB_ARTIFACTS = frozenset({
    "compile-spec.json",
    "factory.ts",
    "model.bundle.js",
    "standalone.html",
})


class _RequestBodyTooLarge(Exception):
    pass


class UploadBodyLimitMiddleware:
    """Cap the /api/jobs request stream before multipart parsing.

    ``UploadFile`` avoids retaining large files in RAM, but without an ASGI
    stream cap a caller could still force the multipart parser to spool an
    unbounded request. Content-Length is a cheap early rejection; counting
    receive chunks enforces the same limit when that header is absent or
    false.
    """

    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if (
            scope["type"] != "http"
            or scope.get("method") != "POST"
            or scope.get("path") != "/api/jobs"
        ):
            await self.app(scope, receive, send)
            return

        body_limit = settings.max_upload_bytes + MULTIPART_OVERHEAD_BYTES
        declared_length: int | None = None
        for name, value in scope.get("headers", []):
            if name.lower() == b"content-length":
                try:
                    declared_length = int(value)
                except (TypeError, ValueError):
                    declared_length = None
                break
        if declared_length is not None and declared_length > body_limit:
            await self._reject(scope, receive, send, body_limit)
            return

        received = 0

        async def limited_receive():
            nonlocal received
            message = await receive()
            if message.get("type") == "http.request":
                received += len(message.get("body", b""))
                if received > body_limit:
                    raise _RequestBodyTooLarge
            return message

        try:
            await self.app(scope, limited_receive, send)
        except _RequestBodyTooLarge:
            await self._reject(scope, receive, send, body_limit)

    @staticmethod
    async def _reject(
        scope: Scope, receive: Receive, send: Send, body_limit: int
    ) -> None:
        response = JSONResponse(
            status_code=413,
            content={
                "error": "request_too_large",
                "detail": (
                    "The multipart request exceeds the configured upload "
                    f"boundary ({body_limit} bytes including form overhead)."
                ),
            },
        )
        await response(scope, receive, send)


# A reservation covers the period between the cheap admission check and
# JobRegistry.create(). It closes the race where many multipart bodies could
# all pass the queue check while awaiting parsing.
_queue_guard = threading.Lock()
_pending_uploads = 0


def _reserve_job_slot() -> tuple[bool, int]:
    global _pending_uploads
    with _queue_guard:
        in_flight = sum(1 for job in registry.jobs.values() if job.status == "running")
        occupied = in_flight + _pending_uploads
        if occupied >= settings.max_in_flight_jobs:
            return False, occupied
        _pending_uploads += 1
        return True, occupied


def _release_job_slot() -> None:
    global _pending_uploads
    with _queue_guard:
        _pending_uploads = max(0, _pending_uploads - 1)


def _create_reserved_job():
    """Atomically convert one upload reservation into a registry job."""
    global _pending_uploads
    with _queue_guard:
        job = registry.create()
        _pending_uploads = max(0, _pending_uploads - 1)
        return job


app = FastAPI(title="img2threejs", docs_url=None, redoc_url=None, openapi_url=None)
app.add_middleware(UploadBodyLimitMiddleware)
# Factory for the per-job LLM client; tests substitute a mock.
app.state.llm_factory = lambda s: LLMClient(s)  # noqa: E731

CSP = (
    "default-src 'self'; script-src 'self' blob:; style-src 'self' 'unsafe-inline'; "
    "img-src 'self' data: blob:; connect-src 'self'; frame-src 'self'; "
    "object-src 'none'; base-uri 'none'; "
    "frame-ancestors 'self' https://huggingface.co"
)

# The sandboxed viewer shell runs LLM-influenced code inside an inline module
# script and blob: imports — its own policy. Sending the app-wide CSP on this
# path would intersect with the document's meta policy and block the shell
# (multiple CSPs intersect; 'self' never matches an opaque origin anyway).
VIEWER_CSP = (
    "default-src 'none'; script-src 'unsafe-inline' blob:; "
    "style-src 'unsafe-inline'; img-src blob: data:; worker-src blob:; "
    "frame-ancestors 'self' https://huggingface.co; sandbox allow-scripts"
)
VIEWER_PATH = "/static/viewer.html"


@app.middleware("http")
async def security_headers(request: Request, call_next):
    response = await call_next(request)
    if request.url.path == VIEWER_PATH:
        response.headers["Content-Security-Policy"] = VIEWER_CSP
    else:
        response.headers.setdefault("Content-Security-Policy", CSP)
    response.headers.setdefault("X-Content-Type-Options", "nosniff")
    response.headers.setdefault("Referrer-Policy", "no-referrer")
    return response


@app.on_event("startup")
async def startup() -> None:
    Path(settings.runs_dir).mkdir(parents=True, exist_ok=True)
    try:
        gallery_store.ensure_ready()
    except GalleryError as exc:
        # Gallery storage is an optional post-generation destination. A
        # transient Bucket mount outage must not take down image conversion;
        # gallery APIs report 503 and publish attempts surface a warning until
        # the mount recovers.
        logger.warning("startup: community gallery unavailable: %s", exc)
    logger.info(
        "startup: llm_configured=%s port=%s space=%s",
        settings.llm_configured, settings.port, settings.space_id or "(local)",
    )

    async def reaper() -> None:
        while True:
            await asyncio.sleep(600)
            doomed = registry.reap(settings.job_ttl_s)
            if doomed:
                logger.info("reaper: evicted %d expired jobs", len(doomed))

    asyncio.create_task(reaper())


# ---------------------------------------------------------------------------
# Health + config
# ---------------------------------------------------------------------------

@app.get("/health")
async def health() -> JSONResponse:
    """Liveness + lightweight readiness. The LLM is deliberately NOT probed:
    a failing provider must not mark the container down — it is surfaced in
    the UI instead."""
    return JSONResponse({
        "status": "ok",
        "llm_configured": settings.llm_configured,
        "space_id": settings.space_id,
        "time": int(time.time()),
    })


@app.get("/api/config")
async def config() -> JSONResponse:
    """Public config view. Never includes the API key or any secret value."""
    return JSONResponse({
        "llm_configured": settings.llm_configured,
        "missing_llm_vars": settings.missing_llm_vars,
        "model": settings.llm_model if settings.llm_configured else None,
        "max_upload_bytes": settings.max_upload_bytes,
        "rate_limit_jobs_per_hour": settings.rate_limit_jobs_per_hour,
        "space_host": settings.space_host,
        "community_gallery": True,
        "share_default": True,
    })


# ---------------------------------------------------------------------------
# Jobs
# ---------------------------------------------------------------------------

_TRUE_FORM_VALUES = frozenset({"1", "true", "yes", "on"})
_FALSE_FORM_VALUES = frozenset({"0", "false", "no", "off"})


def _parse_share_preference(value: object) -> bool | None:
    """Parse the optional multipart share field.

    Missing means the intentionally public-by-default behaviour. ``None`` is
    reserved for malformed values so callers never silently invert a choice.
    """
    if value is None:
        return True
    if not isinstance(value, str):
        return None
    normalized = value.strip().lower()
    if normalized in _TRUE_FORM_VALUES:
        return True
    if normalized in _FALSE_FORM_VALUES:
        return False
    return None


async def _acquire_job_worker(job) -> None:
    """Wait for one worker with truthful elapsed feedback and a hard deadline."""
    deadline = job.created + settings.job_timeout_s
    while True:
        remaining = deadline - time.time()
        if remaining <= 0:
            raise TimeoutError
        interval = min(QUEUE_PROGRESS_INTERVAL_S, remaining)
        try:
            await asyncio.wait_for(job_semaphore.acquire(), timeout=interval)
            return
        except TimeoutError:
            if time.time() >= deadline:
                raise
            elapsed = max(1, round(time.time() - job.created))
            job.emit(
                "queued",
                "progress",
                f"Still queued: waiting {elapsed}s for a conversion worker.",
                elapsedSeconds=elapsed,
                shareRequested=bool(
                    job.events[0].get("data", {}).get("shareRequested", True)
                ),
            )


@app.post("/api/jobs", status_code=202)
async def create_job(request: Request) -> JSONResponse:
    if not settings.llm_configured:
        missing = ", ".join(settings.missing_llm_vars)
        return JSONResponse(status_code=503, content={
            "error": "llm_not_configured",
            "detail": (
                f"This Space needs vision-LLM credentials to author the sculpt spec: set "
                f"{missing} and LLM_BASE_URL as Space Secrets (Settings → Secrets), then "
                "restart the Space. No model was generated — results are never fabricated."
            ),
        })

    # Reserve capacity *before* parsing multipart data. Without this
    # reservation, concurrent requests can all retain/spool their complete
    # images while waiting to create an unbounded number of tasks.
    reserved, occupied = _reserve_job_slot()
    if not reserved:
        return JSONResponse(
            status_code=503,
            headers={"Retry-After": "60"},
            content={
                "error": "queue_full",
                "detail": (
                    f"The Space is busy ({occupied} jobs in flight, max "
                    f"{settings.max_in_flight_jobs}). Try again shortly."
                ),
            },
        )

    reservation_active = True
    try:
        # Uvicorn resolves forwarding headers only from configured trusted
        # proxies; do not reinterpret raw X-Forwarded-For in application code.
        key = client_key(request.client.host if request.client else None)
        retry_after = limiter.check(key)
        if retry_after is not None:
            return JSONResponse(
                status_code=429,
                headers={"Retry-After": str(retry_after)},
                content={
                    "error": "rate_limited",
                    "detail": (
                        f"At most {settings.rate_limit_jobs_per_hour} jobs per hour "
                        f"per client. Try again in {retry_after}s."
                    ),
                },
            )

        request_content_type = request.headers.get("content-type", "").lower()
        if request_content_type and not request_content_type.startswith("multipart/form-data"):
            return JSONResponse(status_code=415, content={
                "error": "unsupported_media_type",
                "detail": "Expected multipart form data containing an image file.",
            })

        # Parse only after admission. max_files/max_fields constrain multipart
        # metadata; UploadBodyLimitMiddleware constrains the complete stream.
        async with request.form(max_files=1, max_fields=3, max_part_size=64 * 1024) as form:
            file = form.get("file")
            if not isinstance(file, UploadFile):
                return JSONResponse(status_code=422, content={
                    "error": "file_required",
                    "detail": "A multipart image field named 'file' is required.",
                })
            declared = (file.content_type or "").lower()
            if declared and not declared.startswith("image/"):
                return JSONResponse(status_code=415, content={
                    "error": "unsupported_media_type",
                    "detail": "Expected an image upload (PNG, JPEG, WebP, GIF or BMP).",
                })
            raw = await file.read(settings.max_upload_bytes + 1)
            if len(raw) > settings.max_upload_bytes:
                return JSONResponse(status_code=413, content={
                    "error": "file_too_large",
                    "detail": (
                        "The upload exceeds the "
                        f"{settings.max_upload_bytes // (1024 * 1024)} MiB limit."
                    ),
                })
            hint_value = form.get("hint")
            hint = hint_value if isinstance(hint_value, str) else None
            share = _parse_share_preference(form.get("share"))
            if share is None:
                return JSONResponse(status_code=422, content={
                    "error": "invalid_share_preference",
                    "detail": (
                        "The optional multipart 'share' field must be true or false."
                    ),
                })

        job = _create_reserved_job()
        reservation_active = False
        object_hint = (hint or "").strip()[:80] or None
        job.emit(
            "queued",
            "started",
            "Job accepted and queued for the next available conversion worker.",
            shareRequested=share,
        )
    finally:
        if reservation_active:
            _release_job_slot()

    async def runner() -> None:
        acquired = False
        try:
            await _acquire_job_worker(job)
            acquired = True
            job.emit(
                "queued",
                "done",
                "Conversion worker acquired; starting image intake.",
                shareRequested=share,
            )
            await run_job(job, raw_upload=raw, object_hint=object_hint,
                          settings=settings,
                          llm=app.state.llm_factory(settings),
                          share=share,
                          gallery_store=gallery_store)
        except TimeoutError:
            fail_job_timeout(job, settings.job_timeout_s)
        finally:
            if acquired:
                job_semaphore.release()

    asyncio.create_task(runner())
    return JSONResponse(status_code=202, content={
        "job_id": job.id,
        "events_url": f"/api/jobs/{job.id}/events",
        "status_url": f"/api/jobs/{job.id}",
        "share_requested": share,
    }, headers={"Cache-Control": "no-store"})


@app.get("/api/jobs/{job_id}")
async def job_status(job_id: str) -> JSONResponse:
    job = registry.get(job_id)
    if job is None:
        return JSONResponse(
            status_code=404,
            content={"error": "job_not_found"},
            headers={"Cache-Control": "no-store"},
        )
    return JSONResponse({
        "job_id": job.id,
        "status": job.status,
        "stage": job.stage,
        "result": job.result,
        "error": job.error,
    }, headers={"Cache-Control": "no-store"})


def _is_terminal(event: dict) -> bool:
    """An event ends the SSE stream: any error, or the final done/done."""
    return event.get("status") == "error" or (
        event.get("status") == "done" and event.get("stage") == "done"
    )


def _sse_resume_seq(request: Request, current_seq: int) -> int:
    """Parse a standard Last-Event-ID or explicit cursor without exceptions."""
    raw = request.headers.get("last-event-id")
    if raw is None:
        raw = request.query_params.get("cursor")
    try:
        parsed = int(raw) if raw is not None else 0
    except (TypeError, ValueError):
        return 0
    if parsed < 0:
        return 0
    return min(parsed, current_seq)


@app.get("/api/jobs/{job_id}/events", response_model=None)
async def job_events(job_id: str, request: Request) -> Response:
    job = registry.get(job_id)
    if job is None:
        return JSONResponse(
            status_code=404,
            content={"error": "job_not_found"},
            headers={"Cache-Control": "no-store"},
        )

    async def stream():
        last_seq = _sse_resume_seq(request, job.seq)
        cursor = next(
            (
                index for index, event in enumerate(job.events)
                if int(event.get("seq", 0)) > last_seq
            ),
            len(job.events),
        )
        while True:
            # Replay any events not yet sent (covers reconnects).
            while cursor < len(job.events):
                event = job.events[cursor]
                cursor += 1
                yield f"id: {event['seq']}\ndata: {json.dumps(event)}\n\n"
                if _is_terminal(event):
                    return
            if job.status != "running":
                return
            if await request.is_disconnected():
                return
            # Clear before the sequence recheck: an event emitted before this
            # clear is observed by the cursor check, while one emitted after it
            # sets the waiter. This avoids the old set+immediate-clear race.
            job.waiter.clear()
            if cursor < len(job.events):
                continue
            try:
                await asyncio.wait_for(job.waiter.wait(), timeout=15)
            except asyncio.TimeoutError:
                yield ": heartbeat\n\n"

    return StreamingResponse(
        stream(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-store", "X-Accel-Buffering": "no"},
    )


@app.get("/api/jobs/{job_id}/artifacts/{name}")
async def job_artifact(job_id: str, name: str, download: bool = False) -> Response:
    job = registry.get(job_id)
    if job is None:
        return JSONResponse(
            status_code=404,
            content={"error": "job_not_found"},
            headers={"Cache-Control": "no-store"},
        )
    if name not in ARTIFACT_NAMES:
        return JSONResponse(
            status_code=404,
            content={"error": "artifact_not_found"},
            headers={"Cache-Control": "no-store"},
        )
    if job.status != "done" and name in SUCCESS_ONLY_JOB_ARTIFACTS:
        return JSONResponse(
            status_code=404,
            content={"error": "artifact_not_found"},
            headers={"Cache-Control": "no-store"},
        )
    path = (job.dir / name).resolve()
    # Containment: the artifact must live inside this job's directory.
    if not path.is_file() or path.parent != job.dir.resolve():
        return JSONResponse(
            status_code=404,
            content={"error": "artifact_not_found"},
            headers={"Cache-Control": "no-store"},
        )
    headers = {"Cache-Control": "no-store"}
    if download or name in {"factory.ts", "standalone.html"}:
        headers["Content-Disposition"] = f'attachment; filename="{name}"'
    media_type = ARTIFACT_NAMES[name].split(";")[0]
    return FileResponse(path, media_type=media_type, headers=headers)


# ---------------------------------------------------------------------------
# Persistent community gallery
# ---------------------------------------------------------------------------

@app.get("/api/gallery")
def gallery_index(
    offset: int = Query(default=0, ge=0, le=1_000_000),
    limit: int = Query(default=24, ge=1, le=100),
) -> JSONResponse:
    try:
        payload = gallery_store.list_items(offset=offset, limit=limit)
    except GalleryError:
        logger.exception("gallery listing failed")
        return JSONResponse(
            status_code=503,
            content={"error": "gallery_unavailable"},
            headers={"Cache-Control": "no-store"},
        )
    return JSONResponse(payload, headers={"Cache-Control": "no-store"})


@app.get("/api/gallery/{item_id}")
def gallery_detail(item_id: str) -> JSONResponse:
    item = gallery_store.get(item_id)
    if item is None:
        return JSONResponse(
            status_code=404,
            content={"error": "gallery_item_not_found"},
            headers={"Cache-Control": "no-store"},
        )
    return JSONResponse(item, headers={"Cache-Control": "public, max-age=60"})


@app.get("/api/gallery/{item_id}/artifacts/{name}")
def gallery_artifact(
    item_id: str,
    name: str,
    download: bool = False,
) -> Response:
    path = gallery_store.artifact_path(item_id, name)
    if path is None:
        return JSONResponse(
            status_code=404,
            content={"error": "gallery_artifact_not_found"},
            headers={"Cache-Control": "no-store"},
        )
    headers = {"Cache-Control": "public, max-age=31536000, immutable"}
    if download or name in {"factory.ts", "standalone.html"}:
        headers["Content-Disposition"] = f'attachment; filename="{name}"'
    media_type = GALLERY_ARTIFACT_NAMES[name].split(";")[0]
    return FileResponse(path, media_type=media_type, headers=headers)


# ---------------------------------------------------------------------------
# Static SPA
# ---------------------------------------------------------------------------

@app.get("/", include_in_schema=False)
async def index() -> FileResponse:
    return FileResponse(STATIC_DIR / "index.html")


@app.get("/gallery", include_in_schema=False)
async def gallery_page() -> FileResponse:
    """Friendly SPA entry for browsing the community gallery."""
    return FileResponse(STATIC_DIR / "index.html")


@app.get("/gallery/{item_id}", include_in_schema=False, response_model=None)
async def gallery_item_page(item_id: str) -> Response:
    """Friendly SPA deep link; reject malformed ids before serving the shell."""
    if not valid_item_id(item_id):
        return JSONResponse(
            status_code=404,
            content={"error": "gallery_item_not_found"},
            headers={"Cache-Control": "no-store"},
        )
    return FileResponse(STATIC_DIR / "index.html")


app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")


@app.exception_handler(ImageRejected)
async def image_rejected_handler(_: Request, exc: ImageRejected) -> JSONResponse:
    return JSONResponse(status_code=422, content={"error": exc.code, "detail": exc.reason})


def main() -> None:
    """Console entry: python -m app.main"""
    import uvicorn

    uvicorn.run(
        "app.main:app",
        host="0.0.0.0",
        port=settings.port,
        workers=1,
        proxy_headers=True,
        log_level="info",
    )


if __name__ == "__main__":
    main()