# -*- coding: utf-8 -*-
"""
图片处理完整链路测试

测试文件: documents/public_kb/三峡公报_16-30页.pdf
MinerU 临时输出: .data/mineru_temp

测试步骤:
1. MinerU 解析 → 检查 chunk.images 和 chunk.image_path
2. 同步到向量库 → 检查 metadata.images_json 和 metadata.image_path
3. RAG 检索 → 检查图片召回结果
"""

import sys
import os
import json
import shutil
from pathlib import Path

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

# 测试文件路径
TEST_PDF = "documents/public_kb/三峡公报_16-30页.pdf"
TEMP_DIR = ".data/mineru_temp"
KB_NAME = "my_kb"


def step1_mineru_parse():
    """步骤1: MinerU 解析，检查图片信息"""
    print("\n" + "=" * 60)
    print("步骤1: MinerU 解析")
    print("=" * 60)

    if not os.path.exists(TEST_PDF):
        print(f"❌ 测试文件不存在: {TEST_PDF}")
        return None

    print(f"📄 解析文件: {TEST_PDF}")

    # 清理临时目录
    if os.path.exists(TEMP_DIR):
        shutil.rmtree(TEMP_DIR)

    # 调用 MinerU 解析
    from parsers.mineru_parser import parse_with_mineru_persistent

    result = parse_with_mineru_persistent(
        TEST_PDF,
        output_base=TEMP_DIR,
        cleanup_after_image_move=False  # 保留临时文件用于检查
    )

    if not result:
        print("❌ MinerU 解析失败")
        return None

    chunks = result.get('chunks', [])
    images = result.get('images', [])

    print(f"\n📊 解析结果:")
    print(f"  - 切片数量: {len(chunks)}")
    print(f"  - 图片数量: {len(images)}")

    # 检查 chunk.images 和 chunk.image_path
    chunks_with_images = 0
    chunks_with_image_path = 0
    image_chunks = 0

    print(f"\n📋 切片图片信息（前 20 个）:")
    print("-" * 40)

    for i, chunk in enumerate(chunks[:20]):
        chunk_type = getattr(chunk, 'chunk_type', 'text')
        has_images = hasattr(chunk, 'images') and chunk.images
        has_image_path = hasattr(chunk, 'image_path') and chunk.image_path

        if has_images:
            chunks_with_images += 1
        if has_image_path:
            chunks_with_image_path += 1
        if chunk_type in ('image', 'chart'):
            image_chunks += 1

        print(f"[{i+1}] type={chunk_type}")
        if has_images:
            print(f"    images: {chunk.images}")
        if has_image_path:
            print(f"    image_path: {chunk.image_path}")

    print("-" * 40)
    print(f"\n📈 统计:")
    print(f"  - 有 images 字段的切片: {chunks_with_images}")
    print(f"  - 有 image_path 字段的切片: {chunks_with_image_path}")
    print(f"  - 图片类型切片: {image_chunks}")

    return result


def step2_sync_to_kb(result):
    """步骤2: 同步到向量库，检查 metadata"""
    print("\n" + "=" * 60)
    print("步骤2: 同步到向量库")
    print("=" * 60)

    from knowledge.manager import get_kb_manager

    kb = get_kb_manager()

    # 检查/创建知识库
    if KB_NAME not in kb.list_collections():
        kb.create_collection(KB_NAME, display_name="测试知识库")
        print(f"✅ 创建知识库: {KB_NAME}")
    else:
        # 清空现有数据
        coll = kb.get_collection(KB_NAME)
        if coll:
            existing = coll.get()
            if existing['ids']:
                coll.delete(ids=existing['ids'])
                print(f"🗑️ 清空知识库: {len(existing['ids'])} 条记录")

    # 同步文件
    print(f"📤 同步文件: {TEST_PDF}")
    count = kb.add_file_to_kb(KB_NAME, TEST_PDF)
    print(f"✅ 同步完成: {count} 个切片")

    # 检查 metadata
    coll = kb.get_collection(KB_NAME)
    data = coll.get(include=['metadatas', 'documents'])

    has_images_json = 0
    has_image_path = 0

    print(f"\n📋 Metadata 检查（前 20 个）:")
    print("-" * 40)

    for i, meta in enumerate(data['metadatas'][:20]):
        images_json = meta.get('images_json')
        image_path = meta.get('image_path')

        if images_json:
            has_images_json += 1
        if image_path:
            has_image_path += 1

        print(f"[{i+1}] type={meta.get('chunk_type')}")
        if images_json:
            print(f"    images_json: {images_json[:80]}...")
        if image_path:
            print(f"    image_path: {image_path}")

    print("-" * 40)
    print(f"\n📈 统计:")
    print(f"  - 有 images_json 的切片: {has_images_json}")
    print(f"  - 有 image_path 的切片: {has_image_path}")

    # 验证结果
    print(f"\n✅ 验证结果:")
    if has_images_json > 0:
        print(f"  ✅ images_json 字段正确写入")
    else:
        print(f"  ❌ images_json 字段未写入")

    if has_image_path > 0:
        print(f"  ✅ image_path 字段正确写入")
    else:
        print(f"  ⚠️ image_path 字段未写入（可能没有独立图片切片）")

    return has_images_json > 0 or has_image_path > 0


def step3_rag_recall():
    """步骤3: RAG 检索，检查图片召回"""
    print("\n" + "=" * 60)
    print("步骤3: RAG 检索图片召回")
    print("=" * 60)

    from core.engine import get_engine

    engine = get_engine()
    if not engine._initialized:
        engine.initialize()

    # 测试查询
    queries = [
        "三峡公报中有哪些图表？",
        "文档中的图片",
        "有哪些数据图表",
    ]

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

        result = engine.search_knowledge(
            query=query,
            top_k=10,
            collections=[KB_NAME]
        )

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

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

        # 统计图片
        images_found = []
        for meta in result.get('metadatas', []):
            images_json = meta.get('images_json')
            image_path = meta.get('image_path')

            if images_json:
                try:
                    imgs = json.loads(images_json)
                    images_found.extend(imgs)
                except:
                    pass
            if image_path:
                images_found.append({'id': image_path})

        if images_found:
            print(f"  ✅ 召回 {len(images_found)} 张图片:")
            for img in images_found:
                print(f"     📷 {img.get('id', img)}")
        else:
            print("  ⚠️ 未召回图片")


def main():
    print("=" * 60)
    print("图片处理完整链路测试")
    print("=" * 60)

    # 步骤1: MinerU 解析
    result = step1_mineru_parse()
    if not result:
        return

    # 步骤2: 同步到向量库
    success = step2_sync_to_kb(result)

    # 步骤3: RAG 检索
    step3_rag_recall()

    print("\n" + "=" * 60)
    print("测试完成")
    print("=" * 60)


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

    main()