from __future__ import annotations

from collections.abc import Awaitable, Callable
from typing import Any

from pydantic import BaseModel


DataSourceFactory = Callable[..., Any | Awaitable[Any]]


_CONFIG_MODELS: dict[str, type[BaseModel]] = {}
_FACTORIES: dict[str, DataSourceFactory] = {}


def _normalize_type_name(type_name: str) -> str:
    normalized = str(type_name).strip()
    if not normalized:
        raise ValueError("datasource type name must not be empty")
    return normalized


def register_datasource_config_model(
    type_name: str,
    config_model: type[BaseModel],
    *,
    replace: bool = False,
) -> None:
    normalized = _normalize_type_name(type_name)
    existing = _CONFIG_MODELS.get(normalized)
    if existing is not None and existing is not config_model and not replace:
        raise ValueError(f"datasource config model already registered for type: {normalized}")
    _CONFIG_MODELS[normalized] = config_model


def register_datasource_factory(
    type_name: str,
    factory: DataSourceFactory,
    *,
    replace: bool = False,
) -> None:
    normalized = _normalize_type_name(type_name)
    existing = _FACTORIES.get(normalized)
    if existing is not None and existing is not factory and not replace:
        raise ValueError(f"datasource factory already registered for type: {normalized}")
    _FACTORIES[normalized] = factory


def register_datasource_type(
    type_name: str,
    *,
    config_model: type[BaseModel] | None = None,
    factory: DataSourceFactory | None = None,
    replace: bool = False,
) -> None:
    if config_model is None and factory is None:
        raise ValueError("register_datasource_type() requires config_model, factory, or both")

    if config_model is not None:
        register_datasource_config_model(type_name, config_model, replace=replace)
    if factory is not None:
        register_datasource_factory(type_name, factory, replace=replace)


def get_datasource_config_model(type_name: str) -> type[BaseModel] | None:
    return _CONFIG_MODELS.get(_normalize_type_name(type_name))


def get_datasource_factory(type_name: str) -> DataSourceFactory | None:
    return _FACTORIES.get(_normalize_type_name(type_name))


def list_registered_datasource_types() -> list[str]:
    return sorted(set(_CONFIG_MODELS) | set(_FACTORIES))


def validate_registered_datasource_config(value: Any) -> BaseModel:
    if isinstance(value, BaseModel):
        payload = value.model_dump(mode="python")
    elif isinstance(value, dict):
        payload = dict(value)
    else:
        raise TypeError(f"datasource config must be a mapping or BaseModel, got {type(value).__name__}")

    type_name = payload.get("type")
    if not isinstance(type_name, str) or not type_name.strip():
        raise ValueError("datasource config must define a non-empty string field 'type'")

    config_model = get_datasource_config_model(type_name)
    if config_model is None:
        available = ", ".join(list_registered_datasource_types()) or "<none>"
        raise ValueError(f"Unsupported datasource type: '{type_name}'. Registered types: {available}")

    return config_model.model_validate(payload)
