from typing import Dict, Optional, List
from emissary_ml_toolkit.storage.base import BaseStorage
from emissary_ml_toolkit.storage.s3 import S3Storage
from emissary_ml_toolkit.storage.local import LocalStorage
from emissary_ml_toolkit.config.storage import StorageConfig, StorageType, StorageConfigs
from emissary_ml_toolkit.utils.exceptions import StorageError


class StorageFactory:
    _storage_classes = {
        StorageType.S3: S3Storage,
        StorageType.LOCAL: LocalStorage,
    }

    @classmethod
    def create_storage(cls, config: StorageConfig) -> BaseStorage:
        storage_class = cls._storage_classes.get(config.type)

        if not storage_class:
            raise StorageError(f"Unsupported storage type: {config.type}")

        return storage_class(config)

    @classmethod
    def create_storage_manager(cls, storage_configs: StorageConfigs) -> "StorageManager":
        return StorageManager(storage_configs)


class StorageManager:
    def __init__(self, storage_configs: StorageConfigs):
        self.configs = storage_configs
        self._storages: Dict[str, BaseStorage] = {}
        self._initialize_storages()

    def _initialize_storages(self):
        for name, config in self.configs.storages.items():
            try:
                self._storages[name] = StorageFactory.create_storage(config)
            except Exception as e:
                raise StorageError(f"Failed to initialize storage {name}: {e}")

    def get_storage(self, name: Optional[str] = None) -> BaseStorage:
        storage_name = name or self.configs.default

        if storage_name not in self._storages:
            raise StorageError(f"Storage {storage_name} not found")

        return self._storages[storage_name]

    def get_default_storage(self) -> BaseStorage:
        return self.get_storage(self.configs.default)

    def list_storages(self) -> List[str]:
        return list(self._storages.keys())