#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
测试MongoDB和BOS存储集成
"""

import os
import logging
from mongodb_manager import get_mongodb_manager
from baidu_bos_manager import get_bos_manager

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def test_mongodb_connection():
    """测试MongoDB连接"""
    try:
        mongodb_mgr = get_mongodb_manager()
        stats = mongodb_mgr.get_stats()
        logger.info(f"✅ MongoDB连接成功，统计信息: {stats}")
        return True
    except Exception as e:
        logger.error(f"❌ MongoDB连接失败: {e}")
        return False

def test_bos_connection():
    """测试BOS连接"""
    try:
        bos_mgr = get_bos_manager()
        objects = bos_mgr.list_objects(max_keys=5)
        logger.info(f"✅ BOS连接成功，找到 {len(objects)} 个对象")
        return True
    except Exception as e:
        logger.error(f"❌ BOS连接失败: {e}")
        return False

def test_file_upload_workflow():
    """测试文件上传工作流"""
    try:
        # 创建测试文件
        test_file = "/tmp/test_storage.txt"
        with open(test_file, 'w', encoding='utf-8') as f:
            f.write("这是一个测试文件，用于验证存储集成功能。")
        
        # 获取管理器
        mongodb_mgr = get_mongodb_manager()
        bos_mgr = get_bos_manager()
        
        # 上传到BOS
        bos_result = bos_mgr.upload_file(test_file)
        logger.info(f"✅ BOS上传成功: {bos_result['bos_key']}")
        
        # 存储元数据到MongoDB
        file_id = mongodb_mgr.store_file_metadata(
            file_path=test_file,
            file_type="text",
            bos_key=bos_result["bos_key"],
            additional_info={
                "test": True,
                "bos_url": bos_result["url"]
            }
        )
        logger.info(f"✅ MongoDB元数据存储成功: {file_id}")
        
        # 存储向量元数据
        mongodb_mgr.store_vector_metadata(
            file_id=file_id,
            vector_type="text_vector",
            vdb_id=bos_result["bos_key"],
            vector_info={"test": True}
        )
        logger.info("✅ 向量元数据存储成功")
        
        # 清理测试文件
        os.remove(test_file)
        
        return True
        
    except Exception as e:
        logger.error(f"❌ 文件上传工作流测试失败: {e}")
        return False

if __name__ == "__main__":
    logger.info("🚀 开始存储集成测试...")
    
    # 测试MongoDB连接
    mongodb_ok = test_mongodb_connection()
    
    # 测试BOS连接
    bos_ok = test_bos_connection()
    
    # 测试完整工作流
    workflow_ok = test_file_upload_workflow()
    
    # 总结
    if mongodb_ok and bos_ok and workflow_ok:
        logger.info("🎉 所有存储集成测试通过！")
    else:
        logger.error("❌ 存储集成测试失败")
        logger.error(f"MongoDB: {'✅' if mongodb_ok else '❌'}")
        logger.error(f"BOS: {'✅' if bos_ok else '❌'}")
        logger.error(f"工作流: {'✅' if workflow_ok else '❌'}")