"""独立重建 KB3，避免交互命令超时中止向量索引的持久化。"""
from __future__ import annotations

import gc
import json
import sys
import time
from pathlib import Path

SERVER_DIR = Path(__file__).resolve().parents[1]
if str(SERVER_DIR) not in sys.path:
    sys.path.insert(0, str(SERVER_DIR))

from sqlmodel import Session

from app.db import engine, init_db
from app.models import KnowledgeBase, now
from app.rag import store
from app.rag.pipeline import rebuild_kb


LOG_PATH = SERVER_DIR / "data" / "kb3-rebuild.log"


def log(payload: dict) -> None:
    payload["at"] = now()
    with LOG_PATH.open("a", encoding="utf-8") as stream:
        stream.write(json.dumps(payload, ensure_ascii=False) + "\n")


def main() -> None:
    init_db()
    with Session(engine) as session:
        kb = session.get(KnowledgeBase, 3)
        if not kb:
            raise RuntimeError("KB3 不存在")
        kb.index_status = "needs_reindex"
        kb.chunk_count = 0
        session.add(kb); session.commit()
        log({"status": "started", "files": kb.file_count})
        completed = 0
        def progress(file_name: str, chunks: int) -> None:
            nonlocal completed
            completed += 1
            log({"status": "progress", "completed_files": completed, "file": file_name, "chunks": chunks})
        result = rebuild_kb(session, 3, progress_cb=progress)
        # 给持久化 HNSW writer 一个明确的落盘窗口，再以同一独立进程验证计数。
        gc.collect()
        time.sleep(5)
        vector_count = store.count(3)
        kb = session.get(KnowledgeBase, 3)
        if vector_count <= 0 or vector_count != kb.chunk_count:
            kb.index_status = "needs_reindex"
            session.add(kb); session.commit()
            log({"status": "verification_failed", "result": result,
                 "database_chunks": kb.chunk_count, "vector_chunks": vector_count})
            raise RuntimeError("KB3 向量计数与数据库不一致")
        kb.index_status = "ready"
        session.add(kb); session.commit()
        log({"status": "completed", "result": result, "vector_chunks": vector_count})
        # Chroma 1.5 的 Rust compactor 在后台写 HNSW；显式停止系统确保索引文件
        # 在进程退出前完整落盘，避免新进程读取到只有 index_metadata 的半成品。
        store.get_client()._system.stop()
        log({"status": "persisted", "vector_chunks": vector_count})


if __name__ == "__main__":
    main()