from typing import Optional, List, Dict
from analogai.foundation.extensions.fallback.repositories.fallback_repository import FallbackRepository
from analogai.foundation.extensions.fallback import Fallback
from analogai.foundation.infrastructure.storage.in_memory import InMemoryStorageAdapterImpl
from analogai.foundation.extensions.fallback.mappers.fallback_storage_mapper import FallbackStorageMapper

class InMemoryFallbackRepositoryImpl(FallbackRepository):
    COLLECTION_NAME = "fallbacks"

    def __init__(self):
        self.storage_adapter = InMemoryStorageAdapterImpl()
        self._fallback_cache: Dict[str, Fallback] = {}

    def create(self, fallback: Fallback) -> Fallback:
        data = FallbackStorageMapper.to_dict(fallback)
        stored_data = self.storage_adapter.store(self.COLLECTION_NAME, data)
        if not stored_data:
            raise ValueError(f"Failed to store fallback {fallback.id}")
        self._fallback_cache[fallback.id] = fallback
        return fallback

    def find_by_id(self, fallback_id: str) -> Optional[Fallback]:
        if fallback_id in self._fallback_cache:
            return self._fallback_cache[fallback_id]
        data = self.storage_adapter.retrieve_by_id(self.COLLECTION_NAME, fallback_id)
        if not data:
            return None
        fallback = FallbackStorageMapper.from_dict(data)
        self._fallback_cache[fallback_id] = fallback
        return fallback

    def find_by_agent_id(self, agent_id: str) -> List[Fallback]:
        fallbacks_data = self.storage_adapter.retrieve_by_attributes(
            self.COLLECTION_NAME,
            {"agent_id": agent_id}
        )
        fallbacks = []
        for data in fallbacks_data:
            fallback_id = data.get("id")
            if fallback_id in self._fallback_cache:
                fallbacks.append(self._fallback_cache[fallback_id])
            else:
                fallback = FallbackStorageMapper.from_dict(data)
                self._fallback_cache[fallback_id] = fallback
                fallbacks.append(fallback)
        return fallbacks
