import time
import threading
import queue
from typing import Optional, Dict, Any, List, Callable, Tuple

# ANSI colors for logging
RED = "\033[91m"
GREEN = "\033[92m"
YELLOW = "\033[93m"
BLUE = "\033[94m"
RESET = "\033[0m"

class AutoModeConfig:
    """Configuration parameters for the automatic mode behavior."""
    def __init__(self):
        # Ping behavior
        self.ping_response_accept = True  # Whether to accept incoming pings
        self.ping_auto_initiate = False   # Whether to initiate pings when connected
        self.ping_retry_count = 3         # Number of ping retries
        self.ping_retry_delay = 5.0       # Seconds between ping retries
        self.ping_timeout = 10.0          # Seconds to wait for ping response
        self.preferred_cipher = 0         # 0=AES-GCM, 1=ChaCha20-Poly1305
        
        # Handshake behavior
        self.handshake_retry_count = 3    # Number of handshake retries
        self.handshake_retry_delay = 5.0  # Seconds between handshake retries
        self.handshake_timeout = 10.0     # Seconds to wait for handshake
        
        # Messaging behavior
        self.auto_message_enabled = False # Whether to auto-send messages
        self.message_interval = 10.0      # Seconds between auto messages
        self.message_content = "Hello, secure world!"  # Default message
        
        # General behavior
        self.active_mode = False          # If true, initiates protocol instead of waiting


class AutoMode:
    """
    Manages automated behavior for the Icing protocol.
    Handles automatic progression through the protocol stages:
    1. Connection setup
    2. Ping/discovery
    3. Key exchange
    4. Encrypted communication
    """
    
    def __init__(self, protocol_interface):
        """
        Initialize the AutoMode manager.
        
        Args:
            protocol_interface: An object implementing the required protocol methods
        """
        self.protocol = protocol_interface
        self.config = AutoModeConfig()
        self.active = False
        self.state = "idle"
        
        # Message queue for automated sending
        self.message_queue = queue.Queue()
        
        # Tracking variables
        self.ping_attempts = 0
        self.handshake_attempts = 0
        self.last_action_time = 0
        self.timer_tasks = []  # List of active timer tasks (for cleanup)
        
    def start(self):
        """Start the automatic mode."""
        if self.active:
            return

        self.active = True
        self.state = "idle"
        self.ping_attempts = 0
        self.handshake_attempts = 0
        self.last_action_time = time.time()
        
        self._log_info("Automatic mode started")
        
        # Start in active mode if configured
        if self.config.active_mode and self.protocol.connections:
            self._start_ping_sequence()
            
    def stop(self):
        """Stop the automatic mode and clean up any pending tasks."""
        if not self.active:
            return
            
        # Cancel any pending timers
        for timer in self.timer_tasks:
            if timer.is_alive():
                timer.cancel()
        self.timer_tasks = []
        
        self.active = False
        self.state = "idle"
        self._log_info("Automatic mode stopped")
    
    def handle_connection_established(self):
        """Called when a new connection is established."""
        if not self.active:
            return
            
        self._log_info("Connection established")
        
        # If in active mode, start pinging
        if self.config.active_mode:
            self._start_ping_sequence()
    
    def handle_ping_received(self, index: int):
        """
        Handle a received ping request.
        
        Args:
            index: Index of the ping request in the protocol's inbound message queue
        """
        if not self.active or not self._is_valid_message_index(index):
            return
            
        self._log_info(f"Ping request received (index={index})")
        
        # Automatically respond to ping if configured to accept
        if self.config.ping_response_accept:
            self._log_info(f"Auto-responding to ping with accept={self.config.ping_response_accept}")
            try:
                # Schedule the response with a small delay to simulate real behavior
                timer = threading.Timer(0.5, self._respond_to_ping, args=[index])
                timer.daemon = True
                timer.start()
                self.timer_tasks.append(timer)
            except Exception as e:
                self._log_error(f"Failed to auto-respond to ping: {e}")
    
    def handle_ping_response_received(self, accepted: bool):
        """
        Handle a received ping response.
        
        Args:
            accepted: Whether the ping was accepted
        """
        if not self.active:
            return
            
        self.ping_attempts = 0  # Reset ping attempts counter
        
        if accepted:
            self._log_info("Ping accepted! Proceeding with handshake")
            # Send handshake if not already done
            if self.state != "handshake_sent":
                self._ensure_ephemeral_keys()
                self._start_handshake_sequence()
        else:
            self._log_info("Ping rejected by peer. Stopping auto-protocol sequence.")
            self.state = "idle"
    
    def handle_handshake_received(self, index: int):
        """
        Handle a received handshake.
        
        Args:
            index: Index of the handshake in the protocol's inbound message queue
        """
        if not self.active or not self._is_valid_message_index(index):
            return
            
        self._log_info(f"Handshake received (index={index})")
        
        try:
            # Ensure we have ephemeral keys
            self._ensure_ephemeral_keys()
            
            # Process the handshake (compute ECDH)
            self.protocol.generate_ecdhe(index)
            
            # Derive HKDF key
            self.protocol.derive_hkdf()
            
            # If we haven't sent our handshake yet, send it
            if self.state != "handshake_sent":
                timer = threading.Timer(0.5, self.protocol.send_handshake)
                timer.daemon = True
                timer.start()
                self.timer_tasks.append(timer)
                self.state = "handshake_sent"
            else:
                self.state = "key_exchange_complete"
                
            # Start sending queued messages if auto messaging is enabled
            if self.config.auto_message_enabled:
                self._start_message_sequence()
                
        except Exception as e:
            self._log_error(f"Failed to process handshake: {e}")
    
    def handle_encrypted_received(self, index: int):
        """
        Handle a received encrypted message.
        
        Args:
            index: Index of the encrypted message in the protocol's inbound message queue
        """
        if not self.active or not self._is_valid_message_index(index):
            return
            
        # Try to decrypt automatically
        try:
            plaintext = self.protocol.decrypt_received_message(index)
            self._log_info(f"Auto-decrypted message: {plaintext}")
        except Exception as e:
            self._log_error(f"Failed to auto-decrypt message: {e}")
    
    def queue_message(self, message: str):
        """
        Add a message to the auto-send queue.
        
        Args:
            message: Message text to send
        """
        self.message_queue.put(message)
        self._log_info(f"Message queued for sending: {message}")
        
        # If we're in the right state, start sending messages
        if self.active and self.state == "key_exchange_complete" and self.config.auto_message_enabled:
            self._process_message_queue()
    
    def _start_ping_sequence(self):
        """Start the ping sequence to discover the peer."""
        if self.ping_attempts >= self.config.ping_retry_count:
            self._log_warning(f"Maximum ping attempts ({self.config.ping_retry_count}) reached")
            self.state = "idle"
            return
            
        self.state = "pinging"
        self.ping_attempts += 1
        
        self._log_info(f"Sending ping request (attempt {self.ping_attempts}/{self.config.ping_retry_count})")
        try:
            self.protocol.send_ping_request(self.config.preferred_cipher)
            self.last_action_time = time.time()
            
            # Schedule next ping attempt if needed
            timer = threading.Timer(
                self.config.ping_retry_delay, 
                self._check_ping_response
            )
            timer.daemon = True
            timer.start()
            self.timer_tasks.append(timer)
            
        except Exception as e:
            self._log_error(f"Failed to send ping: {e}")
    
    def _check_ping_response(self):
        """Check if we got a ping response, retry if not."""
        if not self.active or self.state != "pinging":
            return
            
        # If we've waited long enough for a response, retry
        if time.time() - self.last_action_time >= self.config.ping_timeout:
            self._log_warning("No ping response received, retrying")
            self._start_ping_sequence()
    
    def _respond_to_ping(self, index: int):
        """
        Respond to a ping request.
        
        Args:
            index: Index of the ping request in the inbound messages
        """
        if not self.active or not self._is_valid_message_index(index):
            return
            
        try:
            answer = 1 if self.config.ping_response_accept else 0
            self.protocol.respond_to_ping(index, answer)
            
            if answer == 1:
                # If we accepted, we should expect a handshake
                self.state = "accepted_ping"
                self._ensure_ephemeral_keys()
                
                # Set a timer to send our handshake if we don't receive one
                timer = threading.Timer(
                    self.config.handshake_timeout,
                    self._check_handshake_received
                )
                timer.daemon = True
                timer.start()
                self.timer_tasks.append(timer)
                self.last_action_time = time.time()
                
        except Exception as e:
            self._log_error(f"Failed to respond to ping: {e}")
    
    def _check_handshake_received(self):
        """Check if we've received a handshake after accepting a ping."""
        if not self.active or self.state != "accepted_ping":
            return
            
        # If we've waited long enough and haven't received a handshake, initiate one
        if time.time() - self.last_action_time >= self.config.handshake_timeout:
            self._log_warning("No handshake received after accepting ping, initiating handshake")
            self._start_handshake_sequence()
    
    def _start_handshake_sequence(self):
        """Start the handshake sequence."""
        if self.handshake_attempts >= self.config.handshake_retry_count:
            self._log_warning(f"Maximum handshake attempts ({self.config.handshake_retry_count}) reached")
            self.state = "idle"
            return
            
        self.state = "handshake_sent"
        self.handshake_attempts += 1
        
        self._log_info(f"Sending handshake (attempt {self.handshake_attempts}/{self.config.handshake_retry_count})")
        try:
            self.protocol.send_handshake()
            self.last_action_time = time.time()
            
            # Schedule handshake retry check
            timer = threading.Timer(
                self.config.handshake_retry_delay,
                self._check_handshake_response
            )
            timer.daemon = True
            timer.start()
            self.timer_tasks.append(timer)
            
        except Exception as e:
            self._log_error(f"Failed to send handshake: {e}")
    
    def _check_handshake_response(self):
        """Check if we've completed the key exchange, retry handshake if not."""
        if not self.active or self.state != "handshake_sent":
            return
            
        # If we've waited long enough for a response, retry
        if time.time() - self.last_action_time >= self.config.handshake_timeout:
            self._log_warning("No handshake response received, retrying")
            self._start_handshake_sequence()
    
    def _start_message_sequence(self):
        """Start the automated message sending sequence."""
        if not self.config.auto_message_enabled:
            return
            
        self._log_info("Starting automated message sequence")
        
        # Add the default message if queue is empty
        if self.message_queue.empty():
            self.message_queue.put(self.config.message_content)
            
        # Start processing the queue
        self._process_message_queue()
    
    def _process_message_queue(self):
        """Process messages in the queue and send them."""
        if not self.active or self.state != "key_exchange_complete" or not self.config.auto_message_enabled:
            return
            
        if not self.message_queue.empty():
            message = self.message_queue.get()
            self._log_info(f"Sending queued message: {message}")
            
            try:
                self.protocol.send_encrypted_message(message)
                
                # Schedule next message send
                timer = threading.Timer(
                    self.config.message_interval,
                    self._process_message_queue
                )
                timer.daemon = True
                timer.start()
                self.timer_tasks.append(timer)
                
            except Exception as e:
                self._log_error(f"Failed to send queued message: {e}")
                # Put the message back in the queue
                self.message_queue.put(message)
    
    def _ensure_ephemeral_keys(self):
        """Ensure ephemeral keys are generated if needed."""
        if not hasattr(self.protocol, 'ephemeral_pubkey') or self.protocol.ephemeral_pubkey is None:
            self._log_info("Generating ephemeral keys")
            self.protocol.generate_ephemeral_keys()
            
    def _is_valid_message_index(self, index: int) -> bool:
        """
        Check if a message index is valid in the protocol's inbound_messages queue.
        
        Args:
            index: The index to check
            
        Returns:
            bool: True if the index is valid, False otherwise
        """
        if not hasattr(self.protocol, 'inbound_messages'):
            self._log_error("Protocol has no inbound_messages attribute")
            return False
            
        if index < 0 or index >= len(self.protocol.inbound_messages):
            self._log_error(f"Invalid message index: {index}")
            return False
            
        return True
    
    # Helper methods for logging
    def _log_info(self, message: str):
        print(f"{BLUE}[AUTO]{RESET} {message}")
        if hasattr(self, 'verbose_logging') and self.verbose_logging:
            state_info = f"(state={self.state})"
            if 'pinging' in self.state and hasattr(self, 'ping_attempts'):
                state_info += f", attempts={self.ping_attempts}/{self.config.ping_retry_count}"
            elif 'handshake' in self.state and hasattr(self, 'handshake_attempts'):
                state_info += f", attempts={self.handshake_attempts}/{self.config.handshake_retry_count}"
            print(f"{BLUE}[AUTO-DETAIL]{RESET} {state_info}")
        
    def _log_warning(self, message: str):
        print(f"{YELLOW}[AUTO-WARN]{RESET} {message}")
        if hasattr(self, 'verbose_logging') and self.verbose_logging:
            timer_info = f"Active timers: {len(self.timer_tasks)}"
            print(f"{YELLOW}[AUTO-WARN-DETAIL]{RESET} {timer_info}")
        
    def _log_error(self, message: str):
        print(f"{RED}[AUTO-ERROR]{RESET} {message}")
        if hasattr(self, 'verbose_logging') and self.verbose_logging:
            print(f"{RED}[AUTO-ERROR-DETAIL]{RESET} Current state: {self.state}, Active: {self.active}") 