backend/api/main.py

"""
backend/api/main.py + routes — Full FastAPI application.
"""
from __future__ import annotations
import json
from contextlib import asynccontextmanager
from datetime import datetime, timezone

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

from ..agents.orchestrator import run_workflow, stream_workflow
from ..memory.memory_store import init_db, long_term, short_term
from ..state.graph_state import create_initial_state, TaskStatus
from ..core.config import get_settings
from ..core.logger import setup_logging, get_logger

log = get_logger(__name__)


@asynccontextmanager
async def lifespan(app: FastAPI):
    setup_logging()
    await init_db()
    log.info("Multi-Agent System started")
    yield
    log.info("Multi-Agent System shutting down")


app = FastAPI(
    title="Multi-Agent Workflow System",
    description="LangGraph-based autonomous agent system with Planner, Executor, Critic, and Memory.",
    version="1.0.0",
    lifespan=lifespan,
)

settings = get_settings()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"] if settings.is_dev else ["https://yourdomain.com"],
    allow_methods=["*"],
    allow_headers=["*"],
)


# ── Request/Response models ────────────────────────────────────────────────────

class TaskRequest(BaseModel):
    task: str = Field(..., min_length=5, max_length=2000, description="The task to complete")
    stream: bool = Field(default=False, description="Stream agent events via SSE")


class TaskResponse(BaseModel):
    task_id: str
    task: str
    status: str
    final_output: str | None
    quality_score: float | None
    plan: list[dict]
    events: list[dict]
    total_tokens: int
    error_message: str | None
    created_at: str
    completed_at: str | None


# ── Background task registry ───────────────────────────────────────────────────

_running_tasks: dict[str, dict] = {}


# ── Routes ────────────────────────────────────────────────────────────────────

@app.post("/api/tasks", response_model=TaskResponse)
async def create_task(req: TaskRequest):
    """
    Submit a new task to the multi-agent system.
    Runs the full Planner → Executor → Critic → Memory pipeline.
    """
    state = create_initial_state(req.task)
    task_id = state["task_id"]

    log.info("Task submitted", task_id=task_id, task=req.task[:80])

    final_state = await run_workflow(state)

    return TaskResponse(
        task_id=task_id,
        task=req.task,
        status=final_state.get("status", "unknown"),
        final_output=final_state.get("final_output"),
        quality_score=final_state.get("quality_score"),
        plan=final_state.get("plan", []),
        events=final_state.get("events", []),
        total_tokens=final_state.get("total_tokens", 0),
        error_message=final_state.get("error_message"),
        created_at=final_state.get("created_at", ""),
        completed_at=datetime.now(timezone.utc).isoformat(),
    )


@app.post("/api/tasks/stream")
async def create_task_stream(req: TaskRequest):
    """
    Submit task and stream agent events via Server-Sent Events.
    Frontend receives real-time updates as each agent node runs.
    """
    state = create_initial_state(req.task)

    async def event_generator():
        last_snapshot = None
        try:
            async for snapshot in stream_workflow(state):
                last_snapshot = snapshot
                events = snapshot.get("events", [])
                latest_event = events[-1] if events else {}

                payload = {
                    "task_id": snapshot.get("task_id"),
                    "status": snapshot.get("status"),
                    "iteration": snapshot.get("iteration", 0),
                    "plan": snapshot.get("plan", []),
                    "latest_event": latest_event,
                    "quality_score": snapshot.get("quality_score"),
                    "total_tokens": snapshot.get("total_tokens", 0),
                    "error_message": snapshot.get("error_message"),
                    "final_output": snapshot.get("final_output"),
                }
                yield f"data: {json.dumps(payload, default=str)}\n\n"

            # Emit terminal SSE event based on final status
            final_status = str((last_snapshot or {}).get("status", ""))
            if final_status == "failed":
                err = (last_snapshot or {}).get("error_message") or "Task failed"
                log.error("Task ended in failed state", task_id=state["task_id"], error=err)
                yield f"event: error\ndata: {json.dumps({'error': err, 'task_id': state['task_id']})}\n\n"
            else:
                yield f"event: done\ndata: {json.dumps({'task_id': state['task_id']})}\n\n"

        except Exception as e:
            import traceback
            log.error("Stream workflow crashed", error=str(e), traceback=traceback.format_exc())
            yield f"event: error\ndata: {json.dumps({'error': str(e), 'task_id': state['task_id']})}\n\n"

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


@app.get("/api/tasks/{task_id}")
async def get_task(task_id: str):
    """Get task state — check Redis cache first, then DB."""
    # Try cache
    cached = short_term.get_state(task_id)
    if cached:
        return cached

    # Try DB
    tasks = await long_term.get_recent_tasks(limit=100)
    for t in tasks:
        if t["task_id"] == task_id:
            return t

    raise HTTPException(status_code=404, detail=f"Task {task_id} not found")


@app.get("/api/tasks")
async def list_tasks(limit: int = 20):
    """List recent tasks with status and scores."""
    tasks = await long_term.get_recent_tasks(limit=limit)
    return {"tasks": tasks, "total": len(tasks)}


@app.get("/api/memories")
async def search_memories(q: str = "", memory_type: str | None = None, limit: int = 10):
    """Search the agent's long-term memory."""
    if not q:
        q = "recent"
    memories = await long_term.retrieve(q, memory_type=memory_type, limit=limit)
    return {"memories": memories, "query": q}


@app.get("/api/health")
async def health():
    redis_ok = False
    try:
        r = short_term
        from ..memory.memory_store import get_redis
        rc = get_redis()
        redis_ok = rc is not None and bool(rc.ping())
    except Exception:
        pass

    return {
        "status": "ok",
        "version": "1.0.0",
        "redis": "connected" if redis_ok else "unavailable",
        "agents": ["planner", "executor", "critic", "memory"],
        "tools": ["web_search", "fetch_url", "calculate", "run_python", "write_file", "read_file"],
        "env": settings.app_env,
    }


@app.get("/api/graph")
async def get_graph_definition():
    """Return the agent graph structure for visualization."""
    return {
        "nodes": [
            {"id": "memory_retrieve", "label": "Memory", "role": "memory", "description": "Retrieve relevant past memories"},
            {"id": "planner", "label": "Planner", "role": "planner", "description": "Decompose task into steps"},
            {"id": "executor", "label": "Executor", "role": "executor", "description": "Execute plan steps with tools"},
            {"id": "critic", "label": "Critic", "role": "critic", "description": "Evaluate quality and reflect"},
            {"id": "memory_store", "label": "Memory Store", "role": "memory", "description": "Persist learnings"},
        ],
        "edges": [
            {"from": "START", "to": "memory_retrieve"},
            {"from": "memory_retrieve", "to": "planner"},
            {"from": "planner", "to": "executor", "condition": "plan valid"},
            {"from": "executor", "to": "executor", "condition": "more steps"},
            {"from": "executor", "to": "critic", "condition": "all done"},
            {"from": "critic", "to": "planner", "condition": "needs replan"},
            {"from": "critic", "to": "memory_store", "condition": "approved"},
            {"from": "memory_store", "to": "END"},
        ],
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("backend.api.main:app", host=settings.app_host,
                port=settings.app_port, reload=settings.is_dev)