import requests
import json
import time

API_KEY = "paVrIT+XU1NhwCAOb0X4aYi75QKogK5YNMGvQF1dCyo="
BASE_URL = "https://llama.theaken.com/v1"

MODELS = [
    "gpt-oss-120b",
    "deepseek-r1-671b", 
    "qwen3-embedding-8b"
]

def test_model(model_name):
    """測試單個模型"""
    print(f"\n[測試模型: {model_name}]")
    print("-" * 40)
    
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }
    
    # 測試聊天完成端點
    chat_url = f"{BASE_URL}/chat/completions"
    data = {
        "model": model_name,
        "messages": [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "Say 'Hello, I am working!' if you can see this message."}
        ],
        "temperature": 0.5,
        "max_tokens": 50
    }
    
    try:
        print(f"連接到: {chat_url}")
        response = requests.post(chat_url, headers=headers, json=data, timeout=30)
        
        print(f"HTTP 狀態碼: {response.status_code}")
        
        if response.status_code == 200:
            result = response.json()
            if 'choices' in result and len(result['choices']) > 0:
                content = result['choices'][0]['message']['content']
                print(f"[SUCCESS] AI 回應: {content}")
                return True
            else:
                print("[ERROR] 回應格式異常")
                print(f"回應內容: {json.dumps(result, indent=2)}")
        else:
            print(f"[ERROR] 錯誤回應")
            # 檢查是否是 HTML 錯誤頁面
            if response.text.startswith('<!DOCTYPE'):
                print("收到 HTML 錯誤頁面 (可能是 502 Bad Gateway)")
            else:
                print(f"回應內容: {response.text[:300]}")
                
    except requests.exceptions.Timeout:
        print("[TIMEOUT] 請求超時 (30秒)")
    except requests.exceptions.ConnectionError as e:
        print(f"[CONNECTION ERROR]: {str(e)[:100]}")
    except Exception as e:
        print(f"[UNEXPECTED ERROR]: {str(e)[:100]}")
    
    return False

def test_api_endpoints():
    """測試不同的 API 端點"""
    print("\n[測試 API 端點可用性]")
    print("=" * 50)
    
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }
    
    # 測試不同的可能端點
    endpoints = [
        f"{BASE_URL}/models",
        f"{BASE_URL}/chat/completions",
        BASE_URL
    ]
    
    for endpoint in endpoints:
        try:
            print(f"\n測試端點: {endpoint}")
            response = requests.get(endpoint, headers=headers, timeout=10)
            print(f"  狀態碼: {response.status_code}")
            
            if response.status_code == 200:
                print("  [OK] 端點可訪問")
                # 如果是 JSON 回應，顯示部分內容
                try:
                    data = response.json()
                    print(f"  回應類型: JSON")
                    if 'data' in data:
                        print(f"  包含 {len(data['data'])} 項資料")
                except:
                    print(f"  回應類型: {response.headers.get('content-type', 'unknown')}")
            elif response.status_code == 405:
                print("  [OK] 端點存在 (但不支援 GET 方法)")
            elif response.status_code == 502:
                print("  [ERROR] 502 Bad Gateway - 伺服器暫時無法使用")
            else:
                print(f"  [ERROR] 無法訪問")
                
        except Exception as e:
            print(f"  [ERROR]: {str(e)[:50]}")

def main():
    print("=" * 50)
    print("Llama API 完整測試程式")
    print("=" * 50)
    print(f"API 基礎 URL: {BASE_URL}")
    print(f"API 金鑰: {API_KEY[:10]}...{API_KEY[-5:]}")
    
    # 首先測試端點可用性
    test_api_endpoints()
    
    print("\n" + "=" * 50)
    print("開始測試各個模型")
    print("=" * 50)
    
    success_count = 0
    for model in MODELS:
        if test_model(model):
            success_count += 1
        time.sleep(1)  # 避免請求過快
    
    print("\n" + "=" * 50)
    print(f"測試結果: {success_count}/{len(MODELS)} 個模型成功連接")
    
    if success_count == 0:
        print("\n可能的問題：")
        print("1. API 伺服器暫時離線 (502 錯誤)")
        print("2. API 金鑰可能不正確")
        print("3. 網路連接問題")
        print("4. 防火牆或代理設定")
        print("\n建議稍後再試，或聯繫 API 提供者確認服務狀態。")

if __name__ == "__main__":
    main()