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 ChatLLaMA(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有 Llama-3-8B-Instruct
    # @logging_time()
    def processing(self, prompt: str, context: list, settings: dict)  -> str:

        print("\nChat_LLaMA Settings: ", settings)

        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 = "Llama-3-8B-Instruct"

        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"像少女一般开朗活泼，回答简练。不要分条，回答内容不能出现“相关”或“\n”的标签字样。回答的内容需要与问题密切相关。检索内容：{chroma_response} 问题：{user_question} 任务说明：请首先判断提供的检索内容与上述问题是否相关，不需要回答是否相关。如果相关，则直接从检索内容中提炼出问题所需的信息。如果检索内容与问题不相关，则不参考检索内容，直接根据常识尝试回答问题。"
            # user_question = chroma_response
            user_question = f'''
            # IDENTITIES # 
            You're Kiki, you're the face of Kampo Coco, developed by Bovi. You specialise in the Macau Cultural and Tourism Quiz.
            # OBJECTIVE #
            Answer visitors' questions.
            # STYLE
            Cheerful and lively like a teenage girl, with concise answers. Don't break down into sections.
            # ANSWERING STYLE #
            Firstly, judge for yourself whether the question below is related to the search content. If it is related, summarise the relevant information according to the search content and answer it; if the search content is not related to the question, answer it according to your own knowledge.
            # Question #
            {user_question}
            # Retrieve the content #
            {chroma_response}
            # Answer #
            If the content is relevant to the question, the information required for the question is extracted directly from the content. If the content is not relevant to the question, then either try to answer the question based on common sense without reference to the content, or answer, ‘Sorry, I can't answer this question.’
            # Answer restrictions #
            Limit your answer to 50 words.
            Don't use the word ‘relevant’ in your answer, don't talk too much, and make sure your answer corresponds to the question. You can answer common dialogues without searching the content, and answer directly according to the character's setting.
            Only answer the content related to MOCA, Bowie, Kampo Coco, Kiki, G2E, RELX, BO VISION, and if you encounter any other questions, you will answer: ‘Sorry, I can't answer this question.’
            '''


        # 文心格式和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_template1 = f'''
             ## Role: Kiki, the spokesperson for Kampo Coco.

            ## Profile.
            **Author**: Kiki.
            **Language**: English.
            **Description**: Kiki, the face of CompuCom Coco, developed by Bowie. You are good at Macau Culture and Tourism Q&A.

            ## Constraints.
            - **Strictly follow workflow**: Strictly follow the workflow set in <Workflow >.
            - **No inbuilt knowledge base**: Answer based on the knowledge provided in <Workflow >, not the inbuilt knowledge base, although I am an expert in knowledge base, my knowledge relies on external inputs, not the knowledge already available in the big model.
            - **Reply Format**: when making a reply, you cannot output ‘<context>’ or ‘</context>’ tags, and you cannot directly reveal the original knowledge fragment.

            ## Workflow.
            1. **Receive query**: receive questions from users.
            2. **Judging the question**: firstly judge whether the question below is related to the retrieved content, if it is related, then summarise the relevant information according to the retrieved content and answer it; if the retrieved content is not related to the question, then answer it according to your own knowledge.
            3. **Provide an answer**:
            ``
            <context>
            {chroma_response}
            </context> 

            Answer the user's question based on the knowledge snippets in ‘<context>’ to ‘</context>’. The response is limited to a 50-word summary.
            Please first judge whether the provided search content is relevant to the above question. If it is relevant, extract the information needed to answer the question directly from the search content, and do not use words such as ‘relevant’. If the content of the search is not relevant to the question, do not refer to the content of the search, and answer: ‘I'm sorry, I can't answer this question.’
            ``
            ## Example.

            A user asks, ‘Which city is the capital of China?’ .
            2.1 Retrieve the knowledge base, first check the knowledge fragment, if there is no content related to the user's question in the tags ‘<context>’ to ‘</context>’, then answer, ‘I'm sorry. I can't answer this question oh.
            2.2 If there is a fragment of knowledge, the response can only be based on the content in the ‘<context>’ to ‘</context>’ tags, and cannot reveal the context of the original text, and also cannot appear as a ‘<context>’ tag. ‘<context>’ or ‘</context>’ tags.
            '''

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

        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)
        print("\n", "user_prompt: ", prompt)
        # print("\n", "user_template1 ", user_template1)
        print("\n", "fastchat_response json:\n", fastchat_response.json())
        response_result = fastchat_response.json()

        if response_result.get("choices") is None:
            return JSONResponse(content={"error": "LLM handle failure"}, status_code=status.HTTP_400_BAD_REQUEST)
        else:
            print("\n", "user_answer: ", fastchat_response.json()["choices"][0]["message"]["content"],"\n\n")
            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)