Trixy3/trixy_core/network/server_socket.py

"""
Server Socket Implementation for Trixy Network

This module provides server-side socket handling for the Trixy system.
It implements:

- Multi-threaded server socket management
- Client connection handling and authentication
- Command processing and response handling
- Audio stream server management
- Connection monitoring and health checking
- Integration with satellite registration process

The server supports multiple simultaneous client connections and handles
the full satellite connection workflow as specified in CLAUDE.md.
"""

import socket
import threading
import time
import uuid
import json
from typing import Dict, Any, Optional, List, Tuple, Callable
from dataclasses import dataclass, field
from concurrent.futures import ThreadPoolExecutor
import select

from .protocol import TrixyProtocol, TrixyMessage, MessageFlags
from .cmd import (
    TrixyCommand, 
    SatelliteHelloCommand,
    SatelliteRegistrationCommand,
    SatelliteConnectionCommand,
    create_hello_response,
    create_registration_response,
    create_connection_response
)


def pprint(message: str) -> None:
    """Server socket logging function."""
    print(f"[SERVER_SOCKET] {message}")


class ServerError(Exception):
    """Base exception for server errors."""
    pass


class PortInUseError(ServerError):
    """Raised when a port is already in use."""
    pass


class ClientConnectionError(ServerError):
    """Raised when client connection fails."""
    pass


@dataclass
class ClientConnection:
    """
    Represents a client connection to the server.
    
    Tracks connection state, authentication status, and statistics.
    """
    
    connection_id: str
    socket: socket.socket
    address: Tuple[str, int]
    connected_time: float = field(default_factory=time.time)
    last_activity: float = field(default_factory=time.time)
    
    # Connection state
    authenticated: bool = False
    satellite_id: Optional[str] = None
    satellite_info: Dict[str, Any] = field(default_factory=dict)
    
    # Statistics
    bytes_sent: int = 0
    bytes_received: int = 0
    commands_sent: int = 0
    commands_received: int = 0
    
    # Audio streaming
    audio_streams: Dict[str, int] = field(default_factory=dict)  # stream_type -> port
    
    def update_activity(self, bytes_delta: int = 0):
        """Update last activity time."""
        self.last_activity = time.time()
        if bytes_delta > 0:
            self.bytes_received += bytes_delta
    
    def is_timeout(self, timeout_seconds: float = 300.0) -> bool:
        """Check if connection has timed out."""
        return time.time() - self.last_activity > timeout_seconds
    
    def get_connection_info(self) -> Dict[str, Any]:
        """Get connection information."""
        return {
            'connection_id': self.connection_id,
            'address': self.address,
            'connected_time': self.connected_time,
            'last_activity': self.last_activity,
            'authenticated': self.authenticated,
            'satellite_id': self.satellite_id,
            'satellite_info': self.satellite_info,
            'bytes_sent': self.bytes_sent,
            'bytes_received': self.bytes_received,
            'commands_sent': self.commands_sent,
            'commands_received': self.commands_received,
            'audio_streams': self.audio_streams
        }


class ServerSocketHandler:
    """
    Handles server-side socket operations.
    
    This class manages the command socket server and handles client connections,
    authentication, and message processing.
    """
    
    def __init__(self, application_container):
        """
        Initialize the server socket handler.
        
        Args:
            application_container: Main application container
        """
        self.application = application_container
        self.protocol = TrixyProtocol()
        
        # Server state
        self._running = False
        self._server_socket: Optional[socket.socket] = None
        self._listen_port = 2101
        
        # Connection management
        self._connections: Dict[str, ClientConnection] = {}
        self._connections_lock = threading.RLock()
        
        # Threading
        self._thread_pool = ThreadPoolExecutor(max_workers=50)
        self._accept_thread: Optional[threading.Thread] = None
        self._cleanup_thread: Optional[threading.Thread] = None
        
        # Registration mode
        self._registration_mode = False
        self._registration_timeout = 60.0
        self._registration_timer: Optional[threading.Timer] = None
        
        # Message handlers
        self._message_handlers: Dict[str, Callable] = {}
        self._setup_default_handlers()
        
        # Statistics
        self._stats = {
            'connections_accepted': 0,
            'connections_rejected': 0,
            'messages_processed': 0,
            'errors': 0
        }
        
        pprint("ServerSocketHandler initialized")
    
    def start(self, port: int = 2101) -> None:
        """
        Start the server socket.
        
        Args:
            port: Port to listen on
            
        Raises:
            ServerError: If server cannot be started
        """
        if self._running:
            raise ServerError("Server already running")
        
        self._listen_port = port
        
        try:
            # Create server socket
            self._server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self._server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self._server_socket.bind(('', port))
            self._server_socket.listen(50)
            
            self._running = True
            
            # Start accept thread
            self._accept_thread = threading.Thread(target=self._accept_connections, daemon=True)
            self._accept_thread.start()
            
            # Start cleanup thread
            self._cleanup_thread = threading.Thread(target=self._cleanup_connections, daemon=True)
            self._cleanup_thread.start()
            
            pprint(f"Server started on port {port}")
            
        except OSError as e:
            if e.errno == 98:  # Address already in use
                raise PortInUseError(f"Port {port} is already in use") from e
            raise ServerError(f"Failed to start server: {e}") from e
    
    def stop(self) -> None:
        """Stop the server socket."""
        if not self._running:
            return
        
        pprint("Stopping server...")
        
        self._running = False
        
        # Stop registration mode
        if self._registration_mode:
            self.exit_registration_mode()
        
        # Close server socket
        if self._server_socket:
            self._server_socket.close()
            self._server_socket = None
        
        # Close all client connections
        with self._connections_lock:
            for connection in list(self._connections.values()):
                self._close_connection(connection.connection_id)
        
        # Wait for threads
        if self._accept_thread and self._accept_thread.is_alive():
            self._accept_thread.join(timeout=5.0)
        
        if self._cleanup_thread and self._cleanup_thread.is_alive():
            self._cleanup_thread.join(timeout=5.0)
        
        # Shutdown thread pool
        self._thread_pool.shutdown(wait=True)
        
        pprint("Server stopped")
    
    def enter_registration_mode(self, timeout_seconds: float = 60.0) -> None:
        """
        Enter registration mode to accept new satellites.
        
        Args:
            timeout_seconds: Registration timeout
        """
        if self._registration_mode:
            return
        
        pprint(f"Entering registration mode (timeout: {timeout_seconds}s)")
        
        self._registration_mode = True
        self._registration_timeout = timeout_seconds
        
        # Set timeout timer
        self._registration_timer = threading.Timer(
            timeout_seconds,
            self.exit_registration_mode
        )
        self._registration_timer.start()
        
        # Trigger event
        self._trigger_event("registration_mode_entered", {
            'timeout_seconds': timeout_seconds
        })
    
    def exit_registration_mode(self) -> None:
        """Exit registration mode."""
        if not self._registration_mode:
            return
        
        pprint("Exiting registration mode")
        
        self._registration_mode = False
        
        if self._registration_timer:
            self._registration_timer.cancel()
            self._registration_timer = None
        
        # Trigger event
        self._trigger_event("registration_mode_exited", {})
    
    def is_registration_mode(self) -> bool:
        """Check if in registration mode."""
        return self._registration_mode
    
    def send_message(self, connection_id: str, command: TrixyCommand) -> bool:
        """
        Send a message to a specific client.
        
        Args:
            connection_id: Target connection ID
            command: Command to send
            
        Returns:
            bool: True if sent successfully
        """
        with self._connections_lock:
            connection = self._connections.get(connection_id)
            if not connection:
                return False
            
            try:
                # Create protocol message
                message = self.protocol.create_message(
                    command.get_command_name(),
                    command
                )
                
                # Serialize and send
                data = message.serialize()
                connection.socket.send(data)
                
                # Update statistics
                connection.bytes_sent += len(data)
                connection.commands_sent += 1
                
                pprint(f"Sent {command.get_command_name()} to {connection_id}")
                return True
                
            except Exception as e:
                pprint(f"Error sending message to {connection_id}: {e}")
                return False
    
    def broadcast_message(self, command: TrixyCommand, authenticated_only: bool = False) -> int:
        """
        Broadcast a message to all connected clients.
        
        Args:
            command: Command to broadcast
            authenticated_only: Only send to authenticated clients
            
        Returns:
            int: Number of clients that received the message
        """
        sent_count = 0
        
        with self._connections_lock:
            for connection in self._connections.values():
                if authenticated_only and not connection.authenticated:
                    continue
                
                if self.send_message(connection.connection_id, command):
                    sent_count += 1
        
        return sent_count
    
    def get_connection_info(self, connection_id: str) -> Optional[Dict[str, Any]]:
        """Get information about a connection."""
        with self._connections_lock:
            connection = self._connections.get(connection_id)
            return connection.get_connection_info() if connection else None
    
    def get_all_connections(self) -> Dict[str, Dict[str, Any]]:
        """Get information about all connections."""
        with self._connections_lock:
            return {
                cid: conn.get_connection_info() 
                for cid, conn in self._connections.items()
            }
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get server statistics."""
        with self._connections_lock:
            active_connections = len(self._connections)
            authenticated_connections = sum(
                1 for conn in self._connections.values() 
                if conn.authenticated
            )
        
        return {
            **self._stats,
            'active_connections': active_connections,
            'authenticated_connections': authenticated_connections,
            'registration_mode': self._registration_mode,
            'listen_port': self._listen_port,
            'running': self._running
        }
    
    def _accept_connections(self) -> None:
        """Accept incoming connections (runs in separate thread)."""
        pprint("Started accepting connections")
        
        while self._running and self._server_socket:
            try:
                # Use select with timeout to allow clean shutdown
                ready, _, _ = select.select([self._server_socket], [], [], 1.0)
                
                if ready and self._running:
                    client_socket, address = self._server_socket.accept()
                    
                    # Create connection
                    connection_id = f"client_{address[0]}_{address[1]}_{uuid.uuid4().hex[:8]}"
                    connection = ClientConnection(
                        connection_id=connection_id,
                        socket=client_socket,
                        address=address
                    )
                    
                    # Add to connections
                    with self._connections_lock:
                        self._connections[connection_id] = connection
                    
                    # Handle in thread pool
                    self._thread_pool.submit(self._handle_client, connection)
                    
                    # Update statistics
                    self._stats['connections_accepted'] += 1
                    
                    pprint(f"Accepted connection from {address}: {connection_id}")
                    
            except Exception as e:
                if self._running:
                    pprint(f"Error accepting connection: {e}")
                    self._stats['errors'] += 1
    
    def _handle_client(self, connection: ClientConnection) -> None:
        """Handle a client connection (runs in thread pool)."""
        connection_id = connection.connection_id
        
        try:
            pprint(f"Handling client: {connection_id}")
            
            # Set socket timeout
            connection.socket.settimeout(30.0)
            
            while self._running:
                try:
                    # Receive data
                    data = connection.socket.recv(4096)
                    if not data:
                        break
                    
                    # Update connection activity
                    connection.update_activity(len(data))
                    
                    # Process the data
                    self._process_client_data(connection, data)
                    
                except socket.timeout:
                    # Check if connection timed out
                    if connection.is_timeout():
                        pprint(f"Connection {connection_id} timed out")
                        break
                    continue
                    
                except Exception as e:
                    pprint(f"Error handling client {connection_id}: {e}")
                    self._stats['errors'] += 1
                    break
            
        except Exception as e:
            pprint(f"Client handler error for {connection_id}: {e}")
            self._stats['errors'] += 1
        finally:
            self._close_connection(connection_id)
    
    def _process_client_data(self, connection: ClientConnection, data: bytes) -> None:
        """Process data received from a client."""
        try:
            # Check for hard-coded commands
            if self.protocol.is_hardcoded_command(data):
                command, args = self.protocol.parse_hardcoded_command(data)
                self._handle_hardcoded_command(connection, command, args)
            else:
                # Parse protocol message
                message = self.protocol.deserialize_message(data)
                self._handle_message(connection, message)
            
            # Update statistics
            connection.commands_received += 1
            self._stats['messages_processed'] += 1
            
        except Exception as e:
            pprint(f"Error processing data from {connection.connection_id}: {e}")
            self._stats['errors'] += 1
    
    def _handle_hardcoded_command(self, connection: ClientConnection, command: str, args: str) -> None:
        """Handle hard-coded commands."""
        try:
            if command == "TRXINOOP":
                # Heartbeat - just update activity (already done)
                pass
                
            elif command == "TRXIPING":
                # Respond with pong
                pong_response = f"TRXIPONG {args}".encode('utf-8')
                connection.socket.send(pong_response)
                connection.bytes_sent += len(pong_response)
                
            elif command == "TRXIPRNT":
                # Print command for debugging
                pprint(f"Client {connection.connection_id} print: {args}")
                
            elif command == "TRXYHELO":
                # Hello command for debugging
                pprint(f"Client {connection.connection_id} hello: {args}")
                
        except Exception as e:
            pprint(f"Error handling hard-coded command {command}: {e}")
    
    def _handle_message(self, connection: ClientConnection, message: TrixyMessage) -> None:
        """Handle a protocol message from a client."""
        command_name = message.class_name
        
        try:
            # Check for registered handler
            handler = self._message_handlers.get(command_name)
            if handler:
                handler(connection, message)
            else:
                pprint(f"No handler for command: {command_name}")
                
        except Exception as e:
            pprint(f"Error handling message {command_name}: {e}")
    
    def _setup_default_handlers(self) -> None:
        """Set up default message handlers."""
        self._message_handlers = {
            'SatelliteHelloCommand': self._handle_satellite_hello,
            'SatelliteRegistrationCommand': self._handle_satellite_registration,
            'SatelliteConnectionCommand': self._handle_satellite_connection,
        }
    
    def _handle_satellite_hello(self, connection: ClientConnection, message: TrixyMessage) -> None:
        """Handle satellite hello command."""
        try:
            # Deserialize command data
            hello_cmd = SatelliteHelloCommand.from_dict(message.data.to_dict())
            
            pprint(f"Satellite hello from {hello_cmd.alias} ({hello_cmd.mac_address})")
            
            # Check if satellite is registered
            # This would check against registration files in a real implementation
            registered = True  # Placeholder
            
            if registered:
                # Create positive response
                response = create_hello_response(
                    command_id=hello_cmd.command_id,
                    accepted=True,
                    assigned_satellite_id=f"sat_{hello_cmd.mac_address.replace(':', '')}",
                    registration_required=False
                )
            else:
                # Need registration
                response = create_hello_response(
                    command_id=hello_cmd.command_id,
                    accepted=False,
                    registration_required=True,
                    next_steps=["Enter registration mode on server", "Send registration command"]
                )
            
            # Send response
            self.send_message(connection.connection_id, response)
            
        except Exception as e:
            pprint(f"Error handling satellite hello: {e}")
    
    def _handle_satellite_registration(self, connection: ClientConnection, message: TrixyMessage) -> None:
        """Handle satellite registration command."""
        try:
            # Check if in registration mode
            if not self._registration_mode:
                response = create_registration_response(
                    command_id=message.data.command_id,
                    registration_accepted=False,
                    rejection_reason="Server not in registration mode"
                )
                self.send_message(connection.connection_id, response)
                return
            
            # Deserialize command data
            reg_cmd = SatelliteRegistrationCommand.from_dict(message.data.to_dict())
            
            pprint(f"Satellite registration: {reg_cmd.alias} in {reg_cmd.room_id}")
            
            # Create satellite ID
            satellite_id = f"sat_{reg_cmd.mac_address.replace(':', '')}"
            
            # In a real implementation, this would:
            # 1. Create registration file
            # 2. Validate MAC address format
            # 3. Check for conflicts
            
            # Create positive response
            response = create_registration_response(
                command_id=reg_cmd.command_id,
                registration_accepted=True,
                satellite_id=satellite_id
            )
            
            # Update connection info
            connection.satellite_id = satellite_id
            connection.satellite_info = {
                'room_id': reg_cmd.room_id,
                'alias': reg_cmd.alias,
                'mac_address': reg_cmd.mac_address,
                'device_info': reg_cmd.device_info
            }
            
            # Send response
            self.send_message(connection.connection_id, response)
            
            # Trigger event
            self._trigger_event("satellite_registered", {
                'satellite_id': satellite_id,
                'connection_id': connection.connection_id,
                'satellite_info': connection.satellite_info
            })
            
        except Exception as e:
            pprint(f"Error handling satellite registration: {e}")
    
    def _handle_satellite_connection(self, connection: ClientConnection, message: TrixyMessage) -> None:
        """Handle satellite connection command."""
        try:
            # Deserialize command data
            conn_cmd = SatelliteConnectionCommand.from_dict(message.data.to_dict())
            
            pprint(f"Satellite connection request: {conn_cmd.satellite_id}")
            
            # Validate satellite ID and MAC
            # In a real implementation, check registration files
            
            # Assign audio ports
            assigned_ports = {
                'audio_input': 2102,
                'audio_output': 2103,
                'music': 2104
            }
            
            # Create positive response
            response = create_connection_response(
                command_id=conn_cmd.command_id,
                connection_accepted=True,
                assigned_audio_ports=assigned_ports,
                session_id=str(uuid.uuid4())
            )
            
            # Update connection state
            connection.authenticated = True
            connection.audio_streams = assigned_ports
            
            # Send response
            self.send_message(connection.connection_id, response)
            
            # Trigger event
            self._trigger_event("satellite_connected", {
                'satellite_id': conn_cmd.satellite_id,
                'connection_id': connection.connection_id,
                'assigned_ports': assigned_ports
            })
            
        except Exception as e:
            pprint(f"Error handling satellite connection: {e}")
    
    def _close_connection(self, connection_id: str) -> None:
        """Close a client connection."""
        with self._connections_lock:
            connection = self._connections.pop(connection_id, None)
            
            if connection:
                try:
                    connection.socket.close()
                except:
                    pass
                
                pprint(f"Closed connection: {connection_id}")
                
                # Trigger event if it was an authenticated satellite
                if connection.authenticated and connection.satellite_id:
                    self._trigger_event("satellite_disconnected", {
                        'satellite_id': connection.satellite_id,
                        'connection_id': connection_id
                    })
    
    def _cleanup_connections(self) -> None:
        """Clean up timed out connections (runs in separate thread)."""
        while self._running:
            try:
                time.sleep(30.0)  # Check every 30 seconds
                
                timed_out = []
                with self._connections_lock:
                    for connection_id, connection in self._connections.items():
                        if connection.is_timeout():
                            timed_out.append(connection_id)
                
                # Close timed out connections
                for connection_id in timed_out:
                    pprint(f"Closing timed out connection: {connection_id}")
                    self._close_connection(connection_id)
                    
            except Exception as e:
                if self._running:
                    pprint(f"Error in cleanup thread: {e}")
    
    def _trigger_event(self, event_name: str, event_data: Dict[str, Any]) -> None:
        """Trigger an event through the application event system."""
        try:
            event_handler = self.application.get_event_handler()
            if event_handler:
                event_handler.trigger_event(event_name, event_data)
        except Exception as e:
            pprint(f"Error triggering event {event_name}: {e}")


class TrixyServer:
    """
    High-level Trixy server implementation.
    
    This class provides a simple interface to the server functionality
    and integrates with the application container.
    """
    
    def __init__(self, application_container):
        """Initialize the Trixy server."""
        self.application = application_container
        self.server_handler = ServerSocketHandler(application_container)
        
    def start(self, port: int = 2101) -> None:
        """Start the server."""
        self.server_handler.start(port)
    
    def stop(self) -> None:
        """Stop the server."""
        self.server_handler.stop()
    
    def enter_registration_mode(self, timeout: float = 60.0) -> None:
        """Enter registration mode."""
        self.server_handler.enter_registration_mode(timeout)
    
    def exit_registration_mode(self) -> None:
        """Exit registration mode.""" 
        self.server_handler.exit_registration_mode()
    
    def get_status(self) -> Dict[str, Any]:
        """Get server status."""
        return self.server_handler.get_statistics()


# Module exports
__all__ = [
    "TrixyServer",
    "ServerSocketHandler", 
    "ClientConnection",
    "ServerError",
    "PortInUseError",
    "ClientConnectionError",
]