large_index.py

"""Build a queryable SQLite index from FolderVisualizer scan-result.json."""

from __future__ import annotations

import argparse
import json
import os
import sqlite3
import time
from collections import Counter
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


PROJECT_DIR = Path(__file__).resolve().parent
DEFAULT_INPUT_PATH = PROJECT_DIR / "outputs" / "scan-result.json"
DEFAULT_DATABASE_PATH = PROJECT_DIR / "outputs" / "folder-index.db"
DEFAULT_MANIFEST_PATH = PROJECT_DIR / "outputs" / "manifest.json"
SCHEMA_VERSION = 1
DEFAULT_BATCH_SIZE = 1_000


@dataclass(frozen=True)
class IndexBuildResult:
    database_path: Path
    manifest_path: Path
    elapsed_seconds: float
    database_size: int
    integrity_check: str
    manifest: dict[str, Any]


def _required_nonnegative_int(data: dict[str, Any], key: str) -> int:
    value = data.get(key)
    if not isinstance(value, int) or isinstance(value, bool) or value < 0:
        raise ValueError(f'JSON 字段 "{key}" 必须是非负整数')
    return value


def _node_nonnegative_int(node: dict[str, Any], key: str, path: str) -> int:
    value = node.get(key, 0)
    if not isinstance(value, int) or isinstance(value, bool) or value < 0:
        raise ValueError(f'节点 "{path}" 的字段 "{key}" 必须是非负整数')
    return value


def _node_children(node: dict[str, Any], path: str) -> list[dict[str, Any]]:
    children = node.get("children", [])
    if not isinstance(children, list) or not all(
        isinstance(child, dict) for child in children
    ):
        raise ValueError(f'节点 "{path}" 的 children 必须是节点列表')
    return children


def _remove_sqlite_artifacts(path: Path) -> None:
    for suffix in ("", "-journal", "-wal", "-shm"):
        Path(f"{path}{suffix}").unlink(missing_ok=True)


def _create_schema(connection: sqlite3.Connection) -> None:
    connection.executescript(
        """
        CREATE TABLE metadata (
            key TEXT PRIMARY KEY,
            value TEXT NOT NULL
        );

        CREATE TABLE nodes (
            id INTEGER PRIMARY KEY,
            parent_id INTEGER REFERENCES nodes(id),
            name TEXT NOT NULL,
            normalized_name TEXT NOT NULL,
            path TEXT NOT NULL,
            node_type TEXT NOT NULL CHECK (node_type IN ('file', 'directory')),
            extension TEXT NOT NULL,
            size INTEGER NOT NULL CHECK (size >= 0),
            depth INTEGER NOT NULL CHECK (depth >= 0),
            order_index INTEGER NOT NULL CHECK (order_index >= 0),
            child_count INTEGER NOT NULL CHECK (child_count >= 0),
            scan_error TEXT
        );

        CREATE TABLE extension_statistics (
            extension TEXT PRIMARY KEY,
            file_count INTEGER NOT NULL CHECK (file_count >= 0)
        );
        """
    )


def _create_indexes(connection: sqlite3.Connection) -> None:
    connection.executescript(
        """
        CREATE INDEX idx_nodes_parent_order
            ON nodes(parent_id, order_index);
        CREATE INDEX idx_nodes_normalized_name
            ON nodes(normalized_name);
        CREATE INDEX idx_nodes_extension
            ON nodes(extension);
        CREATE INDEX idx_nodes_node_type
            ON nodes(node_type);
        """
    )


def _validate_database(
    connection: sqlite3.Connection,
    expected_file_count: int,
    expected_directory_count: int,
) -> tuple[int, int, int, str]:
    node_count = connection.execute("SELECT COUNT(*) FROM nodes").fetchone()[0]
    file_count = connection.execute(
        "SELECT COUNT(*) FROM nodes WHERE node_type = 'file'"
    ).fetchone()[0]
    directory_count = connection.execute(
        "SELECT COUNT(*) FROM nodes WHERE node_type = 'directory'"
    ).fetchone()[0]

    if node_count != file_count + directory_count:
        raise ValueError("数据库节点总数不等于文件数与目录数之和")
    if file_count != expected_file_count:
        raise ValueError(
            f"数据库文件数 {file_count} 与 JSON 文件数 {expected_file_count} 不一致"
        )
    if directory_count != expected_directory_count:
        raise ValueError(
            "数据库目录数 "
            f"{directory_count} 与 JSON 目录数 {expected_directory_count} 不一致"
        )

    root_count = connection.execute(
        "SELECT COUNT(*) FROM nodes WHERE parent_id IS NULL"
    ).fetchone()[0]
    if root_count != 1:
        raise ValueError(f"数据库必须且只能包含一个根节点，当前为 {root_count}")

    missing_parent_count = connection.execute(
        """
        SELECT COUNT(*)
        FROM nodes AS child
        LEFT JOIN nodes AS parent ON parent.id = child.parent_id
        WHERE child.parent_id IS NOT NULL AND parent.id IS NULL
        """
    ).fetchone()[0]
    if missing_parent_count:
        raise ValueError(f"数据库存在 {missing_parent_count} 个无有效父节点的节点")

    extension_total = connection.execute(
        "SELECT COALESCE(SUM(file_count), 0) FROM extension_statistics"
    ).fetchone()[0]
    if extension_total != file_count:
        raise ValueError(
            f"扩展名统计合计 {extension_total} 与文件数 {file_count} 不一致"
        )

    integrity_check = connection.execute("PRAGMA integrity_check").fetchone()[0]
    if integrity_check != "ok":
        raise ValueError(f"SQLite 完整性检查失败: {integrity_check}")

    return node_count, file_count, directory_count, integrity_check


def build_large_index(
    input_path: Path = DEFAULT_INPUT_PATH,
    database_path: Path = DEFAULT_DATABASE_PATH,
    manifest_path: Path = DEFAULT_MANIFEST_PATH,
    batch_size: int = DEFAULT_BATCH_SIZE,
) -> IndexBuildResult:
    """Build and validate the SQLite index, then atomically publish its files."""

    if batch_size <= 0:
        raise ValueError("batch_size 必须大于 0")

    input_path = input_path.resolve()
    database_path = database_path.resolve()
    manifest_path = manifest_path.resolve()
    database_path.parent.mkdir(parents=True, exist_ok=True)
    manifest_path.parent.mkdir(parents=True, exist_ok=True)

    temp_database_path = database_path.with_name(f"{database_path.name}.tmp")
    temp_manifest_path = manifest_path.with_name(f"{manifest_path.name}.tmp")
    _remove_sqlite_artifacts(temp_database_path)
    temp_manifest_path.unlink(missing_ok=True)

    started_at = time.perf_counter()
    connection: sqlite3.Connection | None = None

    try:
        source_stat = input_path.stat()
        with input_path.open("r", encoding="utf-8") as input_file:
            data = json.load(input_file)

        if not isinstance(data, dict) or not isinstance(data.get("root"), dict):
            raise ValueError("JSON 中缺少有效的 root 节点")

        root = data["root"]
        expected_file_count = _required_nonnegative_int(data, "file_count")
        expected_directory_count = _required_nonnegative_int(
            data, "directory_count"
        )
        expected_total_size = _required_nonnegative_int(data, "total_size")
        root_path = root.get("path")
        if not isinstance(root_path, str):
            raise ValueError("根节点 path 必须是字符串")
        if _node_nonnegative_int(root, "size", root_path) != expected_total_size:
            raise ValueError("根节点大小与 JSON total_size 不一致")

        connection = sqlite3.connect(temp_database_path)
        connection.execute("PRAGMA foreign_keys = ON")
        connection.execute("PRAGMA journal_mode = MEMORY")
        connection.execute("PRAGMA synchronous = OFF")
        connection.execute("PRAGMA temp_store = MEMORY")
        connection.execute("BEGIN IMMEDIATE")
        _create_schema(connection)

        insert_sql = """
            INSERT INTO nodes (
                id, parent_id, name, normalized_name, path, node_type,
                extension, size, depth, order_index, child_count, scan_error
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """
        pending_rows: list[tuple[Any, ...]] = []
        extension_counts: Counter[str] = Counter()
        file_count = 0
        directory_count = 0
        max_depth = 0
        max_children = 0
        next_id = 1
        stack: list[tuple[dict[str, Any], int | None, int, int]] = [
            (root, None, 0, 0)
        ]

        while stack:
            node, parent_id, depth, order_index = stack.pop()
            node_id = next_id
            next_id += 1

            name = node.get("name")
            path = node.get("path")
            node_type = node.get("type")
            if not isinstance(name, str):
                raise ValueError(f"节点 {node_id} 的 name 必须是字符串")
            if not isinstance(path, str):
                raise ValueError(f"节点 {node_id} 的 path 必须是字符串")
            if node_type not in ("file", "directory"):
                raise ValueError(f'节点 "{path}" 的 type 必须是 file 或 directory')

            children = _node_children(node, path)
            if node_type == "file" and children:
                raise ValueError(f'文件节点 "{path}" 不能包含子节点')

            size = _node_nonnegative_int(node, "size", path)
            child_count = len(children)
            extension = Path(name).suffix.casefold() if node_type == "file" else ""
            scan_error = node.get("error")
            if scan_error is not None and not isinstance(scan_error, str):
                scan_error = str(scan_error)

            pending_rows.append(
                (
                    node_id,
                    parent_id,
                    name,
                    name.casefold(),
                    path,
                    node_type,
                    extension,
                    size,
                    depth,
                    order_index,
                    child_count,
                    scan_error,
                )
            )

            if node_type == "file":
                file_count += 1
                extension_counts[extension] += 1
            else:
                directory_count += 1
            max_depth = max(max_depth, depth)
            max_children = max(max_children, child_count)

            for child_order in range(child_count - 1, -1, -1):
                stack.append(
                    (children[child_order], node_id, depth + 1, child_order)
                )

            if len(pending_rows) >= batch_size:
                connection.executemany(insert_sql, pending_rows)
                pending_rows.clear()

        if pending_rows:
            connection.executemany(insert_sql, pending_rows)

        if file_count != expected_file_count:
            raise ValueError(
                f"树结构文件数 {file_count} 与 JSON 文件数 {expected_file_count} 不一致"
            )
        if directory_count != expected_directory_count:
            raise ValueError(
                "树结构目录数 "
                f"{directory_count} 与 JSON 目录数 {expected_directory_count} 不一致"
            )
        if sum(extension_counts.values()) != file_count:
            raise ValueError("扩展名统计合计与文件数不一致")

        connection.executemany(
            "INSERT INTO extension_statistics (extension, file_count) VALUES (?, ?)",
            sorted(extension_counts.items()),
        )

        generated_at = datetime.now(timezone.utc).isoformat()
        metadata = {
            "schema_version": str(SCHEMA_VERSION),
            "generated_at": generated_at,
            "source_path": str(input_path),
            "source_size": str(source_stat.st_size),
            "source_mtime_ns": str(source_stat.st_mtime_ns),
            "root_id": "1",
        }
        connection.executemany(
            "INSERT INTO metadata (key, value) VALUES (?, ?)", metadata.items()
        )
        _create_indexes(connection)

        node_count, database_file_count, database_directory_count, _ = (
            _validate_database(
                connection, expected_file_count, expected_directory_count
            )
        )
        connection.commit()
        integrity_check = connection.execute(
            "PRAGMA integrity_check"
        ).fetchone()[0]
        if integrity_check != "ok":
            raise ValueError(f"SQLite 完整性检查失败: {integrity_check}")
        connection.close()
        connection = None

        manifest: dict[str, Any] = {
            "schema_version": SCHEMA_VERSION,
            "generated_at": generated_at,
            "source": {
                "path": str(input_path),
                "size": source_stat.st_size,
                "mtime_ns": source_stat.st_mtime_ns,
            },
            "root_id": 1,
            "root_path": root_path,
            "node_count": node_count,
            "file_count": database_file_count,
            "directory_count": database_directory_count,
            "total_size": expected_total_size,
            "max_depth": max_depth,
            "max_children": max_children,
            "extension_counts": dict(sorted(extension_counts.items())),
        }
        with temp_manifest_path.open("w", encoding="utf-8") as output_file:
            json.dump(manifest, output_file, ensure_ascii=False, indent=2)
            output_file.write("\n")

        database_size = temp_database_path.stat().st_size
        os.replace(temp_database_path, database_path)
        os.replace(temp_manifest_path, manifest_path)
        elapsed_seconds = time.perf_counter() - started_at
        return IndexBuildResult(
            database_path=database_path,
            manifest_path=manifest_path,
            elapsed_seconds=elapsed_seconds,
            database_size=database_size,
            integrity_check=integrity_check,
            manifest=manifest,
        )
    except BaseException:
        if connection is not None:
            connection.rollback()
            connection.close()
        _remove_sqlite_artifacts(temp_database_path)
        temp_manifest_path.unlink(missing_ok=True)
        raise


def _format_bytes(size: int) -> str:
    units = ("B", "KB", "MB", "GB", "TB")
    value = float(size)
    for unit in units:
        if value < 1024 or unit == units[-1]:
            return f"{int(value)} {unit}" if unit == "B" else f"{value:.2f} {unit}"
        value /= 1024
    return f"{size} B"


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="根据 scan-result.json 构建大目录 SQLite 索引"
    )
    parser.add_argument("--input", type=Path, default=DEFAULT_INPUT_PATH)
    parser.add_argument("--database", type=Path, default=DEFAULT_DATABASE_PATH)
    parser.add_argument("--manifest", type=Path, default=DEFAULT_MANIFEST_PATH)
    parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE)
    return parser.parse_args()


def main() -> int:
    args = _parse_args()
    try:
        result = build_large_index(
            input_path=args.input,
            database_path=args.database,
            manifest_path=args.manifest,
            batch_size=args.batch_size,
        )
    except KeyboardInterrupt:
        print("\n索引构建已取消。")
        return 130
    except (OSError, ValueError, json.JSONDecodeError, sqlite3.Error) as error:
        print(f"错误: 无法构建大目录索引: {error}")
        return 1

    manifest = result.manifest
    print("大目录索引构建完成")
    print(f"SQLite: {result.database_path}")
    print(f"Manifest: {result.manifest_path}")
    print(f"耗时: {result.elapsed_seconds:.2f} 秒")
    print(
        f"SQLite 大小: {_format_bytes(result.database_size)} "
        f"({result.database_size:,} 字节)"
    )
    print(f"节点总数: {manifest['node_count']:,}")
    print(f"文件数量: {manifest['file_count']:,}")
    print(f"目录数量: {manifest['directory_count']:,}")
    print(f"最大深度: {manifest['max_depth']}")
    print(f"最大直接子节点数: {manifest['max_children']:,}")
    print(f"integrity_check: {result.integrity_check}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())