"""把 kb_build/pages/*.txt 的站点文档写入 nonebot-plugin-wtfllm 知识库。

用法:
    python kb_build/import_kb.py --agent-id <机器人QQ号>

首次运行会自动下载 Qdrant 与 embedding 模型,需要几分钟。
"""
import argparse
import asyncio
import re
import sys
import time
import uuid
from pathlib import Path

# NoneBot 必须先初始化,wtfllm 的配置才可用
import nonebot

nonebot.init()

sys.path.insert(0, str(Path(__file__).parent.parent))

PAGES_DIR = Path(__file__).parent / "pages"

# 每条知识的目标长度(字符),按标题切块后超过再按段落切
CHUNK_SIZE = 600

CATEGORY_BY_DIR = {
    "功能帮助": "功能帮助",
    "脚本进阶": "脚本进阶",
    "简单使用": "简单使用",
    "其它": "其它",
}


def parse_page(path: Path) -> list[dict]:
    """把单页文本按标题切成知识条目。"""
    raw = path.read_text(encoding="utf-8")
    lines = raw.splitlines()
    page_title = lines[0].lstrip("# ").replace("\u200b", "") if lines else path.stem
    source = lines[1].lstrip("> 来源: ").strip() if len(lines) > 1 else ""

    parts: list[dict] = []
    cur_head = page_title
    cur_body: list[str] = []

    def flush() -> None:
        text = "\n".join(cur_body).strip()
        if text:
            parts.append({"title": cur_head, "content": text})

    for line in lines[2:]:
        m = re.match(r"^(#{1,4})\s+(.+)$", line)
        if m:
            flush()
            cur_body = []
            h = m.group(2).replace("\u200b", "").strip()
            cur_head = h if h == page_title else f"{page_title} · {h}"
        else:
            cur_body.append(line)
    flush()

    # 超长条目按段落二次切分
    chunks: list[dict] = []
    for p in parts:
        if len(p["content"]) <= CHUNK_SIZE * 1.5:
            chunks.append(p)
            continue
        paras = p["content"].split("\n\n")
        buf: list[str] = []
        for para in paras:
            if buf and sum(len(b) for b in buf) + len(para) > CHUNK_SIZE:
                chunks.append({"title": p["title"], "content": "\n\n".join(buf)})
                buf = []
            buf.append(para)
        if buf:
            chunks.append({"title": p["title"], "content": "\n\n".join(buf)})

    # 附上页面来源与分类
    first_dir = path.stem.split("_")[0]
    category = CATEGORY_BY_DIR.get(first_dir, "general")
    for c in chunks:
        c["category"] = category
        c["source"] = source
        c["page_title"] = page_title
    return chunks


async def main(agent_id: str, dry_run: bool) -> None:
    entries_spec: list[dict] = []
    for f in sorted(PAGES_DIR.glob("*.txt")):
        entries_spec.extend(parse_page(f))
    print(f"共解析出 {len(entries_spec)} 条知识条目")
    if dry_run:
        for e in entries_spec[:10]:
            print(f"  [{e['category']}] {e['title']} ({len(e['content'])} chars)")
        return

    # 必须先以插件形式加载,让 localstore 能定位数据目录
    nonebot.load_plugin("nonebot_plugin_wtfllm")

    from nonebot_plugin_wtfllm.memory.items.knowledge_base import KnowledgeEntry
    from nonebot_plugin_wtfllm.v_db.lifecycle import on_startup as vdb_startup
    from nonebot_plugin_wtfllm.v_db.repositories import KnowledgeBaseRepository

    print("初始化 Qdrant(首次会自动下载)...")
    await vdb_startup()

    now = int(time.time())
    entries = [
        KnowledgeEntry.create(
            content=e["content"],
            title=e["title"],
            agent_id=agent_id,
            category=e["category"],
            tags=[e["page_title"], "danding.vip"],
            token_count=len(e["content"]),
        )
        for e in entries_spec
    ]
    repo = KnowledgeBaseRepository()
    ids = await repo.save_many_knowledge(entries)
    print(f"已写入 {len(ids)} 条知识")

    # 验证检索
    for q in ["怎么设置定时调度", "积分怎么获得", "下载安装"]:
        results = await repo.search_relevant(agent_id=agent_id, query=q, limit=3)
        print(f"\n检索「{q}」:")
        for r in results:
            print(f"  {r.score:.3f} {r.item.title}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--agent-id", required=True, help="机器人 QQ 号(wtfllm 的 agent_id)")
    parser.add_argument("--dry-run", action="store_true", help="只解析不写入")
    args = parser.parse_args()
    asyncio.run(main(args.agent_id, args.dry_run))