import hashlib
import os
import sqlite3


def GenerateHash(s: str) -> str:
    """
    对输入字符串生成 SHA-256 哈希值。
    用于唯一标识一个路径（PathHash）。
    """
    return hashlib.sha256(s.encode('utf-8')).hexdigest()


def ShouldSkipPath(path: str) -> bool:
    """
    判断是否应跳过该路径（NTFS元文件或系统文件夹）。
    """
    name = os.path.basename(path)
    if name.startswith('$'):
        return True
    if name == "System Volume Information":
        return True
    return False


def ScanVolume(volume_letter: str):
    """
    完整扫描指定磁盘的所有文件和目录，忽略 NTFS 元文件和系统文件夹，
    并为每个节点分配 ParentID。

    返回:
        list of dict：包含文件/目录信息的字典列表
    """
    root_path = f"{volume_letter.upper()}:\\"
    if not os.path.exists(root_path):
        raise ValueError(f"磁盘 {root_path} 不存在")

    result = []
    path_to_id = {}  # 用于记录路径到数据库 ID 的映射
    counter = 1  # 模拟数据库自增 ID

    for root, dirs, files in os.walk(root_path, topdown=True, onerror=None, followlinks=False):
        # 过滤掉需要跳过的目录
        dirs[:] = [d for d in dirs if not ShouldSkipPath(os.path.join(root, d))]

        for entry in files + dirs:
            full_path = os.path.join(root, entry)

            if ShouldSkipPath(full_path):
                continue

            try:
                if os.path.isdir(full_path):
                    is_dir = 1
                    bytes_size = 0
                elif os.path.isfile(full_path):
                    is_dir = 0
                    bytes_size = os.path.getsize(full_path)
                else:
                    continue

                name = entry

                # ✅ 修正点：对 Path 字段进行哈希
                path_hash = GenerateHash(full_path)

                # 计算 ContentSize（KB），小文件至少显示为 1 KB
                content_size = bytes_size // 1024
                if content_size == 0 and bytes_size > 0:
                    content_size = 1

                # 获取父目录路径
                parent_path = os.path.dirname(full_path)
                parent_id = path_to_id.get(parent_path, 0)  # 默认为 0（根目录可能未录入）

                item = {
                    "ID": counter,
                    "Path": full_path,
                    "Name": name,
                    "PathHash": path_hash,
                    "IsDir": is_dir,
                    "ParentID": parent_id,
                    "ContentSize": content_size
                }

                result.append(item)
                path_to_id[full_path] = counter
                counter += 1

            except Exception as e:
                print(f"⚠️ 跳过路径 {full_path}，错误: {e}")

    return result


def InsertPathDataToDB(data, db_path='../src/db_ntfs_info.db', table_name='db_path', batch_size=20):
    """
    批量将扫描结果写入数据库。
    """
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    try:
        # 创建表（如果不存在）
        create_table_sql = f"""
        CREATE TABLE IF NOT EXISTS {table_name} (
            ID INTEGER PRIMARY KEY AUTOINCREMENT,
            Path TEXT NOT NULL,
            Name TEXT NOT NULL,
            PathHash TEXT UNIQUE NOT NULL,
            IsDir INTEGER NOT NULL CHECK(IsDir IN (0, 1)),
            ParentID INTEGER,
            ContentSize INTEGER,

            FOREIGN KEY(ParentID) REFERENCES {table_name}(ID)
        );
        """
        cursor.execute(create_table_sql)

        # 插入语句（忽略重复 PathHash）
        insert_sql = f"""
        INSERT OR IGNORE INTO {table_name} 
        (Path, Name, PathHash, IsDir, ParentID, ContentSize)
        VALUES (?, ?, ?, ?, ?, ?)
        """

        total_inserted = 0
        batch = []

        for item in data:
            batch.append((
                item['Path'],
                item['Name'],
                item['PathHash'],
                item['IsDir'],
                item['ParentID'] or 0,
                item['ContentSize']
            ))

            if len(batch) >= batch_size:
                cursor.executemany(insert_sql, batch)
                conn.commit()
                total_inserted += cursor.rowcount
                print(f"✅ 提交一批 {len(batch)} 条数据")
                batch.clear()

        # 插入剩余数据
        if batch:
            cursor.executemany(insert_sql, batch)
            conn.commit()
            total_inserted += cursor.rowcount
            print(f"✅ 提交最后一批 {len(batch)} 条数据")

        print(f"✅ 总共插入 {total_inserted} 条记录到数据库。")

    except Exception as e:
        print(f"❌ 插入失败: {e}")
        conn.rollback()

    finally:
        conn.close()


# 示例主函数
def main():
    volume_letter = "Z"

    print(f"🔍 开始全盘扫描磁盘 {volume_letter}:\\ ...")
    scanned_data = ScanVolume(volume_letter)

    print(f"📊 共扫描到 {len(scanned_data)} 条有效记录，开始入库...")
    InsertPathDataToDB(scanned_data)

    print("✅ 全盘扫描与入库完成")


if __name__ == "__main__":
    main()