"""Dataset sources for loading dataset assets from local files or S3."""

from __future__ import annotations

import json
from abc import ABC, abstractmethod
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

from agentbyte.dataset.base import DatasetConfig

if TYPE_CHECKING:
    from agentbyte.dataset.config import LocalDatasetConfig, S3DatasetConfig


_REQUIRED_DATASET_SEGMENTS = 2


class DatasetSource(ABC):
    """Abstract base for all dataset sources.

    A source is responsible for reading the three dataset asset files
    (``config.json``, ``data.json``, ``schema.json``) from a storage
    backend and returning a ``DatasetFiles`` container.
    """

    @abstractmethod
    def read_dataset_files(self, dataset_id: str) -> DatasetFiles:
        """Read dataset assets for the given dataset ID."""


@dataclass(slots=True)
class DatasetFiles:
    """Resolved dataset assets ready for dataset construction."""

    dataset_id: str
    config: DatasetConfig
    data_text: str
    schema_text: str | None
    source_uri: str


def validate_dataset_id(dataset_id: str) -> None:
    """Validate the canonical ``category/name`` dataset identifier format."""
    parts = [part for part in dataset_id.split("/") if part]
    if len(parts) != _REQUIRED_DATASET_SEGMENTS:
        raise ValueError(
            f"Dataset ID must be 'category/name', got: {dataset_id}"
        )


class LocalDatasetSource(DatasetSource):
    """Load dataset assets from a local directory.

    Args:
        config: ``LocalDatasetConfig`` containing the validated root path.

    Example:
        >>> from agentbyte.dataset.config import LocalDatasetConfig
        >>> config = LocalDatasetConfig(local_root=Path("/data/datasets"))
        >>> source = LocalDatasetSource(config)
        >>> files = source.read_dataset_files("contracts/asmd")
    """

    def __init__(self, config: LocalDatasetConfig) -> None:
        self.root = config.local_root

    def read_dataset_files(self, dataset_id: str) -> DatasetFiles:
        """Read dataset assets from disk."""
        validate_dataset_id(dataset_id)

        dataset_dir = self.root / dataset_id
        if not dataset_dir.is_dir():
            raise FileNotFoundError(
                f"Dataset not found at: {dataset_dir}\n"
                f"Available datasets at: {self.root}"
            )

        config_path = dataset_dir / "config.json"
        config_text = self._read_required_text(config_path, "Config file not found")
        config = DatasetConfig(**json.loads(config_text))

        data_path = dataset_dir / "data.json"
        data_text = self._read_required_text(data_path, "Data file not found")

        schema_path = dataset_dir / "schema.json"
        schema_text = schema_path.read_text(encoding="utf-8") if schema_path.exists() else None

        return DatasetFiles(
            dataset_id=dataset_id,
            config=config,
            data_text=data_text,
            schema_text=schema_text,
            source_uri=str(dataset_dir),
        )

    @staticmethod
    def _read_required_text(path: Path, not_found_message: str) -> str:
        if not path.exists():
            raise FileNotFoundError(f"{not_found_message}: {path}")
        return path.read_text(encoding="utf-8")


class S3DatasetSource(DatasetSource):
    """Load dataset assets from an S3 bucket.

    Args:
        config: ``S3DatasetConfig`` containing bucket, region and prefix.

    Example:
        >>> from agentbyte.dataset.config import S3DatasetConfig
        >>> config = S3DatasetConfig(s3_bucket="datapsycho-datasets")
        >>> source = S3DatasetSource(config)
        >>> files = source.read_dataset_files("contracts/asmd")
    """

    def __init__(self, config: S3DatasetConfig, client: object | None = None) -> None:
        self.config = config
        self._client = client

    def read_dataset_files(self, dataset_id: str) -> DatasetFiles:
        """Read dataset assets from S3."""
        validate_dataset_id(dataset_id)

        config_key = self._build_key(dataset_id, "config.json")
        config_text = self._read_required_text(config_key)
        config = DatasetConfig(**json.loads(config_text))

        data_key = self._build_key(dataset_id, "data.json")
        data_text = self._read_required_text(data_key)

        schema_key = self._build_key(dataset_id, "schema.json")
        schema_text = self._read_optional_text(schema_key)

        return DatasetFiles(
            dataset_id=dataset_id,
            config=config,
            data_text=data_text,
            schema_text=schema_text,
            source_uri=f"s3://{self.config.s3_bucket}/{self._dataset_prefix(dataset_id)}",
        )

    def _dataset_prefix(self, dataset_id: str) -> str:
        prefix = self.config.s3_prefix.strip("/")
        return f"{prefix}/{dataset_id}" if prefix else dataset_id

    def _build_key(self, dataset_id: str, filename: str) -> str:
        return f"{self._dataset_prefix(dataset_id)}/{filename}"

    def _read_required_text(self, key: str) -> str:
        try:
            return self._read_text(key)
        except FileNotFoundError as exc:
            raise FileNotFoundError(f"Dataset object not found: s3://{self.config.s3_bucket}/{key}") from exc

    def _read_optional_text(self, key: str) -> str | None:
        try:
            return self._read_text(key)
        except FileNotFoundError:
            return None

    def _read_text(self, key: str) -> str:
        client = self._get_client()
        try:
            response = client.get_object(Bucket=self.config.s3_bucket, Key=key)
        except Exception as exc:
            if self._is_not_found_error(exc):
                raise FileNotFoundError(f"s3://{self.config.s3_bucket}/{key}") from exc
            raise

        body = response["Body"].read()
        if isinstance(body, bytes):
            return body.decode("utf-8")
        return body

    @staticmethod
    def _is_not_found_error(exc: Exception) -> bool:
        error_response = getattr(exc, "response", None)
        if not isinstance(error_response, dict):
            return False
        error = error_response.get("Error", {})
        code = str(error.get("Code", "")).lower()
        return code in {"nosuchkey", "notfound", "404"}

    def _get_client(self) -> object:
        if self._client is None:
            try:
                import boto3
            except ImportError as exc:
                raise ImportError(
                    "S3 dataset loading requires boto3. Install the AWS dependency group first."
                ) from exc

            self._client = boto3.client("s3", region_name=self.config.s3_region)
        return self._client


__all__ = [
    "DatasetFiles",
    "DatasetSource",
    "LocalDatasetSource",
    "S3DatasetSource",
    "validate_dataset_id",
]
