import transformers
import torch
import http.client
import json
from accelerate import infer_auto_device_map
from meta_buffer_utilis import meta_distiller_prompt,extract_and_execute_code
from test_templates import game24,checkmate,word_sorting
# 加载模型用
class Pipeline:
    def __init__(self,model_id,api_key=None):
        self.api = False
        self.local = False
        self.model_id = model_id
        # 没提供api
        if api_key is None:
            self.local = True
            # 使用transformer加载
            self.pipeline = transformers.pipeline(
        "text-generation",
        model=self.model_id,
        model_kwargs={"torch_dtype": torch.bfloat16},
        device_map = 'auto'
        )
        else:
            self.api = True
            self.api_key = api_key
    # 获取回复 传入meta_prompt和user_prompt,meta_prompt是作为"system"角色的。
    def get_respond(self,meta_prompt,user_prompt):
        if self.api: # 提供了api
            # 调用api得到结果
            conn = http.client.HTTPSConnection("api.openai.com")
            payload = json.dumps({
            "model": self.model_id,
            "messages": [
                {
                    "role": 'assistant',
                    "content": meta_prompt
                },
                {
                    "role": 'user',
                    "content": user_prompt
                },
            ]
            })
            headers = {
            'Accept': 'application/json',
            'Authorization': f'Bearer {self.api_key}',
            'User-Agent': 'Apifox/1.0.0 (https://apifox.com)',
            'Content-Type': 'application/json'
            }
            conn.request("POST", "/v1/chat/completions", payload, headers)
            res = conn.getresponse()
            data = res.read()
            text = data.decode("utf-8")
            dict_object = json.loads(text)
            respond = dict_object['choices'][0]['message']['content']
            return respond
        # 没提供api，那就是本地模型了。
        else:
            messages = [
            {"role": "system", "content": meta_prompt},
            {"role": "user", "content": user_prompt},
        ]
            # 使用huggingface的tokenizer模板来格式化输入的信息。
            prompt = self.pipeline.tokenizer.apply_chat_template(
                    messages, 
                    tokenize=False, 
                    add_generation_prompt=True
            )
            # 这个字符串指定了要使用的pipeline类型。在这个例子中，它是文本生成，意味着模型将会根据给定的输入文本继续生成后续的文本。

            terminators = [
                self.pipeline.tokenizer.eos_token_id,
                self.pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")
            ]
            # 第一个元素是从tokenizer中获取的默认结束符ID（eos_token_id），第二个元素是通过将特殊标记"<|eot_id|>"转换成对应的token ID。
            # 发送参数。
            outputs = self.pipeline(
                prompt,
                max_new_tokens=2048,
                eos_token_id=terminators, # 设置了终止生成的条件，可以是单个ID或ID列表，这里使用了前面定义的terminators列表。
                do_sample=True,
                temperature=0.4,
                top_p=0.9,
            )
            # 获得response
            respond = outputs[0]["generated_text"][len(prompt):]
            # outputs[0]访问的是第一个批次的结果（通常pipeline只处理一个prompt时只有一个批次），然后从"generated_text"字段获取完整生成的文本。由于我们只对# prompt之后的部分感兴趣，所以使用切片操作[len(prompt):]去除prompt本身的内容，最后返回这部分生成的文本作为响应。
            return respond
            

        
class BoT:
    # 初始化
    def __init__(self, user_input,problem_id,api_key=None,model_id=None):
        self.api_key = api_key
        self.model_id = model_id
        self.pipeline = Pipeline(self.model_id,self.api_key)
        self.user_input = user_input
        # Only for test use, stay tuned for our update
        self.problem_id = problem_id 
        
    def problem_distillation(self):
        print(f'User prompt:{self.user_input}')
        # 传入的meta_distiller_prompt是作为system角色的
        self.distilled_information = self.pipeline.get_respond(meta_distiller_prompt, self.user_input)
        print(f'Distilled information:{self.distilled_information}')
        
    # 这里就是根据问题id来选择模板
    def buffer_retrieve(self):
        # For initial test use, we will later update the embedding retrieval version to support more 
        if self.problem_id == 0:
            self.thought_template = game24
        elif self.problem_id == 1:
            self.thought_template = checkmate
        elif self.problem_id == 2:
            self.thought_template = word_sorting
            
    def reasoner_instantiation(self):
        # Temporay using selection method to select answer extract method
        problem_id_list = [0,1,2]
        self.instantiation_instruct = """
You are an expert in problem analysis and can apply previous problem-solving approaches to new issues. The user will provide a specific task description and a thought template. Your goal is to analyze the user's task and generate a specific solution based on the thought template. If the instantiated solution involves Python code, only provide the code and let the compiler handle it. If the solution does not involve code, provide a final answer that is easy to extract from the text.
        """

        """
        你是问题分析方面的专家，可以将以前解决问题的方法应用于新问题。用户将提供具体的任务描述和思想模板。您的目标是分析用户的任务，并根据思维模板生成特定的解决方案。如果实例化的解决方案涉及Python代码，则只提供代码并让编译器处理。如果解决方案不涉及代码，则提供易于从文本中提取的最终答案。
        """
        # 这里输入了蒸馏问题，还有思维的模板 {self.distilled_information} 用来做系统角色，{self.thought_template}具体问题的模板。
        self.formated_input = f"""
Distilled information:
{self.distilled_information}
Thought template:
{self.thought_template}

Instantiated Solution:
Please analyze the above user task description and thought template, and generate a specific, detailed solution. If the solution involves Python code, only provide the code. If not, provide a clear and extractable final answer.        
        """
        """
        请分析以上用户任务描述和思想模板，并生成具体、详细的解决方案。如果解决方案涉及Python代码，则仅提供代码。如果没有，请提供一个清晰可提取的最终答案。
        """
        # 写入提示词，然后拿到结果。
        self.result = self.pipeline.get_respond(self.instantiation_instruct,self.formated_input)
        print(f'Instantiated reasoning result: {self.result}')
        # 如果问题id在列表中，则执行代码。有可可能有python代码。
        if self.problem_id in problem_id_list:
            self.final_result, code_str = extract_and_execute_code(self.result)
            print(f'The result of code execution: {self.final_result}')
        else:
            self.final_result = self.result 

    
    def bot_run(self):
        # 执行问题蒸馏，缓冲，
        self.problem_distillation()
        self.buffer_retrieve()
        self.reasoner_instantiation()
        return self.final_result