"""按配置构建 Dataset / DataLoader。"""

from __future__ import annotations

from collections.abc import Mapping
from typing import Any

from torch.utils.data import DataLoader, Dataset

from datasets.film_defect import collate_fn
from datasets.registry import DATASETS
from utils.device import pin_memory_supported, resolve_device


def build_dataset(cfg: Mapping[str, Any], **kwargs: Any) -> Dataset:
    return DATASETS.build(dict(cfg), **kwargs)


def build_dataloader(
    dataset: Dataset,
    batch_size: int = 8,
    shuffle: bool = False,
    num_workers: int = 4,
    pin_memory: bool | None = None,
    drop_last: bool = False,
) -> DataLoader:
    """`pin_memory` 缺省时只在 CUDA 上开启（CPU/MPS 上无收益且会告警）。"""
    if pin_memory is None:
        pin_memory = pin_memory_supported(resolve_device("auto"))
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        pin_memory=pin_memory,
        drop_last=drop_last,
        collate_fn=collate_fn,
        persistent_workers=num_workers > 0,
    )