tiku/tiku-server/app/database.py

"""数据库层 — SQLite CRUD + FTS5"""

import hashlib
import json
import sqlite3
import os
from pathlib import Path
from typing import Any


import threading

_db_lock = threading.Lock()

def get_connection(db_path: str | Path) -> sqlite3.Connection:
    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout=10000")
    conn.execute("PRAGMA foreign_keys=ON")
    return conn


def with_lock(fn):
    """装饰器：对数据库写操作加锁，防止多线程并发写入损坏数据库"""
    def wrapper(*args, **kwargs):
        with _db_lock:
            return fn(*args, **kwargs)
    return wrapper


def init_db(db_path: str | Path):
    """建表 + FTS5 + 触发器"""
    conn = get_connection(db_path)
    try:
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS problems (
                id TEXT PRIMARY KEY,
                file_path TEXT NOT NULL UNIQUE,
                mtime REAL NOT NULL,
                checksum TEXT NOT NULL,
                chapters TEXT DEFAULT '[]',
                difficulty REAL DEFAULT 0.5,
                question_type TEXT DEFAULT '',
                import_type TEXT DEFAULT '精选题',
                knowledge_codes TEXT DEFAULT '[]',
                comment TEXT DEFAULT '',
                source TEXT DEFAULT '',
                created_at TEXT DEFAULT (datetime('now','localtime')),
                updated_at TEXT DEFAULT (datetime('now','localtime'))
            );

            CREATE TABLE IF NOT EXISTS knowledge_points (
                code TEXT PRIMARY KEY,
                path TEXT NOT NULL,
                leaf_name TEXT NOT NULL,
                chapter_code TEXT,
                chapter_name TEXT
            );

            CREATE VIRTUAL TABLE IF NOT EXISTS problems_fts USING fts5(
                problem_id UNINDEXED,
                body,
                tokenize='trigram'
            );

            CREATE TABLE IF NOT EXISTS carts (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                created_at TEXT DEFAULT (datetime('now','localtime'))
            );

            CREATE TABLE IF NOT EXISTS cart_items (
                cart_id INTEGER NOT NULL REFERENCES carts(id) ON DELETE CASCADE,
                problem_id TEXT NOT NULL REFERENCES problems(id) ON DELETE CASCADE,
                added_at TEXT DEFAULT (datetime('now','localtime')),
                PRIMARY KEY (cart_id, problem_id)
            );

            CREATE TABLE IF NOT EXISTS import_tasks (
                id TEXT PRIMARY KEY,
                status TEXT DEFAULT 'pending',
                total INTEGER DEFAULT 0,
                completed INTEGER DEFAULT 0,
                message TEXT DEFAULT '',
                result TEXT DEFAULT '{}',
                created_at TEXT DEFAULT (datetime('now','localtime'))
            );
        """)
        # 确保 FTS 触发器存在（重建 FTS 表后需要重建触发器）
        conn.executescript("""
            DROP TRIGGER IF EXISTS problems_ai;
            CREATE TRIGGER problems_ai AFTER INSERT ON problems BEGIN
                INSERT INTO problems_fts(problem_id, body) VALUES (new.id, '');
            END;

            DROP TRIGGER IF EXISTS problems_ad;
            CREATE TRIGGER problems_ad AFTER DELETE ON problems BEGIN
                DELETE FROM problems_fts WHERE problem_id = old.id;
            END;

            DROP TRIGGER IF EXISTS problems_au;
            CREATE TRIGGER problems_au AFTER UPDATE ON problems BEGIN
                DELETE FROM problems_fts WHERE problem_id = old.id;
                INSERT INTO problems_fts(problem_id, body) VALUES (new.id, '');
            END;
        """)
        conn.commit()
    finally:
        conn.close()


# ── Problems CRUD ────────────────────────────────────────

def upsert_problem(conn: sqlite3.Connection, problem: dict):
    conn.execute("""
        INSERT OR REPLACE INTO problems
            (id, file_path, mtime, checksum, chapters, difficulty, question_type,
             import_type, knowledge_codes, comment, source, updated_at)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, datetime('now','localtime'))
    """, (
        problem["id"],
        problem["file_path"],
        problem.get("mtime", 0),
        problem.get("checksum", ""),
        json.dumps(problem.get("chapters", []), ensure_ascii=False),
        problem.get("difficulty", 0.5),
        problem.get("question_type", ""),
        problem.get("import_type", "精选题"),
        json.dumps(problem.get("knowledge_codes", []), ensure_ascii=False),
        problem.get("comment", ""),
        problem.get("source", ""),
    ))
    conn.commit()


def delete_problem(conn: sqlite3.Connection, problem_id: str):
    conn.execute("DELETE FROM problems WHERE id = ?", (problem_id,))
    conn.commit()


def get_problem(conn: sqlite3.Connection, problem_id: str) -> dict | None:
    row = conn.execute("SELECT * FROM problems WHERE id = ?", (problem_id,)).fetchone()
    if row is None:
        return None
    return _row_to_dict(row)


def get_all_problems(conn: sqlite3.Connection) -> list[dict]:
    rows = conn.execute("SELECT * FROM problems ORDER BY updated_at DESC").fetchall()
    return [_row_to_dict(r) for r in rows]


def search_problems(
    conn: sqlite3.Connection,
    chapter: str | None = None,
    knowledge_code: str | None = None,
    question_type: str | None = None,
    import_type: str | None = None,
    annotated: str | None = None,
    page: int = 1,
    page_size: int = 20,
) -> tuple[list[dict], int]:
    """分页查询题目列表"""
    conditions = []
    params = []

    if chapter:
        conditions.append("chapters LIKE ?")
        params.append(f'%"{chapter}"%')
    if knowledge_code:
        # Prefix match: codes like 01.01 → match "01.01.03.02.01"
        # Full leaf code has 5 segments; partial has fewer
        if knowledge_code.count('.') >= 4:  # full leaf code → exact match
            conditions.append("knowledge_codes LIKE ?")
            params.append(f'%"{knowledge_code}"%')
        else:  # partial code → prefix match
            conditions.append("knowledge_codes LIKE ?")
            params.append(f'%"{knowledge_code}.%')
    if question_type:
        conditions.append("question_type = ?")
        params.append(question_type)
    if import_type:
        conditions.append("import_type = ?")
        params.append(import_type)
    if annotated == "yes":
        conditions.append("knowledge_codes != '[]'")
    elif annotated == "no":
        conditions.append("knowledge_codes = '[]'")

    where = ("WHERE " + " AND ".join(conditions)) if conditions else ""

    # 总数
    total_row = conn.execute(f"SELECT COUNT(*) FROM problems {where}", params).fetchone()
    total = total_row[0] if total_row else 0

    # 分页
    offset = (page - 1) * page_size
    rows = conn.execute(
        f"SELECT * FROM problems {where} ORDER BY updated_at DESC LIMIT ? OFFSET ?",
        params + [page_size, offset]
    ).fetchall()

    return [_row_to_dict(r) for r in rows], total


def search_fts(conn: sqlite3.Connection, query: str, limit: int = 50) -> list[dict]:
    """FTS5 全文搜索，短查询回退 LIKE"""
    # 清理查询中的特殊字符
    safe_query = query.strip()
    if len(safe_query) < 3:
        # 短查询 LIKE 回退
        like_pattern = f"%{safe_query}%"
        rows = conn.execute(
            "SELECT * FROM problems WHERE id IN ("
            "  SELECT problem_id FROM problems_fts WHERE body LIKE ?"
            ") LIMIT ?",
            (like_pattern, limit)
        ).fetchall()
    else:
        try:
            rows = conn.execute(
                "SELECT p.* FROM problems p "
                "INNER JOIN problems_fts f ON p.id = f.problem_id "
                "WHERE f.body MATCH ? LIMIT ?",
                (safe_query, limit)
            ).fetchall()
        except sqlite3.OperationalError:
            # FTS5 查询语法错误时回退 LIKE
            like_pattern = f"%{safe_query}%"
            rows = conn.execute(
                "SELECT * FROM problems WHERE id IN ("
                "  SELECT problem_id FROM problems_fts WHERE body LIKE ?"
                ") LIMIT ?",
                (like_pattern, limit)
            ).fetchall()

    return [_row_to_dict(r) for r in rows]


def update_fts_body(conn: sqlite3.Connection, problem_id: str, body: str):
    """更新 FTS 全文索引"""
    conn.execute(
        "UPDATE problems_fts SET body = ? WHERE problem_id = ?",
        (body, problem_id)
    )
    conn.commit()


def update_problem_metadata(conn: sqlite3.Connection, problem_id: str, fields: dict):
    """更新题目元数据"""
    allowed = {"difficulty", "question_type", "import_type", "comment", "source", "knowledge_codes"}
    updates = {}
    for k, v in fields.items():
        if k in allowed:
            if k in ("knowledge_codes",):
                updates[k] = json.dumps(v, ensure_ascii=False) if isinstance(v, list) else v
            else:
                updates[k] = v

    if not updates:
        return

    set_clause = ", ".join(f"{k} = ?" for k in updates)
    values = list(updates.values()) + [problem_id]

    conn.execute(
        f"UPDATE problems SET {set_clause}, updated_at = datetime('now','localtime') WHERE id = ?",
        values
    )
    conn.commit()


def get_next_unannotated(conn: sqlite3.Connection, current_id: str | None, direction: str = "next") -> dict | None:
    """获取下一个未标注题目"""
    if direction == "next":
        if current_id:
            row = conn.execute(
                "SELECT * FROM problems WHERE knowledge_codes = '[]' AND id > ? ORDER BY id LIMIT 1",
                (current_id,)
            ).fetchone()
        else:
            row = conn.execute(
                "SELECT * FROM problems WHERE knowledge_codes = '[]' ORDER BY id LIMIT 1"
            ).fetchone()
    else:
        if current_id:
            row = conn.execute(
                "SELECT * FROM problems WHERE knowledge_codes = '[]' AND id < ? ORDER BY id DESC LIMIT 1",
                (current_id,)
            ).fetchone()
        else:
            row = conn.execute(
                "SELECT * FROM problems WHERE knowledge_codes = '[]' ORDER BY id DESC LIMIT 1"
            ).fetchone()

    return _row_to_dict(row) if row else None


# ── 统计 ─────────────────────────────────────────────────

def count_by_question_type(conn: sqlite3.Connection) -> dict:
    rows = conn.execute(
        "SELECT question_type, COUNT(*) as cnt FROM problems GROUP BY question_type"
    ).fetchall()
    return {r["question_type"] or "未知": r["cnt"] for r in rows}


def count_by_import_type(conn: sqlite3.Connection) -> dict:
    rows = conn.execute(
        "SELECT import_type, COUNT(*) as cnt FROM problems GROUP BY import_type"
    ).fetchall()
    return {r["import_type"] or "未知": r["cnt"] for r in rows}


def count_by_chapter(conn: sqlite3.Connection) -> dict:
    rows = conn.execute("SELECT chapters FROM problems").fetchall()
    chapter_count: dict[str, int] = {}
    for r in rows:
        try:
            chapters = json.loads(r["chapters"])
        except (json.JSONDecodeError, TypeError):
            chapters = []
        for ch in chapters:
            chapter_count[ch] = chapter_count.get(ch, 0) + 1
    return chapter_count


# ── 知识点 ───────────────────────────────────────────────

def sync_knowledge_points(conn: sqlite3.Connection, knowledge_map: dict):
    """同步知识点编码到数据库"""
    conn.execute("DELETE FROM knowledge_points")
    for code, info in knowledge_map.items():
        conn.execute(
            "INSERT INTO knowledge_points (code, path, leaf_name, chapter_code, chapter_name) VALUES (?, ?, ?, ?, ?)",
            (code, info["path"], info["leaf_name"], info.get("chapter_code", ""), info.get("chapter_name", ""))
        )
    conn.commit()


def get_knowledge_tree(conn: sqlite3.Connection) -> list[dict]:
    """获取全部知识点"""
    rows = conn.execute("SELECT * FROM knowledge_points ORDER BY code").fetchall()
    return [_row_to_dict(r) for r in rows]


def search_knowledge(conn: sqlite3.Connection, query: str, limit: int = 30) -> list[dict]:
    """搜索知识点"""
    rows = conn.execute(
        "SELECT * FROM knowledge_points WHERE leaf_name LIKE ? OR path LIKE ? LIMIT ?",
        (f"%{query}%", f"%{query}%", limit)
    ).fetchall()
    return [_row_to_dict(r) for r in rows]


# ── 试题篮 ───────────────────────────────────────────────

def get_carts(conn: sqlite3.Connection) -> list[dict]:
    rows = conn.execute(
        "SELECT c.*, COUNT(ci.problem_id) as item_count "
        "FROM carts c LEFT JOIN cart_items ci ON c.id = ci.cart_id "
        "GROUP BY c.id ORDER BY c.created_at DESC"
    ).fetchall()
    return [_row_to_dict(r) for r in rows]


def create_cart(conn: sqlite3.Connection, name: str) -> dict:
    cur = conn.execute("INSERT INTO carts (name) VALUES (?)", (name,))
    conn.commit()
    row = conn.execute("SELECT * FROM carts WHERE id = ?", (cur.lastrowid,)).fetchone()
    return _row_to_dict(row)


def delete_cart(conn: sqlite3.Connection, cart_id: int):
    conn.execute("DELETE FROM carts WHERE id = ?", (cart_id,))
    conn.commit()


def get_cart(conn: sqlite3.Connection, cart_id: int) -> dict | None:
    row = conn.execute("SELECT * FROM carts WHERE id = ?", (cart_id,)).fetchone()
    if row is None:
        return None
    items = conn.execute(
        "SELECT p.*, ci.added_at as cart_added_at FROM cart_items ci "
        "JOIN problems p ON ci.problem_id = p.id "
        "WHERE ci.cart_id = ? ORDER BY ci.added_at",
        (cart_id,)
    ).fetchall()
    result = _row_to_dict(row)
    item_dicts = []
    for item in items:
        d = _row_to_dict(item)
        # Add body_preview from FTS
        fts_row = conn.execute(
            "SELECT body FROM problems_fts WHERE problem_id = ?",
            (d["id"],)
        ).fetchone()
        if fts_row and fts_row["body"]:
            body = fts_row["body"]
            import re
            body = re.sub(r'^\d+[．.]\s*', '', body)
            d["body_preview"] = body[:80]
        else:
            d["body_preview"] = ""
        item_dicts.append(d)
    result["items"] = item_dicts
    return result


def add_cart_item(conn: sqlite3.Connection, cart_id: int, problem_id: str) -> bool:
    try:
        conn.execute(
            "INSERT OR IGNORE INTO cart_items (cart_id, problem_id) VALUES (?, ?)",
            (cart_id, problem_id)
        )
        conn.commit()
        return True
    except sqlite3.IntegrityError:
        return False


def remove_cart_item(conn: sqlite3.Connection, cart_id: int, problem_id: str):
    conn.execute(
        "DELETE FROM cart_items WHERE cart_id = ? AND problem_id = ?",
        (cart_id, problem_id)
    )
    conn.commit()


# ── 导入任务 ─────────────────────────────────────────────

def create_import_task(conn: sqlite3.Connection, task_id: str) -> dict:
    conn.execute("INSERT INTO import_tasks (id) VALUES (?)", (task_id,))
    conn.commit()
    row = conn.execute("SELECT * FROM import_tasks WHERE id = ?", (task_id,)).fetchone()
    return _row_to_dict(row)


def update_import_task(conn: sqlite3.Connection, task_id: str, **kwargs):
    allowed = {"status", "total", "completed", "message", "result"}
    updates = {k: v for k, v in kwargs.items() if k in allowed}
    if not updates:
        return
    set_clause = ", ".join(f"{k} = ?" for k in updates)
    values = list(updates.values()) + [task_id]
    conn.execute(f"UPDATE import_tasks SET {set_clause} WHERE id = ?", values)
    conn.commit()


def get_import_task(conn: sqlite3.Connection, task_id: str) -> dict | None:
    row = conn.execute("SELECT * FROM import_tasks WHERE id = ?", (task_id,)).fetchone()
    return _row_to_dict(row) if row else None


# ── helpers ───────────────────────────────────────────────

def _row_to_dict(row: sqlite3.Row) -> dict:
    d = dict(row)
    # 将 JSON 字符串字段解析
    for field in ("chapters", "knowledge_codes", "result"):
        if field in d and isinstance(d[field], str):
            try:
                d[field] = json.loads(d[field])
            except (json.JSONDecodeError, TypeError):
                pass
    return d


def compute_checksum(file_path: str | Path) -> str:
    """计算 docx 内 document.xml 的 SHA256"""
    import zipfile
    try:
        with zipfile.ZipFile(file_path, 'r') as zf:
            if 'word/document.xml' in zf.namelist():
                return hashlib.sha256(zf.read('word/document.xml')).hexdigest()
    except Exception:
        pass
    # 回退到全文哈希
    with open(file_path, 'rb') as f:
        return hashlib.sha256(f.read()).hexdigest()


def compute_id_from_docx(file_path: str | Path) -> str:
    """从 docx 的 document.xml 计算题目 ID（SHA256 前 8 位）"""
    return compute_checksum(file_path)[:8]