import json

from openai import OpenAI

from Text2Sql.Util.PostgreSQLUtil import PostgreSQLUtil
from Text2Sql.Util.SaveToExcel import save_to_excel
from Text2Sql.Util.Text2SqlUtil import *

from Util.EchartsUtil import *

'''
经验：
1、尽量使用宽表，少用关联，越少越好
2、应该有一些固定的组合用法预置出来，给出范例，让用户可以简单修改后就能使用
3、应该有类似于 保存为用例，查询历史等功能，让用户方便利旧。
'''
if __name__ == "__main__":
    vn = DeepSeekVanna()

    # 开始训练
    print("开始训练...")
    # 打开AreaSchoolLesson.sql文件内容
    with open("Sql/AreaSchoolLesson.sql", "r", encoding="utf-8") as file:
        ddl = file.read()
    # 训练数据
    vn.train(
        ddl=ddl
    )

    # 自然语言提问
    # 整体情况
    # question = '''
    #     查询:
    #     1、发布时间是2024年度
    #     2、每个行政区每个学校都上传了多少课程数量
    #     3、格式: 行政区划名,学段,排名,学校名称,课程数量
    #     '''

    # 指定行政区域
    # question = '''
    #         查询:
    #         1、发布时间是2024年度
    #         2、二道区每个学校都上传了多少课程数量
    #         3、格式: 行政区划名,学段,排名,学校名称,发布年份,课程数量
    #         '''

    # 指定学段
    question = '''
                查询:
                1、发布时间是2024年度
                2、每个学段，每个科目，上传课程数量，按由多到少排序
                3、字段名: 学段,科目,排名,课程数量
                '''
    common_prompt = '''
    返回的信息要求：
    1、行政区划为NULL 或者是空字符的不参加统计
    2、目标数据库是Postgresql 16
    '''
    question = question + common_prompt
    # 开始查询
    print("开始查询...")
    # 获取完整 SQL
    sql = vn.generate_sql(question)
    print("生成的查询 SQL:\n", sql)

    # 执行SQL查询
    with PostgreSQLUtil() as db:
        _data = db.execute_query(sql)

        # 1、生成柱状图
        generate_bar_chart(
            _data=_data,
            title="学段+科目课程数量柱状图",
            x_columns=["学段", "科目"],  # 动态指定 X 轴列
            y_columns=["课程数量"],  # 动态指定 Y 轴列
            output_file="d:/lesson_bar_chart.html"
        )
        # 2、生成饼状图
        generate_pie_chart(
            _data=_data,
            title="学段+科目分布",
            category_columns=["学段", "科目"],  # 多列组合参数
            value_column="课程数量",
            output_file="d:/lesson_pie_chart.html"
        )

        # 3、生成excel
        filename = "d:/导出信息.xlsx"
        save_to_excel(_data, filename)

        # 4、生成word报告
        # 打印格式化的 JSON 数据
        print(json.dumps(_data, ensure_ascii=False, indent=4))

        prompt='''
        请根据以下 JSON 数据，整理出一段话描述当前数据情况。要求：
        1. 按学段（初中、小学、高中、学前）分别总结课程数量的分布情况。
        2. 指出每个学段中课程数量最多的科目和课程数量最少的科目。
        3. 分析不同学段之间课程数量的差异，并指出哪些科目在不同学段中表现突出。
        4. 语言简洁明了，重点突出数据中的关键信息。
        '''

        # 初始化 OpenAI 客户端
        client = OpenAI(
            api_key=MODEL_API_KEY,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )

        # 调用 OpenAI API 生成总结
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": "你是一个数据分析助手，擅长从 JSON 数据中提取关键信息并生成简洁的总结。"},
                {"role": "user", "content": prompt}
            ],
            max_tokens=1500,  # 控制生成内容的长度
            temperature=0.7  # 控制生成内容的创造性
        )

        # 提取生成的总结
        summary = response.choices[0].message.content
        print("数据总结：\n", summary)