tiku/tiku-server/app/services/indexer.py

"""增量索引 + FTS + 统计模块"""

import os
import shutil
import time
from datetime import datetime
from pathlib import Path

from app.database import (
    get_connection, upsert_problem, delete_problem, update_fts_body,
    get_all_problems, compute_checksum,
)
from app.services.docx_parser import read_custom_props, extract_plain_text
from app.services.knowledge import load_knowledge_map, get_chapter_from_code


def scan_problem_files(problems_dir: Path) -> dict[str, dict]:
    """扫描 problems/ 目录下所有 docx 文件"""
    files = {}
    problems_dir = Path(problems_dir)
    if not problems_dir.exists():
        return files

    for docx_file in problems_dir.glob("*.docx"):
        problem_id = docx_file.stem  # 文件名即 ID
        mtime = docx_file.stat().st_mtime
        checksum = compute_checksum(docx_file)
        files[problem_id] = {
            "file_path": str(docx_file),
            "mtime": mtime,
            "checksum": checksum,
        }
    return files


def parse_docx_metadata(docx_path: str | Path, knowledge_map: dict) -> dict | None:
    """从 docx 自定义属性读取元数据"""
    props = read_custom_props(docx_path)
    if not props:
        return None

    import json
    knowledge_codes_str = props.get("knowledge_codes", "[]")
    try:
        knowledge_codes = json.loads(knowledge_codes_str) if isinstance(knowledge_codes_str, str) else knowledge_codes_str
    except json.JSONDecodeError:
        knowledge_codes = []

    chapters = []
    for code in knowledge_codes:
        ch = get_chapter_from_code(code, knowledge_map)
        if ch:
            chapters.append(ch)

    return {
        "difficulty": float(props.get("difficulty", 0.5)),
        "question_type": str(props.get("question_type", "")),
        "import_type": str(props.get("import_type", "精选题")),
        "knowledge_codes": knowledge_codes,
        "chapters": sorted(set(chapters)),
        "comment": str(props.get("comment", "")),
        "source": str(props.get("source", "")),
    }


def incremental_update(config) -> list[str]:
    """增量索引：扫描 → 对比 → 更新数据库 + FTS"""
    changes = []
    conn = get_connection(config.db_path)
    knowledge_map = load_knowledge_map(config)

    try:
        # 扫描磁盘
        disk_files = scan_problem_files(config.problems_dir)

        # 获取数据库中记录
        db_rows = conn.execute("SELECT id, file_path, mtime, checksum FROM problems").fetchall()
        db_map = {r["id"]: {"file_path": r["file_path"], "mtime": r["mtime"], "checksum": r["checksum"]} for r in db_rows}

        # 新增/修改
        for pid, info in disk_files.items():
            if pid not in db_map:
                # 新增
                meta = parse_docx_metadata(info["file_path"], knowledge_map)
                if meta:
                    problem = {
                        "id": pid,
                        "file_path": info["file_path"],
                        "mtime": info["mtime"],
                        "checksum": info["checksum"],
                        **meta,
                    }
                    upsert_problem(conn, problem)
                    # 更新 FTS
                    body = extract_plain_text(info["file_path"])
                    update_fts_body(conn, pid, body)
                    changes.append(f"+ 新增: {pid}")
            else:
                # 已存在，检查变更
                db_info = db_map[pid]
                if info["checksum"] != db_info["checksum"]:
                    meta = parse_docx_metadata(info["file_path"], knowledge_map)
                    if meta:
                        problem = {
                            "id": pid,
                            "file_path": info["file_path"],
                            "mtime": info["mtime"],
                            "checksum": info["checksum"],
                            **meta,
                        }
                        upsert_problem(conn, problem)
                        body = extract_plain_text(info["file_path"])
                        update_fts_body(conn, pid, body)
                        changes.append(f"~ 修改: {pid}")
                elif info["mtime"] != db_info["mtime"]:
                    # 仅 mtime 变，checksum 不变 → 重索引 FTS
                    body = extract_plain_text(info["file_path"])
                    update_fts_body(conn, pid, body)
                    changes.append(f"* FTS 更新: {pid}")

        # 删除（磁盘不存在但数据库有）
        for pid, db_info in db_map.items():
            if pid not in disk_files:
                changes.append(f"! 磁盘缺失: {pid} ({db_info['file_path']})")
                # 不自动删除，只提示

        conn.commit()
    finally:
        conn.close()

    return changes


def ensure_index(config) -> list[str]:
    """确保索引存在：如果 knowledge_map 或 DB 不存在，则初始化"""
    result = []
    if not config.knowledge_map_path.exists():
        from app.services.knowledge import build_knowledge
        build_knowledge(config)
        result.append("knowledge_map 已构建")
    if not config.db_path.exists():
        from app.database import init_db
        init_db(config.db_path)
        result.append("数据库已初始化")
    return result


def show_stats(config):
    """终端显示统计"""
    from app.database import get_connection, count_by_question_type, count_by_import_type, count_by_chapter

    conn = get_connection(config.db_path)
    try:
        total_row = conn.execute("SELECT COUNT(*) FROM problems").fetchone()
        total = total_row[0] if total_row else 0
        print(f"总题数: {total}")

        print("\n按题型:")
        for t, c in count_by_question_type(conn).items():
            print(f"  {t}: {c}")

        print("\n按入库类型:")
        for t, c in count_by_import_type(conn).items():
            print(f"  {t}: {c}")

        print("\n按章节:")
        for ch, c in sorted(count_by_chapter(conn).items()):
            print(f"  {ch}: {c}")
    finally:
        conn.close()