from typing import Optional, List
from analogai.foundation.extensions.fallback.repositories.fallback_repository import FallbackRepository
from analogai.foundation.extensions.fallback import Fallback
from analogai.foundation.infrastructure.storage.storage_adapter import StorageAdapter
from analogai.foundation.infrastructure.storage.mongodb.mongodb_storage_adapter_impl import MongoDBStorageAdapterImpl
from analogai.foundation.extensions.fallback.mappers.fallback_storage_mapper import FallbackStorageMapper


class MongoDBFallbackRepositoryImpl(FallbackRepository):
    COLLECTION_NAME = "fallbacks"

    def __init__(self, storage_adapter: Optional[StorageAdapter] = None):
        self.storage_adapter = storage_adapter if storage_adapter is not None else MongoDBStorageAdapterImpl()

    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}")
        return fallback

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

    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}
        )
        return [FallbackStorageMapper.from_dict(data) for data in fallbacks_data]
