"""
step1_search.py
接收用户查询，执行 FAISS 向量语义检索，返回最匹配的技能列表（JSON）。

用法：
    python3 step1_search.py --query "用户的问题"
    python3 step1_search.py --query "用户的问题" --top_k 3 --threshold 0.3
"""

import os
import sys
import json
import argparse
import numpy as np

# ===================== 配置 =====================
SKILL_ROUTER_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
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/"
DEFAULT_TOP_K    = 3
DEFAULT_THRESHOLD = 0.3  # 余弦相似度低于此值视为无匹配

# ===================== 懒加载（避免重复初始化）=====================
_model = None

def get_model():
    global _model
    if _model is None:
        from sentence_transformers import SentenceTransformer
        _model = SentenceTransformer(MODEL_PATH)
    return _model

def get_index_and_map():
    import faiss
    if not os.path.exists(INDEX_FILE):
        raise FileNotFoundError(
            f"索引文件不存在：{INDEX_FILE}\n"
            f"请先运行：python3 {SKILL_ROUTER_DIR}/scripts/build_index.py"
        )
    if not os.path.exists(MAP_FILE):
        raise FileNotFoundError(
            f"映射文件不存在：{MAP_FILE}\n"
            f"请先运行：python3 {SKILL_ROUTER_DIR}/scripts/build_index.py"
        )
    index = faiss.read_index(INDEX_FILE)
    with open(MAP_FILE, "r", encoding="utf-8") as f:
        map_data = json.load(f)
    return index, map_data

# ===================== 向量化 =====================
def text_to_vector(text: str) -> np.ndarray:
    vec = get_model().encode(text, normalize_embeddings=True)
    return vec.astype(np.float32)

# ===================== 查询 =====================
def search_skills(query: str, top_k: int = DEFAULT_TOP_K, threshold: float = DEFAULT_THRESHOLD) -> dict:
    index, map_data = get_index_and_map()

    q_vec = text_to_vector(query).reshape(1, -1)
    scores, indices = index.search(q_vec, top_k)

    # 过滤无效结果（faiss 库数量不足时返回 -1）并按阈值筛选
    candidates = []
    for i, idx in enumerate(indices[0]):
        if idx == -1:
            continue
        score = round(float(scores[0][i]), 4)
        if score >= threshold:
            candidates.append({
                "skill":   map_data["names"][idx],
                "score":   score,
                "path":    map_data["paths"][idx],
                "preview": map_data["docs"][idx][:150] + "..."
            })

    if not candidates:
        top_score = round(float(scores[0][0]), 4) if len(scores[0]) > 0 and indices[0][0] != -1 else 0.0
        return {
            "status":    "no_match",
            "message":   f"未找到相似度高于 {threshold} 的匹配技能，最高得分：{top_score}",
            "top_skill": None,
            "top_score": top_score,
            "top_path":  None,
            "candidates": []
        }

    top = candidates[0]
    return {
        "status":    "success",
        "top_skill": top["skill"],
        "top_score": top["score"],
        "top_path":  top["path"],
        "candidates": candidates
    }

# ===================== 主函数 =====================
def main():
    parser = argparse.ArgumentParser(description="技能路由：语义检索最匹配的技能")
    parser.add_argument("--query",     type=str,   required=True,                  help="用户的原始问题")
    parser.add_argument("--top_k",     type=int,   default=DEFAULT_TOP_K,          help=f"返回候选技能数量，默认 {DEFAULT_TOP_K}")
    parser.add_argument("--threshold", type=float, default=DEFAULT_THRESHOLD,      help=f"最低匹配阈值（余弦相似度），默认 {DEFAULT_THRESHOLD}")
    args = parser.parse_args()

    try:
        result = search_skills(args.query, top_k=args.top_k, threshold=args.threshold)
        print(json.dumps(result, ensure_ascii=False, indent=2))
    except FileNotFoundError as e:
        print(json.dumps({"status": "error", "message": str(e)}, ensure_ascii=False, indent=2))
        sys.exit(1)
    except Exception as e:
        print(json.dumps({"status": "error", "message": f"查询失败：{str(e)}"}, ensure_ascii=False, indent=2))
        sys.exit(1)

if __name__ == "__main__":
    main()