from typing import Optional, List, Dict
from analogai.foundation.modules.agent.agent_repository import AgentRepository
from analogai.foundation.modules.agent.agent import Agent
from analogai.foundation.infrastructure.storage.storage_adapter import StorageAdapter
from analogai.foundation.infrastructure.mappers import AgentStorageMapper

class InMemoryAgentRepositoryImpl(AgentRepository):
    COLLECTION_NAME = "agents"

    def __init__(self, storage_adapter=None):
        if storage_adapter is None:
            from analogai.foundation.infrastructure.storage.in_memory.in_memory_storage_adapter_impl import InMemoryStorageAdapterImpl
            self.storage_adapter = InMemoryStorageAdapterImpl()
        else:
            self.storage_adapter = storage_adapter
        self._agent_cache: Dict[str, Agent] = {}

    def create(self, agent: Agent) -> Agent:
        data = AgentStorageMapper.to_dict(agent)
        stored_data = self.storage_adapter.store(self.COLLECTION_NAME, data)
        if not stored_data:
            raise ValueError(f"Failed to store agent {agent.id}")
        self._agent_cache[agent.id] = agent
        return agent

    def find_by_id(self, id: str) -> Optional[Agent]:
        if id in self._agent_cache:
            return self._agent_cache[id]
        
        data = self.storage_adapter.retrieve_by_id(self.COLLECTION_NAME, id)
        if not data:
            return None
        agent = AgentStorageMapper.from_dict(data)
        self._agent_cache[id] = agent
        return agent

    def update(self, agent: Agent) -> Agent:
        data = AgentStorageMapper.to_dict(agent)
        stored_data = self.storage_adapter.store(self.COLLECTION_NAME, data)
        if not stored_data:
            raise ValueError(f"Failed to update agent {agent.id}")
        self._agent_cache[agent.id] = agent
        return agent

    def delete(self, id: str) -> bool:
        self._agent_cache.pop(id, None)
        result = self.storage_adapter.delete(self.COLLECTION_NAME, id)
        return result

    def find_by_creator_id(self, creator_id: str) -> List[Agent]:
        agents_data = self.storage_adapter.retrieve_by_attributes(
            self.COLLECTION_NAME, 
            {"creator_id": creator_id}
        )
        agents = []
        for data in agents_data:
            agent_id = data.get("id")
            if agent_id in self._agent_cache:
                agents.append(self._agent_cache[agent_id])
            else:
                agent = AgentStorageMapper.from_dict(data)
                self._agent_cache[agent_id] = agent
                agents.append(agent)
        return agents

