"""Dataset loading and registry."""

from agentbyte.dataset.base import BaseDataset
from agentbyte.dataset.config import DatasetSourceConfig
from agentbyte.dataset.json import JSONDataset
from agentbyte.dataset.sources import DatasetFiles
from agentbyte.dataset.sqlite import SQLiteDataset

# Engine registry: maps engine name → BaseDataset subclass.
# Third-party engines can register themselves via register_engine().
_ENGINE_REGISTRY: dict[str, type[BaseDataset]] = {}


def register_engine(name: str, cls: type[BaseDataset]) -> None:
    """Register a dataset engine class under a given name.

    Built-in engines (``"sqlite"``, ``"json"``) are pre-registered at import
    time. Call this once — typically at the top of a module — to make a
    custom engine available to ``load_dataset()``.

    Args:
        name: Engine identifier used in ``config.json`` (e.g. ``"duckdb"``).
        cls: A ``BaseDataset`` subclass that implements ``fetch_all()``.

    Example::

        from agentbyte.dataset import register_engine
        from my_package.duckdb_dataset import DuckDBDataset

        register_engine("duckdb", DuckDBDataset)
    """
    _ENGINE_REGISTRY[name] = cls


# Pre-register built-in engines
register_engine("sqlite", SQLiteDataset)
register_engine("json", JSONDataset)


def load_dataset(dataset_id: str, config: DatasetSourceConfig) -> BaseDataset:
    """Load a dataset by ID.

    Dataset IDs follow the pattern: ``"category/name"``.

    The ``config`` object determines the storage backend and holds all
    required parameters. Each config type builds its own source via
    ``config.build_source()``, so ``load_dataset()`` never needs to change
    when new source types are added (OCP).

    Args:
        dataset_id: Dataset identifier, e.g. ``"contracts/asmd"``.
        config: Source configuration — either ``LocalDatasetConfig`` or
            ``S3DatasetConfig`` (or any future ``DatasetSourceConfig``
            subclass).

    Returns:
        Dataset object with ``engine``, ``table_name``, and ``connect()``.

    Raises:
        FileNotFoundError: If dataset or config files not found.
        ValueError: If dataset config is invalid.
        NotImplementedError: If the dataset engine is not yet supported.

    Example:
        >>> from pathlib import Path
        >>> from agentbyte.dataset import load_dataset
        >>> from agentbyte.dataset.config import LocalDatasetConfig, S3DatasetConfig

        >>> # Local — explicit path
        >>> ds = load_dataset("contracts/asmd", LocalDatasetConfig(local_root=Path("/data")))

        >>> # Local — from environment variable AGENTBYTE_DATASET_LOCAL_ROOT
        >>> ds = load_dataset("contracts/asmd", LocalDatasetConfig())

        >>> # S3 — explicit
        >>> ds = load_dataset("contracts/asmd", S3DatasetConfig(s3_bucket="datapsycho-datasets"))

        >>> # S3 — from environment variable AGENTBYTE_DATASET_S3_BUCKET
        >>> ds = load_dataset("contracts/asmd", S3DatasetConfig())
    """
    source = config.build_source()
    dataset_files = source.read_dataset_files(dataset_id)
    return _build_dataset(dataset_id, dataset_files)


def _build_dataset(dataset_id: str, dataset_files: DatasetFiles) -> BaseDataset:
    """Instantiate the correct dataset class via the engine registry."""
    config = dataset_files.config
    cls = _ENGINE_REGISTRY.get(config.engine)
    if cls is None:
        registered = ", ".join(sorted(_ENGINE_REGISTRY))
        raise NotImplementedError(
            f"Engine '{config.engine}' is not registered. "
            f"Registered engines: {registered}. "
            f"Use register_engine() to add a new engine."
        )
    return cls(
        dataset_id=dataset_id,
        config=config,
        data_text=dataset_files.data_text,
        schema_text=dataset_files.schema_text,
        source_uri=dataset_files.source_uri,
    )


__all__ = ["load_dataset", "register_engine"]
