import inspect
import time
from typing import Type, Union, Callable

from apollo_utils.core.parallel import iter_chunk
from apollo_utils.core.utils.result_types import ListResult, DictResult


def request_in_chunks(
    items_kwarg: str,
    chunk_size: int = None,
    result_type: Type[Union[ListResult, DictResult]] = DictResult,
    wait: bool = False,
    wait_time: int = 3
) -> Callable:
    """Execute API requests for chunks of items.
    Executes original request by synchronous chunks

    Args:
        items_kwarg: Args list key.
        chunk_size: A size of a chunk.
        result_type: Wrapped function result type class.
        wait: to wait between requests
        wait_time: wait (sleep) time between requests

    Returns:
        Wrapped function.

    """
    if not items_kwarg:
        raise ValueError(f"'items_kwarg' should be passed. Got item_kwarg={items_kwarg}")

    def wrapper(f: Callable):
        def wrapped(*args, **kwargs):

            result = result_type.result_cls()

            def set_items(_items, _kwargs):
                _kwargs[items_kwarg] = _items
                return _kwargs

            for chunk in iter_chunk(kwargs[items_kwarg], chunk_size=chunk_size):
                chunk_result = f(*args, **set_items(chunk, kwargs))

                result_type.update(result, chunk_result)
                if wait:
                    time.sleep(wait_time)
            return result

        wrapped.__signature__ = inspect.signature(f)
        return wrapped

    return wrapper