"""Utils for breaking up input into batches."""

import collections.abc
import functools

from src.configs import common_config


def execute_in_chunks(max_chunk_size=None):
    """Execute decorated function by splitting its input into chunks.

    Split function input list (first positional argument) into chunks and
    return a result that combines of all the items from individual calls.

    Works with functions that receive a list as the first positional argument
    and return a mutable sequence or a mutable mapping.

    Decorator should be always called, even if you're not overriding the
    default chunk size:

    @chunks.execute_in_chunks()
    def foo(some_input_list):
        ...
        return some_output_list

    Args:
        max_chunk_size (int): Maximum number of elements in each chunk to split
            the input list into.

    """
    # This allows simpler tests, as the config is read during a function call.
    if max_chunk_size is None:
        max_chunk_size = common_config.MAX_CHUNK_SIZE

    def inner(f):
        @functools.wraps(f)
        def wrapped(items, *args, **kwargs):
            result = None
            if not items:
                # The inner loop won't run for empty args list, but we have to
                # return actual result of a certain type.
                return f(items, *args, **kwargs)
            # Split input into chunks and execute wrapped function to create
            # a full result list.
            for i in range(0, len(items), max_chunk_size):
                new_args = (items[i:i + max_chunk_size],) + args
                function_result = f(*new_args, **kwargs)
                if isinstance(
                        function_result, collections.abc.MutableSequence):
                    if result is None:
                        result = []
                    result.extend(function_result)
                elif isinstance(
                        function_result, collections.abc.MutableMapping):
                    if result is None:
                        result = {}
                    result.update(function_result)
                else:
                    raise NotImplementedError
            return result
        return wrapped
    return inner