import json
import os
import time
from functools import lru_cache
from typing import Dict, List

import hubble
from docarray import Document, DocumentArray
from hubble.excepts import AuthenticationRequiredError
from jina import Executor, requests
from jina.enums import LogVerbosity
from jina.logging.logger import JinaLogger

from now.now_dataclasses import UserInput


class SecurityLevel:
    ADMIN = 1
    USER = 2


def secure_request(level: int, on: str = None):
    """decorator to check the authorization of the incoming request"""

    def decorator(func):
        @requests(on=on)
        def wrapper(*args, **kwargs):
            _check_user(
                kwargs,
                level,
                args[0].user_emails,
                args[0].admin_emails,
                args[0].api_keys,
            )

            cls_instance = args[0]
            if 'args' in kwargs:
                cls_instance.logger.debug(f"args: {kwargs['args']}")

            if 'docs' in kwargs and kwargs['docs']:
                if isinstance(kwargs['docs'], DocumentArray):
                    cls_instance.logger.info(f"processing {len(kwargs['docs'])} docs")
                    if cls_instance.logger.logger.level <= LogVerbosity.DEBUG:
                        kwargs['docs'].summary()
                        kwargs['docs'][0].summary()
                elif isinstance(kwargs['docs'], Document):
                    cls_instance.logger.info(f"processing a single doc")
                    if cls_instance.logger.logger.level <= LogVerbosity.DEBUG:
                        kwargs['docs'].summary()

            if 'docs_matrix' in kwargs and kwargs['docs_matrix']:
                cls_instance.logger.info(
                    f"processing {len(kwargs['docs_matrix'])} docs_matrix"
                )
                lengths = [str(len(da)) for da in kwargs['docs_matrix']]
                cls_instance.logger.info(f"lengths: {', '.join(lengths)}")
                if cls_instance.logger.logger.level <= LogVerbosity.DEBUG and len(
                    kwargs['docs_matrix']
                ):
                    kwargs['docs_matrix'][0].summary()
                    kwargs['docs_matrix'][0][0].summary()

            if 'docs_map' in kwargs and kwargs['docs_map']:
                cls_instance.logger.info(
                    f"processing {len(kwargs['docs_map'])} docs_map"
                )
                lengths = [
                    f"{key_encoder}: {len(da)}"
                    for key_encoder, da in kwargs['docs_map'].items()
                ]
                cls_instance.logger.info(f"lengths: {', '.join(lengths)}")

            t0 = time.time()
            res = func(*args, **kwargs)
            t1 = time.time()

            if isinstance(res, DocumentArray):
                cls_instance.logger.info(
                    f"output {len(res)} docs in {(t1 - t0) * 1000:.1f}ms"
                )
                if len(res) and cls_instance.logger.logger.level <= LogVerbosity.DEBUG:
                    res[0].summary()
            elif not res:
                cls_instance.logger.info(f"no output in {(t1 - t0) * 1000:.1f}ms")

            return res

        return wrapper

    return decorator


def _check_user(kwargs, level, user_emails, admin_emails, api_keys):
    if user_emails == [] and admin_emails == [] and api_keys == []:
        return

    if 'parameters' not in kwargs or (
        'api_key' not in kwargs['parameters'] and 'jwt' not in kwargs['parameters']
    ):
        raise PermissionError(
            f'`jwt` or `api_key` needs to be part of the request parameters.'
        )

    if 'api_key' in kwargs['parameters']:
        if kwargs['parameters']['api_key'] in api_keys:
            return
        else:
            raise PermissionError(f'`api_key` is invalid')
    else:
        jwt = kwargs['parameters']['jwt']

    user_info = _get_user_info(jwt['token'])
    for email in admin_emails + user_emails + ['jina.ai']:
        if _valid_user(user_info.get('email'), email):
            if level == SecurityLevel.ADMIN and email not in admin_emails:
                raise PermissionError(f'User {email} is not an admin.')
            return
    raise PermissionError(
        f'User {user_info.get("email") or user_info["_id"]} has no permission'
    )


def _valid_user(user_email, allowed_email):
    if '@' not in allowed_email:
        return user_email.split('@')[1] == allowed_email
    else:
        return user_email == allowed_email


@lru_cache(maxsize=128, typed=False)
def _get_user_info(token):
    try:
        client = hubble.Client(token=token, max_retries=None, jsonify=True)
        user_info = client.get_user_info()
        if user_info['code'] != 200:
            raise PermissionError(
                'Wrong token passed or the token is expired! Please re-login'
            )
        return user_info['data']
    except AuthenticationRequiredError as e:
        raise PermissionError(
            'Token expired or invalid. Please use the updated token', e
        )


def get_auth_executor_class():
    class NOWAuthExecutor(Executor):
        """
        NOWAuthExecutor performs the token check for authorization. It stores the owner ID belonging
        to the authorised user and also the `user_id` of the allowed users with access to the flow
        deployed by the user.

        If no `admin_emals`, `user_emails` and `api_keys` are provided, no checks will be performed.
        """

        def __init__(
            self,
            user_input_dict: Dict = {},
            admin_emails: List[str] = [],
            user_emails: List[str] = [],
            api_keys: List[str] = [],
            *args,
            **kwargs,
        ):
            """
            :param user_input_dict: kwargs to construct UserInput
            :param admin_email: ID of the user deploying this flow. ID is obtained from Hubble
            :param user_emails: Comma separated Email IDs of the allowed users with access to this flow.
                The Email ID from the incoming request to this flow will be verified against this.
            :param pats: List of PATs of the allowed users with access to this flow.
            """
            super().__init__(*args, **kwargs)
            self.logger = JinaLogger(self.__class__.__name__)
            self.admin_emails = admin_emails
            self.user_emails = user_emails
            self.api_keys = api_keys
            self._user = None

            self.user_input = UserInput()
            for attr_name, prev_value in self.user_input.__dict__.items():
                setattr(
                    self.user_input,
                    attr_name,
                    user_input_dict.get(attr_name, prev_value),
                )

            # TODO admin email must be persisted as well
            # TODO also, please remove duplicate code

            self.api_keys_path = (
                os.path.join(self.workspace, 'api_keys.json')
                if self.workspace
                else None
            )
            self.user_emails_path = (
                os.path.join(self.workspace, 'user_emails.json')
                if self.workspace
                else None
            )

            if self.api_keys_path and os.path.exists(self.api_keys_path):
                with open(self.api_keys_path, 'r') as fp:
                    self.api_keys = json.load(fp)
            if self.user_emails_path and os.path.exists(self.user_emails_path):
                with open(self.user_emails_path, 'r') as fp:
                    self.user_emails = json.load(fp)

        @secure_request(on='/admin/updateUserEmails', level=SecurityLevel.ADMIN)
        def update_user_emails(self, parameters: Dict = None, **kwargs):
            """
            Update the email addresses during runtime. That way, we don't have to restart it.
            """
            self.user_emails = parameters['user_emails']
            if self.user_emails_path:
                with open(self.user_emails_path, 'w') as fp:
                    json.dump(self.user_emails, fp)

        @secure_request(on='/admin/updateApiKeys', level=SecurityLevel.ADMIN)
        def update_api_keys(self, parameters: Dict = None, **kwargs):
            """
            Update the api keys during runtime. That way, we don't have to restart it.
            """
            self.api_keys = parameters['api_keys']
            if self.api_keys_path:
                with open(self.api_keys_path, 'w') as fp:
                    json.dump(self.api_keys, fp)

        @secure_request(level=SecurityLevel.USER)
        def check(self, docs: DocumentArray = None, **kwargs):
            """
            Check the authorization for each incoming request. The logic of the function is in the decorator.
            """
            return docs

    return NOWAuthExecutor
