"""Shared split-wrapper helpers for train/eval style pipelines."""

from __future__ import annotations

from typing import TYPE_CHECKING


if TYPE_CHECKING:
    from collections.abc import Callable, Mapping, Sequence


def build_split_loader_map(
    requested_splits: Sequence[str],
    train_loader: object | None,
    val_loader: object | None,
    test_loader: object | None,
) -> dict[str, object]:
    """Build split-to-loader map while skipping unavailable loaders.

    Args:
        requested_splits: Ordered split names requested by caller.
        train_loader: Optional loader for ``train`` split.
        val_loader: Optional loader for ``val`` split.
        test_loader: Optional loader for ``test`` split.

    Returns:
        Dictionary of requested splits that have non-``None`` loaders.
    """
    candidates = {"train": train_loader, "val": val_loader, "test": test_loader}
    return {
        split: candidates[split]
        for split in requested_splits
        if split in candidates and candidates[split] is not None
    }


def evaluate_available_splits(
    requested_splits: Sequence[str],
    split_to_loader: Mapping[str, object],
    evaluate_split: Callable[[str, object], list[dict]],
) -> list[dict]:
    """Evaluate all available requested splits and aggregate payload items.

    Args:
        requested_splits: Ordered split names to evaluate.
        split_to_loader: Mapping from split name to loader object.
        evaluate_split: Callable that evaluates one split and returns payload.

    Returns:
        Concatenated list of payload dictionaries from available splits.
    """
    payload: list[dict] = []
    for split in requested_splits:
        loader = split_to_loader.get(split)
        if loader is None:
            continue
        payload.extend(evaluate_split(split, loader))
    return payload
