# -*- coding: utf-8 -*-
"""
RAG 接口测试：验证图片召回功能

运行方式：
.\venv\Scripts\python.exe scripts/test_rag_image_recall.py
"""

import sys
import os
import json

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from knowledge.manager import get_kb_manager
from core.engine import get_engine


def test_rag_image_recall():
    """测试 RAG 检索的图片召回"""
    print("=" * 60)
    print("RAG 图片召回测试")
    print("=" * 60)

    # 初始化引擎
    engine = get_engine()
    if not engine._initialized:
        engine.initialize()

    # 测试查询
    test_queries = [
        "文档中有哪些图片？",
        "有哪些图表？",
        "表格内容是什么？",
    ]

    for query in test_queries:
        print(f"\n🔍 查询: {query}")
        print("-" * 40)

        # 执行检索
        result = engine.search_knowledge(
            query=query,
            collection_name='public_kb',
            top_k=10
        )

        if not result or not result.get('ids'):
            print("  ❌ 无检索结果")
            continue

        print(f"  📊 检索到 {len(result['ids'])} 个切片")

        # 检查图片信息
        images_found = []
        for i, meta in enumerate(result.get('metadatas', [])):
            images_json = meta.get('images_json')
            image_path = meta.get('image_path')
            chunk_type = meta.get('chunk_type')

            if images_json:
                try:
                    imgs = json.loads(images_json)
                    images_found.extend(imgs)
                    print(f"  [{i+1}] chunk_type={chunk_type}, images_json={len(imgs)}张图片")
                except:
                    pass

            if image_path:
                images_found.append({'id': image_path, 'type': 'image_path'})
                print(f"  [{i+1}] chunk_type={chunk_type}, image_path={image_path}")

        if images_found:
            print(f"\n  ✅ 共召回 {len(images_found)} 张图片")
        else:
            print("\n  ⚠️ 未召回任何图片")


def test_extract_rich_media():
    """测试 _extract_rich_media 函数"""
    print("\n" + "=" * 60)
    print("_extract_rich_media 函数测试")
    print("=" * 60)

    from core.agentic import AgenticRAG

    # 创建 AgenticRAG 实例
    rag = AgenticRAG()

    # 模拟检索结果
    mock_contexts = [
        {
            'doc': '这是一段文本',
            'meta': {
                'chunk_type': 'text',
                'source': 'test.pdf',
                'page': 1,
            }
        },
        {
            'doc': '这是一张图片',
            'meta': {
                'chunk_type': 'image',
                'source': 'test.pdf',
                'page': 2,
                'image_path': 'abc123.jpg',
            }
        },
        {
            'doc': '这是带关联图片的文本',
            'meta': {
                'chunk_type': 'text',
                'source': 'test.pdf',
                'page': 3,
                'images_json': '[{"id": "img1.jpg", "order": 1}, {"id": "img2.jpg", "order": 2}]',
            }
        },
    ]

    # 调用 _extract_rich_media
    rich_media = rag._extract_rich_media(mock_contexts)

    print(f"\n📊 提取结果:")
    print(f"  - 图片数量: {len(rich_media.get('images', []))}")
    print(f"  - 表格数量: {len(rich_media.get('tables', []))}")

    for img in rich_media.get('images', []):
        print(f"  📷 {img}")

    if not rich_media.get('images'):
        print("  ⚠️ 未提取到图片，检查 _extract_rich_media 逻辑")


if __name__ == "__main__":
    if sys.platform == 'win32':
        sys.stdout.reconfigure(encoding='utf-8')

    test_rag_image_recall()
    test_extract_rich_media()