#!/usr/bin/env python3
"""
CN-CLIP 服务测试脚本

用法:
    python scripts/test_cnclip_service.py

选项:
    --url TEXT       服务地址（默认：grpc://localhost:51000）
    --text           只测试文本编码
    --image          只测试图像编码
    --batch-size INT 批处理大小（默认：10）
    --help           显示帮助信息
"""

import sys
import time
import argparse
from pathlib import Path

# 添加项目路径到 sys.path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

# 颜色输出
class Colors:
    GREEN = '\033[0;32m'
    RED = '\033[0;31m'
    YELLOW = '\033[1;33m'
    BLUE = '\033[0;34m'
    NC = '\033[0m'


def print_success(msg):
    print(f"{Colors.GREEN}✓ {msg}{Colors.NC}")


def print_error(msg):
    print(f"{Colors.RED}✗ {msg}{Colors.NC}")


def print_warning(msg):
    print(f"{Colors.YELLOW}⚠ {msg}{Colors.NC}")


def print_info(msg):
    print(f"{Colors.BLUE}ℹ {msg}{Colors.NC}")


def test_imports():
    """测试必要的依赖是否安装"""
    print("\n" + "="*50)
    print("测试 1: 检查依赖")
    print("="*50)

    try:
        import clip_client
        print_success("clip_client 已安装")
    except ImportError as e:
        print_error(f"clip_client 未安装: {e}")
        print_info("请运行: pip install clip-client")
        return False

    try:
        import numpy as np
        print_success("numpy 已安装")
    except ImportError as e:
        print_error(f"numpy 未安装: {e}")
        return False

    return True


def test_connection(url):
    """测试服务连接"""
    print("\n" + "="*50)
    print("测试 2: 连接服务")
    print("="*50)
    print(f"服务地址: {url}")

    try:
        from clip_client import Client

        client = Client(url)
        print_success("客户端创建成功")
        return client
    except Exception as e:
        print_error(f"连接失败: {e}")
        print_info("请确保服务已启动: ./scripts/start_cnclip_service.sh")
        return None


def test_text_encoding(client, batch_size=10):
    """测试文本编码"""
    print("\n" + "="*50)
    print("测试 3: 文本编码")
    print("="*50)

    try:
        # 准备测试数据
        test_texts = [
            '你好，世界',
            'CN-CLIP 图像编码服务',
            '这是一个测试',
            '人工智能',
            '机器学习',
            '深度学习',
            '计算机视觉',
            '自然语言处理',
            '搜索引擎',
            '多模态检索',
        ][:batch_size]

        print(f"测试文本数量: {len(test_texts)}")
        print(f"示例文本: {test_texts[0]}")

        # 执行编码
        start_time = time.time()
        embeddings = client.encode(test_texts)
        elapsed_time = time.time() - start_time

        # 验证结果
        assert embeddings.shape[0] == len(test_texts), "向量数量不匹配"
        assert embeddings.shape[1] == 1024, "向量维度应该是 1024"

        print_success(f"编码成功")
        print(f"  向量形状: {embeddings.shape}")
        print(f"  耗时: {elapsed_time:.2f}秒")
        print(f"  速度: {len(test_texts)/elapsed_time:.2f} 条/秒")
        print(f"  数据类型: {embeddings.dtype}")

        return True

    except Exception as e:
        print_error(f"文本编码失败: {e}")
        return False


def test_image_encoding(client, batch_size=5):
    """测试图像编码"""
    print("\n" + "="*50)
    print("测试 4: 图像编码")
    print("="*50)

    try:
        # 准备测试数据（使用在线图片）
        test_images = [
            'https://picsum.photos/224',
            'https://picsum.photos/224?random=1',
            'https://picsum.photos/224?random=2',
            'https://picsum.photos/224?random=3',
            'https://picsum.photos/224?random=4',
        ][:batch_size]

        print(f"测试图像数量: {len(test_images)}")
        print(f"示例 URL: {test_images[0]}")

        # 执行编码
        start_time = time.time()
        embeddings = client.encode(test_images)
        elapsed_time = time.time() - start_time

        # 验证结果
        assert embeddings.shape[0] == len(test_images), "向量数量不匹配"
        assert embeddings.shape[1] == 1024, "向量维度应该是 1024"

        print_success(f"编码成功")
        print(f"  向量形状: {embeddings.shape}")
        print(f"  耗时: {elapsed_time:.2f}秒")
        print(f"  速度: {len(test_images)/elapsed_time:.2f} 条/秒")
        print(f"  数据类型: {embeddings.dtype}")

        return True

    except Exception as e:
        print_error(f"图像编码失败: {e}")
        print_warning("可能需要网络连接来下载测试图片")
        return False


def test_mixed_encoding(client):
    """测试混合编码（文本+图像）"""
    print("\n" + "="*50)
    print("测试 5: 混合编码")
    print("="*50)

    try:
        # 准备混合数据
        mixed_data = [
            '这是一段测试文本',
            'https://picsum.photos/224?random=10',
            'CN-CLIP 图像编码',
            'https://picsum.photos/224?random=11',
        ]

        print(f"混合数据数量: {len(mixed_data)}")
        print(f"  文本: 2 条")
        print(f"  图像: 2 条")

        # 执行编码
        start_time = time.time()
        embeddings = client.encode(mixed_data)
        elapsed_time = time.time() - start_time

        # 验证结果
        assert embeddings.shape[0] == len(mixed_data), "向量数量不匹配"
        assert embeddings.shape[1] == 1024, "向量维度应该是 1024"

        print_success(f"混合编码成功")
        print(f"  向量形状: {embeddings.shape}")
        print(f"  耗时: {elapsed_time:.2f}秒")

        return True

    except Exception as e:
        print_error(f"混合编码失败: {e}")
        return False


def test_single_encoding(client):
    """测试单个数据编码"""
    print("\n" + "="*50)
    print("测试 6: 单个数据编码")
    print("="*50)

    try:
        # 测试单个文本
        single_text = '测试文本'
        print(f"输入: {single_text}")

        start_time = time.time()
        embedding = client.encode(single_text)
        elapsed_time = time.time() - start_time

        # 注意：单个数据会返回 (1, 1024) 的形状
        if embedding.ndim == 1:
            embedding = embedding.reshape(1, -1)

        assert embedding.shape == (1, 1024), f"向量形状应该是 (1, 1024), 实际是 {embedding.shape}"

        print_success(f"单个文本编码成功")
        print(f"  向量形状: {embedding.shape}")
        print(f"  耗时: {elapsed_time:.2f}秒")

        return True

    except Exception as e:
        print_error(f"单个数据编码失败: {e}")
        return False


def main():
    parser = argparse.ArgumentParser(description='CN-CLIP 服务测试脚本')
    parser.add_argument('--url',
                       default='grpc://localhost:51000',
                       help='服务地址（默认：grpc://localhost:51000）')
    parser.add_argument('--text',
                       action='store_true',
                       help='只测试文本编码')
    parser.add_argument('--image',
                       action='store_true',
                       help='只测试图像编码')
    parser.add_argument('--batch-size',
                       type=int,
                       default=10,
                       help='批处理大小（默认：10）')

    args = parser.parse_args()

    print("\n" + "="*50)
    print("CN-CLIP 服务测试")
    print("="*50)

    # 测试 1: 检查依赖
    if not test_imports():
        sys.exit(1)

    # 测试 2: 连接服务
    client = test_connection(args.url)
    if not client:
        sys.exit(1)

    # 运行测试
    results = []

    if args.text:
        # 只测试文本编码
        results.append(test_text_encoding(client, args.batch_size))
    elif args.image:
        # 只测试图像编码
        results.append(test_image_encoding(client, args.batch_size))
    else:
        # 运行所有测试
        results.append(test_text_encoding(client, args.batch_size))
        results.append(test_image_encoding(client, min(args.batch_size, 5)))
        results.append(test_mixed_encoding(client))
        results.append(test_single_encoding(client))

    # 汇总结果
    print("\n" + "="*50)
    print("测试结果汇总")
    print("="*50)

    total_tests = len(results)
    passed_tests = sum(results)

    print(f"总测试数: {total_tests}")
    print(f"通过: {passed_tests}")
    print(f"失败: {total_tests - passed_tests}")

    if passed_tests == total_tests:
        print_success("\n所有测试通过！")
        sys.exit(0)
    else:
        print_error("\n部分测试失败")
        sys.exit(1)


if __name__ == '__main__':
    main()