from __future__ import annotations

from datasets import DatasetDict, load_dataset

DEFAULT_DATASET_ID = "FlowRank/labeled_emails"
DEFAULT_DATASET_REVISION = "main"


def load_labeled_emails(
    dataset_id: str = DEFAULT_DATASET_ID,
    revision: str | None = None,
    *,
    seed: int = 42,
) -> DatasetDict:
    """Charge le dataset labeled_emails (révision `main` = dernière version publiée)."""
    ds = load_dataset(dataset_id, revision=revision or DEFAULT_DATASET_REVISION)
    if "train" in ds and "test" in ds:
        return ds
    if "train" in ds and "test" not in ds:
        return ds["train"].train_test_split(test_size=0.1, seed=seed)
    return ds