"""
Web Terminal WebSocket Handling
WebSocket broadcast and communication for multi-terminal synchronization
"""

import asyncio
import logging
from typing import Set
import threading

logger = logging.getLogger(__name__)


class WebSocketManager:
    """Manages WebSocket connections and broadcasting"""

    def __init__(self, shared_state):
        """
        Initialize WebSocket manager

        Args:
            shared_state: SharedTerminalState instance
        """
        self.shared_state = shared_state
        self.active_websockets: Set = set()
        self._ws_lock = threading.Lock()
        self._broadcast_task = None

    async def handle_websocket(self, websocket):
        """
        Handle individual WebSocket connection
        FIXED: Properly wait for messages without exiting immediately

        Args:
            websocket: WebSocket connection from client
        """
        # Add to active connections
        with self._ws_lock:
            self.active_websockets.add(websocket)

            # Start broadcast loop if this is the first connection
            if len(self.active_websockets) == 1 and self._broadcast_task is None:
                logger.info("Starting broadcast loop (first WebSocket connection)")
                self._broadcast_task = asyncio.create_task(self._broadcast_output_loop())

        client_id = id(websocket)
        logger.info(f"WebSocket connected: {client_id} (total: {len(self.active_websockets)})")

        try:
            # Send welcome message
            await websocket.send_json({
                'type': 'connection',
                'status': 'connected',
                'message': 'Terminal synchronized'
            })

            # FIXED: Keep connection alive and handle messages
            # Use receive() instead of async for loop
            while True:
                try:
                    # Wait for message from client (with timeout to allow graceful shutdown)
                    message = await asyncio.wait_for(websocket.receive_json(), timeout=1.0)

                    if message.get('type') == 'terminal_input':
                        # User typed in THIS terminal
                        # Send to SSH (will echo back to ALL terminals via broadcast)
                        input_data = message.get('data', '')
                        if self.shared_state.ssh_manager and input_data:
                            self.shared_state.ssh_manager.send_input(input_data)
                            logger.debug(f"Forwarded input to SSH: {repr(input_data[:20])}")

                    elif message.get('type') == 'terminal_resize':
                        # Terminal resized
                        cols = message.get('cols')
                        rows = message.get('rows')
                        if cols and rows and self.shared_state.ssh_manager:
                            self.shared_state.ssh_manager.resize_pty(cols, rows)
                            logger.debug(f"Terminal resized to {cols}x{rows}")

                except asyncio.TimeoutError:
                    # No message received in 1 second - that's fine, keep waiting
                    continue

                except Exception as e:
                    # Connection closed or other error
                    logger.debug(f"WebSocket receive error: {e}")
                    break

        except Exception as e:
            logger.debug(f"WebSocket {client_id} handler error: {e}")

        finally:
            # Remove from active connections
            with self._ws_lock:
                self.active_websockets.discard(websocket)

            logger.info(f"WebSocket disconnected: {client_id} (remaining: {len(self.active_websockets)})")

    async def _broadcast_output_loop(self):
        """
        Background task that broadcasts SSH output to all connected WebSockets
        Runs continuously while server is active
        """
        logger.info("Broadcast loop started successfully")

        while self.shared_state.web_server_running:
            try:
                # Get output from shared state
                output = self.shared_state.get_output()

                if output:
                    # Broadcast to all connected WebSockets
                    message = {
                        'type': 'terminal_output',
                        'data': output
                    }

                    with self._ws_lock:
                        active_count = len(self.active_websockets)
                        if active_count > 0:
                            logger.debug(f"Broadcasting {len(output)} bytes to {active_count} WebSocket(s)")

                        disconnected = set()
                        for ws in self.active_websockets:
                            try:
                                await ws.send_json(message)
                            except Exception as e:
                                logger.debug(f"WebSocket send failed: {e}")
                                disconnected.add(ws)

                        # Clean up disconnected WebSockets
                        self.active_websockets -= disconnected

                # Poll every 50ms (same as original HTTP polling)
                await asyncio.sleep(0.05)

            except Exception as e:
                logger.error(f"Error in broadcast loop: {e}", exc_info=True)
                await asyncio.sleep(0.1)

        logger.info("Output broadcast loop stopped")

    async def broadcast_transfer_update(self, transfer_id: str, progress: dict):
        """
        Broadcast SFTP transfer progress to all connected clients

        Args:
            transfer_id: Transfer identifier
            progress: Progress information dict
        """
        message = {
            "type": "transfer_progress",
            "transfer_id": transfer_id,
            "progress": progress
        }

        with self._ws_lock:
            disconnected = set()
            for ws in self.active_websockets:
                try:
                    await ws.send_json(message)
                except Exception as e:
                    logger.debug(f"Failed to send transfer update: {e}")
                    disconnected.add(ws)

            self.active_websockets -= disconnected

    def close_all(self):
        """Close all active WebSocket connections"""
        with self._ws_lock:
            websocket_count = len(self.active_websockets)
            if websocket_count > 0:
                logger.info(f"Closing {websocket_count} active WebSocket connection(s)...")

                # Create a copy of the set to avoid modification during iteration
                websockets_to_close = list(self.active_websockets)

                for ws in websockets_to_close:
                    try:
                        # Send shutdown message
                        loop = asyncio.new_event_loop()
                        loop.run_until_complete(ws.send_json({
                            'type': 'server_shutdown',
                            'message': 'Server is shutting down'
                        }))
                        loop.run_until_complete(ws.close())
                        loop.close()
                    except Exception as e:
                        logger.debug(f"Error closing WebSocket: {e}")

                self.active_websockets.clear()
                logger.info(f"Closed {websocket_count} WebSocket connection(s)")
