# -*- coding: utf-8 -*-
# Copyright (c) 2017 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

from typing import Callable, Iterable

from toolz.dicttoolz import get_in
from toolz.itertoolz import partition_all
from voluptuous import ALLOW_EXTRA, Invalid, Optional, Schema

from czc.string import vacuum_pack, vacuum_unpack


class PaginatedApiError(Exception):
    pass


def iterable(xs):
    if isinstance(xs, Iterable):
        return xs
    else:
        raise Invalid('expected Iterable')


AUTO_CURSOR_KEY = 'auto_cursor'
CURSOR_KEY = 'cursor'
DATA_KEY = 'data'
ORIGINAL_EVENT_KEY = 'original'
PAGINATED_RESULT = Schema({
    DATA_KEY: iterable,
    Optional(CURSOR_KEY): dict,
}, required=True, extra=ALLOW_EXTRA)


def _validate(result):
    try:
        return PAGINATED_RESULT(result)
    except Invalid as e:
        raise PaginatedApiError(f'Expected valid paginated response: {e}')


def _unpack(event):
    if CURSOR_KEY in event:
        try:
            cursor = vacuum_unpack(event[CURSOR_KEY].encode())
            original = _get_original_event(cursor)
            return {
                **original,
                CURSOR_KEY: cursor,
            }
        except Exception as e:
            raise PaginatedApiError(f'Failed to unpack cursor: {e}')
    return event


def _get_auto_cursor(unpacked):
    return get_in([CURSOR_KEY, AUTO_CURSOR_KEY], unpacked, default={})


def _get_original_event(cursor):
    return cursor.get(ORIGINAL_EVENT_KEY, {})


def _next_chunk(xs, n, o):
    chunks = partition_all(n, xs)
    r = []

    try:
        for x in range(o + 1):
            r = list(next(chunks))
        next(chunks)
        return (r, True)
    except StopIteration:
        return (r, False)


def _paginate_data(data, cursor, default_limit):
    try:
        limit = cursor.get('limit', default_limit)
        offset = cursor.get('offset', 0)
        (chunk, more) = _next_chunk(data, limit, offset)
        next_cursor = {'limit': limit, 'offset': (offset + 1)} if more else None
        return (chunk, next_cursor)
    except Exception as e:
        raise PaginatedApiError(f'Failed to pack API response cursor: {e}')


def _next_cursor(auto_cursor, unpacked):
    cursor = {
        AUTO_CURSOR_KEY: auto_cursor,
        ORIGINAL_EVENT_KEY: unpacked,
    }
    return {CURSOR_KEY: vacuum_pack(cursor).decode()} if auto_cursor else {}


def _pack_auto_cursor(data, auto_cursor, unpacked):
    try:
        return {
            **{DATA_KEY: data},
            **_next_cursor(auto_cursor, unpacked)
        }
    except Exception as e:
        raise PaginatedApiError(f'Failed to pack API response cursor: {e}')


def _pack_handler_cursor(result):
    try:
        if CURSOR_KEY in result:
            return {**result, CURSOR_KEY: vacuum_pack(result[CURSOR_KEY]).decode()}
        return result
    except Exception as e:
        raise PaginatedApiError(f'Failed to pack API response cursor: {e}')


def paginated(limit: int) -> Callable:
    """Lambda-to-lambda server decorator.

    This decorator paginates a lambda handler function. The lambda handler must return a response that meets the
    API schema: {'data': Any, Optional('cursor'): str}.

    If *limit* is > 0, this decorator auto paginates 'data' and manages the 'cursor'. In this case of auto-pagination,
    this decorator will re-invoke the lambda handler with its initial arguments on each successive call, returning a
    slice of the 'data' managed by 'cursor'.

    If *limit* is <= 0, this decorator simply packs up the existing 'cursor' key, leaving all pagination and cursor
    management to the lambda handler.

    Example usage:

    @paginated(10)
    def handler(event, context, **kwargs):
        return {'data': list(range(20))}

    The decorator will return a lamda handler that will automatically return the first 10 elements on the first
    invocation, followed by the last 10 elements on second invocation, given the auto generated cursor is passed
    back to the function.
    """
    def decorator(handler: Callable) -> Callable:

        def handler_wrapper(event: dict, context, **kwargs) -> dict:
            unpacked = _unpack(event)
            result = handler(unpacked, context, **kwargs)
            valid = _validate(result)

            if limit > 0:
                (paginated, auto_cursor) = _paginate_data(valid[DATA_KEY], _get_auto_cursor(unpacked), limit)
                return _pack_auto_cursor(paginated, auto_cursor, unpacked)
            return _pack_handler_cursor(valid)

        return handler_wrapper

    return decorator
