import argparse
import asyncio
import os
import sys
from typing import List, Dict, Any, Optional

# 强制使用 utf-8，避免控制台编码报错影响 Embedding 加载
try:
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8")
    if hasattr(sys.stderr, "reconfigure"):
        sys.stderr.reconfigure(encoding="utf-8")
except Exception:
    pass

# 确保能导入 src.*
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from src.common.logger import get_logger
from src.config.config import global_config
from src.chat.knowledge import lpmm_start_up
from src.memory_system.retrieval_tools.query_lpmm_knowledge import query_lpmm_knowledge

logger = get_logger("test_lpmm_retrieval")


DEFAULT_TEST_CASES: List[Dict[str, Any]] = [
    {
        "name": "回滚一批知识",
        "query": "LPMM是什么?",
        "expect_keywords": ["哈希列表", "删除脚本", "OpenIE"],
    },
    {
        "name": "调整 LPMM 检索参数",
        "query": "不同用词习惯带来的检索偏差该如何解决",
        "expect_keywords": ["bot_config.toml", "lpmm_knowledge", "qa_paragraph_search_top_k"],
    },
]


async def run_tests(test_cases: Optional[List[Dict[str, Any]]] = None) -> None:
    """简单测试 LPMM 知识库检索能力"""
    if not global_config.lpmm_knowledge.enable:
        logger.warning("当前配置中 lpmm_knowledge.enable 为 False，检索测试可能直接返回“未启用”。")

    logger.info("开始初始化 LPMM 知识库...")
    lpmm_start_up()
    logger.info("LPMM 知识库初始化完成，开始执行测试用例。")

    cases = test_cases if test_cases is not None else DEFAULT_TEST_CASES

    for case in cases:
        name = case["name"]
        query = case["query"]
        expect_keywords: List[str] = case.get("expect_keywords", [])

        print("\n" + "=" * 60)
        print(f"[TEST] {name}")
        print(f"[Q] {query}")

        result = await query_lpmm_knowledge(query, limit=3)

        print("\n[RAW RESULT]")
        print(result)

        status = "UNKNOWN"
        hit_keywords: List[str] = []

        if isinstance(result, str):
            if "未启用" in result or "未初始化" in result or "查询失败" in result:
                status = "ERROR"
            elif "未找到与" in result:
                status = "NO_HIT"
            else:
                if expect_keywords:
                    hit_keywords = [kw for kw in expect_keywords if kw in result]
                    status = "PASS" if hit_keywords else "WARN"
                else:
                    status = "PASS"

        print("\n[CHECK]")
        print(f"Status: {status}")
        if expect_keywords:
            print(f"Expected keywords: {expect_keywords}")
            print(f"Hit keywords: {hit_keywords}")

    print("\n" + "=" * 60)
    print("LPMM 检索测试完成。请根据每条用例的 Status 和命中关键词判断检索效果是否符合预期。")


def main() -> None:
    parser = argparse.ArgumentParser(
        description=(
            "测试 LPMM 知识库检索能力。\n"
            "如不提供参数，则执行内置的默认用例；\n"
            "也可以通过 --query 与 --expect-keyword 自定义一条测试用例。"
        )
    )
    parser.add_argument(
        "--query",
        help="自定义测试问题（单条）。提供该参数时，将仅运行这一条用例。",
    )
    parser.add_argument(
        "--expect-keyword",
        action="append",
        help="期望在检索结果中出现的关键字，可重复多次指定；仅在提供 --query 时生效。",
    )
    args = parser.parse_args()

    if args.query:
        custom_case = {
            "name": "custom",
            "query": args.query,
            "expect_keywords": args.expect_keyword or [],
        }
        asyncio.run(run_tests([custom_case]))
    else:
        asyncio.run(run_tests())


if __name__ == "__main__":
    main()