import random
from typing import Callable, List
from ding.data.buffer.buffer import BufferedData


def group_sample(size_in_group: int, ordered_in_group: bool = True, max_use_in_group: bool = True) -> Callable:
    """
    Overview:
        The middleware is designed to process the data in each group after sampling from the buffer.
    Arguments:
        - size_in_group (:obj:`int`): Sample size in each group.
        - ordered_in_group (:obj:`bool`): Whether to keep the original order of records, default is true.
        - max_use_in_group (:obj:`bool`): Whether to use as much data in each group as possible, default is true.
    """

    def sample(chain: Callable, *args, **kwargs) -> List[List[BufferedData]]:
        if not kwargs.get("groupby"):
            raise Exception("Group sample must be used when the `groupby` parameter is specified.")
        sampled_data = chain(*args, **kwargs)
        for i, grouped_data in enumerate(sampled_data):
            if ordered_in_group:
                if max_use_in_group:
                    end = max(0, len(grouped_data) - size_in_group) + 1
                else:
                    end = len(grouped_data)
                start_idx = random.choice(range(end))
                sampled_data[i] = grouped_data[start_idx:start_idx + size_in_group]
            else:
                sampled_data[i] = random.sample(grouped_data, k=size_in_group)
        return sampled_data

    def _group_sample(action: str, chain: Callable, *args, **kwargs):
        if action == "sample":
            return sample(chain, *args, **kwargs)
        return chain(*args, **kwargs)

    return _group_sample