"""
build_index.py
扫描 skill-router/skills/ 目录下所有子技能的 SKILL.md，
使用 bge-small-zh-v1.5 生成向量并构建 FAISS 索引。

用法：
    python3 build_index.py
"""

import os
import json
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer

# ===================== 配置 =====================
SKILL_ROUTER_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SKILLS_DIR       = os.path.join(SKILL_ROUTER_DIR, "skills")
INDEX_FILE       = os.path.join(SKILL_ROUTER_DIR, "skills_faiss.index")
MAP_FILE         = os.path.join(SKILL_ROUTER_DIR, "skills_map.json")
MODEL_PATH       = "/home/admin/bge-small-zh-v1.5/"
VECTOR_DIM       = 512  # bge-small-zh-v1.5 输出维度

# ===================== 加载模型 =====================
print("正在加载 embedding 模型...")
model = SentenceTransformer(MODEL_PATH)
print("✅ 模型加载完成\n")

# ===================== 扫描 SKILL.md =====================
def load_all_skill_files():
    documents   = []
    skill_names = []
    file_paths  = []

    print(f"正在扫描目录：{SKILLS_DIR}")
    if not os.path.exists(SKILLS_DIR):
        print(f"❌ skills 目录不存在：{SKILLS_DIR}")
        return documents, skill_names, file_paths

    for root, dirs, files in os.walk(SKILLS_DIR):
        for file in files:
            if file == "SKILL.md":
                full_path = os.path.join(root, file)
                try:
                    with open(full_path, "r", encoding="utf-8") as f:
                        content = f.read().strip()
                    skill_name = os.path.basename(root)
                    documents.append(content)
                    skill_names.append(skill_name)
                    file_paths.append(full_path)
                    print(f"  ✅ 已加载：{skill_name}")
                except Exception as e:
                    print(f"  ❌ 读取失败：{full_path} — {e}")

    print(f"\n✅ 总共加载技能数：{len(documents)}")
    return documents, skill_names, file_paths

# ===================== 构建并保存索引 =====================
def build_index():
    docs, names, paths = load_all_skill_files()

    if not docs:
        print("❌ 没有找到任何 SKILL.md，索引未构建。")
        return

    print("\n正在生成向量...")
    vecs = np.array(
        [model.encode(d, normalize_embeddings=True).astype(np.float32) for d in docs],
        dtype=np.float32
    )

    # IndexFlatIP + normalize = 余弦相似度，score 越大越相关
    index = faiss.IndexFlatIP(VECTOR_DIM)
    index.add(vecs)
    faiss.write_index(index, INDEX_FILE)
    print(f"✅ FAISS 索引已保存：{INDEX_FILE}")

    with open(MAP_FILE, "w", encoding="utf-8") as f:
        json.dump(
            {"docs": docs, "names": names, "paths": paths},
            f,
            ensure_ascii=False
        )
    print(f"✅ 映射文件已保存：{MAP_FILE}")
    print(f"\n🎉 索引构建完成，共收录 {len(docs)} 个技能。")

if __name__ == "__main__":
    build_index()