"""Dataset loaders for SWE-bench and SWE-bench_Lite."""

from __future__ import annotations

from typing import Optional

from .schema import SWEbenchExample

try:
    from datasets import load_dataset
except ImportError:  # pragma: no cover - optional dependency
    load_dataset = None


def _require_datasets():
    if load_dataset is None:
        raise ImportError(
            "datasets is required to load SWE-bench. Install with `uv run pip install datasets`."
        )


def load_swebench_examples(
    dataset_name: str,
    split: str,
    sample_size: Optional[int] = None,
    *,
    shuffle: bool = False,
    seed: int = 0,
) -> list[SWEbenchExample]:
    """
    Load SWE-bench examples and return them as typed objects.

    Args:
        dataset_name: HF dataset id (e.g., 'princeton-nlp/SWE-bench_Lite').
        split: dataset split to load (train/dev/test).
        sample_size: optional cap on number of rows returned.
        shuffle: whether to shuffle before sampling.
        seed: seed for deterministic sampling.

    Returns:
        List of SWEbenchExample.
    """
    _require_datasets()
    dataset = load_dataset(dataset_name, split=split)
    if shuffle:
        dataset = dataset.shuffle(seed=seed)
    if sample_size is not None:
        dataset = dataset.select(range(sample_size))

    examples: list[SWEbenchExample] = []
    for row in dataset:
        try:
            repo = row["repo"]
            instance_id = row["instance_id"]
            task_id = f"{repo}:{instance_id}"
        except KeyError as exc:  # pragma: no cover - dataset schema guarantees keys
            raise KeyError(f"Missing expected key {exc!r} in SWE-bench row") from exc

        examples.append(
            SWEbenchExample(
                task_id=task_id,
                problem_statement=row["problem_statement"],
                hints_text=row.get("hints_text"),
                repo=repo,
                instance_id=instance_id,
                base_commit=row["base_commit"],
                fail_to_pass=row.get("FAIL_TO_PASS"),
                pass_to_pass=row.get("PASS_TO_PASS"),
            )
        )
    return examples
