#!/usr/bin/env python3
"""
BIZ-26 限流器测试脚本

测试场景：
1. 令牌桶限流功能
2. 优先级队列调度
3. 缓存管理器
4. 重试机制
5. 429 错误模拟

运行方式：
    python3 scripts/test_rate_limiter.py
"""

import sys
import time
import threading
from datetime import datetime

# 添加脚本目录到路径
sys.path.insert(0, "/home/vincent/.openclaw/workspace/costcodev/EnterpriseArchitect/scripts")

from rate_limiter import (
    TokenBucket,
    CacheManager,
    RequestScheduler,
    Priority,
    retry_with_backoff,
    CoordinatedPoller,
    is_nvidia_gateway,
)


def test_token_bucket():
    """测试令牌桶限流器"""
    print("=" * 60)
    print("测试 1: 令牌桶限流器")
    print("=" * 60)
    
    # 创建限流器：40 RPM = 0.67 令牌/秒
    bucket = TokenBucket(rate=40/60, capacity=40)
    
    print(f"\n初始状态：{bucket.get_status()}")
    
    # 快速消费 10 个令牌
    print("\n快速消费 10 个令牌...")
    success_count = 0
    for i in range(10):
        if bucket.consume():
            success_count += 1
    
    print(f"成功消费：{success_count}/10")
    print(f"消费后状态：{bucket.get_status()}")
    
    # 测试等待获取令牌
    print("\n测试等待获取令牌...")
    start = time.time()
    got_token = bucket.wait_for_token(timeout=2.0)
    elapsed = time.time() - start
    
    print(f"等待耗时：{elapsed:.3f}s, 获取成功：{got_token}")
    print(f"等待后状态：{bucket.get_status()}")
    
    print("\n✅ 令牌桶测试完成\n")


def test_cache_manager():
    """测试缓存管理器"""
    print("=" * 60)
    print("测试 2: 缓存管理器")
    print("=" * 60)
    
    cache = CacheManager()
    
    # 测试 WorkBoard 缓存（TTL 5 分钟）
    print("\n1. 设置 WorkBoard 缓存（TTL 5 分钟）")
    cache.set("workboard", {"query": "status=todo"}, [{"id": "card1", "title": "Test"}])
    
    # 立即读取
    result = cache.get("workboard", {"query": "status=todo"})
    print(f"   立即读取：{result is not None}")
    
    # 测试配置缓存（TTL 1 小时）
    print("\n2. 设置配置缓存（TTL 1 小时）")
    cache.set("config", "agent_list", ["costcodev", "secretary", "coo"])
    result = cache.get("config", "agent_list")
    print(f"   读取配置：{result}")
    
    # 测试缓存统计
    print("\n3. 缓存统计")
    stats = cache.get_stats()
    print(f"   总条目数：{stats['total_entries']}")
    print(f"   按类别：{stats['by_category']}")
    
    # 测试缓存删除
    print("\n4. 删除缓存")
    deleted = cache.delete("workboard", {"query": "status=todo"})
    print(f"   删除成功：{deleted}")
    result = cache.get("workboard", {"query": "status=todo"})
    print(f"   删除后读取：{result is None}")
    
    print("\n✅ 缓存管理器测试完成\n")


def test_priority_queue():
    """测试优先级队列调度"""
    print("=" * 60)
    print("测试 3: 优先级队列调度（简化版，不启动工作线程）")
    print("=" * 60)
    
    scheduler = RequestScheduler(rate=40/60, capacity=40, enable_cache=True)
    
    # 模拟请求处理结果
    results = []
    
    def record_result(data):
        results.append((time.time(), data))
        return data
    
    # 提交不同优先级的请求（不启动工作线程，只测试队列）
    print("\n提交请求（按顺序）:")
    scheduler.submit(
        payload={"task": "normal_1"},
        priority=Priority.NORMAL,
        callback=record_result
    )
    print("  1. 正常优先级：normal_1")
    
    scheduler.submit(
        payload={"task": "urgent_1"},
        priority=Priority.URGENT,
        callback=record_result
    )
    print("  2. 紧急优先级：urgent_1")
    
    scheduler.submit(
        payload={"task": "low_1"},
        priority=Priority.LOW,
        callback=record_result
    )
    print("  3. 低优先级：low_1")
    
    scheduler.submit(
        payload={"task": "high_1"},
        priority=Priority.HIGH,
        callback=record_result
    )
    print("  4. 高优先级：high_1")
    
    # 查看队列大小
    print(f"\n队列大小：{scheduler.get_queue_size()}")
    
    # 查看状态
    status = scheduler.get_status()
    print(f"初始令牌数：{status['token_bucket']['tokens']}")
    
    print("\n✅ 优先级队列测试完成（仅提交，未处理）\n")


def test_retry_decorator():
    """测试重试装饰器"""
    print("=" * 60)
    print("测试 4: 重试装饰器")
    print("=" * 60)
    
    attempt_count = [0]
    
    @retry_with_backoff(max_retries=3, base_delay=0.1, jitter=False)
    def flaky_function():
        attempt_count[0] += 1
        if attempt_count[0] < 3:
            raise Exception(f"模拟失败 (尝试 {attempt_count[0]})")
        return f"成功 (尝试 {attempt_count[0]})"
    
    print("\n调用易失败函数（前 2 次失败，第 3 次成功）...")
    start = time.time()
    result = flaky_function()
    elapsed = time.time() - start
    
    print(f"结果：{result}")
    print(f"总尝试次数：{attempt_count[0]}")
    print(f"总耗时：{elapsed:.3f}s")
    
    print("\n✅ 重试装饰器测试完成\n")


def test_coordinated_poller():
    """测试统一轮询器"""
    print("=" * 60)
    print("测试 5: COO 统一轮询器（简化版，短间隔测试）")
    print("=" * 60)
    
    scheduler = RequestScheduler(rate=40/60, capacity=40)
    poller = CoordinatedPoller(scheduler, poll_interval=2)  # 2 秒轮询一次（测试用）
    
    received_results = []
    
    def on_poll_result(result):
        received_results.append((datetime.now().strftime("%H:%M:%S"), result))
        print(f"   [{datetime.now().strftime('%H:%M:%S')}] 收到轮询结果")
    
    poller.subscribe(on_poll_result)
    
    print("\n启动轮询器（轮询间隔 2 秒，运行 5 秒后停止）...")
    poller.start()
    
    # 等待 5 秒
    time.sleep(5)
    
    poller.stop()
    
    print(f"\n收到结果次数：{len(received_results)}")
    for ts, result in received_results:
        print(f"  {ts}: {result['timestamp'][:19]}")
    
    print("\n✅ 统一轮询器测试完成\n")


def test_rate_limit_stress():
    """压力测试：快速提交大量请求"""
    print("=" * 60)
    print("测试 6: 压力测试（40 RPM 限制下提交 50 个请求）")
    print("=" * 60)
    
    scheduler = RequestScheduler(rate=40/60, capacity=40, enable_cache=True)
    scheduler.start()
    
    completed = []
    failed = []
    lock = threading.Lock()
    
    def callback(data):
        with lock:
            completed.append(data)
        return data
    
    print("\n快速提交 50 个请求...")
    start_time = time.time()
    
    for i in range(50):
        priority = Priority.NORMAL if i % 10 != 0 else Priority.URGENT
        scheduler.submit(
            payload={"index": i, "provider": "nvidia"},
            priority=priority,
            callback=callback,
            gateway="nvidia"
        )
    
    print("提交完成，等待处理...")
    
    # 等待 10 秒
    time.sleep(10)
    
    elapsed = time.time() - start_time
    
    # 查看统计
    status = scheduler.get_status()
    print(f"\n耗时：{elapsed:.2f}s")
    print(f"队列大小：{status['queue_size']}")
    print(f"已完成：{status['stats']['completed_requests']}")
    print(f"失败：{status['stats']['failed_requests']}")
    print(f"降级：{status['stats']['fallback_requests']}")
    print(f"令牌桶状态：{status['token_bucket']}")
    
    scheduler.stop()
    
    print("\n✅ 压力测试完成\n")


def test_gateway_scope():
    """测试限流范围：只对 NVIDIA 网关生效"""
    print("=" * 60)
    print("测试 7: 网关范围识别（只限 NVIDIA）")
    print("=" * 60)

    assert is_nvidia_gateway("nvidia") is True
    assert is_nvidia_gateway("nvidiavx18088980513/deepseek-ai/deepseek-v4-pro") is True
    assert is_nvidia_gateway("volcengine-plan/ark-code-latest") is False
    assert is_nvidia_gateway("siliconflow/Qwen/Qwen3") is False
    assert is_nvidia_gateway("deepseek/deepseek-chat") is False
    assert is_nvidia_gateway(None) is False

    scheduler = RequestScheduler(rate=1/60, capacity=1, enable_cache=True)
    # 先耗尽 NVIDIA 桶
    scheduler.submit(payload={"provider": "nvidia", "i": 1}, priority=Priority.NORMAL, callback=lambda x: x, gateway="nvidia")
    # 非 NVIDIA 请求应直接执行，不受桶状态影响
    non_nv = {"provider": "volcengine-plan", "i": 2}
    assert scheduler._should_rate_limit(type("R", (), {"gateway": "volcengine-plan", "model": None, "payload": non_nv})()) is False

    print("✅ 网关范围识别测试完成：volcengine-plan/siliconflow/DeepSeek 不限流，NVIDIA 限流\n")


def main():
    """运行所有测试"""
    print("\n")
    print("╔" + "=" * 58 + "╗")
    print("║" + " " * 58 + "║")
    print("║" + "  BIZ-26 限流器测试套件".center(58) + "║")
    print("║" + "  API 请求优先级队列 + 令牌桶限流".center(58) + "║")
    print("║" + " " * 58 + "║")
    print("╚" + "=" * 58 + "╝")
    print()
    
    try:
        test_token_bucket()
        test_cache_manager()
        test_priority_queue()
        test_retry_decorator()
        test_coordinated_poller()
        test_rate_limit_stress()
        test_gateway_scope()
        
        print("\n")
        print("╔" + "=" * 58 + "╗")
        print("║" + " " * 58 + "║")
        print("║" + "  ✅ 所有测试完成".center(58) + "║")
        print("║" + " " * 58 + "║")
        print("╚" + "=" * 58 + "╝")
        print()
        
    except KeyboardInterrupt:
        print("\n\n⚠️  测试被用户中断\n")
    except Exception as e:
        print(f"\n\n❌ 测试出错：{e}\n")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()