import sqlite3

import psutil


def ScanSpecialVolumes(volume_letter):
    """
    扫描指定的单个磁盘卷，返回其基本信息字典。

    参数:
        volume_letter: str，磁盘盘符（例如 "C", "Z"）

    返回:
        dict: 包含 Path, Type, Option 字段的字典
    """
    # 简单校验盘符格式（去除可能的冒号）
    if len(volume_letter) >= 1 and volume_letter[-1] == ":":
        volume_letter = volume_letter[:-1]

    if not volume_letter or len(volume_letter) != 1 or not volume_letter.isalpha():
        raise ValueError("无效的磁盘盘符，应为单个字母，如 'C' 或 'Z'")

    return {
        "Path": volume_letter.upper(),
        "Type": "磁盘",
        "Option": None
    }


def ScanNTFSVolumes():
    """
    扫描当前系统中所有 NTFS 格式的磁盘卷。

    返回:
        list of dict: 包含盘符等信息的字典列表
    """
    ntfs_volumes = []

    for partition in psutil.disk_partitions():
        if partition.fstype.upper() == 'NTFS':
            # 提取盘符（去掉冒号）
            drive_letter = partition.device[0] if len(partition.device) >= 2 and partition.device[1] == ':' else None
            if drive_letter:
                ntfs_volumes.append({
                    "Path": drive_letter,
                    "Type": "磁盘",
                    "Option": None
                })

    return ntfs_volumes


def InsertVolumesToDB(data, db_path='../src/db_ntfs_info.db', table_name='db_device'):
    """
    将 NTFS 磁盘信息写入数据库表，并防止重复插入。

    参数:
        data: list of dict，包含 Path, Type, Option 字段的数据
        db_path: str，SQLite 数据库路径
        table_name: str，目标表名

    返回:
        int: 成功插入的记录数
    """
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    try:
        # 创建表（如果不存在），并添加 Path 的唯一性约束
        create_table_sql = f"""
        CREATE TABLE IF NOT EXISTS {table_name} (
            ID INTEGER PRIMARY KEY AUTOINCREMENT,
            Path TEXT NOT NULL UNIQUE,
            Type TEXT NOT NULL CHECK(Type IN ('磁盘', '文件', '文件夹')),
            Option TEXT
        );
        """
        cursor.execute(create_table_sql)

        inserted_count = 0
        insert_sql = f"""
        INSERT OR IGNORE INTO {table_name} (Path, Type, Option)
        VALUES (?, ?, ?)
        """

        for item in data:
            cursor.execute(insert_sql, (
                item['Path'],
                item['Type'],
                item['Option']
            ))
            if cursor.rowcount > 0:
                inserted_count += 1

        conn.commit()
        print(f"✅ 成功插入 {inserted_count} 条 NTFS 磁盘信息")
        return inserted_count

    except Exception as e:
        print(f"❌ 数据库操作失败: {e}")
        conn.rollback()
        return 0

    finally:
        conn.close()


def main():
    # 扫描系统下所有 NTFS 磁盘
    # volumes = ScanNTFSVolumes()
    # print("🔍 找到以下 NTFS 磁盘：")
    # for vol in volumes:
    #     print(vol)
    #
    # success_count = InsertVolumesToDB(volumes)
    # print(f"共插入 {success_count} 条记录到数据库。")

    # 扫描单个磁盘
    volume_latter = "Z"
    device_data = ScanSpecialVolumes(volume_latter)
    InsertVolumesToDB([device_data])


if __name__ == "__main__":
    main()