State Layer (Level-1 FSM)

class AgentStatus(Enum):
    """Enumeration of agent states"""
    ERROR = "ERROR"            # Critical error occurred
    FINISH = "FINISH"          # Task completed successfully
    CONTINUE = "CONTINUE"      # Normal execution, continue processing
    FAIL = "FAIL"              # Task failed, cannot proceed
    PENDING = "PENDING"        # Waiting for external event (user input, async operation)
    CONFIRM = "CONFIRM"        # Awaiting user confirmation before proceeding
    SCREENSHOT = "SCREENSHOT"  # Capture observation data (screenshot, UI tree)

class AgentStateManager(ABC, metaclass=SingletonABCMeta):
    """
    Singleton state manager for agent states.

    Responsibilities:
    - Register state classes via decorator
    - Lazy load state instances on demand
    - Provide centralized state access
    """

    _state_mapping: Dict[str, Type[AgentState]] = {}  # Class registry

    def __init__(self):
        self._state_instance_mapping: Dict[str, AgentState] = {}  # Instance cache

    def get_state(self, status: str) -> AgentState:
        """
        Get state instance for the given status string.

        :param status: The status string (e.g., "CONTINUE")
        :return: The state instance

        Note: Uses lazy loading - instances created on first access
        """
        # Lazy load: create instance only when first requested
        if status not in self._state_instance_mapping:
            state_class = self._state_mapping.get(status)
            if state_class:
                self._state_instance_mapping[status] = state_class()
            else:
                # Fallback to none_state if status not registered
                self._state_instance_mapping[status] = self.none_state

        return self._state_instance_mapping.get(status, self.none_state)

    def add_state(self, status: str, state: AgentState) -> None:
        """
        Add a state instance at runtime (advanced usage).

        :param status: The status string
        :param state: The state instance
        """
        self._state_instance_mapping[status] = state

    @property
    def state_map(self) -> Dict[str, AgentState]:
        """
        The state mapping of status to state.
        :return: The state mapping.
        """
        return self._state_instance_mapping

    @classmethod
    def register(cls, state_class: Type[AgentState]) -> Type[AgentState]:
        """
        Decorator to register state class.

        Usage:
            @AgentStateManager.register
            class ContinueAppAgentState(AgentState):
                @staticmethod
                def name() -> str:
                    return AgentStatus.CONTINUE.value

        :param state_class: The state class to register
        :return: The state class (unchanged)
        """
        cls._state_mapping[state_class.name()] = state_class
        return state_class

    @property
    @abstractmethod
    def none_state(self) -> AgentState:
        """
        Fallback state when requested state not found.

        :return: Default/fallback state instance
        """
        pass

# AppAgent states
class AppAgentStateManager(AgentStateManager):
    @property
    def none_state(self):
        return NoneAppAgentState()

@AppAgentStateManager.register
class ContinueAppAgentState(AgentState):
    @classmethod
    def name(cls):
        return AgentStatus.CONTINUE.value

class AgentState(ABC):
    """
    Abstract base class for agent states.
    """

    @abstractmethod
    async def handle(
        self, agent: BasicAgent, context: Optional[Context] = None
    ) -> None:
        """
        Handle the agent for the current step.
        :param agent: The agent to handle.
        :param context: The context for the agent and session.
        """
        pass

    @abstractmethod
    def next_agent(self, agent: BasicAgent) -> BasicAgent:
        """
        Get the agent for the next step.
        :param agent: The agent for the current step.
        :return: The agent for the next step.
        """
        return agent

    @abstractmethod
    def next_state(self, agent: BasicAgent) -> AgentState:
        """
        Get the state for the next step.
        :param agent: The agent for the current step.
        :return: The state for the next step.
        """
        pass

    @abstractmethod
    def is_round_end(self) -> bool:
        """
        Check if the round ends.
        :return: True if the round ends, False otherwise.
        """
        pass

    @abstractmethod
    def is_subtask_end(self) -> bool:
        """
        Check if the subtask ends.
        :return: True if the subtask ends, False otherwise.
        """
        pass

    @classmethod
    @abstractmethod
    def agent_class(cls) -> Type[BasicAgent]:
        """
        The class of the agent.
        :return: The class of the agent.
        """
        pass

    @classmethod
    @abstractmethod
    def name(cls) -> str:
        """
        The class name of the state.
        :return: The class name of the state.
        """
        return ""

@AppAgentStateManager.register
class ContinueAppAgentState(AgentState):
    """
    Continue state for AppAgent - normal execution flow.
    """

    async def handle(self, agent: AppAgent, context: Context):
        """Execute AppAgent processor strategies."""
        # Get processor (Level-2 Strategy Layer)
        processor = agent.processor

        # Execute all strategies in sequence
        await processor.process(agent, context)

        # Processor updates agent.status based on LLM response
        # Possible status: CONTINUE, FINISH, FAIL, ERROR, CONFIRM, etc.

    def next_state(self, agent: AppAgent) -> AgentState:
        """Transition to next state based on agent status."""
        state_manager = AppAgentStateManager()
        return state_manager.get_state(agent.status)

    def next_agent(self, agent: AppAgent) -> BasicAgent:
        """For AppAgent, typically stays on same agent."""
        # AppAgent continues executing unless delegating back to HostAgent
        if agent.status == AgentStatus.FINISH:
            return agent.host  # Return to HostAgent
        return agent  # Continue with current agent

    @classmethod
    def name(cls) -> str:
        """State name for registration"""
        return AgentStatus.CONTINUE.value  # "CONTINUE"

async def handle(self, agent, context):
    """Execute full processor workflow"""
    # Run all four strategy phases
    await agent.processor.process(agent, context)
    # Status updated by LLM response parsing

async def handle(self, agent, context):
    """Wait for external event"""
    # Do not execute processor
    # Wait for callback, user input, or timeout
    event = await wait_for_event(timeout=60)
    if event:
        agent.status = AgentStatus.CONTINUE
    else:
        agent.status = AgentStatus.FAIL

async def handle(self, agent, context):
    """Request user confirmation"""
    # Execute DATA_COLLECTION to show current state
    await agent.processor.execute_phase(ProcessingPhase.DATA_COLLECTION)

    # Prompt user for confirmation
    confirmed = await prompt_user_confirmation()

    if confirmed:
        agent.status = AgentStatus.CONTINUE
    else:
        agent.status = AgentStatus.FAIL

async def handle(self, agent, context):
    """No action - state is terminal"""
    # Terminal states do not execute processor
    # Round manager will detect terminal status and end execution
    pass

class ContinueHostAgentState(AgentState):
    def next_agent(self, agent: HostAgent) -> BasicAgent:
        """Delegate to AppAgent when task decomposed"""
        if agent.status == "DELEGATE_TO_APP":
            # Create AppAgent for selected application
            app_agent = AgentFactory.create_agent(
                agent_type="app",
                name=f"AppAgent/{agent.selected_app}",
                process_name=agent.selected_process,
                app_root_name=agent.selected_app,
                is_visual=True,
                main_prompt=config.appagent_prompt,
                example_prompt=config.appagent_example_prompt
            )

            # Set HostAgent as host (for returning)
            app_agent.host = agent

            # Transfer context via blackboard
            app_agent.blackboard = agent.blackboard

            return app_agent

        # No delegation, continue with HostAgent
        return agent

class FinishAppAgentState(AgentState):
    def next_agent(self, agent: AppAgent) -> BasicAgent:
        """Return to HostAgent when app task complete"""
        if agent.host:
            # Update HostAgent's blackboard with results
            agent.host.blackboard = agent.blackboard

            # Set HostAgent status to continue
            agent.host.status = AgentStatus.CONTINUE

            return agent.host

        # No host, AppAgent finishes independently
        return agent

@AppAgentStateManager.register
class ContinueAppAgentState(AgentState):
    """Continue state for Windows AppAgent"""
    # Implements UI Automation-specific logic

@AppAgentStateManager.register
class ScreenshotAppAgentState(AgentState):
    """Screenshot state for Windows AppAgent"""
    # Captures Windows UI tree + screenshot

@LinuxAgentStateManager.register
class ContinueLinuxAgentState(AgentState):
    """Continue state for Linux Agent"""
    # Implements shell command execution logic

@LinuxAgentStateManager.register
class FinishLinuxAgentState(AgentState):
    """Finish state for Linux Agent"""
    # Terminal state for Linux workflow