#!/usr/bin/env python3
"""
尝试获取 Amazon Q API 支持的所有模型
"""

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

BASE_URL = "http://localhost:8000"

def try_models_endpoint():
    """尝试调用 /v1/models 端点（如果存在）"""
    try:
        response = requests.get(f"{BASE_URL}/v1/models", timeout=10)
        if response.status_code == 200:
            data = response.json()
            print("✅ 发现 /v1/models 端点")
            return data
        else:
            print(f"❌ /v1/models 端点返回: {response.status_code}")
            return None
    except Exception as e:
        print(f"❌ /v1/models 端点异常: {e}")
        return None

def extract_model_from_error(error_text: str) -> List[str]:
    """从错误信息中提取可能的模型名称"""
    models = []
    
    # 常见的模型名称模式
    patterns = [
        r'claude-[\w\.-]+',
        r'anthropic\.[\w\.-]+',
        r'model\s+["\']?([a-zA-Z0-9_\.-]+)["\']?',
        r'available\s+models?[:\s]+([^\n\r]+)',
        r'supported\s+models?[:\s]+([^\n\r]+)',
    ]
    
    for pattern in patterns:
        matches = re.findall(pattern, error_text, re.IGNORECASE)
        models.extend(matches)
    
    return list(set(models))  # 去重

def try_common_amazon_q_models() -> List[str]:
    """尝试 Amazon Q 可能支持的其他模型名称"""
    # 基于 AWS Bedrock 和 Amazon Q 的常见模型命名
    amazon_q_models = [
        # Claude 模型 (AWS Bedrock 格式)
        "anthropic.claude-v2:1",
        "anthropic.claude-v2",
        "anthropic.claude-instant-v1",
        "anthropic.claude-3-sonnet-20240229-v1:0",
        "anthropic.claude-3-sonnet-20240229-v1",
        "anthropic.claude-3-haiku-20240307-v1:0", 
        "anthropic.claude-3-haiku-20240307-v1",
        "anthropic.claude-3-opus-20240229-v1:0",
        "anthropic.claude-3-opus-20240229-v1",
        
        # Amazon Q 特定模型
        "amazon.q",
        "amazon.q-turbo",
        "amazon.q-pro",
        "amazon.q-max",
        "q.amazon",
        "q-turbo.amazon",
        "q-pro.amazon",
        "q-max.amazon",
        
        # 其他可能的命名
        "claude",
        "claude-v2",
        "claude-instant",
        "claude-3-sonnet",
        "claude-3-haiku", 
        "claude-3-opus",
        "sonnet",
        "haiku",
        "opus",
        
        # 通用模型名
        "default",
        "base",
        "latest",
        "text-davinci-003",  # 测试是否支持OpenAI模型名
        "gpt-3.5-turbo",
        "gpt-4",
    ]
    
    return amazon_q_models

def test_model_with_details(model_name: str) -> Dict[str, any]:
    """详细测试单个模型"""
    url = f"{BASE_URL}/v1/chat/completions"
    
    payload = {
        "model": model_name,
        "messages": [
            {"role": "user", "content": "test"}
        ],
        "stream": False,
        "max_tokens": 5
    }
    
    try:
        response = requests.post(url, json=payload, timeout=15)
        
        result = {
            "model": model_name,
            "status_code": response.status_code,
            "available": response.status_code == 200,
            "error": None,
            "error_details": None
        }
        
        if response.status_code != 200:
            result["error"] = response.text[:300]
            # 尝试从错误中提取模型信息
            extracted_models = extract_model_from_error(response.text)
            if extracted_models:
                result["extracted_models"] = extracted_models
        
        return result
        
    except Exception as e:
        return {
            "model": model_name,
            "status_code": None,
            "available": False,
            "error": str(e),
            "error_details": None
        }

def comprehensive_model_search():
    """全面搜索可用模型"""
    print("🔍 全面搜索 Amazon Q API 支持的模型")
    print("=" * 60)
    
    # 1. 尝试标准 models 端点
    print("\n1️⃣ 尝试标准 /v1/models 端点...")
    models_data = try_models_endpoint()
    if models_data:
        print("发现模型列表:")
        if isinstance(models_data, dict) and "data" in models_data:
            for model in models_data["data"]:
                print(f"  • {model.get('id', 'unknown')}")
        elif isinstance(models_data, list):
            for model in models_data:
                if isinstance(model, dict):
                    print(f"  • {model.get('id', model.get('model', 'unknown'))}")
                else:
                    print(f"  • {model}")
        return models_data
    
    # 2. 测试更多模型名称
    print("\n2️⃣ 测试扩展的模型名称列表...")
    additional_models = try_common_amazon_q_models()
    
    print(f"测试 {len(additional_models)} 个额外的模型名称...")
    
    available_models = []
    error_models = []
    extracted_from_errors = set()
    
    for i, model in enumerate(additional_models, 1):
        print(f"[{i}/{len(additional_models)}] {model}...", end=" ")
        
        result = test_model_with_details(model)
        
        if result["available"]:
            print("✅")
            available_models.append(model)
        else:
            print("❌")
            error_models.append((model, result["error"]))
            
            # 从错误中提取模型信息
            if "extracted_models" in result:
                extracted_from_errors.update(result["extracted_models"])
        
        # 避免请求过快
        import time
        time.sleep(0.3)
    
    # 3. 输出结果
    print("\n" + "=" * 60)
    print("📊 搜索结果")
    print("=" * 60)
    
    if available_models:
        print(f"\n✅ 发现可用模型 ({len(available_models)} 个):")
        for model in available_models:
            print(f"  • {model}")
    else:
        print("\n❌ 没有发现新的可用模型")
    
    if extracted_from_errors:
        print(f"\n🔍 从错误信息中提取的可能模型 ({len(extracted_from_errors)} 个):")
        for model in sorted(extracted_from_errors):
            print(f"  • {model}")
    
    # 4. 保存结果
    search_results = {
        "available_models": available_models,
        "error_models": [{"model": m, "error": e} for m, e in error_models[:10]],  # 只保存前10个错误
        "extracted_from_errors": list(extracted_from_errors),
        "timestamp": str(time.time())
    }
    
    with open("model_search_results.json", "w", encoding="utf-8") as f:
        json.dump(search_results, f, ensure_ascii=False, indent=2)
    
    print(f"\n📁 详细结果已保存到: model_search_results.json")
    
    return search_results

def inspect_api_info():
    """检查API的其他信息端点"""
    endpoints_to_try = [
        "/",
        "/docs", 
        "/openapi.json",
        "/info",
        "/v1",
        "/status",
        "/health"
    ]
    
    print("\n🔍 检查API信息端点...")
    
    for endpoint in endpoints_to_try:
        try:
            response = requests.get(f"{BASE_URL}{endpoint}", timeout=5)
            if response.status_code == 200:
                print(f"✅ {endpoint} - 可访问")
                # 检查是否包含模型信息
                content = response.text.lower()
                if any(keyword in content for keyword in ["model", "claude", "anthropic"]):
                    print(f"   📄 可能包含模型相关信息")
            else:
                print(f"❌ {endpoint} - {response.status_code}")
        except:
            print(f"❌ {endpoint} - 连接失败")

if __name__ == "__main__":
    import time
    
    # 检查API信息端点
    inspect_api_info()
    
    # 全面模型搜索
    comprehensive_model_search()
    
    print("\n✨ 搜索完成!")