#!/usr/bin/env python3
"""
VNC WebSocket 프록시
브라우저와 Proxmox VNC 서버 간 WebSocket 연결을 중계하며
PVE 인증을 자동으로 처리합니다.
"""

import asyncio
import websockets
import ssl
import logging
from utils.proxmox_client import ProxmoxClient

# 로깅 설정
logger = logging.getLogger(__name__)

class VNCWebSocketProxy:
    def __init__(self, proxmox_host, proxmox_username, proxmox_password):
        self.proxmox_host = proxmox_host
        self.proxmox_username = proxmox_username
        self.proxmox_password = proxmox_password
        self.proxmox_client = None
        
    async def create_proxmox_client(self):
        """Proxmox 클라이언트 생성 및 로그인"""
        if not self.proxmox_client:
            self.proxmox_client = ProxmoxClient(
                self.proxmox_host, 
                self.proxmox_username, 
                self.proxmox_password
            )
            
        if not self.proxmox_client.login():
            logger.error("Proxmox 로그인 실패")
            return False
            
        logger.info("Proxmox 로그인 성공")
        return True
        
    async def get_vnc_connection_info(self, node, vm_id):
        """VNC 연결 정보 생성"""
        if not await self.create_proxmox_client():
            return None
            
        # VM 상태 확인
        vm_status = self.proxmox_client.get_vm_status(node, vm_id)
        if vm_status.get('status') != 'running':
            logger.error(f"VM {vm_id}가 실행 중이 아님: {vm_status.get('status')}")
            return None
            
        # VNC 티켓 생성
        vnc_data = self.proxmox_client.get_vnc_ticket(node, vm_id)
        if not vnc_data:
            logger.error(f"VM {vm_id} VNC 티켓 생성 실패")
            return None
            
        # WebSocket 연결 정보 준비
        connection_info = {
            'websocket_url': vnc_data['websocket_url'],
            'password': vnc_data['password'],
            'auth_headers': {
                'Cookie': f'PVEAuthCookie={self.proxmox_client.ticket}'
            }
        }
        
        logger.info(f"VM {vm_id} VNC 연결 정보 생성 완료")
        return connection_info
        
    async def create_proxmox_websocket(self, connection_info):
        """Proxmox VNC WebSocket 연결 생성"""
        ssl_context = ssl.create_default_context()
        ssl_context.check_hostname = False
        ssl_context.verify_mode = ssl.CERT_NONE
        
        try:
            websocket = await websockets.connect(
                connection_info['websocket_url'],
                ssl=ssl_context,
                additional_headers=connection_info['auth_headers']
            )
            logger.info("Proxmox VNC WebSocket 연결 성공")
            return websocket
            
        except Exception as e:
            logger.error(f"Proxmox VNC WebSocket 연결 실패: {e}")
            return None
            
    async def proxy_data(self, browser_ws, proxmox_ws):
        """브라우저와 Proxmox 간 WebSocket 데이터 양방향 중계"""
        async def forward_browser_to_proxmox():
            """브라우저 → Proxmox 데이터 전달"""
            try:
                async for message in browser_ws:
                    await proxmox_ws.send(message)
                    logger.debug(f"브라우저 → Proxmox: {len(message)} bytes")
            except websockets.exceptions.ConnectionClosed:
                logger.info("브라우저 연결 종료")
            except Exception as e:
                logger.error(f"브라우저 → Proxmox 전달 오류: {e}")
                
        async def forward_proxmox_to_browser():
            """Proxmox → 브라우저 데이터 전달"""
            try:
                async for message in proxmox_ws:
                    await browser_ws.send(message)
                    logger.debug(f"Proxmox → 브라우저: {len(message)} bytes")
            except websockets.exceptions.ConnectionClosed:
                logger.info("Proxmox 연결 종료")
            except Exception as e:
                logger.error(f"Proxmox → 브라우저 전달 오류: {e}")
                
        # 양방향 데이터 전달을 병렬로 실행
        await asyncio.gather(
            forward_browser_to_proxmox(),
            forward_proxmox_to_browser(),
            return_exceptions=True
        )
        
    async def handle_vnc_proxy(self, browser_websocket, node, vm_id):
        """VNC 프록시 메인 핸들러"""
        logger.info(f"VNC 프록시 시작: VM {vm_id}")
        
        try:
            # 1. Proxmox VNC 연결 정보 생성
            connection_info = await self.get_vnc_connection_info(node, vm_id)
            if not connection_info:
                await browser_websocket.send("ERROR: VNC 연결 정보 생성 실패")
                return False
                
            # 2. Proxmox VNC WebSocket 연결
            proxmox_websocket = await self.create_proxmox_websocket(connection_info)
            if not proxmox_websocket:
                await browser_websocket.send("ERROR: Proxmox VNC 연결 실패")
                return False
                
            # 3. 연결 성공 알림
            logger.info(f"VM {vm_id} VNC 프록시 연결 완료")
            
            # 4. 데이터 중계 시작
            await self.proxy_data(browser_websocket, proxmox_websocket)
            
            return True
            
        except Exception as e:
            logger.error(f"VNC 프록시 처리 오류: {e}")
            try:
                await browser_websocket.send(f"ERROR: {str(e)}")
            except:
                pass
            return False
            
        finally:
            logger.info(f"VNC 프록시 종료: VM {vm_id}")


# 전역 VNC 프록시 인스턴스 (설정값은 app.py에서 주입)
vnc_proxy_instance = None

def init_vnc_proxy(proxmox_host, proxmox_username, proxmox_password):
    """VNC 프록시 인스턴스 초기화"""
    global vnc_proxy_instance
    vnc_proxy_instance = VNCWebSocketProxy(proxmox_host, proxmox_username, proxmox_password)
    logger.info("VNC WebSocket 프록시 초기화 완료")
    
def get_vnc_proxy():
    """VNC 프록시 인스턴스 반환"""
    global vnc_proxy_instance
    return vnc_proxy_instance