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)

        # 生成柱状图
        generate_bar_chart(
            _data,
            x_columns=["学段", "科目"],  # 动态指定 X 轴列
            y_columns=["课程数量"],  # 动态指定 Y 轴列
            output_file="d:/lesson_bar_chart.html"
        )

        generate_pie_chart(
            _data=_data,
            category_columns=["学段", "科目"],  # 多列组合参数
            value_column="课程数量",
            output_file="d:/lesson_pie_chart.html"
        )

        filename = "d:/导出信息.xlsx"
        save_to_excel(_data, filename)
        os.startfile(filename)