#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
测试LiteLLM自定义API端点
"""

import requests
import json
import os

def test_litellm_api():
    """测试LiteLLM API端点"""
    api_url = "http://master.tailnet-68f9.ts.net:40012"
    
    print(f"🔍 测试LiteLLM API端点: {api_url}")
    
    # 获取用户的API密钥
    gemini_key = os.getenv('GEMINI_API_KEY', '')
    
    # 尝试不同的API密钥格式
    test_keys = [
        f"sk-{gemini_key}",  # 添加sk-前缀
        gemini_key,  # 原始密钥
        "sk-test",  # 测试密钥
        "test-key",  # 简单测试
    ]
    
    for api_key in test_keys:
        if not api_key or api_key == "sk-":
            continue
            
        print(f"\n🔑 测试API密钥: {api_key[:10]}...")
        
        # 测试模型列表
        try:
            headers = {"x-litellm-api-key": api_key}
            response = requests.get(f"{api_url}/v1/models", headers=headers, timeout=10)
            print(f"模型列表状态码: {response.status_code}")
            
            if response.status_code == 200:
                models = response.json()
                print(f"✅ 找到 {len(models.get('data', []))} 个可用模型")
                for model in models.get('data', [])[:3]:  # 显示前3个模型
                    print(f"  - {model.get('id', 'unknown')}")
                    
                # 测试聊天完成
                test_payload = {
                    "model": "gemini-2.5-flash",
                    "messages": [
                        {"role": "user", "content": "Hello, this is a test message. Please respond briefly."}
                    ],
                    "max_tokens": 50
                }
                
                chat_response = requests.post(
                    f"{api_url}/v1/chat/completions", 
                    json=test_payload, 
                    headers=headers, 
                    timeout=30
                )
                
                print(f"聊天完成状态码: {chat_response.status_code}")
                
                if chat_response.status_code == 200:
                    result = chat_response.json()
                    content = result.get('choices', [{}])[0].get('message', {}).get('content', '')
                    print(f"✅ API测试成功！响应: {content[:100]}...")
                    print(f"\n🎉 可用的API配置:")
                    print(f"  - 端点: {api_url}/v1/chat/completions")
                    print(f"  - 头部: x-litellm-api-key: {api_key}")
                    print(f"  - 模型: gemini-2.5-flash")
                    return True
                else:
                    print(f"❌ 聊天测试失败: {chat_response.text[:200]}...")
                    
            elif response.status_code == 401:
                print(f"❌ 认证失败: {response.text[:100]}...")
            else:
                print(f"❌ 请求失败: {response.text[:100]}...")
                
        except requests.exceptions.RequestException as e:
            print(f"❌ 连接失败: {e}")
            
    return False

def test_environment_setup():
    """测试环境变量设置"""
    print("\n🔧 当前环境变量:")
    gemini_key = os.getenv('GEMINI_API_KEY', '')
    google_key = os.getenv('GOOGLE_API_KEY', '')
    
    print(f"GEMINI_API_KEY: {'已设置' if gemini_key else '未设置'} ({gemini_key[:10]}... 如果已设置)")
    print(f"GOOGLE_API_KEY: {'已设置' if google_key else '未设置'} ({google_key[:10]}... 如果已设置)")
    
    return gemini_key

if __name__ == "__main__":
    print("🚀 开始测试LiteLLM自定义API端点...")
    
    # 检查环境
    api_key = test_environment_setup()
    
    if not api_key:
        print("\n⚠️  警告: 未找到GEMINI_API_KEY环境变量")
    
    # 测试API
    success = test_litellm_api()
    
    if success:
        print("\n✅ LiteLLM API端点测试成功！")
        print("\n💡 建议: 可以使用这个端点替代Google官方API")
    else:
        print("\n❌ LiteLLM API端点测试失败")
        print("\n🔍 可能的解决方案:")
        print("  1. 检查LiteLLM服务器配置")
        print("  2. 确认API密钥格式")
        print("  3. 检查网络连接")