"""活动 Agent Run 注册表：为 steer、follow-up、abort 提供进程内控制面。"""

from __future__ import annotations

import asyncio

from app.application.agent.agent import Agent


class AgentRunRegistry:
    """跟踪当前进程中的活动 Agent；持久状态仍以 Run/EventStore 为准。"""

    def __init__(self) -> None:
        self._agents: dict[str, Agent] = {}
        self._lock = asyncio.Lock()

    async def register(self, run_id: str, agent: Agent) -> None:
        async with self._lock:
            self._agents[run_id] = agent

    async def unregister(self, run_id: str, agent: Agent) -> None:
        async with self._lock:
            if self._agents.get(run_id) is agent:
                self._agents.pop(run_id, None)

    def get(self, run_id: str) -> Agent | None:
        return self._agents.get(run_id)

    def steer(self, run_id: str, content: str) -> bool:
        agent = self.get(run_id)
        if agent is None:
            return False
        return agent.steer(content)

    def follow_up(self, run_id: str, content: str) -> bool:
        agent = self.get(run_id)
        if agent is None:
            return False
        return agent.follow_up(content)

    async def abort(self, run_id: str) -> bool:
        agent = self.get(run_id)
        if agent is None or not await agent.abort():
            return False
        await agent.wait_for_idle()
        return True


__all__ = ["AgentRunRegistry"]