"""Factory function for DataSource."""

from inspect import isabstract
from typing import Any, Dict, Union

from hydromt.data_catalog.sources import (
    DataFrameSource,
    DatasetSource,
    DataSource,
    GeoDataFrameSource,
    GeoDatasetSource,
    RasterDatasetSource,
)
from hydromt.typing.type_def import DataType

# Map DataType to DataSource, need to add here when implementing a new Type
available_sources: Dict[DataType, DataSource] = {
    "DataFrame": DataFrameSource,
    "Dataset": DatasetSource,
    "RasterDataset": RasterDatasetSource,
    "GeoDataFrame": GeoDataFrameSource,
    "GeoDataset": GeoDatasetSource,
}


def create_source(data: Union[Dict[str, Any], DataSource]) -> DataSource:
    """Create a DataSource.

    Create a datasource from a dictionary, or another DataSource.
    """
    if isinstance(data, DataSource):
        if isabstract(DataSource):
            raise ValueError("DataSource is an Abstract Class")
        else:
            # Already is a subclass of DataSource
            return data

    elif isinstance(data, dict):
        if data_type := data.get("data_type", None):
            if target_source := available_sources.get(data_type):
                return target_source.model_validate(data)

            raise ValueError(f"Unknown 'data_type': '{data_type}'")
        else:
            raise ValueError("DataSource needs 'data_type'.")
    else:
        raise ValueError(f"Invalid argument for creating DataSource: {data}")
