#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
百度图像搜索API测试脚本
简单测试API功能，打印结果
支持实际API和模拟模式
"""

import os
import json
import argparse
from dotenv import load_dotenv
from app.api.baidu_image_search import BaiduImageSearch
from app.api.image_utils import ImageUtils

# 导入模拟类
from mock_baidu_api import MockBaiduImageSearch

# 全局变量，用于存储测试过程中的图片ID
test_cont_sign = None

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='百度图像搜索API测试脚本')
    parser.add_argument('--mock', action='store_true', help='使用模拟模式，不连接实际API')
    return parser.parse_args()

def check_api_keys():
    """检查API密钥是否已配置"""
    api_key = os.getenv('BAIDU_API_KEY')
    secret_key = os.getenv('BAIDU_SECRET_KEY')
    
    if not api_key or not secret_key:
        print("\n\033[91m错误: API密钥未配置\033[0m")
        print("\n请在.env文件中添加以下内容:")
        print("BAIDU_API_KEY=你的API密钥")
        print("BAIDU_SECRET_KEY=你的密钥")
        print("\n您可以从百度AI开放平台获取密钥: https://ai.baidu.com/")
        print("\n或者使用 --mock 参数运行模拟模式: python test_baidu_api.py --mock")
        return False
    
    print(f"API密钥: {api_key[:4]}...")
    print(f"Secret密钥: {secret_key[:4]}...")
    return True

def main():
    # 解析命令行参数
    args = parse_arguments()
    
    # 加载环境变量
    print("加载环境变量...")
    load_dotenv()
    
    # 判断是否使用模拟模式
    if args.mock:
        print("\n\033[93m注意: 使用模拟模式，不连接实际API\033[0m")
        api = MockBaiduImageSearch()
        
        # 测试添加图片
        test_add_image(api)
        
        # 测试搜索图片
        test_search_image(api)
        
        # 测试更新图片
        test_update_image(api)
        
        # 测试删除图片
        test_delete_image(api)
        return
    
    # 如果不是模拟模式，检查API密钥
    if not check_api_keys():
        return
    
    # 初始化实际API
    print("\n初始化百度图像搜索API...")
    try:
        api = BaiduImageSearch()
        if api.access_token:
            print(f"获取到的Access Token: {api.access_token[:10]}...（已截断）")
            
            # 测试添加图片
            test_add_image(api)
            
            # 测试搜索图片
            test_search_image(api)
            
            # 测试更新图片
            test_update_image(api)
            
            # 测试删除图片
            test_delete_image(api)
        else:
            print("\033[91m错误: 无法获取Access Token\033[0m")
            print("\n尝试使用模拟模式: python test_baidu_api.py --mock")
    except Exception as e:
        print(f"\033[91m错误: {e}\033[0m")
        print("\n尝试使用模拟模式: python test_baidu_api.py --mock")

def test_add_image(api):
    """测试添加图片功能"""
    print("\n===== 测试添加图片 =====")
    
    # 使用URL方式添加图片
    test_image_url = "https://www.baidu.com/img/PCtm_d9c8750bed0b3c7d089fa7d55720d6cf.png"  # 百度logo
    brief = json.dumps({"name": "测试图片", "id": "test001"}, ensure_ascii=False)
    tags = "测试,图片"
    
    try:
        result = api.add_image(
            url=test_image_url,
            brief=brief,
            tags=tags
        )
        print("添加图片结果:")
        print(json.dumps(result, indent=2, ensure_ascii=False))
        
        # 保存contSign用于后续测试
        if result.get("cont_sign"):
            global test_cont_sign
            test_cont_sign = result.get("cont_sign")
            print(f"已保存图片ID: {test_cont_sign} 用于后续测试")
    except Exception as e:
        print(f"添加图片失败: {e}")

def test_search_image(api):
    """测试搜索图片功能"""
    print("\n===== 测试搜索图片 =====")
    
    # 使用URL方式搜索图片
    test_image_url = "https://www.baidu.com/img/PCtm_d9c8750bed0b3c7d089fa7d55720d6cf.png"  # 百度logo
    
    try:
        result = api.search_image(
            url=test_image_url,
            tags="测试"
        )
        print("搜索图片结果:")
        print(f"找到 {len(result.get('result', []))} 个匹配项")
        
        # 只打印前2个结果
        for i, item in enumerate(result.get("result", [])[:2]):
            print(f"结果 {i+1}:")
            print(f"  相似度: {item.get('score', 0)}")
            print(f"  简介: {item.get('brief', '')}")
    except Exception as e:
        print(f"搜索图片失败: {e}")

def test_update_image(api):
    """测试更新图片功能"""
    print("\n===== 测试更新图片 =====")
    
    global test_cont_sign
    if not test_cont_sign:
        print("没有可用的图片ID，跳过更新图片测试")
        return
    
    try:
        new_brief = json.dumps({"name": "更新后的测试图片", "id": "test001"}, ensure_ascii=False)
        result = api.update_image(test_cont_sign, new_brief)
        print("更新图片结果:")
        print(json.dumps(result, indent=2, ensure_ascii=False))
    except Exception as e:
        print(f"更新图片失败: {e}")

def test_delete_image(api):
    """测试删除图片功能"""
    print("\n===== 测试删除图片 =====")
    
    global test_cont_sign
    if not test_cont_sign:
        print("没有可用的图片ID，跳过删除图片测试")
        return
    
    try:
        result = api.delete_image(test_cont_sign)
        print("删除图片结果:")
        print(json.dumps(result, indent=2, ensure_ascii=False))
        # 清除已删除的图片ID
        test_cont_sign = None
    except Exception as e:
        print(f"删除图片失败: {e}")

if __name__ == "__main__":
    main()