"""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


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:
    """Route dataset files to the correct engine implementation."""
    config = dataset_files.config

    if config.engine == "sqlite":
        return SQLiteDataset(
            dataset_id=dataset_id,
            config=config,
            data_text=dataset_files.data_text,
            schema_text=dataset_files.schema_text,
            source_uri=dataset_files.source_uri,
        )
    if config.engine == "json":
        return JSONDataset(
            dataset_id=dataset_id,
            config=config,
            data_text=dataset_files.data_text,
            source_uri=dataset_files.source_uri,
        )

    raise NotImplementedError(
        f"Engine '{config.engine}' not yet supported.\n"
        f"Currently supported engines: sqlite, json\n"
        f"Future support planned for: duckdb, lance, qdrant"
    )


__all__ = ["load_dataset"]
