api/main.py

from __future__ import annotations

import logging
import time
from collections import defaultdict, deque
from pathlib import Path
from uuid import uuid4

from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel

from datapilot.config import get_settings
from datapilot.data import load_sample, read_dataset
from datapilot.insights import answer_follow_up
from datapilot.jobs import JobManager
from datapilot.persistence import RunStore
from datapilot.workflow import run_analysis

settings = get_settings()
logger = logging.getLogger(__name__)
jobs = JobManager(workers=2)
request_windows: dict[str, deque[float]] = defaultdict(deque)
app = FastAPI(
    title="DataPilot AI API",
    version="1.0.0",
    description="Evidence-grounded autonomous data science with LangGraph.",
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=[item.strip() for item in settings.cors_origins.split(",")],
    allow_methods=["GET", "POST"],
    allow_headers=["*"],
)


@app.middleware("http")
async def security_middleware(request: Request, call_next):
    correlation_id = request.headers.get("X-Correlation-ID") or uuid4().hex
    request.state.correlation_id = correlation_id
    if request.url.path.startswith("/v1/"):
        if settings.api_key and request.headers.get("X-API-Key") != settings.api_key:
            return _error_response(401, "Authentication required.", correlation_id)
        client = request.client.host if request.client else "unknown"
        now = time.monotonic()
        window = request_windows[client]
        while window and now - window[0] > 60:
            window.popleft()
        if len(window) >= settings.requests_per_minute:
            return _error_response(429, "Rate limit exceeded.", correlation_id)
        window.append(now)
    response = await call_next(request)
    response.headers["X-Correlation-ID"] = correlation_id
    return response


def _error_response(status_code: int, message: str, correlation_id: str):
    from fastapi.responses import JSONResponse

    return JSONResponse(
        status_code=status_code,
        content={"detail": message, "correlation_id": correlation_id},
        headers={"X-Correlation-ID": correlation_id},
    )


def _safe_failure(exc: Exception, correlation_id: str) -> HTTPException:
    logger.exception(
        "Analysis failure correlation_id=%s category=%s", correlation_id, type(exc).__name__
    )
    return HTTPException(
        status_code=500,
        detail={
            "message": "Analysis failed. Use the correlation ID when contacting support.",
            "correlation_id": correlation_id,
        },
    )


class SampleRequest(BaseModel):
    sample: str


class ChatRequest(BaseModel):
    question: str


@app.get("/")
def root() -> dict[str, str]:
    return {"name": settings.app_name, "status": "ready", "docs": "/docs"}


@app.get("/health")
def health() -> dict[str, object]:
    return {
        "status": "healthy",
        "environment": settings.environment,
        "limits": {
            "max_upload_mb": settings.max_upload_mb,
            "max_rows": settings.max_rows,
            "max_columns": settings.max_columns,
        },
    }


@app.post("/v1/analyze/sample", status_code=202)
def analyze_sample(request: SampleRequest, http_request: Request):
    try:
        frame, target, dataset_name = load_sample(request.sample)
        job = jobs.submit(lambda: run_analysis(frame, target, dataset_name, settings))
        return {"job_id": job.job_id, "status": job.status, "status_url": f"/v1/jobs/{job.job_id}"}
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    except Exception as exc:
        raise _safe_failure(exc, http_request.state.correlation_id) from exc


@app.post("/v1/analyze/upload")
async def analyze_upload(
    request: Request,
    file: UploadFile = File(...),
    target: str = Form(...),
):
    try:
        content = await file.read(settings.max_upload_mb * 1024 * 1024 + 1)
        frame = read_dataset(content, file.filename or "dataset.csv", settings)
        job = jobs.submit(
            lambda: run_analysis(frame, target, file.filename or "uploaded_dataset", settings)
        )
        return {"job_id": job.job_id, "status": job.status, "status_url": f"/v1/jobs/{job.job_id}"}
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    except Exception as exc:
        raise _safe_failure(exc, request.state.correlation_id) from exc


@app.get("/v1/jobs/{job_id}")
def get_job(job_id: str):
    job = jobs.get(job_id)
    if job is None:
        raise HTTPException(status_code=404, detail="Job not found.")
    return job.public()


@app.delete("/v1/jobs/{job_id}")
def cancel_job(job_id: str):
    if jobs.get(job_id) is None:
        raise HTTPException(status_code=404, detail="Job not found.")
    return {"job_id": job_id, "cancelled": jobs.cancel(job_id)}


@app.get("/v1/runs")
def recent_runs():
    return RunStore(settings).list_recent()


@app.get("/v1/runs/{run_id}")
def get_run(run_id: str):
    run = RunStore(settings).get(run_id)
    if run is None:
        raise HTTPException(status_code=404, detail="Run not found.")
    return run


@app.post("/v1/runs/{run_id}/chat")
def chat_with_run(run_id: str, request: ChatRequest):
    run = RunStore(settings).get(run_id)
    if run is None:
        raise HTTPException(status_code=404, detail="Run not found.")
    return {"answer": answer_follow_up(run, request.question)}


@app.get("/v1/runs/{run_id}/artifacts/{artifact_name}")
def download_artifact(run_id: str, artifact_name: str):
    run = RunStore(settings).get(run_id)
    if run is None:
        raise HTTPException(status_code=404, detail="Run not found.")
    path_string = run.get("artifacts", {}).get(artifact_name)
    if not path_string:
        raise HTTPException(status_code=404, detail="Artifact not found.")
    path = Path(path_string).resolve()
    artifact_root = settings.artifact_root.resolve()
    if artifact_root not in path.parents or not path.is_file():
        raise HTTPException(status_code=404, detail="Artifact not found.")
    return FileResponse(path, filename=path.name)