from typing import Dict, Type, Any

from atomingest.normalization.sources.base import DataSource
from atomingest.normalization.sources.csv import CSVDataSource
from atomingest.normalization.sources.json import JSONDataSource
from atomingest.normalization.sources.excel import ExcelDataSource
from atomingest.normalization.sources.parquet import ParquetDataSource


class DataSourceFactory:
    _REGISTRY: Dict[str, Type[DataSource]] = {
        "csv": CSVDataSource,
        "json": JSONDataSource,
        "excel": ExcelDataSource,
        "parquet": ParquetDataSource,
    }

    @classmethod
    def create(cls, source_type: str, config: Dict[str, Any]) -> DataSource:
        source_class = cls._REGISTRY.get(source_type.lower())
        if not source_class:
            valid_keys = ", ".join(sorted(cls._REGISTRY.keys()))
            raise ValueError(
                f"Unsupported source type '{source_type}' is not registered. Valid types are: {valid_keys}"
            )
        return source_class(**config)

    @classmethod
    def register(cls, name: str, source_cls: Type[DataSource]):
        """Allow users to register custom data sources"""
        cls._REGISTRY[name.lower()] = source_cls
