"""MSK common functions."""

from concurrent.futures.thread import ThreadPoolExecutor
from functools import wraps
import typing

from lambdacommon.common_config import logger, MSK_HANDLER_CONCURRENCY
from lambdacommon.models.lambda_msk_event_source_message import (
    LambdaMSKEventSourceMessage,
    MSKMessageException
)


def msk_handler(
    process_exceptions: bool = False,
    concurrent: bool = False
) -> typing.Callable[
    [typing.Callable[[typing.Dict[str, typing.Any], typing.Any], typing.Any]],
    typing.Callable[[typing.Dict[str, typing.Any], typing.Any], typing.Any]
]:
    """
    Unwrap MSK Kafka event and call the decorated function with payloads.

    Executes the decorated func with events from the msk event records
    or with the event itself in case of manual lambda execution
    from the AWS console.

    Example:
        @msk_handler(process_exceptions=True, concurrent=True)
        def handler(
            event: Dict[str, Any],
            context: Any
        ) -> None:
            ...
    """

    def decorator(
        func: typing.Callable[[typing.Dict[str, typing.Any], typing.Any], typing.Any]
    ) -> typing.Callable[[typing.Dict[str, typing.Any], typing.Any], typing.Any]:

        @wraps(func)
        def wrapper(
            event: typing.Dict[str, typing.Any],
            context: typing.Any
        ) -> typing.Dict[str, typing.Any]:
            try:
                if concurrent:
                    with ThreadPoolExecutor(
                        max_workers=MSK_HANDLER_CONCURRENCY
                    ) as executor:
                        for key, message in LambdaMSKEventSourceMessage(event):
                            executor.submit(
                                _process_payload,
                                func, message.value, context, process_exceptions
                            )
                else:
                    for key, message in LambdaMSKEventSourceMessage(event):
                        _process_payload(func, message.value, context, process_exceptions)
            except MSKMessageException:
                _process_payload(func, event, context, process_exceptions)
            return {'status': 'success'}

        return wrapper

    return decorator


def _process_payload(
    func: typing.Callable[[typing.Dict[str, typing.Any], typing.Any], typing.Any],
    record: typing.Dict[str, typing.Any],
    context: typing.Any,
    process_exceptions: bool = False
):
    if not process_exceptions:
        func(_unwrap_payload(record), context)
        return
    try:
        func(_unwrap_payload(record), context)
    except Exception as e:
        logger.exception(f'Exception on record processing: {e}')


def _unwrap_payload(
    payload: typing.Dict[str, typing.Any]
) -> typing.Dict[str, typing.Any]:
    return payload.get('payload') or payload