"""Dataset loading and storage module for Agentbyte.

This module provides schema-driven dataset loading from JSON into engine-specific
backends. Currently supports:
    - SQLite: Structured data with strict schema and SQL queries
    - JSON: Semi-structured data with flexible JSON format

Architecture:
    Source Configs (config.py):
        - DatasetSourceConfig: Abstract base for all source configs
        - LocalDatasetConfig: Local filesystem datasets with explicit root path
        - S3DatasetConfig: S3-hosted datasets with bucket/region/prefix

    Generic Classes (base.py):
        - Column: Column definition (engine-agnostic)
        - Schema: Schema definition (engine-agnostic)
        - DataRecord: Generic record type
        - DataStore: Abstract storage interface
        - BaseDataset: Abstract interface for all dataset types
        - DatasetConfig: Configuration (engine, table name, etc.)
        - SchemaFormat: Display format enum

    SQLite-Specific Classes (sqlite.py):
        - SqliteColumn: Column with SQL constraints
        - SqliteSchema: Schema with SQL generation and formatted display
        - SQLiteDataset: SQLite implementation with JSON loading

    JSON-Specific Classes (json.py):
        - JSONDataset: JSON file-based dataset storage

    Loader:
        - load_dataset(): Factory function. Pass a config object to select
          source and parameters.

Public API:
    Source configs:
        DatasetSourceConfig: Abstract base (extend for custom sources)
        LocalDatasetConfig: Load datasets from a local root directory
        S3DatasetConfig: Load datasets from an S3 bucket

    From base.py:
        Column, Schema, DataRecord, DataStore
        DatasetConfig, SchemaFormat, BaseDataset

    From sqlite.py:
        SqliteColumn, SqliteSchema, SQLiteDataset

    From json.py:
        JSONDataset

    Loader:
        load_dataset(): Routes to correct implementation based on config.engine

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

    >>> # Local dataset
    >>> ds = load_dataset("contracts/asmd", LocalDatasetConfig(local_root=Path("/data")))
    >>> conn = ds.connect()
    >>> rows = conn.execute("SELECT * FROM contracts_asmd").fetchall()

    >>> # S3 dataset
    >>> ds = load_dataset("contracts/asmd", S3DatasetConfig(s3_bucket="datapsycho-datasets"))
    >>> conn = ds.connect()
"""

from agentbyte.dataset.base import (
    BaseDataset,
    Column,
    DataRecord,
    DataStore,
    DatasetConfig,
    Schema,
    SchemaFormat,
)
from agentbyte.dataset.config import DatasetSourceConfig, LocalDatasetConfig, S3DatasetConfig
from agentbyte.dataset.json import JSONDataset
from agentbyte.dataset.loader import load_dataset
from agentbyte.dataset.sqlite import (
    SQLiteDataset,
    SqliteColumn,
    SqliteSchema,
)

__all__ = [
    # Source configuration
    "DatasetSourceConfig",
    "LocalDatasetConfig",
    "S3DatasetConfig",
    # Schema definitions
    "Column",
    "SqliteColumn",
    "Schema",
    "SqliteSchema",
    "SchemaFormat",
    # Records and storage
    "DataRecord",
    "DataStore",
    # Configuration
    "DatasetConfig",
    # Base interface
    "BaseDataset",
    # Implementations
    "SQLiteDataset",
    "JSONDataset",
    # Loader
    "load_dataset",
]

