from abc import ABC, abstractmethod
from typing import List, Optional
from .types import AgentInfo


class AgentRegistry(ABC):
    @abstractmethod
    async def add(self, agent_info: AgentInfo):
        pass

    @abstractmethod
    async def get(self, agent_id: str) -> Optional[AgentInfo]:
        pass

    @abstractmethod
    async def remove(self, agent_id: str):
        pass

    @abstractmethod
    async def list(self) -> List[AgentInfo]:
        pass


class InMemoryRegistry(AgentRegistry):
    def __init__(self):
        self._agents: dict[str, AgentInfo] = {}

    async def add(self, agent_info: AgentInfo):
        self._agents[agent_info.id] = agent_info

    async def get(self, agent_id: str) -> Optional[AgentInfo]:
        return self._agents.get(agent_id)

    async def remove(self, agent_id: str):
        if agent_id in self._agents:
            del self._agents[agent_id]

    async def list(self) -> List[AgentInfo]:
        return list(self._agents.values())
