from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Any, Callable

from .base import PipelineComponent
from .factory import make_adapter

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class UtilitySpec:
    """Specification for curated sklearn utility functions."""

    func: Callable[..., Any]


ADAPTER_REGISTRY: dict[str, type[PipelineComponent]] = {}
UTILITY_REGISTRY: dict[str, UtilitySpec] = {}
_INITIALIZED = False


def _register_estimators() -> None:
    try:
        from sklearn.utils import all_estimators
    except Exception as exc:
        logger.warning("Could not import sklearn estimator index: %s", exc)
        return

    for estimator_name, estimator_class in all_estimators():
        try:
            ADAPTER_REGISTRY[estimator_name] = make_adapter(estimator_class)
        except Exception as exc:
            logger.debug("Skipping sklearn estimator '%s': %s", estimator_name, exc)


def _register_utilities() -> None:
    try:
        from sklearn.model_selection import train_test_split

        UTILITY_REGISTRY["train_test_split"] = UtilitySpec(func=train_test_split)
    except Exception as exc:
        logger.warning("Could not register sklearn utility train_test_split: %s", exc)


def initialize_registry() -> None:
    global _INITIALIZED
    if _INITIALIZED:
        return
    _register_estimators()
    _register_utilities()
    _INITIALIZED = True


def get_adapter(component_name: str) -> type[PipelineComponent] | None:
    initialize_registry()
    return ADAPTER_REGISTRY.get(component_name)


def get_utility(component_name: str) -> UtilitySpec | None:
    initialize_registry()
    return UTILITY_REGISTRY.get(component_name)


def list_components() -> list[str]:
    initialize_registry()
    names = set(ADAPTER_REGISTRY.keys()) | set(UTILITY_REGISTRY.keys())
    return sorted(names)

