#!/usr/bin/env python3
"""
直接查询 Amazon Q 服务来获取支持的模型
"""

import requests
import json
import time
from replicate import send_chat_request

def get_available_models_from_amazon_q():
    """尝试通过特殊查询获取 Amazon Q 支持的模型"""
    
    # 尝试不同的查询方式来获取模型信息
    test_queries = [
        "你支持哪些模型？",
        "有哪些可用的模型？", 
        "列出所有支持的模型名称",
        "What models are available?",
        "List all supported models",
        "Show model options",
        "model list",
        "models",
        "help"
    ]
    
    print("🔍 通过查询尝试获取 Amazon Q 支持的模型信息")
    print("=" * 60)
    
    # 这里需要一个有效的 access token，我们先尝试从现有账号获取
    try:
        # 检查是否有可用账号
        accounts_response = requests.get("http://localhost:8000/v2/accounts", timeout=10)
        if accounts_response.status_code != 200:
            print("❌ 无法获取账号信息")
            return
        
        accounts = accounts_response.json()
        enabled_accounts = [acc for acc in accounts if acc.get("enabled", False)]
        
        if not enabled_accounts:
            print("❌ 没有可用的账号")
            return
        
        account = enabled_accounts[0]
        access_token = account.get("accessToken")
        
        if not access_token:
            # 尝试刷新 token
            print("🔄 刷新 access token...")
            refresh_response = requests.post(f"http://localhost:8000/v2/accounts/{account['id']}/refresh", timeout=10)
            if refresh_response.status_code == 200:
                updated_account = refresh_response.json()
                access_token = updated_account.get("accessToken")
        
        if not access_token:
            print("❌ 无法获取 access token")
            return
        
        print(f"✅ 获取到 access token，开始查询...")
        
        # 对每个查询进行测试
        for i, query in enumerate(test_queries, 1):
            print(f"\n[{i}/{len(test_queries)}] 查询: {query}")
            print("-" * 40)
            
            try:
                text, _, tracker = send_chat_request(
                    access_token, 
                    [{"role": "user", "content": query}], 
                    model=None,  # 不指定模型，使用默认
                    stream=False,
                    timeout=(10, 30)
                )
                
                if text:
                    print(f"响应: {text[:200]}...")
                    
                    # 检查响应中是否包含模型信息
                    model_keywords = ["model", "claude", "anthropic", "sonnet", "haiku", "opus"]
                    found_models = []
                    
                    for keyword in model_keywords:
                        if keyword.lower() in text.lower():
                            found_models.append(keyword)
                    
                    if found_models:
                        print(f"🔍 发现模型相关关键词: {', '.join(found_models)}")
                else:
                    print("❌ 无响应")
                
            except Exception as e:
                print(f"❌ 查询失败: {e}")
            
            time.sleep(1)  # 避免请求过快
    
    except Exception as e:
        print(f"❌ 获取模型信息失败: {e}")

def test_model_variations():
    """测试模型名称的变体"""
    # 基于已知可用的 claude-sonnet-4，测试可能的变体
    variations = [
        "claude-sonnet-4",
        "Claude-Sonnet-4", 
        "CLAUDE-SONNET-4",
        "sonnet-4",
        "claude-4-sonnet",
        "claude-4",
        "sonnet4",
        "anthropic-sonnet-4",
        "anthropic-claude-sonnet-4",
    ]
    
    print("\n🧪 测试模型名称变体")
    print("=" * 40)
    
    for model in variations:
        try:
            text, _, tracker = send_chat_request(
                "dummy_token",  # 这里会被替换为真实 token
                [{"role": "user", "content": "test"}], 
                model=model,
                stream=False,
                timeout=(5, 15)
            )
            
            if text:
                print(f"✅ {model} - 可用")
            else:
                print(f"❌ {model} - 不可用")
                
        except Exception as e:
            print(f"❌ {model} - 错误: {str(e)[:50]}")
        
        time.sleep(0.5)

if __name__ == "__main__":
    get_available_models_from_amazon_q()
    # test_model_variations()  # 暂时不运行，因为需要有效token