Trixy3/trixy_core/network/client_socket.py

"""
Client Socket Implementation for Trixy Network

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

- Client connection management with auto-retry
- Satellite authentication and registration
- Command sending and response handling
- Audio stream client management
- Connection monitoring and heartbeat
- Integration with wakeword detection workflow

The client handles the full satellite connection process as specified
in CLAUDE.md including registration and connection establishment.
"""

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

from .protocol import TrixyProtocol, TrixyMessage, MessageFlags
from .cmd import (
    TrixyCommand,
    CommandResponse, 
    SatelliteHelloCommand,
    SatelliteRegistrationCommand,
    SatelliteConnectionCommand,
    create_satellite_hello,
    create_satellite_registration,
    create_satellite_connection,
    create_heartbeat,
    create_ping
)


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


class ClientError(Exception):
    """Base exception for client errors."""
    pass


class ConnectionFailedError(ClientError):
    """Raised when connection to server fails."""
    pass


class AuthenticationError(ClientError):
    """Raised when authentication fails."""
    pass


@dataclass
class ConnectionAttempt:
    """Information about a connection attempt."""
    attempt_number: int
    timestamp: float
    success: bool
    error_message: Optional[str] = None
    response_time: Optional[float] = None


class ClientSocketHandler:
    """
    Handles client-side socket operations.
    
    This class manages connection to the Trixy server, handles authentication,
    and processes messages according to the satellite workflow.
    """
    
    def __init__(self, application_container):
        """
        Initialize the client socket handler.
        
        Args:
            application_container: Main application container
        """
        self.application = application_container
        self.protocol = TrixyProtocol()
        
        # Connection state
        self._connected = False
        self._authenticated = False
        self._socket: Optional[socket.socket] = None
        self._server_address: Optional[Tuple[str, int]] = None
        
        # Client identity
        self._room_id = "unknown"
        self._alias = "satellite"
        self._mac_address = "00:00:00:00:00:00"
        self._satellite_id: Optional[str] = None
        self._session_id: Optional[str] = None
        
        # Connection management
        self._auto_reconnect = True
        self._reconnect_delay = 5.0
        self._max_reconnect_attempts = 0  # 0 = infinite
        self._connection_attempts: list[ConnectionAttempt] = []
        
        # Threading
        self._thread_pool = ThreadPoolExecutor(max_workers=10)
        self._running = True
        self._receive_thread: Optional[threading.Thread] = None
        self._heartbeat_thread: Optional[threading.Thread] = None
        
        # Message handling
        self._message_queue = queue.Queue()
        self._message_handlers: Dict[str, Callable] = {}
        self._pending_responses: Dict[str, queue.Queue] = {}
        self._setup_default_handlers()
        
        # Audio streams
        self._audio_streams: Dict[str, int] = {}  # stream_type -> port
        
        # Statistics
        self._stats = {
            'connection_attempts': 0,
            'successful_connections': 0,
            'messages_sent': 0,
            'messages_received': 0,
            'bytes_sent': 0,
            'bytes_received': 0,
            'last_ping_time': None,
            'last_pong_time': None
        }
        
        pprint("ClientSocketHandler initialized")
    
    def configure_satellite(self, room_id: str, alias: str, mac_address: str) -> None:
        """
        Configure satellite identity.
        
        Args:
            room_id: Room identifier
            alias: Satellite alias/name
            mac_address: MAC address
        """
        self._room_id = room_id
        self._alias = alias
        self._mac_address = mac_address
        
        pprint(f"Satellite configured: {alias} in {room_id} ({mac_address})")
    
    def connect(self, host: str, port: int = 2101, timeout: float = 10.0) -> bool:
        """
        Connect to the Trixy server.
        
        Args:
            host: Server hostname or IP
            port: Server port
            timeout: Connection timeout
            
        Returns:
            bool: True if connected successfully
        """
        if self._connected:
            return True
        
        pprint(f"Connecting to {host}:{port}...")
        
        attempt = ConnectionAttempt(
            attempt_number=len(self._connection_attempts) + 1,
            timestamp=time.time(),
            success=False
        )
        
        try:
            # Create socket
            self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self._socket.settimeout(timeout)
            
            # Connect
            start_time = time.time()
            self._socket.connect((host, port))
            attempt.response_time = time.time() - start_time
            
            self._server_address = (host, port)
            self._connected = True
            attempt.success = True
            
            # Start message processing
            self._start_message_processing()
            
            # Start authentication process
            self._thread_pool.submit(self._authenticate)
            
            # Start heartbeat
            self._start_heartbeat()
            
            # Update statistics
            self._stats['connection_attempts'] += 1
            self._stats['successful_connections'] += 1
            
            pprint(f"Connected to server successfully")
            
            # Trigger event
            self._trigger_event("client_connected", {
                'server_address': self._server_address,
                'response_time': attempt.response_time
            })
            
            return True
            
        except Exception as e:
            attempt.error_message = str(e)
            pprint(f"Failed to connect: {e}")
            
            if self._socket:
                self._socket.close()
                self._socket = None
            
            self._stats['connection_attempts'] += 1
            
            return False
            
        finally:
            self._connection_attempts.append(attempt)
    
    def disconnect(self) -> None:
        """Disconnect from server."""
        if not self._connected:
            return
        
        pprint("Disconnecting from server...")
        
        self._connected = False
        self._authenticated = False
        self._auto_reconnect = False
        
        # Close socket
        if self._socket:
            self._socket.close()
            self._socket = None
        
        # Wait for threads
        if self._receive_thread and self._receive_thread.is_alive():
            self._receive_thread.join(timeout=5.0)
        
        if self._heartbeat_thread and self._heartbeat_thread.is_alive():
            self._heartbeat_thread.join(timeout=5.0)
        
        # Clear state
        self._server_address = None
        self._satellite_id = None
        self._session_id = None
        self._audio_streams.clear()
        
        pprint("Disconnected from server")
        
        # Trigger event
        self._trigger_event("client_disconnected", {})
    
    def send_command(self, command: TrixyCommand, wait_for_response: bool = False, timeout: float = 30.0) -> Optional[CommandResponse]:
        """
        Send a command to the server.
        
        Args:
            command: Command to send
            wait_for_response: Whether to wait for response
            timeout: Response timeout
            
        Returns:
            Optional[CommandResponse]: Response if wait_for_response is True
        """
        if not self._connected or not self._socket:
            raise ClientError("Not connected to server")
        
        try:
            # Create protocol message
            message = self.protocol.create_message(
                command.get_command_name(),
                command
            )
            
            # Set up response waiting if needed
            response_queue = None
            if wait_for_response:
                response_queue = queue.Queue()
                self._pending_responses[command.command_id] = response_queue
            
            # Serialize and send
            data = message.serialize()
            self._socket.send(data)
            
            # Update statistics
            self._stats['messages_sent'] += 1
            self._stats['bytes_sent'] += len(data)
            
            pprint(f"Sent command: {command.get_command_name()}")
            
            # Wait for response if requested
            if wait_for_response and response_queue:
                try:
                    response = response_queue.get(timeout=timeout)
                    return response
                except queue.Empty:
                    pprint(f"Timeout waiting for response to {command.get_command_name()}")
                    return None
                finally:
                    self._pending_responses.pop(command.command_id, None)
            
            return None
            
        except Exception as e:
            pprint(f"Error sending command: {e}")
            if wait_for_response:
                self._pending_responses.pop(command.command_id, None)
            raise ClientError(f"Failed to send command: {e}") from e
    
    def send_hardcoded_command(self, command_text: str) -> bool:
        """
        Send a hard-coded command.
        
        Args:
            command_text: Hard-coded command string
            
        Returns:
            bool: True if sent successfully
        """
        if not self._connected or not self._socket:
            return False
        
        try:
            data = command_text.encode('utf-8')
            self._socket.send(data)
            
            self._stats['messages_sent'] += 1
            self._stats['bytes_sent'] += len(data)
            
            pprint(f"Sent hard-coded command: {command_text}")
            return True
            
        except Exception as e:
            pprint(f"Error sending hard-coded command: {e}")
            return False
    
    def ping_server(self) -> Optional[float]:
        """
        Ping the server and return round-trip time.
        
        Returns:
            Optional[float]: RTT in seconds or None if failed
        """
        if not self._connected:
            return None
        
        ping_id = str(int(time.time() * 1000))
        ping_cmd = create_ping(ping_id)
        
        start_time = time.time()
        self._stats['last_ping_time'] = start_time
        
        response = self.send_command(ping_cmd, wait_for_response=True, timeout=10.0)
        
        if response and hasattr(response, 'ping_id') and response.ping_id == ping_id:
            rtt = time.time() - start_time
            self._stats['last_pong_time'] = time.time()
            return rtt
        
        return None
    
    def enable_auto_reconnect(self, enabled: bool = True, delay: float = 5.0, max_attempts: int = 0) -> None:
        """
        Enable or disable auto-reconnect.
        
        Args:
            enabled: Whether to enable auto-reconnect
            delay: Delay between reconnect attempts
            max_attempts: Maximum attempts (0 = infinite)
        """
        self._auto_reconnect = enabled
        self._reconnect_delay = delay
        self._max_reconnect_attempts = max_attempts
        
        pprint(f"Auto-reconnect: {enabled} (delay={delay}s, max_attempts={max_attempts})")
    
    def get_status(self) -> Dict[str, Any]:
        """Get client status information."""
        return {
            'connected': self._connected,
            'authenticated': self._authenticated,
            'server_address': self._server_address,
            'satellite_id': self._satellite_id,
            'session_id': self._session_id,
            'room_id': self._room_id,
            'alias': self._alias,
            'mac_address': self._mac_address,
            'audio_streams': self._audio_streams,
            'auto_reconnect': self._auto_reconnect,
            'statistics': self._stats,
            'connection_attempts': len(self._connection_attempts)
        }
    
    def _start_message_processing(self) -> None:
        """Start message processing threads."""
        self._receive_thread = threading.Thread(target=self._receive_messages, daemon=True)
        self._receive_thread.start()
        
        self._thread_pool.submit(self._process_message_queue)
    
    def _receive_messages(self) -> None:
        """Receive messages from server (runs in separate thread)."""
        pprint("Started receiving messages")
        
        while self._connected and self._socket:
            try:
                self._socket.settimeout(1.0)
                data = self._socket.recv(4096)
                
                if not data:
                    pprint("Server closed connection")
                    break
                
                # Update statistics
                self._stats['bytes_received'] += len(data)
                self._stats['messages_received'] += 1
                
                # Queue for processing
                self._message_queue.put(data)
                
            except socket.timeout:
                continue
            except Exception as e:
                if self._connected:
                    pprint(f"Error receiving messages: {e}")
                break
        
        # Connection lost - attempt reconnect if enabled
        if self._auto_reconnect and self._connected:
            self._connected = False
            self._authenticated = False
            self._trigger_event("connection_lost", {})
            self._thread_pool.submit(self._attempt_reconnect)
    
    def _process_message_queue(self) -> None:
        """Process queued messages."""
        while self._running:
            try:
                data = self._message_queue.get(timeout=1.0)
                self._process_received_data(data)
                self._message_queue.task_done()
            except queue.Empty:
                continue
            except Exception as e:
                pprint(f"Error processing message: {e}")
    
    def _process_received_data(self, data: bytes) -> None:
        """Process received data from server."""
        try:
            # Check for hard-coded commands/responses
            if self.protocol.is_hardcoded_command(data):
                command, args = self.protocol.parse_hardcoded_command(data)
                self._handle_hardcoded_response(command, args)
            else:
                # Parse protocol message
                message = self.protocol.deserialize_message(data)
                self._handle_message(message)
                
        except Exception as e:
            pprint(f"Error processing received data: {e}")
    
    def _handle_hardcoded_response(self, command: str, args: str) -> None:
        """Handle hard-coded command responses."""
        if command == "TRXIPONG":
            # Pong response to ping
            self._stats['last_pong_time'] = time.time()
            pprint(f"Received pong: {args}")
    
    def _handle_message(self, message: TrixyMessage) -> None:
        """Handle a protocol message."""
        command_name = message.class_name
        
        # Check for response to pending command
        if hasattr(message.data, 'command_id'):
            response_queue = self._pending_responses.get(message.data.command_id)
            if response_queue:
                response_queue.put(message.data)
                return
        
        # Check for registered handler
        handler = self._message_handlers.get(command_name)
        if handler:
            handler(message)
        else:
            # Trigger event for unhandled messages
            self._trigger_event("message_received", {
                'message': message,
                'command_name': command_name
            })
    
    def _setup_default_handlers(self) -> None:
        """Set up default message handlers."""
        # Default handlers can be added here
        pass
    
    def _authenticate(self) -> None:
        """Perform authentication sequence with server."""
        try:
            pprint("Starting authentication...")
            
            # Step 1: Send hello command
            hello_cmd = create_satellite_hello(
                room_id=self._room_id,
                alias=self._alias,
                mac_address=self._mac_address,
                version="1.0.0",
                capabilities=["wakeword", "audio_input", "audio_output"]
            )
            
            response = self.send_command(hello_cmd, wait_for_response=True, timeout=15.0)
            
            if not response or not response.is_success():
                if response and hasattr(response, 'registration_required') and response.registration_required:
                    pprint("Registration required")
                    self._trigger_event("registration_required", {})
                    return
                else:
                    raise AuthenticationError("Hello command failed")
            
            # Extract satellite ID if provided
            if hasattr(response, 'assigned_satellite_id') and response.assigned_satellite_id:
                self._satellite_id = response.assigned_satellite_id
            
            # Step 2: If registered, send connection command
            if self._satellite_id:
                conn_cmd = create_satellite_connection(
                    satellite_id=self._satellite_id,
                    mac_address=self._mac_address,
                    requested_audio_ports={
                        'input': 0,  # Let server assign
                        'output': 0,
                        'music': 0
                    }
                )
                
                conn_response = self.send_command(conn_cmd, wait_for_response=True, timeout=30.0)
                
                if conn_response and conn_response.is_success():
                    self._authenticated = True
                    
                    # Extract session info
                    if hasattr(conn_response, 'session_id'):
                        self._session_id = conn_response.session_id
                    
                    if hasattr(conn_response, 'assigned_audio_ports'):
                        self._audio_streams = conn_response.assigned_audio_ports
                    
                    pprint(f"Authentication successful: {self._satellite_id}")
                    
                    # Trigger event
                    self._trigger_event("authentication_successful", {
                        'satellite_id': self._satellite_id,
                        'session_id': self._session_id,
                        'audio_streams': self._audio_streams
                    })
                else:
                    raise AuthenticationError("Connection command failed")
            else:
                raise AuthenticationError("No satellite ID assigned")
                
        except Exception as e:
            pprint(f"Authentication failed: {e}")
            self._trigger_event("authentication_failed", {'error': str(e)})
    
    def _start_heartbeat(self) -> None:
        """Start heartbeat thread."""
        self._heartbeat_thread = threading.Thread(target=self._heartbeat_loop, daemon=True)
        self._heartbeat_thread.start()
    
    def _heartbeat_loop(self) -> None:
        """Send periodic heartbeat messages."""
        heartbeat_interval = 30.0  # seconds
        
        while self._connected:
            try:
                time.sleep(heartbeat_interval)
                
                if self._connected and self._socket:
                    # Send heartbeat
                    heartbeat = create_heartbeat()
                    self.send_command(heartbeat)
                    
            except Exception as e:
                if self._connected:
                    pprint(f"Error sending heartbeat: {e}")
    
    def _attempt_reconnect(self) -> None:
        """Attempt to reconnect to server."""
        if not self._auto_reconnect or not self._server_address:
            return
        
        pprint("Attempting to reconnect...")
        
        attempt_count = 0
        
        while self._auto_reconnect and not self._connected:
            attempt_count += 1
            
            if self._max_reconnect_attempts > 0 and attempt_count > self._max_reconnect_attempts:
                pprint("Max reconnect attempts reached")
                break
            
            pprint(f"Reconnect attempt {attempt_count}")
            
            if self.connect(self._server_address[0], self._server_address[1]):
                pprint("Reconnected successfully")
                break
            
            time.sleep(self._reconnect_delay)
        
        if not self._connected:
            pprint("Failed to reconnect")
            self._trigger_event("reconnect_failed", {
                'attempts': attempt_count
            })
    
    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(f"client_{event_name}", event_data)
        except Exception as e:
            pprint(f"Error triggering event {event_name}: {e}")
    
    def shutdown(self) -> None:
        """Shutdown the client socket handler."""
        pprint("Shutting down client socket handler...")
        
        self._running = False
        self._auto_reconnect = False
        
        self.disconnect()
        
        # Shutdown thread pool
        self._thread_pool.shutdown(wait=True)
        
        pprint("Client socket handler shutdown complete")


class TrixyClient:
    """
    High-level Trixy client implementation.
    
    This class provides a simple interface for satellite clients
    and integrates with the application container.
    """
    
    def __init__(self, application_container):
        """Initialize the Trixy client."""
        self.application = application_container
        self.client_handler = ClientSocketHandler(application_container)
    
    def configure_satellite(self, room_id: str, alias: str, mac_address: str) -> None:
        """Configure satellite identity."""
        self.client_handler.configure_satellite(room_id, alias, mac_address)
    
    def connect(self, host: str, port: int = 2101) -> bool:
        """Connect to server."""
        return self.client_handler.connect(host, port)
    
    def disconnect(self) -> None:
        """Disconnect from server."""
        self.client_handler.disconnect()
    
    def send_command(self, command: TrixyCommand) -> bool:
        """Send a command to server."""
        try:
            self.client_handler.send_command(command)
            return True
        except:
            return False
    
    def ping_server(self) -> Optional[float]:
        """Ping the server."""
        return self.client_handler.ping_server()
    
    def get_status(self) -> Dict[str, Any]:
        """Get client status."""
        return self.client_handler.get_status()
    
    def shutdown(self) -> None:
        """Shutdown the client."""
        self.client_handler.shutdown()


# Module exports
__all__ = [
    "TrixyClient",
    "ClientSocketHandler",
    "ConnectionAttempt",
    "ClientError",
    "ConnectionFailedError",
    "AuthenticationError",
]