api.py

import os
import uuid
import tempfile
import asyncio
import logging
from datetime import datetime, timedelta
from typing import Dict, Any

import pandas as pd
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
import uvicorn

from core.pipeline import AgenticPipeline

logger = logging.getLogger(__name__)

# ============================================================================
# GLOBAL STATE
# ============================================================================

active_runs: Dict[str, Dict[str, Any]] = {}
completed_runs: Dict[str, Dict[str, Any]] = {}

UPLOAD_DIR = tempfile.mkdtemp(prefix="agentic_uploads_")
RESULTS_DIR = tempfile.mkdtemp(prefix="agentic_results_")

# FIX #8: cap how many completed runs we keep in memory
MAX_COMPLETED_RUNS = 50
# FIX #8: delete result files older than this
RESULT_TTL_HOURS = 2

APP_VERSION = "1.1.0"

print(f"📁 Upload dir:  {UPLOAD_DIR}")
print(f"📁 Results dir: {RESULTS_DIR}")

# ============================================================================
# FASTAPI APP
# ============================================================================

app = FastAPI(
    title="Agentic Data Scientist API",
    description="Autonomous Multi-Agent Data Science Pipeline",
    version=APP_VERSION,
)

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


# ============================================================================
# HELPERS
# ============================================================================

def _add_log(run_state: dict, message: str, level: str = "info"):
    """Append a timestamped log entry to a run's state."""
    run_state["logs"].append({
        "timestamp": datetime.now().isoformat(),
        "message": message,
        "level": level,
    })
    run_state["message"] = message
    run_state["current_stage"] = message[:100]


def _evict_old_runs():
    """
    FIX #8: Remove oldest completed runs when the dict exceeds MAX_COMPLETED_RUNS.
    Also deletes associated result files so temp dirs don't grow unboundedly.
    """
    if len(completed_runs) <= MAX_COMPLETED_RUNS:
        return

    # Sort by completion time, remove oldest surplus entries
    sorted_runs = sorted(
        completed_runs.items(),
        key=lambda kv: kv[1].get("completed_at") or "",
    )
    to_evict = sorted_runs[: len(completed_runs) - MAX_COMPLETED_RUNS]

    for run_id, run_state in to_evict:
        _delete_run_files(run_id, run_state)
        del completed_runs[run_id]
        logger.info(f"Evicted old run {run_id[:8]}")


def _delete_run_files(run_id: str, run_state: dict):
    """Delete all files on disk associated with a completed run."""
    result = run_state.get("result") or {}

    for key in ("model_path", "cleaned_data_path"):
        path = result.get(key)
        if path and os.path.exists(path):
            try:
                os.remove(path)
            except OSError:
                pass

    # Report PDF may be stored separately
    report_path = result.get("report_path")
    if report_path and os.path.exists(report_path):
        try:
            os.remove(report_path)
        except OSError:
            pass

    # Uploaded CSV
    upload_path = os.path.join(UPLOAD_DIR, f"{run_id}.csv")
    if os.path.exists(upload_path):
        try:
            os.remove(upload_path)
        except OSError:
            pass


# ============================================================================
# HEALTH CHECK
# ============================================================================

@app.get("/")
@app.get("/health")
async def health_check():
    return {
        "status": "healthy",
        "version": APP_VERSION,
        "timestamp": datetime.now().isoformat(),
        "active_runs": len(active_runs),
        "completed_runs": len(completed_runs),
    }


# ============================================================================
# UPLOAD
# ============================================================================

@app.post("/upload")
async def upload_file(file: UploadFile = File(...)):
    if not file.filename or not file.filename.endswith(".csv"):
        raise HTTPException(status_code=400, detail="Only CSV files supported")

    file_id = str(uuid.uuid4())
    file_path = os.path.join(UPLOAD_DIR, f"{file_id}.csv")

    try:
        contents = await file.read()
        with open(file_path, "wb") as f:
            f.write(contents)

        df = pd.read_csv(file_path)

        return {
            "file_id": file_id,
            "filename": file.filename,
            "size_bytes": len(contents),
            "rows": len(df),
            "columns": len(df.columns),
            "column_names": df.columns.tolist(),
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ============================================================================
# PIPELINE RUN
# ============================================================================

@app.post("/pipeline/run")
async def run_pipeline(
    background_tasks: BackgroundTasks,
    file_id: str = Form(...),
    target_column: str = Form(...),
    problem_type: str = Form(default="classification"),
):
    file_path = os.path.join(UPLOAD_DIR, f"{file_id}.csv")

    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="File not found")

    if problem_type not in ["classification", "regression"]:
        raise HTTPException(status_code=400, detail="problem_type must be classification or regression")

    run_id = str(uuid.uuid4())

    active_runs[run_id] = {
        "status": "initializing",
        "progress": 5.0,
        "current_stage": "Initializing pipeline...",
        "message": "Pipeline starting",
        "logs": [
            {"timestamp": datetime.now().isoformat(), "message": "🚀 Pipeline initialized", "level": "info"},
            {"timestamp": datetime.now().isoformat(), "message": "📊 Loading data...", "level": "info"},
        ],
        "started_at": datetime.now().isoformat(),
        "completed_at": None,
        "result": None,
    }

    background_tasks.add_task(
        _execute_pipeline,
        run_id=run_id,
        file_path=file_path,
        target_column=target_column,
        problem_type=problem_type,
    )

    return {"run_id": run_id, "status": "initializing", "message": "Pipeline started"}


async def _execute_pipeline(run_id: str, file_path: str, target_column: str, problem_type: str):
    """
    Execute the pipeline in a background task.

    FIX #2: The CPU-bound pipeline.run() call is wrapped in asyncio.to_thread()
    so it executes in a thread-pool worker and never blocks FastAPI's event loop.
    Other requests (status polls, health checks) remain responsive during training.
    """
    run_state = active_runs[run_id]

    try:
        run_state["status"] = "running"
        _add_log(run_state, "📊 Stage 1/6: Profiling data...", "info")
        run_state["progress"] = 10

        df = pd.read_csv(file_path)
        _add_log(run_state, f"✅ Data loaded: {df.shape[0]} rows, {df.shape[1]} columns", "success")

        def log_callback(message: str, level: str = "info"):
            _add_log(run_state, message, level)
            # Update progress based on stage markers in the message
            if "Stage 1" in message:
                run_state["progress"] = 15
            elif "Stage 2" in message:
                run_state["progress"] = 30
            elif "Stage 3" in message:
                run_state["progress"] = 45
            elif "Stage 4" in message:
                run_state["progress"] = 60
            elif "Stage 5" in message:
                run_state["progress"] = 75
            elif "Stage 6" in message:
                run_state["progress"] = 90
            elif "completed" in message.lower():
                run_state["progress"] = 100

        pipeline = AgenticPipeline(
            df=df,
            target_col=target_column,
            problem_type=problem_type,
            log_callback=log_callback,
        )

        _add_log(run_state, "🚀 Starting pipeline execution...", "info")

        # FIX #2: run CPU-bound work in a thread so the event loop stays free
        results = await asyncio.to_thread(pipeline.run)

        run_state["progress"] = 100
        run_state["current_stage"] = "Complete"

        if results.get("success"):
            run_state["status"] = "completed"
            _add_log(run_state, "✅ Pipeline completed successfully!", "success")
            _add_log(
                run_state,
                f"🏆 Best model: {results.get('metrics', {}).get('best_model', 'N/A')}",
                "success",
            )

            run_state["result"] = {
                "success": True,
                "pipeline_summary": results.get("pipeline_summary", {}),
                "metrics": results.get("metrics", {}),
                "report_path": results.get("report_path"),
            }

            # Save model
            if results.get("models") and "best" in results["models"]:
                import joblib
                model_path = os.path.join(RESULTS_DIR, f"{run_id}_model.pkl")
                joblib.dump(results["models"]["best"], model_path)
                run_state["result"]["model_path"] = model_path
                _add_log(run_state, "💾 Model saved successfully", "success")

            # Save cleaned data
            if results.get("cleaned_df") is not None:
                cleaned_path = os.path.join(RESULTS_DIR, f"{run_id}_cleaned.csv")
                results["cleaned_df"].to_csv(cleaned_path, index=False)
                run_state["result"]["cleaned_data_path"] = cleaned_path
                _add_log(run_state, "💾 Cleaned data saved", "success")

        else:
            run_state["status"] = "failed"
            _add_log(run_state, f"❌ Pipeline failed: {results.get('error', 'Unknown error')}", "error")
            run_state["result"] = {"success": False, "error": results.get("error", "Unknown error")}

    except Exception as e:
        run_state["status"] = "failed"
        _add_log(run_state, f"❌ Error: {str(e)}", "error")
        run_state["result"] = {"success": False, "error": str(e)}

    finally:
        run_state["completed_at"] = datetime.now().isoformat()
        completed_runs[run_id] = run_state
        if run_id in active_runs:
            del active_runs[run_id]
        # FIX #8: evict old runs to prevent unbounded memory growth
        _evict_old_runs()


# ============================================================================
# STATUS & RESULTS
# ============================================================================

@app.get("/pipeline/status/{run_id}")
async def get_status(run_id: str):
    run = active_runs.get(run_id) or completed_runs.get(run_id)
    if run is None:
        raise HTTPException(status_code=404, detail="Run ID not found")

    return {
        "run_id": run_id,
        "status": run["status"],
        "progress": run["progress"],
        "current_stage": run["current_stage"],
        "message": run["message"],
        "started_at": run["started_at"],
        "completed_at": run.get("completed_at"),
        "logs": run.get("logs", []),
        "result": run.get("result"),
    }


@app.get("/pipeline/result/{run_id}")
async def get_result(run_id: str):
    if run_id in active_runs:
        raise HTTPException(status_code=400, detail="Pipeline still running")

    run = completed_runs.get(run_id)
    if run is None:
        raise HTTPException(status_code=404, detail="Run ID not found")

    if run["status"] == "completed":
        return run["result"]

    raise HTTPException(status_code=400, detail="Pipeline failed")


# ============================================================================
# DOWNLOADS
# ============================================================================

@app.get("/download/model/{run_id}")
async def download_model(run_id: str):
    model_path = os.path.join(RESULTS_DIR, f"{run_id}_model.pkl")
    if not os.path.exists(model_path):
        raise HTTPException(status_code=404, detail="Model not found")
    return FileResponse(model_path, filename=f"model_{run_id[:8]}.pkl")


@app.get("/download/data/{run_id}")
async def download_data(run_id: str):
    data_path = os.path.join(RESULTS_DIR, f"{run_id}_cleaned.csv")
    if not os.path.exists(data_path):
        raise HTTPException(status_code=404, detail="Data not found")
    return FileResponse(data_path, filename=f"cleaned_{run_id[:8]}.csv")


@app.get("/download/report/{run_id}")
async def download_report(run_id: str):
    run = completed_runs.get(run_id)
    if run:
        report_path = run.get("result", {}).get("report_path")
        if report_path and os.path.exists(report_path):
            return FileResponse(report_path, filename="report.pdf")
    raise HTTPException(status_code=404, detail="Report not found")


# ============================================================================
# MAIN
# ============================================================================

if __name__ == "__main__":
    print("=" * 50)
    print("🚀 Starting Agentic Data Scientist API")
    print("📍 http://localhost:8000")
    print("📖 http://localhost:8000/docs")
    print("=" * 50)
    uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=True)