import struct
import time
import zlib
from auth0 import Auth0Error
from flask import Response, stream_with_context
from functools import wraps
from marshmallow import Schema

from src.config.auth import Auth0ManagementConfig
from src.constants.http_status import HTTP_401_UNAUTHORIZED, HTTP_429_TOO_MANY_REQUESTS
from src.legacy.auth.exceptions import ManagementAPIInvalidResponse


auth0_config = Auth0ManagementConfig()


def streaming(output_schema: Schema, use_compression: bool = True) -> Response:
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            def compress(compressor, item, crc, length):
                item = item.encode()
                chunk = compressor.compress(item)
                crc = zlib.crc32(item, crc) & 0xFFFFFFFF
                length += len(item)

                return chunk, crc, length

            def get_values(item: dict):
                return "," + ",".join(
                    list(
                        map(
                            lambda s: str(s) if isinstance(s, int) else '"' + str(s) + '"',
                            output_schema.dump(item).values(),
                        )
                    )
                )

            def streamer(data):
                if use_compression:
                    yield b"\037\213\010\000" + struct.pack("<L", int(time.time())) + b"\002\377"

                    compressor = zlib.compressobj(1, zlib.DEFLATED, -zlib.MAX_WBITS, zlib.DEF_MEM_LEVEL, 0)
                    crc = zlib.crc32(b"")
                    length = 0

                item = "["

                if use_compression:
                    item, crc, length = compress(compressor, item, crc, length)
                yield item

                item = next(data, None)

                if item:
                    _item_keys = list(map(lambda s: '"' + str(s) + '"', output_schema.dump(item).keys()))

                    item = str(len(_item_keys)) + "," + ",".join(_item_keys) + get_values(item)

                    if use_compression:
                        item, crc, length = compress(compressor, item, crc, length)
                    yield item

                    for item in data:
                        item = get_values(item)
                        if use_compression:
                            item, crc, length = compress(compressor, item, crc, length)
                        yield item

                item = "]"
                if use_compression:
                    item, crc, length = compress(compressor, item, crc, length)
                yield item

                if use_compression:
                    yield compressor.flush()
                    yield struct.pack("<2L", crc, length & 0xFFFFFFFF)

            mime_type = "application/json"
            if use_compression:
                mime_type = "application/gzip"

            response = Response(stream_with_context(streamer(data=func(*args, **kwargs))), mimetype=mime_type)
            return response

        return wrapper

    return decorator


def handle_auth0_too_many_requests(func):
    """Decorate function to handle Auth0 too many requests .

    Arguments:
        func (function): Management API class method

    Raises:
        ManagementAPIInvalidResponse: unhandled Auth0 response or error

    Returns:
        function: decorated function
    """

    @wraps(func)
    def wrapper(self, *args, **kwargs):
        for t in range(1, auth0_config.RATE_LIMIT_RETRY_COUNT + 1):
            try:
                return func(self, *args, **kwargs)
            except Auth0Error as err:
                if err.status_code == HTTP_429_TOO_MANY_REQUESTS:
                    time.sleep(t * auth0_config.RATE_LIMIT_WAIT_MULTIPLIER)
                else:
                    raise ManagementAPIInvalidResponse(
                        original_status_code=err.status_code,
                        original_response=err.message,
                    )

    return wrapper


def handle_auth0_unauthorized(func):
    """Decorate function to handle Auth0 unauthorized error.

    Arguments:
        func (function): Management API class method

    Raises:
        Auth0Error: Auth0 exception

    Returns:
        function: decorated function
    """

    @wraps(func)
    def wrapper(self, *args, **kwargs):
        try:
            return func(self, *args, **kwargs)
        except Auth0Error as err:
            if err.status_code == HTTP_401_UNAUTHORIZED:
                token = self.get_management_token()
                self.set_token(token)
                return func(self, *args, **kwargs)
            else:
                raise

    return wrapper