# -*- coding: utf-8 -*-
import os
import psycopg2


class AlphaTxtImporter:
    def __init__(self, txt_file='alpha.txt'):
        self.database_name = "alpha"
        self.txt_file = txt_file
        self.connection = None

    def get_connection(self):
        """获取数据库连接"""
        if self.connection is None or self.connection.closed:
            self.connection = psycopg2.connect(
                host="localhost",
                port="5432",
                database=self.database_name,
                user="jack",
                password="aaaAAA111"
            )
        return self.connection

    def close_connection(self):
        """关闭数据库连接"""
        if self.connection and not self.connection.closed:
            self.connection.close()

    def load_alpha_from_txt(self):
        """从txt文件加载alpha表达式"""
        if not os.path.exists(self.txt_file):
            print(f"未找到 {self.txt_file} 文件")
            return []

        alpha_list = []
        with open(self.txt_file, 'r', encoding='utf-8') as file:
            for line_num, line in enumerate(file, 1):
                line = line.strip()
                if line and not line.startswith('#'):  # 跳过空行和注释行
                    alpha_list.append(line)
                    print(f"第 {line_num} 行: {line}")

        return alpha_list

    def import_to_database(self, alpha_list):
        """导入alpha表达式到数据库"""
        if not alpha_list:
            print("没有要导入的alpha表达式")
            return 0

        conn = self.get_connection()
        cursor = conn.cursor()

        imported_count = 0
        duplicate_count = 0

        for alpha in alpha_list:
            try:
                # 使用 ON CONFLICT 避免重复插入
                cursor.execute('''
                               INSERT INTO alpha_prepare (alpha, unused)
                               VALUES (%s, %s) ON CONFLICT (alpha) DO NOTHING
                               ''', (alpha, True))

                if cursor.rowcount > 0:
                    imported_count += 1
                    print(f"✅ 导入: {alpha}")
                else:
                    duplicate_count += 1
                    print(f"⏭️  跳过(重复): {alpha}")

            except Exception as e:
                print(f"❌ 导入失败: {alpha} - 错误: {e}")
                continue

        conn.commit()
        cursor.close()

        return imported_count, duplicate_count

    def check_existing_data(self):
        """检查数据库中已存在的数据"""
        conn = self.get_connection()
        cursor = conn.cursor()

        cursor.execute("SELECT COUNT(*) FROM alpha_prepare")
        total_count = cursor.fetchone()[0]

        cursor.execute("SELECT COUNT(*) FROM alpha_prepare WHERE unused = TRUE")
        unused_count = cursor.fetchone()[0]

        cursor.close()

        print(f"数据库统计:")
        print(f"  - 总alpha数量: {total_count}")
        print(f"  - 未使用的alpha数量: {unused_count}")

        return total_count, unused_count

    def run_import(self):
        """执行导入过程"""
        print(f"开始从 {self.txt_file} 导入alpha表达式到 PostgreSQL...")

        # 检查现有数据
        print("\n=== 数据库当前状态 ===")
        total_count, unused_count = self.check_existing_data()

        # 加载txt文件
        print(f"\n=== 加载 {self.txt_file} ===")
        alpha_list = self.load_alpha_from_txt()
        print(f"从文件加载了 {len(alpha_list)} 个alpha表达式")

        if not alpha_list:
            print("没有找到有效的alpha表达式，导入终止")
            return

        # 导入到数据库
        print(f"\n=== 导入到数据库 ===")
        imported_count, duplicate_count = self.import_to_database(alpha_list)

        # 显示结果
        print(f"\n=== 导入结果 ===")
        print(f"文件中的表达式数量: {len(alpha_list)}")
        print(f"成功导入: {imported_count} 条")
        print(f"跳过重复: {duplicate_count} 条")

        # 再次检查数据库状态
        print(f"\n=== 导入后数据库状态 ===")
        new_total_count, new_unused_count = self.check_existing_data()
        print(f"新增记录: {new_total_count - total_count} 条")
        print(f"新增未使用记录: {new_unused_count - unused_count} 条")

        print("\n导入完成！")


def main():
    # 可以指定不同的txt文件
    txt_file = 'alpha.txt'  # 可以改为其他文件名

    importer = AlphaTxtImporter(txt_file)
    try:
        importer.run_import()
    finally:
        importer.close_connection()


if __name__ == "__main__":
    main()