"""Dataset source configuration contracts.

Each config class holds all parameters needed for a specific source type and
builds its own source object. This follows OCP: adding a new source type
only requires a new config class — ``load_dataset()`` never changes.

Usage:
    from pathlib import Path
    from agentbyte.dataset.config import LocalDatasetConfig, S3DatasetConfig

    # Local — explicit path
    config = LocalDatasetConfig(local_root=Path("/data/datasets"))

    # Local — from environment variable AGENTBYTE_DATASET_LOCAL_ROOT
    config = LocalDatasetConfig()

    # S3 — explicit
    config = S3DatasetConfig(s3_bucket="datapsycho-datasets")

    # S3 — from environment variable AGENTBYTE_DATASET_S3_BUCKET
    config = S3DatasetConfig()
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from pathlib import Path
from typing import TYPE_CHECKING

from pydantic import model_validator
from pydantic_settings import BaseSettings, SettingsConfigDict

if TYPE_CHECKING:
    from agentbyte.dataset.sources import DatasetSource


class DatasetSourceConfig(ABC):
    """Abstract base for all dataset source configurations.

    Subclasses hold the parameters needed for a specific storage backend
    (local filesystem, S3, etc.) and are responsible for constructing
    the matching ``DatasetSource`` via ``build_source()``.

    All concrete subclasses extend both this class and ``BaseSettings``
    so they can be populated from environment variables.
    """

    @abstractmethod
    def build_source(self) -> DatasetSource:
        """Build and return the dataset source for this configuration."""


class LocalDatasetConfig(DatasetSourceConfig, BaseSettings):
    """Configuration for loading datasets from a local directory.

    Reads from environment variables with prefix ``AGENTBYTE_DATASET_``:
    - ``AGENTBYTE_DATASET_LOCAL_ROOT`` → ``local_root``

    Attributes:
        local_root: Root directory containing datasets organized as
            ``{category}/{name}/`` subfolders. Must exist on disk.

    Example:
        >>> config = LocalDatasetConfig(local_root=Path("/data/datasets"))
        >>> # Or from environment: AGENTBYTE_DATASET_LOCAL_ROOT=/data/datasets
        >>> config = LocalDatasetConfig()
    """

    model_config = SettingsConfigDict(env_prefix="AGENTBYTE_DATASET_", extra="ignore")

    local_root: Path

    @model_validator(mode="after")
    def _validate_local_root(self) -> "LocalDatasetConfig":
        if not self.local_root.is_dir():
            raise ValueError(f"local_root directory does not exist: {self.local_root}")
        return self

    def build_source(self) -> DatasetSource:
        from agentbyte.dataset.sources import LocalDatasetSource
        return LocalDatasetSource(self)


class S3DatasetConfig(DatasetSourceConfig, BaseSettings):
    """Configuration for loading datasets from an S3 bucket.

    Reads from environment variables with prefix ``AGENTBYTE_DATASET_``:
    - ``AGENTBYTE_DATASET_S3_BUCKET`` → ``s3_bucket``
    - ``AGENTBYTE_DATASET_S3_REGION`` → ``s3_region`` (default: "eu-west-1")
    - ``AGENTBYTE_DATASET_S3_PREFIX`` → ``s3_prefix`` (default: "")

    Attributes:
        s3_bucket: S3 bucket name (e.g., "datapsycho-datasets").
        s3_region: AWS region (default: "eu-west-1").
        s3_prefix: Optional shared root prefix inside the bucket.
            Example: ``"datasets"`` resolves ``contracts/asmd``
            to ``s3://bucket/datasets/contracts/asmd/``.

    Example:
        >>> config = S3DatasetConfig(s3_bucket="datapsycho-datasets")
        >>> # Or from environment: AGENTBYTE_DATASET_S3_BUCKET=datapsycho-datasets
        >>> config = S3DatasetConfig()
    """

    model_config = SettingsConfigDict(env_prefix="AGENTBYTE_DATASET_", extra="ignore")

    s3_bucket: str
    s3_region: str = "eu-west-1"
    s3_prefix: str = ""

    def build_source(self) -> DatasetSource:
        from agentbyte.dataset.sources import S3DatasetSource
        return S3DatasetSource(self)


__all__ = ["DatasetSourceConfig", "LocalDatasetConfig", "S3DatasetConfig"]
