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

if __name__ == "__main__":
    vn = DeepSeekVanna()

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

    # 自然语言提问
    #question = "查询小学语文有哪些册,返回册信息的所有相关属性。"
    question = "查询每个区每个校都上传了多少课程数量，需要返回行政区名称，学校名称，上传课程数量等属性.先按行政区划排序，再按课程数量由高到低排序，只要行政区划是二道区的"
    # 开始查询
    print("开始查询...")
    # 获取完整 SQL
    sql = vn.generate_sql(question)
    print("生成的查询 SQL:\n", sql)

    # 执行SQL查询
    with PostgreSQLUtil() as db:
        sample_data = db.execute_query(sql)
        filename = "d:/导出信息.xlsx"
        save_to_excel(sample_data, filename)
        # 在代码最后添加自动打开逻辑
        if platform.system() == "Windows":
            try:
                full_path = os.path.abspath(filename)
                print(f"\n✅ 文件已保存到：{full_path}")
                os.startfile(full_path)  # 关键代码
            except Exception as e:
                print(f"\n⚠️ 自动打开失败: {str(e)}，请手动打开文件")
        else:
            print("\n⚠️ 非Windows系统，请手动打开文件")