Trixy3/trixy_core/network/network_manager.py

"""
Network Manager for Trixy System

This module provides the main NetworkManager class that handles:

- Command socket management (Port 2101)
- Audio streaming sockets (Ports 2102, 2103, 2104) 
- Multi-client support with thread safety
- Connection management and monitoring
- Protocol serialization/deserialization
- Integration with event system and application container

The NetworkManager serves as the central hub for all network communication
in the Trixy system, supporting both server and client modes.
"""

import socket
import threading
import time
import queue
import weakref
from typing import Dict, Any, Optional, List, Callable, Union, Tuple
from dataclasses import dataclass, field
from enum import Enum, IntEnum
from concurrent.futures import ThreadPoolExecutor
import select

from .protocol import TrixyProtocol, TrixyMessage, MessageFlags, ProtocolError
from .cmd import TrixyCommand, CommandResponse


def pprint(message: str) -> None:
    """Network manager logging function."""
    print(f"[NETWORK_MANAGER] {message}")


class ConnectionState(Enum):
    """Connection state enumeration."""
    DISCONNECTED = "disconnected"
    CONNECTING = "connecting" 
    CONNECTED = "connected"
    AUTHENTICATING = "authenticating"
    AUTHENTICATED = "authenticated"
    DISCONNECTING = "disconnecting"
    ERROR = "error"


class AudioStreamType(Enum):
    """Audio stream type enumeration."""
    INPUT = "input"      # Port 2102 - Raw audio from satellites
    OUTPUT = "output"    # Port 2103 - TTS/response audio  
    MUSIC = "music"      # Port 2104 - Music/media playback


@dataclass
class AudioStreamConfig:
    """Configuration for audio streams."""
    sample_rate: int = 16000
    channels: int = 1
    sample_width: int = 2  # bytes
    format: str = "PCM"
    buffer_size: int = 4096


@dataclass  
class ConnectionInfo:
    """Information about a network connection."""
    connection_id: str
    remote_address: Tuple[str, int]
    local_address: Tuple[str, int]
    state: ConnectionState
    connected_time: float
    last_activity: float
    bytes_sent: int = 0
    bytes_received: int = 0
    commands_sent: int = 0
    commands_received: int = 0
    satellite_id: Optional[str] = None
    session_id: Optional[str] = None


class NetworkManagerError(Exception):
    """Base exception for network manager errors."""
    pass


class ConnectionManager:
    """Manages individual connections and their state."""
    
    def __init__(self):
        self._connections: Dict[str, ConnectionInfo] = {}
        self._sockets: Dict[str, socket.socket] = {}
        self._lock = threading.RLock()
    
    def add_connection(self, connection_id: str, sock: socket.socket, remote_addr: Tuple[str, int]) -> ConnectionInfo:
        """Add a new connection."""
        with self._lock:
            local_addr = sock.getsockname()
            connection = ConnectionInfo(
                connection_id=connection_id,
                remote_address=remote_addr,
                local_address=local_addr,
                state=ConnectionState.CONNECTED,
                connected_time=time.time(),
                last_activity=time.time()
            )
            
            self._connections[connection_id] = connection
            self._sockets[connection_id] = sock
            
            pprint(f"Connection added: {connection_id} from {remote_addr}")
            return connection
    
    def remove_connection(self, connection_id: str) -> Optional[ConnectionInfo]:
        """Remove a connection."""
        with self._lock:
            connection = self._connections.pop(connection_id, None)
            sock = self._sockets.pop(connection_id, None)
            
            if sock:
                try:
                    sock.close()
                except:
                    pass
            
            if connection:
                pprint(f"Connection removed: {connection_id}")
            
            return connection
    
    def get_connection(self, connection_id: str) -> Optional[ConnectionInfo]:
        """Get connection info."""
        with self._lock:
            return self._connections.get(connection_id)
    
    def get_socket(self, connection_id: str) -> Optional[socket.socket]:
        """Get socket for connection."""
        with self._lock:
            return self._sockets.get(connection_id)
    
    def update_activity(self, connection_id: str, bytes_delta: int = 0, command_delta: int = 0):
        """Update connection activity."""
        with self._lock:
            connection = self._connections.get(connection_id)
            if connection:
                connection.last_activity = time.time()
                if bytes_delta > 0:
                    connection.bytes_received += bytes_delta
                if command_delta > 0:
                    connection.commands_received += command_delta
    
    def get_all_connections(self) -> Dict[str, ConnectionInfo]:
        """Get all connections."""
        with self._lock:
            return self._connections.copy()


class AudioStreamManager:
    """Manages audio streaming sockets."""
    
    def __init__(self):
        self._streams: Dict[AudioStreamType, Dict[str, Any]] = {}
        self._configs: Dict[AudioStreamType, AudioStreamConfig] = {}
        self._lock = threading.RLock()
        
        # Initialize default configs
        for stream_type in AudioStreamType:
            if stream_type == AudioStreamType.MUSIC:
                self._configs[stream_type] = AudioStreamConfig(
                    sample_rate=48000, 
                    channels=2,
                    sample_width=2
                )
            else:
                self._configs[stream_type] = AudioStreamConfig()
    
    def start_stream(self, stream_type: AudioStreamType, port: int, config: Optional[AudioStreamConfig] = None):
        """Start an audio stream."""
        with self._lock:
            if stream_type in self._streams:
                raise NetworkManagerError(f"Audio stream {stream_type.value} already active")
            
            if config:
                self._configs[stream_type] = config
            
            # Create server socket for audio stream
            server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            server_sock.bind(('', port))
            server_sock.listen(10)
            
            self._streams[stream_type] = {
                'server_socket': server_sock,
                'port': port,
                'connections': {},
                'active': True
            }
            
            pprint(f"Audio stream {stream_type.value} started on port {port}")
    
    def stop_stream(self, stream_type: AudioStreamType):
        """Stop an audio stream."""
        with self._lock:
            stream_info = self._streams.pop(stream_type, None)
            if stream_info:
                stream_info['active'] = False
                server_sock = stream_info.get('server_socket')
                if server_sock:
                    server_sock.close()
                
                # Close all connections
                for conn_sock in stream_info.get('connections', {}).values():
                    try:
                        conn_sock.close()
                    except:
                        pass
                
                pprint(f"Audio stream {stream_type.value} stopped")
    
    def get_stream_info(self, stream_type: AudioStreamType) -> Optional[Dict[str, Any]]:
        """Get stream information."""
        with self._lock:
            return self._streams.get(stream_type)


class NetworkManager:
    """
    Main network manager for the Trixy system.
    
    Handles all network communication including command sockets and audio streams.
    Integrates with the event system and application container.
    """
    
    def __init__(self, application_container):
        """
        Initialize the network manager.
        
        Args:
            application_container: The main application container
        """
        self.application = application_container
        self.application_ref = weakref.ref(application_container)
        
        # Core components
        self.protocol = TrixyProtocol()
        self.connection_manager = ConnectionManager()
        self.audio_stream_manager = AudioStreamManager()
        
        # Server state
        self._server_running = False
        self._command_server_socket: Optional[socket.socket] = None
        self._command_port = 2101
        
        # Client state
        self._client_connected = False
        self._client_socket: Optional[socket.socket] = None
        self._server_address: Optional[Tuple[str, int]] = None
        
        # Threading
        self._thread_pool = ThreadPoolExecutor(max_workers=20)
        self._running = True
        self._lock = threading.RLock()
        
        # Message handling
        self._message_queue = queue.Queue()
        self._message_handlers: Dict[str, Callable] = {}
        self._response_handlers: Dict[str, Callable] = {}
        
        # Event system integration
        self._event_handler = None
        
        pprint("NetworkManager initialized")
    
    def start_server(self, command_port: int = 2101, audio_ports: Optional[Dict[str, int]] = None) -> None:
        """
        Start the network manager in server mode.
        
        Args:
            command_port: Port for command socket
            audio_ports: Ports for audio streams
        """
        if self._server_running:
            raise NetworkManagerError("Server already running")
        
        pprint(f"Starting network server on port {command_port}")
        
        self._command_port = command_port
        
        try:
            # Create command server socket
            self._command_server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self._command_server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self._command_server_socket.bind(('', command_port))
            self._command_server_socket.listen(50)
            
            # Start audio streams
            if audio_ports is None:
                audio_ports = {
                    'input': 2102,
                    'output': 2103,
                    'music': 2104
                }
            
            for stream_name, port in audio_ports.items():
                if stream_name in ['input', 'output', 'music']:
                    stream_type = AudioStreamType(stream_name)
                    self.audio_stream_manager.start_stream(stream_type, port)
            
            self._server_running = True
            
            # Start accepting connections
            self._thread_pool.submit(self._accept_connections)
            self._thread_pool.submit(self._process_messages)
            
            # Trigger server started event
            self._trigger_event("network_server_started", {
                'command_port': command_port,
                'audio_ports': audio_ports
            })
            
            pprint(f"Network server started successfully")
            
        except Exception as e:
            self._cleanup_server()
            raise NetworkManagerError(f"Failed to start server: {e}") from e
    
    def stop_server(self) -> None:
        """Stop the network server."""
        if not self._server_running:
            return
        
        pprint("Stopping network server...")
        self._server_running = False
        
        self._cleanup_server()
        
        self._trigger_event("network_server_stopped", {})
        pprint("Network server stopped")
    
    def connect_to_server(self, host: str, port: int = 2101) -> bool:
        """
        Connect to a Trixy server as a client.
        
        Args:
            host: Server hostname or IP
            port: Server command port
            
        Returns:
            bool: True if connected successfully
        """
        if self._client_connected:
            raise NetworkManagerError("Client already connected")
        
        pprint(f"Connecting to server {host}:{port}")
        
        try:
            self._client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self._client_socket.settimeout(10.0)
            self._client_socket.connect((host, port))
            
            self._server_address = (host, port)
            self._client_connected = True
            
            # Start client message processing
            connection_id = f"client_{int(time.time())}"
            self.connection_manager.add_connection(
                connection_id, 
                self._client_socket, 
                self._server_address
            )
            
            self._thread_pool.submit(self._handle_client_connection, connection_id)
            self._thread_pool.submit(self._process_messages)
            
            self._trigger_event("network_client_connected", {
                'server_address': self._server_address
            })
            
            pprint(f"Connected to server successfully")
            return True
            
        except Exception as e:
            self._cleanup_client()
            pprint(f"Failed to connect to server: {e}")
            return False
    
    def disconnect_from_server(self) -> None:
        """Disconnect from server."""
        if not self._client_connected:
            return
        
        pprint("Disconnecting from server...")
        self._client_connected = False
        
        self._cleanup_client()
        
        self._trigger_event("network_client_disconnected", {})
        pprint("Disconnected from server")
    
    def send_command(self, command: TrixyCommand, connection_id: Optional[str] = None) -> bool:
        """
        Send a command over the network.
        
        Args:
            command: Command to send
            connection_id: Target connection (None for client mode)
            
        Returns:
            bool: True if sent successfully
        """
        try:
            # Create protocol message
            message = self.protocol.create_message(
                command.get_command_name(),
                command
            )
            
            # Serialize message
            data = message.serialize()
            
            # Send based on mode
            if self._client_connected:
                # Client mode - send to server
                sock = self._client_socket
                if sock:
                    sock.send(data)
                    pprint(f"Sent command: {command.get_command_name()}")
                    return True
            else:
                # Server mode - send to specific connection
                if connection_id:
                    sock = self.connection_manager.get_socket(connection_id)
                    if sock:
                        sock.send(data)
                        pprint(f"Sent command to {connection_id}: {command.get_command_name()}")
                        return True
            
            return False
            
        except Exception as e:
            pprint(f"Failed to send command: {e}")
            return False
    
    def broadcast_command(self, command: TrixyCommand) -> int:
        """
        Broadcast a command to all connected clients.
        
        Args:
            command: Command to broadcast
            
        Returns:
            int: Number of clients that received the command
        """
        if not self._server_running:
            return 0
        
        connections = self.connection_manager.get_all_connections()
        sent_count = 0
        
        for connection_id in connections:
            if self.send_command(command, connection_id):
                sent_count += 1
        
        return sent_count
    
    def register_message_handler(self, command_name: str, handler: Callable[[TrixyMessage, str], None]) -> None:
        """
        Register a handler for incoming messages.
        
        Args:
            command_name: Name of command to handle
            handler: Handler function (message, connection_id) -> None
        """
        self._message_handlers[command_name] = handler
        pprint(f"Registered message handler for: {command_name}")
    
    def _accept_connections(self) -> None:
        """Accept incoming connections (server thread)."""
        pprint("Started accepting connections")
        
        while self._server_running and self._command_server_socket:
            try:
                # Use select to check for new connections with timeout
                ready, _, _ = select.select([self._command_server_socket], [], [], 1.0)
                
                if ready:
                    client_sock, addr = self._command_server_socket.accept()
                    connection_id = f"client_{addr[0]}_{addr[1]}_{int(time.time())}"
                    
                    self.connection_manager.add_connection(connection_id, client_sock, addr)
                    
                    # Handle connection in separate thread
                    self._thread_pool.submit(self._handle_client_connection, connection_id)
                    
                    self._trigger_event("network_client_connected", {
                        'connection_id': connection_id,
                        'remote_address': addr
                    })
                    
            except Exception as e:
                # Only log errors if server is still supposed to be running
                if self._server_running and self._running:
                    pprint(f"Error accepting connection: {e}")
    
    def _handle_client_connection(self, connection_id: str) -> None:
        """Handle messages from a client connection."""
        sock = self.connection_manager.get_socket(connection_id)
        if not sock:
            return
        
        pprint(f"Handling connection: {connection_id}")
        
        try:
            while self._running and sock:
                # Receive data with timeout
                sock.settimeout(1.0)
                
                try:
                    data = sock.recv(4096)
                    if not data:
                        break
                    
                    # Update connection activity
                    self.connection_manager.update_activity(connection_id, len(data))
                    
                    # Process the received data
                    self._process_received_data(data, connection_id)
                    
                except socket.timeout:
                    continue
                except Exception as e:
                    pprint(f"Error receiving data from {connection_id}: {e}")
                    break
                    
        except Exception as e:
            pprint(f"Error handling connection {connection_id}: {e}")
        finally:
            self.connection_manager.remove_connection(connection_id)
            self._trigger_event("network_client_disconnected", {
                'connection_id': connection_id
            })
    
    def _process_received_data(self, data: bytes, connection_id: str) -> None:
        """Process received data and extract messages."""
        try:
            # Check if it's a hard-coded command
            if self.protocol.is_hardcoded_command(data):
                command, args = self.protocol.parse_hardcoded_command(data)
                pprint(f"Received hard-coded command from {connection_id}: {command}")
                
                # Handle hard-coded commands immediately
                self._handle_hardcoded_command(command, args, connection_id)
            else:
                # Parse as protocol message
                message = self.protocol.deserialize_message(data)
                
                # Queue for processing
                self._message_queue.put((message, connection_id))
                
        except Exception as e:
            pprint(f"Error processing received data: {e}")
    
    def _process_messages(self) -> None:
        """Process queued messages."""
        pprint("Started message processing")
        
        while self._running:
            try:
                # Get message with timeout
                message, connection_id = self._message_queue.get(timeout=1.0)
                
                # Handle the message
                self._handle_message(message, connection_id)
                
                self._message_queue.task_done()
                
            except queue.Empty:
                continue
            except Exception as e:
                pprint(f"Error processing message: {e}")
    
    def _handle_message(self, message: TrixyMessage, connection_id: str) -> None:
        """Handle a received message."""
        try:
            command_name = message.class_name
            
            # Check for registered handler
            handler = self._message_handlers.get(command_name)
            if handler:
                handler(message, connection_id)
            else:
                # Default handling - trigger event
                self._trigger_event("network_message_received", {
                    'message': message,
                    'connection_id': connection_id,
                    'command_name': command_name
                })
            
            pprint(f"Handled message: {command_name} from {connection_id}")
            
        except Exception as e:
            pprint(f"Error handling message: {e}")
    
    def _handle_hardcoded_command(self, command: str, args: str, connection_id: str) -> None:
        """Handle hard-coded commands."""
        try:
            if command == "TRXINOOP":
                # No-op/heartbeat - just update activity
                self.connection_manager.update_activity(connection_id)
                
            elif command == "TRXIPING":
                # Respond with pong
                pong_data = f"TRXIPONG {args}".encode('utf-8')
                sock = self.connection_manager.get_socket(connection_id)
                if sock:
                    sock.send(pong_data)
                
            elif command == "TRXIPRNT":
                # Print command
                pprint(f"Print from {connection_id}: {args}")
                
            elif command == "TRXYHELO":
                # Hello command
                pprint(f"Hello from {connection_id}: {args}")
                
        except Exception as e:
            pprint(f"Error handling hard-coded command {command}: {e}")
    
    def _trigger_event(self, event_name: str, event_data: Dict[str, Any]) -> None:
        """Trigger an event through the event system."""
        try:
            # Don't trigger events if shutting down
            if not self._running:
                return
                
            if not self._event_handler:
                app = self.application_ref()
                if app:
                    self._event_handler = app.get_event_handler()
            
            if self._event_handler:
                self._event_handler.trigger_event(event_name, event_data)
                
        except Exception as e:
            # Only log errors if still running (avoid shutdown-related errors)
            if self._running:
                pprint(f"Error triggering event {event_name}: {e}")
    
    def _cleanup_server(self) -> None:
        """Clean up server resources."""
        if self._command_server_socket:
            self._command_server_socket.close()
            self._command_server_socket = None
        
        # Stop all audio streams
        for stream_type in AudioStreamType:
            self.audio_stream_manager.stop_stream(stream_type)
        
        # Close all connections
        for connection_id in list(self.connection_manager.get_all_connections().keys()):
            self.connection_manager.remove_connection(connection_id)
    
    def _cleanup_client(self) -> None:
        """Clean up client resources."""
        if self._client_socket:
            self._client_socket.close()
            self._client_socket = None
        
        self._server_address = None
    
    def shutdown(self) -> None:
        """Shutdown the network manager."""
        pprint("Shutting down network manager...")
        
        self._running = False
        
        if self._server_running:
            self.stop_server()
        
        if self._client_connected:
            self.disconnect_from_server()
        
        self._thread_pool.shutdown(wait=True)
        
        pprint("Network manager shutdown complete")
    
    def get_status(self) -> Dict[str, Any]:
        """Get network manager status."""
        connections = self.connection_manager.get_all_connections()
        
        return {
            'server_running': self._server_running,
            'client_connected': self._client_connected,
            'command_port': self._command_port,
            'server_address': self._server_address,
            'active_connections': len(connections),
            'connection_details': {cid: {
                'remote_address': conn.remote_address,
                'state': conn.state.value,
                'connected_time': conn.connected_time,
                'bytes_sent': conn.bytes_sent,
                'bytes_received': conn.bytes_received
            } for cid, conn in connections.items()}
        }


# Module exports
__all__ = [
    "NetworkManager",
    "NetworkManagerError",
    "ConnectionManager", 
    "ConnectionState",
    "ConnectionInfo",
    "AudioStreamManager",
    "AudioStreamType",
    "AudioStreamConfig",
]