"""Hezor2 SDK knowledge_retrieve 使用示例。

本示例演示通过 Hezor2 SDK 调用知识检索接口：

1. 基础检索（使用默认参数）
2. 自定义参数检索（top_k / score_threshold）
3. 结果详情展示（chunks、entities、communities、pictures）

Usage
-----
运行所有示例:
    uv run examples/transfer/hezor2_sdk/knowledge_retrieve_example.py

运行指定示例:
    uv run examples/transfer/hezor2_sdk/knowledge_retrieve_example.py 1
    uv run examples/transfer/hezor2_sdk/knowledge_retrieve_example.py 2
    uv run examples/transfer/hezor2_sdk/knowledge_retrieve_example.py 3

列出所有示例:
    uv run examples/transfer/hezor2_sdk/knowledge_retrieve_example.py --list
"""

import asyncio
import logging
import sys
from pathlib import Path

from dotenv import load_dotenv

# ⚠️ 重要：必须在导入 SDK 之前加载环境变量
# 因为 constants.py 在模块导入时会读取环境变量创建默认值
env_file = Path(__file__).parent.parent.parent / ".env"
load_dotenv(dotenv_path=env_file)

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

from hezor_common.transfer.hezor2_sdk import Hezor2SDK, MetaInfo  # noqa: E402
from hezor_common.transfer.hezor2_sdk.env_config import Hezor2EnvConfig  # noqa: E402

logging.basicConfig(
    level=logging.WARNING,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)

print(f"✓ 已加载环境变量: {env_file}")
print()

config = Hezor2EnvConfig()
print(f"✓ API Base URL: {config.hezor2_api_base_url}")
print()


def _build_sdk() -> Hezor2SDK:
    """构建 SDK 实例（带可选的认证配置）。"""
    kwargs: dict = {
        "base_url": config.hezor2_api_base_url,
        "api_key": config.hezor2_api_key,
    }

    if config.hezor2_app_name:
        kwargs["app_name"] = config.hezor2_app_name

    kwargs["meta_info"] = MetaInfo(
        caller_id="example/knowledge_retrieve",
        subject="example",
        subject_code="example_001",
        creation_name="示例创建",
        creation_slug="example_creation",
        data_coverage="20240101-20241231",
    )

    if config.hezor2_header_pk_filepath:
        private_key_path = Path(config.hezor2_header_pk_filepath)
        if not private_key_path.is_absolute():
            examples_dir = Path(__file__).parent.parent.parent
            private_key_path = examples_dir / private_key_path
        if private_key_path.exists():
            kwargs["private_key_path"] = str(private_key_path)
            if config.hezor2_header_pk_password:
                kwargs["password"] = config.hezor2_header_pk_password.encode()
            print(f"  ✓ 使用私钥认证: {private_key_path}")
        else:
            print(f"  ⚠ 私钥文件不存在: {private_key_path}，将使用匿名密钥签名")

    return Hezor2SDK(**kwargs)


def example_1_basic_retrieve():
    """Example 1: 基础检索（默认参数）。"""
    print("=== Example 1: 基础知识检索 ===")
    print()

    query = "什么是稻飞虱？"
    print(f"  查询: {query}")
    print()

    async def run():
        async with _build_sdk() as sdk:
            result = await sdk.knowledge_retrieve(query=query)

            print("  ✓ 检索完成")
            print(f"  Chunks:       {len(result.chunks.items)} 条")
            print(f"  Entities:     {len(result.entities.items)} 条")
            print(f"  Communities:  {len(result.communities.items)} 条")
            print(f"  Pictures:     {len(result.pictures.items)} 条")
            print()

            if result.chunks.items:
                top = result.chunks.items[0]
                print(f"  Top Chunk (rank={top.rank}, score={top.score:.4f}):")
                print(f"    doc_name: {top.chunk.doc_name}")
                print(f"    text:     {top.chunk.text[:100]}...")

    try:
        asyncio.run(run())
    except Exception as e:
        print(f"  调用失败（需要连接 Hezor2 API）: {e}")

    print()


def example_2_custom_params():
    """Example 2: 自定义 top_k 和 score_threshold。"""
    print("=== Example 2: 自定义检索参数 ===")
    print()

    query = "稻纵卷叶螟发生分析"
    top_k = 5
    score_threshold = 0.5

    print(f"  查询:             {query}")
    print(f"  top_k:            {top_k}")
    print(f"  score_threshold:  {score_threshold}")
    print()

    async def run():
        async with _build_sdk() as sdk:
            result = await sdk.knowledge_retrieve(
                query=query,
                top_k=top_k,
                score_threshold=score_threshold,
            )

            print("  ✓ 检索完成")
            print(f"  Chunks:       {len(result.chunks.items)} 条")
            print(f"  Entities:     {len(result.entities.items)} 条")
            print(f"  Communities:  {len(result.communities.items)} 条")
            print(f"  Pictures:     {len(result.pictures.items)} 条")

    try:
        asyncio.run(run())
    except Exception as e:
        print(f"  调用失败（需要连接 Hezor2 API）: {e}")

    print()


def example_3_result_details():
    """Example 3: 展示结构化结果详情。"""
    print("=== Example 3: 结果详情展示 ===")
    print()

    query = "棉花病虫害有哪些？"
    print(f"  查询: {query}")
    print()

    async def run():
        async with _build_sdk() as sdk:
            result = await sdk.knowledge_retrieve(query=query, top_k=3)

            # Chunks
            print(f"  【Chunks】 (collection={result.chunks.collection})")
            for item in result.chunks.items:
                print(
                    f"    #{item.rank} score={item.score:.4f}  doc={item.chunk.doc_name}  page={item.chunk.page}"
                )
                print(f"       {item.chunk.text[:80].strip()}...")
            print()

            # Entities
            print(f"  【Entities】 (collection={result.entities.collection})")
            for item in result.entities.items:
                print(
                    f"    #{item.rank} score={item.score:.4f}  [{item.entity.entity_type}] {item.entity.name}"
                )
                if item.entity.description:
                    print(f"       {item.entity.description[:80]}...")
            print()

            # Communities
            print(f"  【Communities】 (collection={result.communities.collection})")
            for item in result.communities.items:
                print(f"    #{item.rank} score={item.score:.4f}  {item.community.title}")
                print(f"       {item.community.summary[:80]}...")
                if item.community.findings:
                    print(f"       findings: {len(item.community.findings)} 条")
            print()

            # Pictures
            print(f"  【Pictures】 (collection={result.pictures.collection})")
            for item in result.pictures.items:
                print(f"    #{item.rank} score={item.score:.4f}  {item.picture.filename}")
                if item.picture.description:
                    print(f"       {item.picture.description[:80]}...")
            if not result.pictures.items:
                print("    (无匹配图片)")

    try:
        asyncio.run(run())
    except Exception as e:
        print(f"  调用失败（需要连接 Hezor2 API）: {e}")

    print()


def list_examples():
    """列出所有可用示例。"""
    examples = {
        1: ("基础检索（默认参数）", example_1_basic_retrieve),
        2: ("自定义 top_k / score_threshold", example_2_custom_params),
        3: ("结果详情展示（chunks/entities/communities/pictures）", example_3_result_details),
    }

    print("Available examples:")
    for num, (name, _) in examples.items():
        print(f"  {num}. {name}")
    print()
    print("Usage:")
    print(f"  uv run {sys.argv[0]} [example_number]")
    print(f"  uv run {sys.argv[0]}              # 运行所有示例")
    print(f"  uv run {sys.argv[0]} --list       # 列出示例")

    return examples


def main():
    """根据命令行参数运行示例。"""
    examples = {
        1: ("基础检索（默认参数）", example_1_basic_retrieve),
        2: ("自定义 top_k / score_threshold", example_2_custom_params),
        3: ("结果详情展示（chunks/entities/communities/pictures）", example_3_result_details),
    }

    if len(sys.argv) > 1:
        arg = sys.argv[1]

        if arg in ["--list", "-l", "list"]:
            list_examples()
            return

        try:
            example_num = int(arg)
            if example_num not in examples:
                print(f"Error: Example {example_num} not found.")
                print()
                list_examples()
                sys.exit(1)

            name, func = examples[example_num]
            print("=" * 60)
            print(f"knowledge_retrieve - Example {example_num}")
            print("=" * 60)
            print()
            func()
            print("=" * 60)
            print(f"Example {example_num} completed successfully!")
            print("=" * 60)
            return

        except ValueError:
            print(f"Error: Invalid argument '{arg}'")
            print()
            list_examples()
            sys.exit(1)

    # 运行所有示例
    print("=" * 60)
    print("knowledge_retrieve - Usage Examples")
    print("=" * 60)
    print()

    for func in [example_1_basic_retrieve, example_2_custom_params, example_3_result_details]:
        func()

    print("=" * 60)
    print("All examples completed successfully!")
    print("=" * 60)


if __name__ == "__main__":
    main()
