api/agent_fsm.py

"""
backend/api/agent_fsm.py — AgentLoop come Macchina a Stati (ARCH-I4.5)

Rifactorizza il flusso dell'AgentLoop come FSM (Finite State Machine) esplicita.
Stati: IDLE → PLAN → THINK → TOOL → WAIT → OBSERVE → DECIDE → DONE | FAILED

Vantaggi rispetto all'implementazione implicita in agent.py:
  - Testabilità: ogni transizione è una funzione pura
  - Osservabilità: stato corrente sempre visibile via API
  - Idempotenza: ogni stato ha guard di entry/exit
  - Debuggabilità: history completa delle transizioni

Integrazione con agent.py:
  - AgentFSM NON sostituisce agent.py (troppo rischio regressione)
  - Viene usato come orchestratore esterno per nuovi task creati via Kernel
  - agent.py esistente continua a funzionare invariato

ADR: S10 S16 S21 S27
"""
from __future__ import annotations

import asyncio
import logging
import time
import uuid
from enum import Enum
from typing import Any, Callable, Coroutine

from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel, Field

from .auth_guard import AuthRole, require_role

_logger = logging.getLogger("api.agent_fsm")

# ── Guards ──────────────────────────────────────────────────────────────────────
try:
    from .brain_planner import planner as _planner, PlanRequest as _PlanReq
    _PLANNER_AVAILABLE = True
except Exception:
    _planner = None; _PlanReq = None; _PLANNER_AVAILABLE = False  # type: ignore

try:
    from .tool_engine import tool_executor as _tool_exec, ToolExecuteRequest as _TExecReq
    _TOOL_ENGINE_AVAILABLE = True
except Exception:
    _tool_exec = None; _TExecReq = None; _TOOL_ENGINE_AVAILABLE = False  # type: ignore

try:
    from .kernel import kernel as _kernel
    _KERNEL_AVAILABLE = True
except Exception:
    _kernel = None; _KERNEL_AVAILABLE = False  # type: ignore

# ── States ──────────────────────────────────────────────────────────────────────

class AgentState(str, Enum):
    IDLE    = "idle"      # in attesa di un goal
    PLAN    = "plan"      # BrainPlanner genera WorkflowPlan
    THINK   = "think"     # LLM reasoning: analisi contesto, scelta prossimo tool
    TOOL    = "tool"      # esecuzione tool via ToolEngine
    WAIT    = "wait"      # attesa risultato asincrono (polling)
    OBSERVE = "observe"   # elaborazione risultato tool, aggiornamento contesto
    DECIDE  = "decide"    # decide se DONE, se re-THINK, o se FAILED
    DONE    = "done"      # goal raggiunto
    FAILED  = "failed"    # errore non recuperabile

# Transizioni valide: stato → [stati raggiungibili]
_TRANSITIONS: dict[AgentState, list[AgentState]] = {
    AgentState.IDLE:    [AgentState.PLAN],
    AgentState.PLAN:    [AgentState.THINK, AgentState.FAILED],
    AgentState.THINK:   [AgentState.TOOL, AgentState.DECIDE, AgentState.FAILED],
    AgentState.TOOL:    [AgentState.WAIT, AgentState.OBSERVE, AgentState.FAILED],
    AgentState.WAIT:    [AgentState.OBSERVE, AgentState.FAILED],
    AgentState.OBSERVE: [AgentState.DECIDE, AgentState.FAILED],
    AgentState.DECIDE:  [AgentState.THINK, AgentState.DONE, AgentState.FAILED],
    AgentState.DONE:    [],
    AgentState.FAILED:  [],
}

# ── Models ──────────────────────────────────────────────────────────────────────

class StateTransition(BaseModel):
    from_state:  AgentState
    to_state:    AgentState
    reason:      str           = ""
    ts:          float         = Field(default_factory=time.time)
    metadata:    dict[str, Any] = Field(default_factory=dict)


class AgentContext(BaseModel):
    """Contesto condiviso tra tutti gli stati della FSM."""
    session_id:     str            = Field(default_factory=lambda: str(uuid.uuid4()))
    goal:           str            = ""
    plan_id:        str | None     = None
    messages:       list[dict]     = Field(default_factory=list)
    tool_results:   list[dict]     = Field(default_factory=list)
    current_step:   int            = 0
    max_steps:      int            = 10
    last_tool:      str | None     = None
    last_result:    Any            = None
    error:          str | None     = None
    metadata:       dict[str, Any] = Field(default_factory=dict)


class FSMRunRequest(BaseModel):
    goal:           str
    session_id:     str | None     = None
    max_steps:      int            = 10
    hints:          list[str]      = Field(default_factory=list,
                                           description="Capability/tool suggeriti")
    correlation_id: str            = Field(default_factory=lambda: str(uuid.uuid4()))


class FSMStatus(BaseModel):
    fsm_id:      str
    state:       AgentState
    context:     AgentContext
    history:     list[StateTransition]
    started_at:  float
    updated_at:  float
    done:        bool


# ── AgentFSM ────────────────────────────────────────────────────────────────────

class AgentFSM:
    """
    Istanza di FSM per un singolo run dell'agente.
    Ogni run ha il proprio fsm_id, contesto e history.
    """

    def __init__(self, fsm_id: str, goal: str, max_steps: int = 10,
                 hints: list[str] | None = None) -> None:
        self.fsm_id     = fsm_id
        self.state      = AgentState.IDLE
        self.context    = AgentContext(goal=goal, max_steps=max_steps,
                                       metadata={"hints": hints or []})
        self.history:   list[StateTransition] = []
        self.started_at = time.time()
        self.updated_at = time.time()

    def can_transition(self, to: AgentState) -> bool:
        return to in _TRANSITIONS.get(self.state, [])

    def transition(self, to: AgentState, reason: str = "", metadata: dict | None = None) -> None:
        if not self.can_transition(to):
            raise ValueError(f"Transizione non valida: {self.state} → {to}")
        t = StateTransition(from_state=self.state, to_state=to,
                            reason=reason, metadata=metadata or {})
        self.history.append(t)
        _logger.debug("[fsm:%s] %s → %s (%s)", self.fsm_id, self.state.value, to.value, reason)
        self.state      = to
        self.updated_at = time.time()

    @property
    def done(self) -> bool:
        return self.state in (AgentState.DONE, AgentState.FAILED)

    def to_status(self) -> FSMStatus:
        return FSMStatus(
            fsm_id=self.fsm_id, state=self.state, context=self.context,
            history=self.history, started_at=self.started_at,
            updated_at=self.updated_at, done=self.done)


# ── AgentFSMRunner ──────────────────────────────────────────────────────────────

class AgentFSMRunner:
    """
    Esegue la FSM step-by-step.
    Ogni stato chiama il componente appropriato (Planner, LLM, ToolEngine).
    """

    async def run(self, fsm: AgentFSM) -> FSMStatus:
        """Esegue la FSM fino a DONE o FAILED."""
        try:
            # IDLE → PLAN
            fsm.transition(AgentState.PLAN, "starting run")
            await self._state_plan(fsm)

            while not fsm.done:
                if fsm.context.current_step >= fsm.context.max_steps:
                    fsm.transition(AgentState.FAILED, f"max_steps={fsm.context.max_steps} raggiunto")
                    break

                if fsm.state == AgentState.THINK:
                    await self._state_think(fsm)
                elif fsm.state == AgentState.TOOL:
                    await self._state_tool(fsm)
                elif fsm.state == AgentState.WAIT:
                    await self._state_wait(fsm)
                elif fsm.state == AgentState.OBSERVE:
                    await self._state_observe(fsm)
                elif fsm.state == AgentState.DECIDE:
                    await self._state_decide(fsm)
                else:
                    fsm.transition(AgentState.FAILED, f"stato sconosciuto: {fsm.state}")
                    break
        except Exception as exc:
            _logger.error("[fsm:%s] unhandled error: %s", fsm.fsm_id, exc)
            if not fsm.done:
                fsm.context.error = str(exc)
                try:
                    fsm.transition(AgentState.FAILED, f"unhandled: {exc}")
                except Exception as _te:
                    _logger.debug("[fsm:%s] transition→FAILED fallito: %s", fsm.fsm_id, _te)
        return fsm.to_status()

    async def _state_plan(self, fsm: AgentFSM) -> None:
        """PLAN: chiama BrainPlanner per generare il piano."""
        if _PLANNER_AVAILABLE and _planner is not None and _PlanReq is not None:
            try:
                plan = await _planner.plan(_PlanReq(
                    goal   = fsm.context.goal,
                    hints  = fsm.context.metadata.get("hints", []),
                    max_steps = fsm.context.max_steps,
                ))
                fsm.context.plan_id = plan.plan_id
                # Carica capabilities come "tool queue" nel contesto
                fsm.context.metadata["tool_queue"] = [
                    {"capability": s.capability, "payload": s.payload, "step_id": s.step_id}
                    for s in plan.steps
                ]
                fsm.transition(AgentState.THINK, f"piano generato: {len(plan.steps)} step")
                return
            except Exception as exc:
                _logger.warning("[fsm:%s] planner error: %s", fsm.fsm_id, exc)
        # Fallback: THINK diretto senza piano
        fsm.transition(AgentState.THINK, "planner non disponibile, THINK diretto")

    async def _state_think(self, fsm: AgentFSM) -> None:
        """THINK: seleziona il prossimo tool da eseguire."""
        tool_queue = fsm.context.metadata.get("tool_queue", [])
        if not tool_queue:
            fsm.transition(AgentState.DECIDE, "tool queue vuota → DECIDE")
            return
        next_tool = tool_queue.pop(0)
        fsm.context.metadata["current_tool"] = next_tool
        fsm.context.last_tool = next_tool.get("capability")
        fsm.context.current_step += 1
        fsm.transition(AgentState.TOOL, f"eseguo tool: {next_tool.get('capability')}")

    async def _state_tool(self, fsm: AgentFSM) -> None:
        """TOOL: esegui il tool selezionato in THINK."""
        current = fsm.context.metadata.get("current_tool", {})
        cap     = current.get("capability", "llm")
        payload = current.get("payload", {})

        if _TOOL_ENGINE_AVAILABLE and _tool_exec is not None and _TExecReq is not None:
            try:
                result = await _tool_exec.execute(_TExecReq(
                    tool_name = cap, payload = payload,
                    correlation_id = fsm.fsm_id,
                ))
                fsm.context.last_result = result.model_dump()
                fsm.context.tool_results.append({"tool": cap, "result": fsm.context.last_result})
                fsm.transition(AgentState.OBSERVE, f"tool {cap} eseguito: {result.status}")
                return
            except Exception as exc:
                fsm.context.error = str(exc)
                _logger.warning("[fsm:%s] tool %s error: %s", fsm.fsm_id, cap, exc)
        # Fallback: submit via Kernel
        if _KERNEL_AVAILABLE and _kernel is not None:
            result = await _kernel.submit_task({"capability": cap, **payload})
            fsm.context.last_result = result.model_dump() if hasattr(result, "model_dump") else str(result)
            fsm.context.tool_results.append({"tool": cap, "result": fsm.context.last_result})
            fsm.transition(AgentState.OBSERVE, f"kernel submit: {cap}")
        else:
            fsm.transition(AgentState.OBSERVE, f"tool {cap} skipped (no executor)")

    async def _state_wait(self, fsm: AgentFSM) -> None:
        """WAIT: attende risultato asincrono (polling con backoff)."""
        await asyncio.sleep(1)
        fsm.transition(AgentState.OBSERVE, "wait complete")

    async def _state_observe(self, fsm: AgentFSM) -> None:
        """OBSERVE: elabora risultato, aggiorna messages."""
        result = fsm.context.last_result
        if result:
            fsm.context.messages.append({
                "role": "tool", "content": str(result)[:2000],  # truncate
                "tool": fsm.context.last_tool,
            })
        fsm.transition(AgentState.DECIDE, "osservazione completata")

    async def _state_decide(self, fsm: AgentFSM) -> None:
        """DECIDE: goal raggiunto? Continua o termina."""
        tool_queue = fsm.context.metadata.get("tool_queue", [])
        if tool_queue:
            fsm.transition(AgentState.THINK, f"coda non vuota: {len(tool_queue)} step rimanenti")
        elif fsm.context.error:
            fsm.transition(AgentState.FAILED, fsm.context.error)
        else:
            fsm.transition(AgentState.DONE, "tutti gli step completati")


# ── AgentFSMManager (registry delle FSM attive) ─────────────────────────────────

class AgentFSMManager:
    def __init__(self) -> None:
        self._instances: dict[str, AgentFSM]   = {}
        self._runner     = AgentFSMRunner()
        self._tasks:     dict[str, asyncio.Task] = {}
        self._lock       = asyncio.Lock()
        self._MAX_STORE  = 200

    async def start(self, req: FSMRunRequest) -> FSMStatus:
        fsm_id = str(uuid.uuid4())
        fsm    = AgentFSM(fsm_id, req.goal, req.max_steps, req.hints)
        if req.session_id:
            fsm.context.session_id = req.session_id
        async with self._lock:
            self._instances[fsm_id] = fsm
            if len(self._instances) > self._MAX_STORE:
                oldest = sorted(self._instances, key=lambda k: self._instances[k].started_at)
                for k in oldest[:10]:
                    self._instances.pop(k, None)

        task = asyncio.create_task(self._runner.run(fsm))
        self._tasks[fsm_id] = task
        _logger.info("[fsm-manager] started fsm_id=%s goal=%s", fsm_id, req.goal[:60])
        return fsm.to_status()

    def get(self, fsm_id: str) -> FSMStatus | None:
        fsm = self._instances.get(fsm_id)
        return fsm.to_status() if fsm else None

    def list_active(self, limit: int = 20) -> list[FSMStatus]:
        all_fsm = sorted(self._instances.values(), key=lambda f: f.started_at, reverse=True)
        return [f.to_status() for f in all_fsm[:limit]]

    def status(self) -> dict:
        instances = list(self._instances.values())
        return {
            "total": len(instances),
            "running": sum(1 for f in instances if not f.done),
            "done":    sum(1 for f in instances if f.state == AgentState.DONE),
            "failed":  sum(1 for f in instances if f.state == AgentState.FAILED),
        }


# ── Singletons ──────────────────────────────────────────────────────────────────
fsm_manager = AgentFSMManager()

# ── HTTP Router ──────────────────────────────────────────────────────────────────
router = APIRouter(
    prefix="/api/agent-fsm",
    tags=["agent-fsm"],
    dependencies=[Depends(require_role(AuthRole.MACHINE))],
)


@router.post("/run", summary="Avvia un AgentFSM run (goal → DONE|FAILED in background)")
async def route_run(req: FSMRunRequest) -> FSMStatus:
    return await fsm_manager.start(req)


@router.get("/runs/{fsm_id}", summary="Stato di un AgentFSM run")
async def route_get(fsm_id: str) -> FSMStatus:
    status = fsm_manager.get(fsm_id)
    if not status:
        raise HTTPException(404, f"FSM '{fsm_id}' non trovata")
    return status


@router.get("/runs", summary="Lista AgentFSM runs attivi")
async def route_list(limit: int = 20) -> dict:
    runs = fsm_manager.list_active(limit)
    return {"count": len(runs), "runs": [r.model_dump() for r in runs]}


@router.get("/status", summary="Stato aggregato AgentFSM manager")
async def route_status() -> dict:
    return fsm_manager.status()