#!/usr/bin/env python3
"""
VNC WebSocket 연결 테스트 스크립트
Claude Code 환경에서 직접 테스트
"""

import asyncio
import websockets
import ssl
import json
import sys
import os

# Flask app.py와 동일한 경로에서 import
sys.path.append('/srv/headscale-setup/farmq-admin')
from utils.proxmox_client import ProxmoxClient

# 설정
PROXMOX_HOST = "pve7.0bin.in"
PROXMOX_USERNAME = "root@pam" 
PROXMOX_PASSWORD = "trajet6640"
VM_ID = 102
NODE_NAME = 'pve7'

async def test_vnc_websocket():
    """VNC WebSocket 연결 테스트"""
    print("=" * 60)
    print("🧪 Claude Code에서 VNC WebSocket 연결 테스트")
    print("=" * 60)
    
    # 1. Proxmox 클라이언트 생성 및 로그인
    print("1️⃣ Proxmox 로그인 중...")
    client = ProxmoxClient(PROXMOX_HOST, PROXMOX_USERNAME, PROXMOX_PASSWORD)
    
    if not client.login():
        print("❌ Proxmox 로그인 실패")
        return
    print("✅ Proxmox 로그인 성공")
    
    # 2. VM 상태 확인
    print("2️⃣ VM 상태 확인 중...")
    vm_status = client.get_vm_status(NODE_NAME, VM_ID)
    print(f"🔍 VM {VM_ID} 상태: {vm_status.get('status', 'unknown')}")
    
    if vm_status.get('status') != 'running':
        print(f"❌ VM이 실행 중이 아닙니다: {vm_status.get('status')}")
        return
    
    # 3. VNC 티켓 생성
    print("3️⃣ VNC 티켓 생성 중...")
    vnc_data = client.get_vnc_ticket(NODE_NAME, VM_ID)
    
    if not vnc_data:
        print("❌ VNC 티켓 생성 실패")
        return
        
    print("✅ VNC 티켓 생성 성공!")
    print(f"   - WebSocket URL: {vnc_data['websocket_url']}")
    print(f"   - VNC 패스워드: {vnc_data.get('password', 'N/A')}")
    print(f"   - 포트: {vnc_data.get('port', 'N/A')}")
    
    # 4. WebSocket 연결 테스트
    print("4️⃣ WebSocket 연결 테스트...")
    websocket_url = vnc_data['websocket_url']
    
    # SSL 컨텍스트 설정 (자체 서명 인증서 허용)
    ssl_context = ssl.create_default_context()
    ssl_context.check_hostname = False
    ssl_context.verify_mode = ssl.CERT_NONE
    
    try:
        # WebSocket 연결 시도
        print(f"🔌 연결 시도: {websocket_url}")
        
        # Proxmox 인증 쿠키를 헤더로 추가
        headers = {
            'Cookie': f'PVEAuthCookie={client.ticket}'
        }
        
        print(f"🔐 인증 쿠키 추가: PVEAuthCookie={client.ticket[:50]}...")
        
        # WebSocket 연결 시 인증 헤더 추가 (다른 방식)
        async with websockets.connect(
            websocket_url,
            ssl=ssl_context,
            additional_headers=headers
        ) as websocket:
            print("✅ WebSocket 연결 성공!")
            
            # VNC 프로토콜 초기 메시지 받기
            try:
                initial_message = await asyncio.wait_for(websocket.recv(), timeout=5.0)
                print(f"📨 초기 메시지 수신 ({len(initial_message)} bytes)")
                
                # VNC 프로토콜 버전 확인
                if isinstance(initial_message, bytes) and initial_message.startswith(b'RFB'):
                    version = initial_message.decode('ascii').strip()
                    print(f"🔗 VNC 프로토콜 버전: {version}")
                    
                    # 클라이언트 버전 응답
                    await websocket.send(b"RFB 003.008\n")
                    print("📤 클라이언트 버전 응답 완료")
                    
                    print("🎉 VNC WebSocket 연결 및 프로토콜 핸드셰이크 성공!")
                    return True
                else:
                    print(f"❓ 예상과 다른 초기 메시지: {initial_message[:50]}...")
                    
            except asyncio.TimeoutError:
                print("⏰ 초기 메시지 수신 타임아웃 - 연결은 성공했지만 VNC 서버 응답 없음")
                return True  # 연결 자체는 성공
                
    except websockets.exceptions.ConnectionClosed as e:
        print(f"❌ WebSocket 연결 종료: 코드={e.code}, 이유={e.reason}")
        return False
        
    except websockets.exceptions.WebSocketException as e:
        print(f"❌ WebSocket 예외: {e}")
        return False
        
    except asyncio.TimeoutError:
        print("❌ WebSocket 연결 타임아웃")
        return False
        
    except Exception as e:
        print(f"❌ 예상치 못한 오류: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    result = asyncio.run(test_vnc_websocket())
    
    print("=" * 60)
    if result:
        print("🎉 테스트 성공! WebSocket 연결이 Claude Code 환경에서 작동합니다.")
        print("💡 브라우저에서 문제가 있다면 브라우저 보안 정책 문제일 가능성이 높습니다.")
    else:
        print("❌ 테스트 실패! WebSocket 연결에 문제가 있습니다.")
        print("🔍 Proxmox 서버나 네트워크 설정을 확인해야 합니다.")
    print("=" * 60)