# -*- coding: utf-8 -*-
"""
功能验证测试套件

验证 MCP 接口向后兼容性，确保单 KEY 用户无感升级，多 KEY 功能正常工作。

测试覆盖:
1. MCP 接口签名验证
2. 单 KEY 兼容性
3. 负载均衡功能验证
"""

import asyncio
import os
import sys
import inspect
import pytest
from unittest.mock import patch, AsyncMock, MagicMock

# 导入待测试模块
from server import search_web
from load_balancer import is_key_failure, ApiKeyManager
from config import load_api_keys


# ============================================================================
# Task 1: MCP 接口签名验证
# ============================================================================

def test_mcp_interface_signature() -> None:
    """
    验证 search_web 函数签名与 Phase 1 之前一致。

    验证内容:
    1. 函数参数: query, search_type, num_results
    2. 参数类型: str, str, Optional[int | str]
    3. 默认值: search_type="search", num_results=4
    4. 返回值类型: str
    5. 文档字符串存在且描述一致
    6. Gradio API 注册: api_name="search_web"
    """
    print("=" * 60)
    print("Task 1: MCP 接口签名验证")
    print("=" * 60)

    # 1. 验证函数签名
    sig = inspect.signature(search_web)
    print(f"函数签名: {sig}")

    # 1.1 验证参数名称
    params = list(sig.parameters.keys())
    expected_params = ["query", "search_type", "num_results"]
    assert params == expected_params, f"参数名称不匹配: {params} != {expected_params}"
    print("✓ 参数名称验证通过")

    # 1.2 验证参数类型
    query_param = sig.parameters["query"]
    search_type_param = sig.parameters["search_type"]
    num_results_param = sig.parameters["num_results"]

    assert query_param.annotation is str, "query 参数类型应为 str"
    assert search_type_param.annotation is str, "search_type 参数类型应为 str"
    # num_results 可以是 int 或 str 或 None
    print("✓ 参数类型验证通过")

    # 1.3 验证默认值
    assert search_type_param.default == "search", "search_type 默认值应为 'search'"
    assert num_results_param.default == 4, "num_results 默认值应为 4"
    print("✓ 默认值验证通过")

    # 1.4 验证返回值类型
    # search_web 是 async 函数，返回 Coroutine，但实际返回值是 str
    return_annotation = str(search_web.__annotations__.get("return", ""))
    # 检查是否包含 "str"（可能是 "str" 或 "Coroutine[str]" 等形式）
    assert "str" in return_annotation or "Coroutine" in return_annotation, f"返回值应为 str 或 Coroutine: {return_annotation}"
    print("✓ 返回值类型验证通过")

    # 1.5 验证文档字符串
    assert search_web.__doc__ is not None, "函数应包含文档字符串"
    assert "Search the web" in search_web.__doc__, "文档应包含 'Search the web' 描述"
    print("✓ 文档字符串验证通过")

    # 2. 验证 FastMCP API 注册
    # 检查 server.py 中是否有 @mcp.tool 装饰器
    with open("server.py", "r", encoding="utf-8") as f:
        server_content = f.read()
    assert '@mcp.tool' in server_content and 'async def search_web' in server_content, "FastMCP API 注册应为 @mcp.tool"
    print("✓ FastMCP API 注册验证通过")

    print("=" * 60)
    print("Task 1: MCP 接口签名验证 - 全部通过")
    print("=" * 60)


# ============================================================================
# Task 2: 单 KEY 兼容性验证
# ============================================================================

def test_single_key_compatibility() -> None:
    """
    验证单 KEY 配置正常工作，用户无需修改配置。

    测试内容:
    1. 单 KEY 环境变量加载
    2. 单 KEY 场景下 load_api_keys() 返回单元素列表
    3. 单 KEY 场景下 KEY 解析正确
    """
    print("=" * 60)
    print("Task 2: 单 KEY 兼容性验证")
    print("=" * 60)

    # 1. 测试单 KEY 配置加载
    test_key = "test_single_key_12345"

    # 备份原始环境变量
    original_env = os.environ.get("SERPER_API_KEYS")

    try:
        # 设置单 KEY 环境变量
        os.environ["SERPER_API_KEYS"] = test_key

        # 2. 验证 load_api_keys() 返回单元素列表
        # 注意: 由于 get_manager() 会缓存实例，我们需要重新加载模块或清除缓存
        # 这里我们直接测试 load_api_keys 函数
        keys = load_api_keys()
        assert len(keys) == 1, f"单 KEY 配置应返回 1 个 KEY，实际返回 {len(keys)} 个"
        assert keys[0] == test_key, f"KEY 解析错误: {keys[0]} != {test_key}"
        print(f"✓ 单 KEY 配置加载验证通过: {keys}")

        # 3. 验证单 KEY 场景下的负载均衡器
        manager = ApiKeyManager(keys)
        next_key = manager.get_next_key()
        assert next_key == test_key, f"单 KEY 场景下应返回唯一的 KEY: {next_key} != {test_key}"
        print("✓ 单 KEY 场景负载均衡验证通过")

        # 4. 验证多次调用返回同一个 KEY
        next_key2 = manager.get_next_key()
        assert next_key2 == test_key, "单 KEY 场景下多次调用应返回同一个 KEY"
        print("✓ 单 KEY 场景轮询验证通过")

    finally:
        # 恢复原始环境变量
        if original_env is not None:
            os.environ["SERPER_API_KEYS"] = original_env
        elif "SERPER_API_KEYS" in os.environ:
            del os.environ["SERPER_API_KEYS"]

    print("=" * 60)
    print("Task 2: 单 KEY 兼容性验证 - 全部通过")
    print("=" * 60)


@pytest.mark.asyncio
async def test_single_key_search_web_mock() -> None:
    """
    使用 mock 模拟成功的 API 响应，验证单 KEY 场景下的 search_web 调用。

    注意: 此测试需要 mock HTTP 请求，仅验证接口行为不变。
    """
    print("=" * 60)
    print("Task 2: 单 KEY search_web Mock 测试")
    print("=" * 60)

    # 备份原始环境变量
    original_env = os.environ.get("SERPER_API_KEYS")

    try:
        # 设置单 KEY 环境变量
        os.environ["SERPER_API_KEYS"] = "test_single_key_12345"

        # Mock HTTP 响应
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "organic": [
                {
                    "title": "Test Result",
                    "link": "https://example.com/test",
                    "snippet": "Test snippet"
                }
            ]
        }

        # Mock httpx.AsyncClient
        mock_client = AsyncMock()
        mock_client.__aenter__.return_value = mock_client
        mock_client.post.return_value = mock_response
        mock_client.get.return_value = MagicMock(
            status_code=200,
            text="<html><body>Test content</body></html>"
        )

        with patch("server.httpx.AsyncClient", return_value=mock_client):
            with patch("server.limiter.hit", new_callable=AsyncMock, return_value=True):
                    result = await search_web("test query", "search", 1)

                    # 验证返回值格式
                    assert isinstance(result, str), "返回值应为字符串"
                    assert "Successfully extracted content" in result, "成功响应应包含确认信息"
                    print("✓ 单 KEY search_web 调用验证通过")
                    print("✓ 返回值类型和格式验证通过")

    finally:
        # 恢复原始环境变量
        if original_env is not None:
            os.environ["SERPER_API_KEYS"] = original_env
        elif "SERPER_API_KEYS" in os.environ:
            del os.environ["SERPER_API_KEYS"]

    print("=" * 60)
    print("Task 2: 单 KEY search_web Mock 测试 - 全部通过")
    print("=" * 60)


# ============================================================================
# Task 3: 负载均衡功能验证
# ============================================================================

def test_multi_key_roundrobin() -> None:
    """
    验证多 KEY Round Robin 轮询逻辑。

    测试内容:
    1. KEY 按顺序轮换
    2. 所有 KEY 都被使用
    3. 轮询循环正确
    """
    print("=" * 60)
    print("Task 3: 多 KEY Round Robin 轮询验证")
    print("=" * 60)

    # 创建多 KEY 测试场景
    keys = ["key1", "key2", "key3"]
    manager = ApiKeyManager(keys)

    # 1. 验证 KEY 按顺序轮换
    # 注意: get_next_key() 成功后会移动到下一个位置
    key_sequence = []
    for i in range(6):  # 轮询 2 圈
        key = manager.get_next_key()
        key_sequence.append(key)

    print(f"KEY 轮换序列: {key_sequence}")

    # 验证所有 KEY 都被使用
    unique_keys = set(key_sequence)
    assert unique_keys == set(keys), f"所有 KEY 都应被使用: {unique_keys} != {set(keys)}"
    print(f"✓ 所有 KEY 都被使用: {unique_keys}")

    # 验证轮询顺序
    # Round Robin 逻辑: index 从 0 开始，每次成功后移动到下一个位置
    # 初始 current_index = 0
    # 第一次: keys[0] -> key1, 然后 current_index = 1
    # 第二次: keys[1] -> key2, 然后 current_index = 2
    # 第三次: keys[2] -> key3, 然后 current_index = 0
    # 第四次: keys[0] -> key1, 然后 current_index = 1
    # ...
    expected_sequence = ["key1", "key2", "key3", "key1", "key2", "key3"]
    assert key_sequence == expected_sequence, f"轮询顺序错误: {key_sequence} != {expected_sequence}"
    print("✓ 轮询顺序验证通过")

    print("=" * 60)
    print("Task 3: 多 KEY Round Robin 轮询验证 - 全部通过")
    print("=" * 60)


def test_key_failure_handling() -> None:
    """
    验证 KEY 失效处理逻辑。

    测试内容:
    1. mark_key_failed() 将 KEY 标记为失效
    2. 失效的 KEY 进入冷却期
    3. get_next_key() 跳过冷却中的 KEY
    4. mark_key_success() 清除 KEY 的冷却状态
    """
    print("=" * 60)
    print("Task 3: KEY 失效处理验证")
    print("=" * 60)

    # 创建多 KEY 测试场景
    keys = ["key1", "key2", "key3"]
    manager = ApiKeyManager(keys, cooldown_seconds=60)  # 60 秒冷却期

    # 1. 标记 key1 为失效（使用实例方法而非全局函数）
    manager.mark_key_failed("key1")
    print("✓ 标记 key1 为失效")

    # 2. 验证失效的 KEY 在冷却中
    assert "key1" in manager.key_cooldown_until, "key1 应在冷却中"
    print(f"✓ key1 进入冷却期: {manager.key_cooldown_until['key1']}")

    # 3. 验证 get_next_key() 跳过冷却中的 KEY
    next_key = manager.get_next_key()
    assert next_key != "key1", "应跳过冷却中的 key1"
    print(f"✓ 跳过冷却中的 KEY，返回: {next_key}")

    # 4. 验证其他 KEY 仍可使用
    next_keys = []
    for _ in range(5):
        key = manager.get_next_key()
        next_keys.append(key)

    assert "key1" not in next_keys, "冷却中的 key1 不应被返回"
    assert all(k in ["key2", "key3"] for k in next_keys), "应只使用 key2 和 key3"
    print(f"✓ 只使用可用 KEY: {set(next_keys)}")

    # 5. 验证 mark_key_success() 清除冷却状态（使用实例方法）
    manager.mark_key_success("key1")
    assert "key1" not in manager.key_cooldown_until, "key1 应从冷却中清除"
    print("✓ 清除 key1 冷却状态")

    # 6. 验证清除后 KEY 可用
    next_key = manager.get_next_key()
    assert next_key in keys, f"清除后 KEY 应可用: {next_key}"
    print(f"✓ 清除后 KEY 可用: {next_key}")

    print("=" * 60)
    print("Task 3: KEY 失效处理验证 - 全部通过")
    print("=" * 60)


def test_key_failure_detection() -> None:
    """
    验证失效检测函数 is_key_failure()。

    测试内容:
    1. 401 状态码返回 True (认证失败)
    2. 403 状态码返回 True (授权失败)
    3. 429 状态码返回 True (速率限制)
    4. 5xx 状态码返回 True (服务器错误)
    5. 200 状态码返回 False (成功)
    6. 400 状态码返回 False (客户端错误，非 KEY 问题)
    7. 超时异常返回 True
    8. 连接异常返回 True
    """
    print("=" * 60)
    print("Task 3: 失效检测函数验证")
    print("=" * 60)

    # 1. 验证状态码检测
    test_cases = [
        (401, True, "401 认证失败应返回 True"),
        (403, True, "403 授权失败应返回 True"),
        (429, True, "429 速率限制应返回 True"),
        (500, True, "500 服务器错误应返回 True"),
        (503, True, "503 服务不可用应返回 True"),
        (200, False, "200 成功应返回 False"),
        (400, False, "400 客户端错误应返回 False"),
        (404, False, "404 未找到应返回 False"),
    ]

    for status_code, expected, description in test_cases:
        result = is_key_failure(status_code)
        assert result == expected, f"{description}: {result} != {expected}"
        print(f"✓ {description}: {result}")

    # 2. 验证异常检测
    from httpx import TimeoutException, ConnectError

    assert is_key_failure(200, TimeoutException("timeout")), "超时异常应返回 True"
    print("✓ 超时异常检测: True")

    assert is_key_failure(200, ConnectError("connection failed")), "连接异常应返回 True"
    print("✓ 连接异常检测: True")

    assert not is_key_failure(200, ValueError("other error")), "其他异常应返回 False"
    print("✓ 其他异常检测: False")

    print("=" * 60)
    print("Task 3: 失效检测函数验证 - 全部通过")
    print("=" * 60)


# ============================================================================
# 手动测试指南
# ============================================================================

def print_manual_test_guide() -> None:
    """
    打印手动测试指南，供用户自行验证功能。

    指南内容:
    1. 如何配置环境变量
    2. 如何运行实际 API 测试
    3. 单 KEY 测试步骤
    4. 多 KEY 测试步骤
    5. 轮询验证方法
    """
    print("=" * 60)
    print("手动测试指南")
    print("=" * 60)

    guide = """
# 手动测试指南

## 环境配置

### 单 KEY 配置
```bash
export SERPER_API_KEYS="your_single_api_key_here"
```

### 多 KEY 配置
```bash
export SERPER_API_KEYS="key1,key2,key3"
```

## 运行测试

### 1. 单 KEY 测试
```bash
# 配置单 KEY
export SERPER_API_KEYS="your_api_key_here"

# 启动 MCP 服务器
python server.py

# 使用 MCP 客户端调用 search_web 工具
# 验证接口签名和功能正常
```

### 2. 多 KEY 测试
```bash
# 配置多 KEY (逗号分隔)
export SERPER_API_KEYS="key1,key2,key3"

# 启动 MCP 服务器
python server.py

# 观察控制台输出，验证 KEY 轮换
# 多次调用 search_web 观察不同 KEY 被使用
```

### 3. 轮询验证
```bash
# 在 MCP 客户端中连续调用 search_web 多次
# 观察控制台输出中的 KEY 使用情况
# 验证 KEY 按顺序轮换: key1 -> key2 -> key3 -> key1 -> ...

# 查看负载均衡器内部状态
python -c "
from load_balancer import get_manager
manager = get_manager()
print(f'Available keys: {manager.keys}')
print(f'Current index: {manager.current_index}')
print(f'Next key: {manager.get_next_key()}')
"
```

### 4. 失效处理验证
```bash
# 1. 配置无效 KEY 测试失效检测
export SERPER_API_KEYS="invalid_key_1,invalid_key_2,valid_key_3"

# 2. 观察控制台输出，验证失败检测和 KEY 切换
# 3. 验证最终使用有效 KEY 完成

# 手动标记 KEY 失效测试
python -c "
from load_balancer import get_manager, mark_key_failed, mark_key_success

manager = get_manager()
print('Available keys:', manager.keys)

# 标记 key1 失效
mark_key_failed(manager.keys[0])
print(f'Marked {manager.keys[0]} as failed')

# 查看下一个 KEY
next_key = manager.get_next_key()
print(f'Next key: {next_key}')

# 恢复 key1
mark_key_success(manager.keys[0])
print(f'Restored {manager.keys[0]}')
"
```

## 验证标准

- [ ] 单 KEY 配置正常工作，无需修改配置
- [ ] MCP 接口签名完全不变
- [ ] 多 KEY 场景下 KEY 按顺序轮换
- [ ] 失效的 KEY 进入冷却期并被跳过
- [ ] 所有 KEY 失效时返回友好错误信息
"""

    print(guide)
    print("=" * 60)


# ============================================================================
# 主测试运行器
# ============================================================================

def main() -> None:
    """运行所有测试。"""
    print("\n" + "=" * 60)
    print("开始执行功能验证测试套件")
    print("=" * 60 + "\n")

    try:
        # Task 1: MCP 接口签名验证
        test_mcp_interface_signature()
        print()

        # Task 2: 单 KEY 兼容性验证
        test_single_key_compatibility()
        print()

        # 运行异步测试
        asyncio.run(test_single_key_search_web_mock())
        print()

        # Task 3: 负载均衡功能验证
        test_multi_key_roundrobin()
        print()

        test_key_failure_handling()
        print()

        test_key_failure_detection()
        print()

        # 打印手动测试指南
        print_manual_test_guide()

        print("\n" + "=" * 60)
        print("所有自动化测试通过！")
        print("=" * 60 + "\n")

    except AssertionError as e:
        print(f"\n❌ 测试失败: {e}")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ 测试出错: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()
