from __future__ import annotations

from typing import Any

from pydantic import BaseModel


class FlowState(BaseModel):
    stage_trace: list[str] = []
    current_stage: str | None = None
    message: str = ""
    context: dict[str, Any] = {}


class AgentFlow:
    def __init__(self) -> None:
        self.state = FlowState()

    async def run(self) -> dict[str, Any]:
        result = await self.intent_recognition()
        result = await self.task_execution(result)
        result = await self.result_reporting(result)
        return result

    async def intent_recognition(self) -> dict[str, Any]:
        self.state.current_stage = "intent"
        self.state.stage_trace.append("intent")
        return {"stage": "intent", "result": "intent recognized"}

    async def task_execution(self, prev_result: dict[str, Any]) -> dict[str, Any]:
        self.state.current_stage = "execution"
        self.state.stage_trace.append("execution")
        return {"stage": "execution", "result": "task executed"}

    async def result_reporting(self, prev_result: dict[str, Any]) -> dict[str, Any]:
        self.state.current_stage = "reporting"
        self.state.stage_trace.append("reporting")
        return {"stage": "reporting", "result": "reported"}
