from __future__ import annotations

from dataclasses import dataclass
from typing import Any


class DatasetInfoError(RuntimeError):
    """Raised when dataset metadata cannot be retrieved."""

@dataclass
class DatasetDetails:
    dataset_id: str
    config_name: str | None
    description: str
    homepage: str
    license: str
    citation: str
    features: dict[str, Any]
    splits: dict[str, dict[str, Any]]
    builder_configs: list[str]
    size_in_bytes: int | None
    download_size: int | None


def describe_dataset(dataset_id: str, config_name: str | None = None) -> DatasetDetails:
    try:
        builder = _load_builder(dataset_id, config_name)
    except DatasetInfoError:
        raise
    except Exception as exc:  # pragma: no cover - propagated as DatasetInfoError
        raise DatasetInfoError(str(exc)) from exc
    info = builder.info

    features = info.features.to_dict() if info.features else {}
    splits = {}
    if info.splits:
        for name, split in info.splits.items():  # type: ignore[attr-defined]
            splits[name] = {
                "num_examples": getattr(split, "num_examples", None),
                "dataset_size": getattr(split, "num_bytes", None)
                or getattr(split, "dataset_size", None),
            }

    builder_configs = sorted(builder.builder_configs.keys()) if builder.builder_configs else []

    return DatasetDetails(
        dataset_id=dataset_id,
        config_name=config_name,
        description=info.description or "",
        homepage=info.homepage or "",
        license=info.license or "",
        citation=info.citation or "",
        features=features,
        splits=splits,
        builder_configs=builder_configs,
        size_in_bytes=getattr(info, "dataset_size", None) or getattr(info, "size_in_bytes", None),
        download_size=getattr(info, "download_size", None),
    )


def _load_builder(dataset_id: str, config_name: str | None):
    from datasets import load_dataset_builder

    try:
        return load_dataset_builder(dataset_id, name=config_name)
    except ValueError as exc:
        if "trust_remote_code" not in str(exc):
            raise DatasetInfoError(str(exc)) from exc
    except Exception as exc:
        raise DatasetInfoError(str(exc)) from exc
    try:
        return load_dataset_builder(dataset_id, name=config_name, trust_remote_code=True)
    except Exception as exc:
        raise DatasetInfoError(str(exc)) from exc


def summarize_feature(spec: Any) -> str:
    if isinstance(spec, list):
        if not spec:
            return "[]"
        return f"[{summarize_feature(spec[0])}]"
    if not isinstance(spec, dict):
        return str(spec)

    feature_type = spec.get("_type")
    if feature_type == "Value":
        return spec.get("dtype", "Value")
    if feature_type == "ClassLabel":
        names = spec.get("names")
        if names and len(names) <= 6:
            return f"ClassLabel[{', '.join(map(str, names))}]"
        return f"ClassLabel(num_classes={spec.get('num_classes')})"
    if feature_type == "Sequence":
        inner = summarize_feature(spec.get("feature"))
        return f"Sequence[{inner}]"
    if feature_type == "Audio":
        return "Audio"
    if feature_type == "Image":
        return "Image"
    if feature_type == "Translation":
        return f"Translation[{', '.join(spec.get('languages') or [])}]"
    if feature_type:
        return feature_type
    if "feature" in spec:
        return summarize_feature(spec["feature"])
    return "Struct"
