"""
DataLoader factory for training and validation pipelines.
"""

from typing import List, Tuple, Optional
import torch
from torch.utils.data import DataLoader, Dataset

from slm.dataset.dataset import CausalLMDataset
from slm.tokenizer.bpe import BPETokenizer


def causal_collate_fn(batch: List[Tuple[torch.Tensor, torch.Tensor]]) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Collate function combining sequence samples into padded batch tensors.

    Args:
        batch: List of (input_ids, target_ids) tuple pairs.

    Returns:
        Tuple of (input_batch [batch_size, seq_len], target_batch [batch_size, seq_len]).
    """
    inputs = [item[0] for item in batch]
    targets = [item[1] for item in batch]

    input_batch = torch.stack(inputs, dim=0)
    target_batch = torch.stack(targets, dim=0)

    return input_batch, target_batch


def create_dataloader(
    dataset: Dataset,
    batch_size: int = 16,
    shuffle: bool = True,
    num_workers: int = 0,
    drop_last: bool = True
) -> DataLoader:
    """
    Constructs PyTorch DataLoader configured for causal LM training.

    Args:
        dataset: Target dataset instance.
        batch_size: Batch capacity.
        shuffle: Whether to shuffle data samples.
        num_workers: Multi-processing data worker count.
        drop_last: Whether to drop final incomplete batch.

    Returns:
        DataLoader instance.
    """
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        collate_fn=causal_collate_fn,
        drop_last=drop_last
    )