#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
WebSocket连接诊断工具
快速定位502错误的原因
"""

import asyncio
import websockets
import requests
import socket
from urllib.parse import urlparse
import ssl
import sys
from config import BACKEND_HOST, BACKEND_WS_URL, get_gui_ws_url, get_saved_token

def test_basic_connectivity():
    """测试基础网络连通性"""
    print("=== 1. 基础连通性测试 ===")
    
    # DNS解析测试
    try:
        ip = socket.gethostbyname(BACKEND_HOST)
        print(f"✅ DNS解析成功: {BACKEND_HOST} -> {ip}")
    except Exception as e:
        print(f"❌ DNS解析失败: {e}")
        return False
    
    # HTTP连通性测试
    try:
        response = requests.get(f"https://{BACKEND_HOST}", timeout=10)
        print(f"✅ HTTPS连接成功: 状态码 {response.status_code}")
    except requests.exceptions.SSLError as e:
        print(f"⚠️ SSL证书问题: {e}")
        return False
    except requests.exceptions.ConnectionError as e:
        print(f"❌ 连接失败: {e}")
        return False
    except requests.exceptions.Timeout:
        print(f"❌ 连接超时")
        return False
    except Exception as e:
        print(f"⚠️ HTTPS测试异常: {e}")
    
    return True

def test_websocket_endpoint():
    """测试WebSocket端点"""
    print("\n=== 2. WebSocket端点测试 ===")
    
    # 获取token
    token = get_saved_token()
    if not token:
        print("❌ 未找到有效token")
        return False
    
    print(f"📝 使用token: {token[:20]}...")
    
    # 构建WebSocket URL
    ws_url = get_gui_ws_url(token)
    print(f"🔗 WebSocket URL: {ws_url}")
    
    return ws_url

async def test_websocket_connection(ws_url):
    """异步测试WebSocket连接"""
    print("\n=== 3. WebSocket连接测试 ===")
    
    try:
        print("🔄 尝试连接WebSocket...")
        
        # 添加连接超时
        websocket = await asyncio.wait_for(
            websockets.connect(ws_url),
            timeout=15.0
        )
        
        print("✅ WebSocket连接成功!")
        
        # 测试发送消息
        test_message = {"type": "ping", "data": "connection_test"}
        await websocket.send(str(test_message))
        print("✅ 消息发送成功")
        
        # 等待响应
        try:
            response = await asyncio.wait_for(websocket.recv(), timeout=5.0)
            print(f"✅ 收到响应: {response}")
        except asyncio.TimeoutError:
            print("⚠️ 未收到响应(超时5秒)")
        
        await websocket.close()
        print("✅ 连接正常关闭")
        return True
        
    except websockets.exceptions.InvalidStatusCode as e:
        if e.status_code == 502:
            print("❌ HTTP 502 Bad Gateway - 后端服务器问题")
            print("💡 可能原因:")
            print("   - 后端服务未启动")
            print("   - Nginx配置错误")
            print("   - 服务器资源不足")
        else:
            print(f"❌ HTTP状态码错误: {e.status_code}")
        return False
        
    except websockets.exceptions.ConnectionClosedError as e:
        print(f"❌ 连接被关闭: {e}")
        return False
        
    except websockets.exceptions.WebSocketException as e:
        print(f"❌ WebSocket异常: {e}")
        return False
        
    except asyncio.TimeoutError:
        print("❌ 连接超时 - 服务器可能未响应")
        return False
        
    except Exception as e:
        print(f"❌ 未知错误: {type(e).__name__}: {e}")
        return False

def test_alternative_endpoints():
    """测试备用连接方式"""
    print("\n=== 4. 备用端点测试 ===")
    
    # 测试HTTP API端点
    try:
        api_url = f"https://{BACKEND_HOST}/api/health"
        response = requests.get(api_url, timeout=10)
        print(f"✅ API端点响应: {response.status_code}")
    except Exception as e:
        print(f"⚠️ API端点不可用: {e}")
    
    # 测试WebSocket根路径
    try:
        root_ws_url = f"wss://{BACKEND_HOST}/ws/"
        print(f"🔗 测试根WebSocket路径: {root_ws_url}")
    except Exception as e:
        print(f"⚠️ 根路径测试失败: {e}")

async def main():
    """主诊断流程"""
    print("🔍 WebSocket连接诊断工具")
    print("=" * 50)
    
    # 1. 基础连通性
    if not test_basic_connectivity():
        print("\n❌ 基础连通性测试失败，请检查网络连接")
        return
    
    # 2. WebSocket端点配置
    ws_url = test_websocket_endpoint()
    if not ws_url:
        print("\n❌ WebSocket端点配置错误")
        return
    
    # 3. WebSocket连接测试
    success = await test_websocket_connection(ws_url)
    
    # 4. 备用测试
    test_alternative_endpoints()
    
    # 5. 总结和建议
    print("\n" + "=" * 50)
    if success:
        print("🎉 诊断完成: WebSocket连接正常!")
    else:
        print("❌ 诊断完成: 发现连接问题")
        print("\n🔧 建议解决方案:")
        print("1. 联系服务器管理员检查后端服务状态")
        print("2. 检查Nginx WebSocket配置")
        print("3. 确认防火墙和安全组设置")
        print("4. 检查服务器资源使用情况")
        print("5. 验证SSL证书有效性")

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n\n⏹️ 诊断被用户中断")
    except Exception as e:
        print(f"\n\n💥 诊断工具异常: {e}")