import time
from pipeline.claim_generate import * 
from pipeline.query_generate import *
from pipeline.tool_execute import *
from pipeline.judge import *
from pipeline.openai_wrapper import *

class Pipeline:
    def __init__(self, type, api_key, base_url=None):
        # 全局只实例化一个对象 会不会干扰prompt的结果
        self.syncchat = SyncChat(model="gpt-4-1106-preview", api_key=api_key, base_url=base_url)
        self.asyncchat = AsyncChat(model="gpt-4-1106-preview", api_key=api_key, base_url=base_url)
        self.visionchat = VisionChat(model="gpt-4-vision-preview", api_key=api_key, base_url=base_url)

        self.claim_generator = ClaimGenerator(prompt_path="pipeline/prompts/claim_generate.yaml",chat=self.syncchat)
        self.query_generator = QueryGenerator(prompt_path="pipeline/prompts/query_generate.yaml",chat=self.asyncchat)
        self.tool = Tool()
        self.judger = Judger(prompt_path="pipeline/prompts/verification.yaml", chat=self.visionchat)

    def run(self, text, image_path, type, use_attribue=False):
        time1 = time.time()
        response, claim_list = self.claim_generator.get_response(text=text)
        time2 = time.time()
        if use_attribue:
            objects, attribute_ques_list, scenetext_ques_list, fact_ques_list = self.query_generator.get_response(claim_list=claim_list)
        else:
            objects, scenetext_ques_list, fact_ques_list = self.query_generator.get_response(claim_list=claim_list,type=type)
            attribute_ques_list = None
        time3 = time.time()
        print(objects)
        print(attribute_ques_list)
        print(scenetext_ques_list)
        print(fact_ques_list)
        """
        需要定时清除文件
        """
        object_res, attribue_res, text_res, fact_res = self.tool.execute(image_path=image_path,
                                                            new_path="pipeline/cache_files/", 
                                                            objects=objects,
                                                            attribute_list=attribute_ques_list, 
                                                            scenetext_list=scenetext_ques_list, 
                                                            fact_list=fact_ques_list)
        time4 = time.time()
        response = self.judger.get_response(type, object_res, attribue_res, text_res, fact_res, claim_list, image_path)
        time5 = time.time()
        print("claim generate time:" + str(time2-time1))
        print("query generate time:" + str(time3-time2))
        print("tool execute time:" + str(time4-time3))
        print("judge time:" + str(time5-time4))
        return response,claim_list