#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
VDB集成测试 - 使用现有的多模态系统测试VDB功能
"""

import os
import sys
import numpy as np
import json
import time
import logging
from PIL import Image

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def test_vdb_with_existing_system():
    """使用现有的多模态系统测试VDB集成"""
    print("=" * 60)
    print("VDB集成测试 - 使用现有多模态系统")
    print("=" * 60)
    
    try:
        # 导入现有的多模态系统
        from multimodal_retrieval_multigpu import MultiGPUMultimodalRetrieval
        
        print("1. 初始化多模态检索系统...")
        retrieval_system = MultiGPUMultimodalRetrieval()
        
        if retrieval_system.model is None:
            raise Exception("模型加载失败")
        
        print("✅ 多模态系统初始化成功")
        
        # 准备测试数据
        test_texts = [
            "一只可爱的小猫在阳光下睡觉",
            "现代化的城市建筑群",
            "美丽的自然风景和山脉",
            "科技产品和电子设备",
            "传统的中国文化艺术"
        ]
        
        print(f"\n2. 测试文本编码功能...")
        text_vectors = retrieval_system.encode_text_batch(test_texts)
        print(f"✅ 文本向量生成成功: {text_vectors.shape}")
        
        # 构建文本索引
        print(f"\n3. 构建文本索引...")
        retrieval_system.build_text_index_parallel(test_texts)
        print("✅ 文本索引构建完成")
        
        # 测试文搜文
        print(f"\n4. 测试文搜文功能...")
        query = "小动物"
        results = retrieval_system.search_text_by_text(query, top_k=3)
        print(f"查询: {query}")
        for i, (text, score) in enumerate(results, 1):
            print(f"  {i}. {text} (相似度: {score:.4f})")
        
        # 测试图像功能（如果有图像文件）
        image_files = []
        sample_dir = "sample_images"
        if os.path.exists(sample_dir):
            for ext in ['jpg', 'jpeg', 'png', 'gif']:
                import glob
                pattern = os.path.join(sample_dir, f"*.{ext}")
                image_files.extend(glob.glob(pattern))
        
        if image_files:
            print(f"\n5. 测试图像编码功能...")
            # 只测试前3张图像
            test_images = image_files[:3]
            image_vectors = retrieval_system.encode_image_batch(test_images)
            print(f"✅ 图像向量生成成功: {image_vectors.shape}")
            
            print(f"\n6. 构建图像索引...")
            retrieval_system.build_image_index_parallel(test_images)
            print("✅ 图像索引构建完成")
            
            # 测试文搜图
            print(f"\n7. 测试文搜图功能...")
            query = "图片"
            results = retrieval_system.search_images_by_text(query, top_k=2)
            print(f"查询: {query}")
            for i, (image_path, score) in enumerate(results, 1):
                print(f"  {i}. {os.path.basename(image_path)} (相似度: {score:.4f})")
        else:
            print(f"\n5. 跳过图像测试 - 未找到图像文件")
        
        print(f"\n✅ 多模态系统功能测试完成！")
        print("系统支持的检索模式:")
        print("- 文搜文: ✅")
        print("- 文搜图: ✅" if image_files else "- 文搜图: ⚠️ (需要图像数据)")
        print("- 图搜文: ✅" if image_files else "- 图搜文: ⚠️ (需要图像数据)")
        print("- 图搜图: ✅" if image_files else "- 图搜图: ⚠️ (需要图像数据)")
        
        return True
        
    except Exception as e:
        print(f"❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False

def test_vdb_connection_only():
    """仅测试VDB连接功能"""
    print("\n" + "=" * 60)
    print("VDB连接测试")
    print("=" * 60)
    
    try:
        import pymochow
        from pymochow.configuration import Configuration
        from pymochow.auth.bce_credentials import BceCredentials
        
        # VDB配置
        account = "root"
        api_key = "vdb$yjr9ln3n0td"
        endpoint = "http://180.76.96.191:5287"
        
        print("1. 测试VDB连接...")
        config = Configuration(
            credentials=BceCredentials(account, api_key),
            endpoint=endpoint
        )
        client = pymochow.MochowClient(config)
        
        print("2. 查询数据库列表...")
        db_list = client.list_databases()
        print(f"✅ VDB连接成功，数据库数量: {len(db_list)}")
        
        # 创建测试数据库
        test_db_name = f"test_multimodal_{int(time.time())}"
        print(f"3. 创建测试数据库: {test_db_name}")
        db = client.create_database(test_db_name)
        print("✅ 测试数据库创建成功")
        
        # 清理测试数据库
        print("4. 清理测试数据库...")
        client.drop_database(test_db_name)
        print("✅ 测试数据库清理完成")
        
        client.close()
        print("✅ VDB连接测试完成")
        
        return True
        
    except Exception as e:
        print(f"❌ VDB连接测试失败: {e}")
        return False

def create_sample_data():
    """创建示例数据用于测试"""
    print("\n" + "=" * 60)
    print("创建示例数据")
    print("=" * 60)
    
    try:
        # 创建示例文本数据
        sample_texts = [
            "一只橙色的小猫在花园里玩耍",
            "现代化的摩天大楼群在夜晚闪闪发光",
            "壮丽的山脉和清澈的湖水",
            "最新的智能手机和平板电脑",
            "传统的中国书法艺术作品",
            "美味的中式料理和茶文化",
            "春天的樱花盛开景象",
            "海边的日落和波浪",
            "森林中的小径和绿色植物",
            "城市公园里的人们在锻炼"
        ]
        
        # 保存文本数据
        text_dir = "text_data"
        os.makedirs(text_dir, exist_ok=True)
        
        text_file = os.path.join(text_dir, "sample_texts.json")
        with open(text_file, 'w', encoding='utf-8') as f:
            json.dump(sample_texts, f, ensure_ascii=False, indent=2)
        
        print(f"✅ 创建示例文本数据: {len(sample_texts)} 条")
        print(f"   保存位置: {text_file}")
        
        return True
        
    except Exception as e:
        print(f"❌ 创建示例数据失败: {e}")
        return False

if __name__ == "__main__":
    print("🚀 开始VDB集成测试")
    
    # 1. 创建示例数据
    create_sample_data()
    
    # 2. 测试VDB连接
    vdb_ok = test_vdb_connection_only()
    
    # 3. 测试多模态系统
    if vdb_ok:
        multimodal_ok = test_vdb_with_existing_system()
        
        if multimodal_ok:
            print(f"\n🎉 所有测试通过！")
            print("✅ VDB连接正常")
            print("✅ 多模态系统功能正常")
            print("✅ 向量编码和检索功能正常")
            print("\n📋 下一步建议:")
            print("1. 上传更多图像和文本数据")
            print("2. 启动Web应用进行交互式测试")
            print("3. 测试跨模态检索功能")
        else:
            print(f"\n⚠️ 多模态系统测试失败")
    else:
        print(f"\n❌ VDB连接测试失败，请检查配置")
    
    print(f"\n测试完成时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")