import os
from openai import OpenAI
from pathlib import Path
import re
import json

API_KEY = 'sk-15ecf7e273ad4b729c7f7f42b542749e'

msg = """
你是一名经验丰富的通信协议分析师。你拥有成熟的协议文档阅读能力和协议文档分析能力。 1.你需要运用自己的阅读能力理解上面这些文档，这些文档都属于型号为XA-5D的无人机系统； 2.文档中描述的协议整体结构划分为：型号下包含设备，设备下包含数据流，数据流下包含探测帧，探测帧包含数据包，数据包下面包含具体的参数结构； 3.所有信息来源请依据文档进行分析，不要使用网络内容和自主生成； 4.使用树形结构的JSON将你理解的结构进行输出，多条信息使用数组包含对象的格式，子信息包含在key为child的数组中，子层结构与父层相同； 5.不要输出其他文字，直接输出对应的json；请输出文档中描述的探测帧结构；
"""


class LangFlow:
    files = []
    file_objects = []

    def __init__(self, doc_files):
        self.client = OpenAI(
            api_key=API_KEY,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
        if doc_files:
            self.files = doc_files
        self.load_file_objs()
        self.delete_all_files()
        self.upload_files()

    def load_file_objs(self):
        file_stk = self.client.files.list()
        self.file_objects = file_stk.data

    def delete_all_files(self):
        for file_object in self.file_objects:
            self.client.files.delete(file_object.id)

    def upload_file(self, file_path):
        file_object = self.client.files.create(file=Path(file_path), purpose="file-extract")
        return file_object

    def upload_files(self):
        self.file_objects = []
        for file_path in self.files:
            file_object = self.upload_file(file_path)
            self.file_objects.append(file_object)

    def run(self):
        result = self.tm_frame_gen()
        result += self.tm_pkts_gen()
        return result

    def _gen(self, msgs, msg):
        messages = [] if msgs is None else msgs
        if len(messages) == 0:
            # 如果是第一次提问加入文档
            for file_object in self.file_objects:
                messages.append({'role': 'system', 'content': 'fileid://' + file_object.id})
        messages.append({'role': 'user', 'content': msg})
        completion = self.client.chat.completions.create(
            model="qwen-long",
            messages=messages,
            stream=False,
            temperature=0.0,
            top_p=0,
            # stream_options={"include_usage": True}
        )
        return completion

    def tm_frame_gen(self):
        result = ""
        messages = []
        _msg = """
你是一名经验丰富的通信协议分析师。你拥有成熟的协议文档阅读能力和协议文档分析能力。
1.你需要运用自己的阅读能力理解上面这些文档，这些文档都属于型号为XA-5D的无人机系统；
2.文档中描述的协议整体结构划分为：型号下包含设备，设备下包含数据流，数据流下包含探测帧，探测帧包含数据包，数据包下面包含具体的参数结构；
3.所有信息来源请依据文档进行分析，不要使用网络内容和自主生成；
4.使用树形结构的JSON将你理解的结构进行输出，多条信息使用数组包含对象的格式，子信息包含在key为child的数组中，子层结构与父层相同；
5.不要输出其他文字，直接输出对应的json；
请输出文档中描述的探测帧结构
        """
        completion = self._gen(messages, _msg)
        full_content = self.remove_markdown(completion.choices[0].message.content)
        print(full_content)
        messages.append({'role': 'assistant', 'content': full_content})
        _msg2 = """
请阅读文档，给上面json加入以下字段：
1.type，例如设备层type为device，
2.length，例如8bit，
仅仅输出json
        """
        completion = self._gen(messages, _msg2)
        messages.append({'role': 'user', 'content': _msg2})
        full_content = self.remove_markdown(completion.choices[0].message.content)
        print('探测帧结构：'+full_content)
        result += '探测帧结构：'
        result += full_content
        return result

    def tm_pkts_gen(self):
        result = ""
        messages = []
        _msg = """
你是一名经验丰富的通信协议分析师。你拥有成熟的协议文档阅读能力和协议文档分析能力。
1.你需要运用自己的阅读能力理解上面这些文档，这些文档都属于型号为XA-5D的无人机系统；
2.所有信息来源请依据文档进行分析，不要使用网络内容和自主生成；
3.不要输出其他文字，直接输出对应的json；
4.请输出探测帧下探测参数源包列表，仅输出包名称数组。
        """
        completion = self._gen(messages, _msg)
        full_content = completion.choices[0].message.content
        full_content = self.remove_markdown(full_content)
        print(f'所有探测帧参数源包：')
        print(full_content)
        result += '所有探测帧参数源包：'
        result += full_content
        # 解析json
        pkts = json.loads(full_content)
        for pkt in pkts:
            pkt_json = self.tm_pkt_gen(pkt)
            print(f'数据包：{pkt}')
            print(pkt_json)
            result += f'数据包：{pkt}\n'
            result += pkt_json
        return result

    def tm_pkt_gen(self, pkt_name):
        messages = []
        _msg = f"""
你是一名经验丰富的通信协议分析师。你拥有成熟的协议文档阅读能力和协议文档分析能力。 
1.你需要运用自己的阅读能力理解上面这些文档，这些文档都属于型号为XA-5D的无人机系统； 
2.所有信息来源请依据文档进行分析，不要使用网络内容和自主生成； 
3.使用树形结构的JSON将你理解的结构进行输出，多条信息使用数组包含对象的格式，子信息包含在key为child的数组中，子层结构与父层相同； 
4.不要输出其他文字，直接输出对应的json；
请输出文档中描述的{pkt_name}的结构的各个层级
        """
        completion = self._gen(messages, _msg)
        return self.remove_markdown(completion.choices[0].message.content)

    def remove_markdown(self, text):
        # 去掉开头的```json
        text = re.sub(r'^```json', '', text)
        # 去掉结尾的```json
        text = re.sub(r'```$', '', text)
        return text
# if __name__ == '__main__':
#     text = ChatFlow().run()
#     print(text)