app.py

"""
ONNX Grid Vision API
=====================
FastAPI-based image grid classification API with web automation.
Uses ONNX Runtime for inference, patchright for browser automation.

Deploy to cloud platforms with Docker support.
"""

import asyncio
import logging
import uuid
import time
from contextlib import asynccontextmanager
from typing import Optional

from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

from browser_engine import VisionEngine

# ── Logging ──────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logger = logging.getLogger("gridvision")

# ── Globals ──────────────────────────────────────────────────────────────────
engine: Optional[VisionEngine] = None
active_sessions: dict = {}  # task_id -> status info


# ── Request / Response Models ────────────────────────────────────────────────
class ProcessRequest(BaseModel):
    sitekey: str = Field(..., description="Widget site key")
    url: str = Field(..., description="Page URL for domain context")
    is_enterprise: bool = Field(False, description="Enterprise widget mode")
    is_invisible: bool = Field(False, description="Invisible widget mode")
    proxy: Optional[str] = Field(None, description="Proxy URL (socks5://user:pass@host:port)")
    headless: bool = Field(True, description="Run browser headless")
    timeout: int = Field(120, ge=10, le=300, description="Processing timeout in seconds")


class ProcessResponse(BaseModel):
    success: bool
    token: Optional[str] = None
    error: Optional[str] = None
    elapsed: Optional[float] = None


class StatusResponse(BaseModel):
    task_id: str
    status: str  # pending, processing, completed, failed
    token: Optional[str] = None
    error: Optional[str] = None
    elapsed: float


# ── Lifespan ─────────────────────────────────────────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
    global engine
    logger.info("Loading ONNX models...")
    engine = VisionEngine(models_dir="models")
    engine.load_models()
    logger.info("ONNX Grid Vision API ready")
    yield
    logger.info("Shutting down")


# ── FastAPI App ──────────────────────────────────────────────────────────────
app = FastAPI(
    title="ONNX Grid Vision API",
    description=(
        "Image grid classification service using YOLOv5 segmentation "
        "and per-class binary classifiers with web automation."
    ),
    version="1.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ── Background Processor ────────────────────────────────────────────────────
async def _process_background(task_id: str, req: ProcessRequest):
    active_sessions[task_id]["status"] = "processing"
    start = time.time()
    try:
        token = await engine.process(
            sitekey=req.sitekey,
            url=req.url,
            proxy=req.proxy,
            headless=req.headless,
            timeout=req.timeout,
            is_enterprise=req.is_enterprise,
            is_invisible=req.is_invisible,
        )
        active_sessions[task_id]["status"] = "completed"
        active_sessions[task_id]["token"] = token
        active_sessions[task_id]["elapsed"] = round(time.time() - start, 2)
        logger.info(f"Task {task_id} completed in {active_sessions[task_id]['elapsed']}s")
    except Exception as e:
        active_sessions[task_id]["status"] = "failed"
        active_sessions[task_id]["error"] = str(e)
        active_sessions[task_id]["elapsed"] = round(time.time() - start, 2)
        logger.error(f"Task {task_id} failed: {e}")


# ── Endpoints ────────────────────────────────────────────────────────────────
@app.get("/", tags=["Info"])
async def root():
    return {
        "name": "ONNX Grid Vision API",
        "version": "1.0.0",
        "endpoints": {
            "POST /process": "Process verification synchronously (blocking)",
            "POST /process/async": "Process verification asynchronously",
            "GET /status/{task_id}": "Check async task status",
            "GET /health": "Health check",
        },
    }


@app.get("/health", tags=["Info"])
async def health():
    return {"status": "ok", "models_loaded": engine is not None and engine.loaded}


@app.post("/process", response_model=ProcessResponse, tags=["Processing"])
async def process(req: ProcessRequest):
    """
    Process verification synchronously.
    Returns the token directly (blocks until done or timeout).
    """
    if not engine or not engine.loaded:
        raise HTTPException(503, "Engine not ready (models still loading)")

    start = time.time()
    try:
        token = await engine.process(
            sitekey=req.sitekey,
            url=req.url,
            proxy=req.proxy,
            headless=req.headless,
            timeout=req.timeout,
            is_enterprise=req.is_enterprise,
            is_invisible=req.is_invisible,
        )
        elapsed = round(time.time() - start, 2)
        if token:
            return ProcessResponse(success=True, token=token, elapsed=elapsed)
        else:
            return ProcessResponse(success=False, error="Timeout: could not complete", elapsed=elapsed)
    except Exception as e:
        elapsed = round(time.time() - start, 2)
        logger.error(f"Sync process failed: {e}")
        return ProcessResponse(success=False, error=str(e), elapsed=elapsed)


@app.post("/process/async", tags=["Processing"])
async def process_async(req: ProcessRequest, background_tasks: BackgroundTasks):
    """
    Process verification asynchronously.
    Returns a task_id immediately. Poll /status/{task_id} for results.
    """
    if not engine or not engine.loaded:
        raise HTTPException(503, "Engine not ready (models still loading)")

    task_id = str(uuid.uuid4())
    active_sessions[task_id] = {
        "status": "pending",
        "token": None,
        "error": None,
        "elapsed": 0.0,
    }

    background_tasks.add_task(_process_background, task_id, req)
    return {"task_id": task_id, "status": "pending"}


@app.get("/status/{task_id}", response_model=StatusResponse, tags=["Processing"])
async def status(task_id: str):
    """Check the status of an async processing task."""
    if task_id not in active_sessions:
        raise HTTPException(404, "Task not found")

    session = active_sessions[task_id]
    return StatusResponse(
        task_id=task_id,
        status=session["status"],
        token=session.get("token"),
        error=session.get("error"),
        elapsed=session["elapsed"],
    )


# ── Entry Point ──────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)