from enum import Enum
from json import dumps, loads
from requests import post, Response
from typing import Any, Dict, Union, Generator, List
from core.model_runtime.model_providers.wenxin.llm.ernie_bot_errors import BadRequestError, InvalidAPIKeyError, \
    InternalServerError, RateLimitReachedError, InvalidAuthenticationError
from core.model_runtime.entities.message_entities import PromptMessageTool
from datetime import datetime, timedelta
from threading import Lock

# map api_key to access_token
baidu_access_tokens: Dict[str, 'BaiduAccessToken'] = {}
baidu_access_tokens_lock = Lock()

class BaiduAccessToken:
    api_key: str
    access_token: str
    expires: datetime

    def __init__(self, api_key: str) -> None:
        self.api_key = api_key
        self.access_token = ''
        self.expires = datetime.now() + timedelta(days=3)

    def _get_access_token(api_key: str, secret_key: str) -> str:
        """
            request access token from Baidu
        """
        try:
            response = post(
                url=f'https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={api_key}&client_secret={secret_key}',
                headers={
                    'Content-Type': 'application/json',
                    'Accept': 'application/json'
                },
            )
        except Exception as e:
            raise InvalidAuthenticationError(f'Failed to get access token from Baidu: {e}')

        resp = response.json()
        if 'error' in resp:
            if resp['error'] == 'invalid_client':
                raise InvalidAPIKeyError(f'Invalid API key or secret key: {resp["error_description"]}')
            elif resp['error'] == 'unknown_error':
                raise InternalServerError(f'Internal server error: {resp["error_description"]}')
            elif resp['error'] == 'invalid_request':
                raise BadRequestError(f'Bad request: {resp["error_description"]}')
            elif resp['error'] == 'rate_limit_exceeded':
                raise RateLimitReachedError(f'Rate limit reached: {resp["error_description"]}')
            else:
                raise Exception(f'Unknown error: {resp["error_description"]}')
                
        return resp['access_token']

    @staticmethod
    def get_access_token(api_key: str, secret_key: str) -> 'BaiduAccessToken':
        """
            LLM from Baidu requires access token to invoke the API.
            however, we have api_key and secret_key, and access token is valid for 30 days.
            so we can cache the access token for 3 days. (avoid memory leak)

            it may be more efficient to use a ticker to refresh access token, but it will cause
            more complexity, so we just refresh access tokens when get_access_token is called.
        """

        # loop up cache, remove expired access token
        baidu_access_tokens_lock.acquire()
        now = datetime.now()
        for key in list(baidu_access_tokens.keys()):
            token = baidu_access_tokens[key]
            if token.expires < now:
                baidu_access_tokens.pop(key)

        if api_key not in baidu_access_tokens:
            # if access token not in cache, request it
            token = BaiduAccessToken(api_key)
            baidu_access_tokens[api_key] = token
            # release it to enhance performance
            # btw, _get_access_token will raise exception if failed, release lock here to avoid deadlock
            baidu_access_tokens_lock.release()
            # try to get access token
            token_str = BaiduAccessToken._get_access_token(api_key, secret_key)
            token.access_token = token_str
            token.expires = now + timedelta(days=3)
            return token
        else:
            # if access token in cache, return it
            token = baidu_access_tokens[api_key]
            baidu_access_tokens_lock.release()
            return token

class ErnieMessage:
    class Role(Enum):
        USER = 'user'
        ASSISTANT = 'assistant'
        FUNCTION = 'function'
        SYSTEM = 'system'

    role: str = Role.USER.value
    content: str
    usage: Dict[str, int] = None
    stop_reason: str = ''

    def to_dict(self) -> Dict[str, Any]:
        return {
            'role': self.role,
            'content': self.content,
        }
    
    def __init__(self, content: str, role: str = 'user') -> None:
        self.content = content
        self.role = role

class ErnieBotModel(object):
    api_bases = {
        'ernie-bot': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions',
        'ernie-bot-4': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro',
        'ernie-bot-8k': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie_bot_8k',
        'ernie-bot-turbo': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant',
    }

    function_calling_supports = [
        'ernie-bot',
        'ernie-bot-8k',
    ]

    api_key: str = ''
    secret_key: str = ''

    def __init__(self, api_key: str, secret_key: str):
        self.api_key = api_key
        self.secret_key = secret_key

    def generate(self, model: str, stream: bool, messages: List[ErnieMessage], 
                 parameters: Dict[str, Any], timeout: int, tools: List[PromptMessageTool], \
                 stop: List[str], user: str) \
        -> Union[Generator[ErnieMessage, None, None], ErnieMessage]:

        # check parameters
        self._check_parameters(model, parameters, tools, stop)

        # get access token
        access_token = self._get_access_token()

        # generate request body
        url = f'{self.api_bases[model]}?access_token={access_token}'

        # clone messages
        messages_cloned = self._copy_messages(messages=messages)

        # build body
        body = self._build_request_body(model, messages=messages_cloned, stream=stream,
                                        parameters=parameters, tools=tools, stop=stop, user=user)
        headers = {
            'Content-Type': 'application/json',
        }

        resp = post(url=url, data=dumps(body), headers=headers, stream=stream)

        if resp.status_code != 200:
            raise InternalServerError(f'Failed to invoke ernie bot: {resp.text}')

        if stream:
            return self._handle_chat_stream_generate_response(resp)
        return self._handle_chat_generate_response(resp)
    
    def _handle_error(self, code: int, msg: str):
        error_map = {
            1: InternalServerError,
            2: InternalServerError,
            3: BadRequestError,
            4: RateLimitReachedError,
            6: InvalidAuthenticationError,
            13: InvalidAPIKeyError,
            14: InvalidAPIKeyError,
            15: InvalidAPIKeyError,
            17: RateLimitReachedError,
            18: RateLimitReachedError,
            19: RateLimitReachedError,
            100: InvalidAPIKeyError,
            111: InvalidAPIKeyError,
            200: InternalServerError,
            336000: InternalServerError,
            336001: BadRequestError,
            336002: BadRequestError,
            336003: BadRequestError,
            336004: InvalidAuthenticationError,
            336005: InvalidAPIKeyError,
            336006: BadRequestError,
            336007: BadRequestError,
            336008: BadRequestError,
            336100: InternalServerError,
            336101: BadRequestError,
            336102: BadRequestError,
            336103: BadRequestError,
            336104: BadRequestError,
            336105: BadRequestError,
            336200: InternalServerError,
            336303: BadRequestError,
            337006: BadRequestError
        }

        if code in error_map:
            raise error_map[code](msg)
        else:
            raise InternalServerError(f'Unknown error: {msg}')

    def _get_access_token(self) -> str:
        token = BaiduAccessToken.get_access_token(self.api_key, self.secret_key)
        return token.access_token
    
    def _copy_messages(self, messages: List[ErnieMessage]) -> List[ErnieMessage]:
        return [ErnieMessage(message.content, message.role) for message in messages]

    def _check_parameters(self, model: str, parameters: Dict[str, Any], 
                          tools: List[PromptMessageTool], stop: List[str]) -> None:
        if model not in self.api_bases:
            raise BadRequestError(f'Invalid model: {model}')
        
        # if model not in self.function_calling_supports and tools is not None and len(tools) > 0:
        #     raise BadRequestError(f'Model {model} does not support calling function.')
        # ErnieBot supports function calling, however, there is lots of limitations.
        # such as, the messages should be ordered as user by assistant or function...
        # so, we just disable function calling for now.

        if tools is not None and len(tools) > 0:
            raise BadRequestError(f'function calling is not supported yet.')

        if stop is not None:
            if len(stop) > 4:
                raise BadRequestError(f'stop list should not exceed 4 items.')

            for s in stop:
                if len(s) > 20:
                    raise BadRequestError(f'stop item should not exceed 20 characters.')
        
    def _build_request_body(self, model: str, messages: List[ErnieMessage], stream: bool, parameters: Dict[str, Any],
                            tools: List[PromptMessageTool], stop: List[str], user: str) -> Dict[str, Any]:
        # if model in self.function_calling_supports:
        #     return self._build_function_calling_request_body(model, messages, parameters, tools, stop, user)
        return self._build_chat_request_body(model, messages, stream, parameters, stop, user)
        
    def _build_function_calling_request_body(self, model: str, messages: List[ErnieMessage], stream: bool,
                                                parameters: Dict[str, Any], tools: List[PromptMessageTool], 
                                                stop: List[str], user: str) \
        -> Dict[str, Any]:
        if len(messages) % 2 == 0:
            raise BadRequestError(f'The number of messages should be odd.')
        if messages[0].role == 'function':
            raise BadRequestError(f'The first message should be user message.')
        
        """
        TODO: implement function calling
        """

    def _build_chat_request_body(self, model: str, messages: List[ErnieMessage], stream: bool, 
                                 parameters: Dict[str, Any], stop: List[str], user: str) \
        -> Dict[str, Any]:
        if len(messages) == 0:
            raise BadRequestError(f'The number of messages should not be zero.')
        
        # check if the first element is system, shift it
        system_message = ''
        if messages[0].role == 'system':
            message = messages.pop(0)
            system_message = message.content

        if len(messages) % 2 == 0:
            raise BadRequestError(f'The number of messages should be odd.')
        if messages[0].role != 'user':
            raise BadRequestError(f'The first message should be user message.')
        body = {
            'messages': [message.to_dict() for message in messages],
            'stream': stream,
            'stop': stop,
            'user_id': user,
            **parameters
        }

        if system_message:
            body['system'] = system_message

        return body
        
    def _handle_chat_generate_response(self, response: Response) -> ErnieMessage:
        data = response.json()
        if 'error_code' in data:
            code = data['error_code']
            msg = data['error_msg']
            # raise error
            self._handle_error(code, msg)

        result = data['result']
        usage = data['usage']

        message = ErnieMessage(content=result, role='assistant')
        message.usage = {
            'prompt_tokens': usage['prompt_tokens'],
            'completion_tokens': usage['completion_tokens'],
            'total_tokens': usage['total_tokens']
        }

        return message

    def _handle_chat_stream_generate_response(self, response: Response) -> Generator[ErnieMessage, None, None]:
        for line in response.iter_lines():
            if len(line) == 0:
                continue
            line = line.decode('utf-8')
            if line[0] == '{':
                try:
                    data = loads(line)
                    if 'error_code' in data:
                        code = data['error_code']
                        msg = data['error_msg']
                        # raise error
                        self._handle_error(code, msg)
                except Exception as e:
                    raise InternalServerError(f'Failed to parse response: {e}')
            
            if line.startswith('data:'):
                line = line[5:].strip()
            else:
                continue

            if not line:
                continue
            try:
                data = loads(line)
            except Exception as e:
                raise InternalServerError(f'Failed to parse response: {e}')
            
            result = data['result']
            is_end = data['is_end']

            if is_end:
                usage = data['usage']
                finish_reason = data.get('finish_reason', None)
                message = ErnieMessage(content=result, role='assistant')
                message.usage = {
                    'prompt_tokens': usage['prompt_tokens'],
                    'completion_tokens': usage['completion_tokens'],
                    'total_tokens': usage['total_tokens']
                }
                message.stop_reason = finish_reason

                yield message
            else:
                message = ErnieMessage(content=result, role='assistant')
                yield message