#!/usr/bin/env python3
"""
K8s 蜂群系统测试脚本
测试 Orchestrator API 和系统功能
"""

import requests
import time
import sys

# Orchestrator 地址
ORCHESTRATOR_URL = "http://52.139.240.116:8000"

def print_section(title):
    """打印分隔线"""
    print(f"\n{'='*60}")
    print(f"  {title}")
    print(f"{'='*60}\n")

def test_health():
    """测试健康检查"""
    print_section("1. 健康检查")
    try:
        response = requests.get(f"{ORCHESTRATOR_URL}/health", timeout=5)
        data = response.json()
        print(f"✅ 状态: {data['status']}")
        print(f"✅ Redis: {data['redis']}")
        print(f"✅ 活跃连接: {data['active_connections']}")
        return True
    except Exception as e:
        print(f"❌ 健康检查失败: {e}")
        return False

def test_agents():
    """测试 Agent 列表"""
    print_section("2. Agent 列表")
    try:
        response = requests.get(f"{ORCHESTRATOR_URL}/agents", timeout=5)
        agents = response.json().get("agents", [])
        print(f"当前 Agent 数量: {len(agents)}")
        if agents:
            for agent in agents:
                print(f"  - Agent ID: {agent['agent_id']}")
                print(f"    状态: {agent['status']}")
                print(f"    最后心跳: {agent.get('last_heartbeat', 'N/A')}")
        else:
            print("  暂无 Agent 运行")
        return True
    except Exception as e:
        print(f"❌ 获取 Agent 列表失败: {e}")
        return False

def test_tasks():
    """测试任务列表"""
    print_section("3. 任务列表")
    try:
        response = requests.get(f"{ORCHESTRATOR_URL}/tasks", timeout=5)
        tasks = response.json().get("tasks", [])
        print(f"任务总数: {len(tasks)}")
        if tasks:
            for task in tasks:
                print(f"  - 任务 ID: {task['task_id']}")
                print(f"    状态: {task['status']}")
                print(f"    描述: {task.get('description', 'N/A')[:50]}...")
        else:
            print("  暂无任务")
        return True
    except Exception as e:
        print(f"❌ 获取任务列表失败: {e}")
        return False

def test_create_task():
    """测试创建任务"""
    print_section("4. 创建测试任务")
    try:
        task_data = {
            "description": "测试任务：修复一个简单的语法错误",
            "context": {
                "workspace_id": "test-workspace-001",
                "required_agents": 1
            },
            "max_retries": 3
        }
        response = requests.post(
            f"{ORCHESTRATOR_URL}/tasks",
            json=task_data,
            timeout=5
        )
        if response.status_code == 200:
            task = response.json()
            print(f"✅ 任务创建成功")
            print(f"  任务 ID: {task['task_id']}")
            print(f"  状态: {task['status']}")
            return task['task_id']
        else:
            print(f"❌ 任务创建失败: {response.status_code}")
            print(f"  响应: {response.text}")
            return None
    except Exception as e:
        print(f"❌ 创建任务失败: {e}")
        return None

def test_metrics():
    """测试 Prometheus 指标"""
    print_section("5. Prometheus 指标")
    try:
        response = requests.get(f"{ORCHESTRATOR_URL}/metrics", timeout=5)
        metrics = response.text

        # 提取关键指标
        lines = metrics.split('\n')
        key_metrics = [
            'swarm_agents_created_total',
            'swarm_agents_active',
            'swarm_tasks_created_total',
            'swarm_tasks_completed_total'
        ]

        print("关键指标:")
        for line in lines:
            for metric in key_metrics:
                if line.startswith(metric) and not line.startswith('#'):
                    print(f"  {line}")
        return True
    except Exception as e:
        print(f"❌ 获取指标失败: {e}")
        return False

def test_redis_connection():
    """测试 Redis 连接"""
    print_section("6. Redis 连接测试")
    try:
        # 通过健康检查验证 Redis
        response = requests.get(f"{ORCHESTRATOR_URL}/health", timeout=5)
        data = response.json()
        if data['redis'] == 'connected':
            print("✅ Redis 连接正常")
            return True
        else:
            print(f"❌ Redis 状态: {data['redis']}")
            return False
    except Exception as e:
        print(f"❌ Redis 连接测试失败: {e}")
        return False

def main():
    """主测试流程"""
    print("\n" + "="*60)
    print("  K8s 蜂群系统 - 集成测试")
    print("="*60)
    print(f"\nOrchestrator URL: {ORCHESTRATOR_URL}")
    print(f"测试时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")

    results = []

    # 运行所有测试
    results.append(("健康检查", test_health()))
    results.append(("Redis 连接", test_redis_connection()))
    results.append(("Agent 列表", test_agents()))
    results.append(("任务列表", test_tasks()))
    results.append(("创建任务", test_create_task() is not None))
    results.append(("Prometheus 指标", test_metrics()))

    # 总结
    print_section("测试总结")
    passed = sum(1 for _, result in results if result)
    total = len(results)

    for name, result in results:
        status = "✅ 通过" if result else "❌ 失败"
        print(f"{status} - {name}")

    print(f"\n总计: {passed}/{total} 测试通过")

    if passed == total:
        print("\n🎉 所有测试通过！系统运行正常。")
        return 0
    else:
        print(f"\n⚠️  {total - passed} 个测试失败，请检查日志。")
        return 1

if __name__ == "__main__":
    sys.exit(main())