#!/usr/bin/env python3
"""
测试 Amazon Q API 可用模型的脚本
"""

import requests
import json
import time
from typing import List, Dict, Optional

# API 配置
BASE_URL = "http://localhost:8000"
API_KEY = ""  # 根据你的 .env 配置设置，如果为空则使用开发模式

# 常见的 Claude 模型列表（基于 Amazon Q 可能支持的模型）
TEST_MODELS = [
    "claude-sonnet-4",
    "claude-sonnet-3.5", 
    "claude-haiku-3.5",
    "claude-opus-3",
    "claude-sonnet-3",
    "claude-haiku-3",
    "claude-instant-1.2",
    "anthropic.claude-3-sonnet-20240229-v1:0",
    "anthropic.claude-3-haiku-20240307-v1:0",
    "anthropic.claude-3-opus-20240229-v1:0",
]

def test_model(model_name: str, api_key: str = "") -> Dict[str, any]:
    """测试单个模型是否可用"""
    url = f"{BASE_URL}/v1/chat/completions"
    
    headers = {
        "Content-Type": "application/json",
    }
    
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    
    payload = {
        "model": model_name,
        "messages": [
            {"role": "user", "content": "请简单回复：你好"}
        ],
        "stream": False,
        "max_tokens": 10
    }
    
    try:
        response = requests.post(url, headers=headers, json=payload, timeout=30)
        
        result = {
            "model": model_name,
            "status_code": response.status_code,
            "available": response.status_code == 200,
            "error": None,
            "response_preview": None
        }
        
        if response.status_code == 200:
            try:
                data = response.json()
                content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
                result["response_preview"] = content[:100]  # 只显示前100个字符
            except:
                result["response_preview"] = "解析响应失败"
        else:
            result["error"] = response.text[:200]  # 只显示前200个字符的错误信息
            
    except requests.exceptions.Timeout:
        result = {
            "model": model_name,
            "status_code": None,
            "available": False,
            "error": "请求超时",
            "response_preview": None
        }
    except Exception as e:
        result = {
            "model": model_name,
            "status_code": None,
            "available": False,
            "error": str(e)[:200],
            "response_preview": None
        }
    
    return result

def check_api_health() -> bool:
    """检查 API 服务是否健康"""
    try:
        response = requests.get(f"{BASE_URL}/healthz", timeout=10)
        return response.status_code == 200
    except:
        return False

def check_accounts() -> bool:
    """检查是否有可用账号"""
    try:
        headers = {"Authorization": f"Bearer {API_KEY}"} if API_KEY else {}
        response = requests.get(f"{BASE_URL}/v2/accounts", headers=headers, timeout=10)
        if response.status_code == 200:
            accounts = response.json()
            enabled_accounts = [acc for acc in accounts if acc.get("enabled", False)]
            return len(enabled_accounts) > 0
        return False
    except:
        return False

def main():
    """主函数"""
    print("🔍 Amazon Q API 模型测试脚本")
    print("=" * 50)
    
    # 检查 API 服务状态
    print("📡 检查 API 服务状态...")
    if not check_api_health():
        print("❌ API 服务不可用，请确保服务正在运行在 http://localhost:8000")
        return
    
    print("✅ API 服务正常")
    
    # 检查账号状态
    print("👤 检查可用账号...")
    if not check_accounts():
        print("❌ 没有可用的账号，请先添加并启用 Amazon Q 账号")
        return
    
    print("✅ 发现可用账号")
    
    # 测试模型
    print(f"\n🧪 开始测试 {len(TEST_MODELS)} 个模型...")
    print("-" * 80)
    
    available_models = []
    failed_models = []
    
    for i, model in enumerate(TEST_MODELS, 1):
        print(f"[{i}/{len(TEST_MODELS)}] 测试: {model}...", end=" ")
        
        result = test_model(model, API_KEY)
        
        if result["available"]:
            print("✅ 可用")
            available_models.append(model)
            if result["response_preview"]:
                print(f"    响应预览: {result['response_preview']}")
        else:
            print("❌ 不可用")
            failed_models.append((model, result["error"]))
        
        # 避免请求过于频繁
        time.sleep(0.5)
    
    # 输出结果总结
    print("\n" + "=" * 80)
    print("📊 测试结果总结")
    print("=" * 80)
    
    if available_models:
        print(f"\n✅ 可用模型 ({len(available_models)} 个):")
        for model in available_models:
            print(f"  • {model}")
    else:
        print("\n❌ 没有发现可用的模型")
    
    if failed_models:
        print(f"\n❌ 不可用模型 ({len(failed_models)} 个):")
        for model, error in failed_models[:5]:  # 只显示前5个错误
            print(f"  • {model}: {error}")
        if len(failed_models) > 5:
            print(f"  ... 还有 {len(failed_models) - 5} 个模型不可用")
    
    # 生成配置建议
    if available_models:
        print(f"\n💡 建议配置:")
        print(f"推荐使用: {available_models[0]}")
        print(f"在代码中使用: model = \"{available_models[0]}\"")
        
        # 保存可用模型到文件
        with open("available_models.txt", "w", encoding="utf-8") as f:
            f.write("可用模型列表:\n")
            for model in available_models:
                f.write(f"{model}\n")
        print(f"📁 可用模型已保存到: available_models.txt")

if __name__ == "__main__":
    main()