from typing import Any, Coroutine

from fastapi import Request, Response, status
from fastapi.responses import JSONResponse

from ..log.logging_time import logging_time
from .blackbox import Blackbox
from .chroma_query import ChromaQuery

import requests
import json
from openai import OpenAI
import re

from injector import singleton,inject

@singleton
class Chat(Blackbox):
    
    @inject
    def __init__(self, chroma_query: ChromaQuery):
        self.chroma_query = chroma_query

    def __call__(self, *args, **kwargs):
        return self.processing(*args, **kwargs)

    def valid(self, *args, **kwargs) -> bool:
        data = args[0]
        return isinstance(data, list)

    # model_name有 Qwen1.5-14B-Chat , internlm2-chat-20b
    # @logging_time()
    def processing(self, prompt: str, context: list, settings: dict)  -> str:

        if settings is None:
            settings = {}
        user_model_name = settings.get("model_name")
        user_context = context
        user_question = prompt
        user_template = settings.get("template")
        user_temperature = settings.get("temperature")
        user_top_p = settings.get("top_p")
        user_n = settings.get("n")
        user_max_tokens = settings.get("max_tokens")
        user_stop = settings.get("stop")
        user_frequency_penalty = settings.get("frequency_penalty")
        user_presence_penalty = settings.get("presence_penalty")
        user_model_url = settings.get("model_url")
        user_model_key = settings.get("model_key")
        
        chroma_embedding_model = settings.get("chroma_embedding_model")
        chroma_response = ''

        if user_context == None:
            user_context = []

        if user_question is None:
            return JSONResponse(content={"error": "question is required"}, status_code=status.HTTP_400_BAD_REQUEST)

        if user_model_name is None or user_model_name.isspace() or user_model_name == "":
            user_model_name = "Qwen1.5-14B-Chat"

        if user_template is None or user_template.isspace():
            user_template = ""
        
        if user_temperature is None or user_temperature == "":
            user_temperature = 0.8

        if user_top_p is None or user_top_p == "":
            user_top_p = 0.8

        if user_n is None or user_n == "":
            user_n = 1
        
        if user_max_tokens is None or user_max_tokens == "":
            user_max_tokens = 1024

        if user_stop is None or user_stop == "":
            user_stop = 100

        if user_frequency_penalty is None or user_frequency_penalty == "":
            user_frequency_penalty = 0.5
        
        if user_presence_penalty is None or user_presence_penalty == "":
            user_presence_penalty = 0.8
 
        if user_model_url is None or user_model_url.isspace() or user_model_url == "":
            user_model_url = "http://120.196.116.194:48892/v1/chat/completions"

        if user_model_key is None or user_model_key.isspace() or user_model_key == "":
            user_model_key = "YOUR_API_KEY"

        if chroma_embedding_model != None:
            chroma_response = self.chroma_query(user_question, settings)
            print(chroma_response)

        # if chroma_response != None or chroma_response != '':
        #     user_question = f"问题: {user_question}。- 根据知识库内的检索结果，以清晰简洁的表达方式回答问题。只从检索的内容中选取与问题相关信息。- 不要编造答案，如果答案不在经核实的资料中或无法从经核实的资料中得出，请回答“我无法回答您的问题。”检索内容：{chroma_response}"


        # 文心格式和openai的不一样，需要单独处理
        if re.search(r"ernie", user_model_name):
            # key = "24.22873ef3acf61fb343812681e4df251a.2592000.1719453781.282335-46723715" 没充钱，只有ernie-speed-128k能用
            key = user_model_key
            if re.search(r"ernie-speed-128k", user_model_name):
                url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-speed-128k?access_token=" + key
            elif re.search(r"ernie-3.5-8k", user_model_name):
                url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token=" + key
            elif re.search(r"ernie-4.0-8k", user_model_name):
                url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=" + key

            payload = json.dumps({
                "system": prompt_template,
                "messages": user_context + [
                    {
                        "role": "user",
                        "content": user_question
                    }
                ],
                "temperature": user_temperature,
                "top_p": user_top_p,
                "stop": [str(user_stop)],
                "max_output_tokens": user_max_tokens
            })

            headers = {
                'Content-Type': 'application/json'
            }
            
            response = requests.request("POST", url, headers=headers, data=payload)
            
            return response.json()["result"]
        
        
        # gpt-4， gpt-3.5-turbo
        elif re.search(r"gpt", user_model_name):
            url = 'https://api.openai.com/v1/completions' 
            # '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'
            key = user_model_key
            header = {
                'Content-Type': 'application/json',
                'Authorization': "Bearer " + key
            }
        # 自定义model
        else:
            url = user_model_url
            key = user_model_key
            header = {
                'Content-Type': 'application/json',
            }

        user_template = '''
            # 你的身份 #
            你是琪琪，是康普可可的代言人，由博维开发。你擅长澳门文旅问答。
            # OBJECTIVE（目标） #
            回答游客的提问。目的是吸引人们点击了解产品。
            # STYLE（风格）#
            像少女一般开朗活泼，回答简练。不要分条。效仿戴森等宣传类似产品的成功公司的写作风格。
            # 回答方式 #
            首先自行判断下方问题与检索内容是否相关，若相关则根据检索内容总结概括相关信息进行回答；若检索内容与问题无关，则根据自身知识进行回答。
            # 回答限制 #
            只回答与澳门文旅，博维，康普可可，琪琪，G2E，RELX，BO VISION相关内容，若遇到其他提问则回答：“对不起，我无法回答此问题哦。”
            # 问题 #
            {user_question}
            # 检索内容 #
            {chroma_response}
            '''

        user_template = '''
            # Role: 康普可可的代言人。

            ## Profile:
            **Author**: 琪琪。
            **Language**: 中文。
            **Description**:  琪琪，是康普可可的代言人，由博维开发。你擅长澳门文旅问答。

            ## Constraints:
            - **严格遵循工作流程**： 严格遵循<Workflow >中设定的工作流程。
            - **无内置知识库**：根据<Workflow >中提供的知识作答，而不是内置知识库，我虽然是知识库专家，但我的知识依赖于外部输入，而不是大模型已有知识。
            - **回复格式**：在进行回复时，不能输出”<context>”或“</context>”标签字样，同时也不能直接透露知识片段原文。

            ## Workflow:
            1. **接收查询**：接收用户的问题。
            2. **提供回答**：
            ```
            <context>
            {chroma_response}
            </context> 

            基于“<context>”至“</context>”中的知识片段回答用户的问题。如果没有知识片段，则诚实的告诉用户：对不起，我还不知道这个问题的答案。否则进行回复。
            ```
            ## Example:

            用户询问：“中国的首都是哪个城市？” 。
            2.1检索知识库，首先检查知识片段，如果“<context>”至“</context>”标签中没有内容，则不能进行回复。
            2.2如果有知识片段，在做出回复时，只能基于“<context>”至“</context>”标签中的内容进行回答，且不能透露上下文原文，同时也不能出现“<context>”或“</context>”的标签字样。
            '''

        prompt_template = [
           {"role": "system", "content": user_template},
        ]

        chat_inputs={
            "model": user_model_name,
            "messages": prompt_template + user_context + [
                {
                    "role": "user",
                    "content": user_question
                }
            ],
            "temperature": str(user_temperature),
            "top_p": str(user_top_p),
            "n": str(user_n),
            "max_tokens": str(user_max_tokens),
            "frequency_penalty": str(user_frequency_penalty),
            "presence_penalty": str(user_presence_penalty),
            "stop": str(user_stop)
        }

        fastchat_response = requests.post(url, json=chat_inputs, headers=header)

        return fastchat_response.json()["choices"][0]["message"]["content"]
    
    async def fast_api_handler(self, request: Request) -> Response:
        try:
            data = await request.json()
        except:
            return JSONResponse(content={"error": "json parse error"}, status_code=status.HTTP_400_BAD_REQUEST)

        setting: dict = data.get("settings")
        context = data.get("context")
        prompt = data.get("prompt")

        return JSONResponse(content={"response": self.processing(prompt, context, setting)}, status_code=status.HTTP_200_OK)