import random
import os
import time
import threading
from typing import List, Dict, Any, Optional, Tuple

from crypto_utils import (
    generate_identity_keys,
    load_peer_identity_key,
    sign_data,
    verify_signature,
    get_ephemeral_keypair,
    compute_ecdh_shared_key,
    der_to_raw,
    raw_signature_to_der
)
from messages import (
    PingRequest, PingResponse, Handshake,
    compute_pfs_hash
)
import transmission
from encryption import (
    EncryptedMessage, MessageHeader,
    generate_iv, encrypt_message, decrypt_message
)
from auto_mode import AutoMode, AutoModeConfig

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


class IcingProtocol:
    def __init__(self):
        # Identity keys (each 512 bits when printed as hex of 64 bytes)
        self.identity_privkey, self.identity_pubkey = generate_identity_keys()

        # Peer identity for verifying ephemeral signatures
        self.peer_identity_pubkey_obj = None
        self.peer_identity_pubkey_bytes = None

        # Ephemeral keys (our side)
        self.ephemeral_privkey = None
        self.ephemeral_pubkey = None

        # Last computed shared secret (hex string)
        self.shared_secret = None

        # Derived HKDF key (hex string, 256 bits)
        self.hkdf_key = None
        
        # Negotiated cipher (0 = AES-256-GCM, 1 = ChaCha20-Poly1305)
        self.cipher_type = 0

        # For PFS: track per-peer session info (session number and last shared secret)
        self.pfs_history: Dict[bytes, Tuple[int, str]] = {}

        # Protocol flags
        self.state = {
            "ping_sent": False,
            "ping_received": False,
            "handshake_sent": False,
            "handshake_received": False,
            "key_exchange_complete": False
        }

        # Auto mode for automated protocol operation
        self.auto_mode = AutoMode(self)
        
        # Legacy auto-responder toggle (kept for backward compatibility)
        self.auto_responder = False

        # Active connections list
        self.connections = []

        # Inbound messages (each message is a dict with keys: type, raw, parsed, connection)
        self.inbound_messages: List[Dict[str, Any]] = []

        # Store the session nonce (17 bytes but only 129 bits are valid) from first sent or received PING
        self.session_nonce: bytes = None
        
        # Last used IV for encrypted messages
        self.last_iv: bytes = None

        self.local_port = random.randint(30000, 40000)
        self.server_listener = transmission.ServerListener(
            host="127.0.0.1",
            port=self.local_port,
            on_new_connection=self.on_new_connection,
            on_data_received=self.on_data_received
        )
        self.server_listener.start()

    # -------------------------------------------------------------------------
    # Transport callbacks
    # -------------------------------------------------------------------------

    def on_new_connection(self, conn: transmission.PeerConnection):
        print(f"{GREEN}[IcingProtocol]{RESET} New incoming connection.")
        self.connections.append(conn)
        
        # Notify auto mode
        self.auto_mode.handle_connection_established()

    def on_data_received(self, conn: transmission.PeerConnection, data: bytes):
        bits_count = len(data) * 8
        print(
            f"{GREEN}[RECV]{RESET} {bits_count} bits from peer: {data.hex()[:60]}{'...' if len(data.hex()) > 60 else ''}")

        # PING REQUEST (22 bytes)
        if len(data) == 22:
            ping_request = PingRequest.deserialize(data)
            if ping_request:
                self.state["ping_received"] = True
                
                # If received cipher is not supported, force to 0 (AES-256-GCM)
                if ping_request.cipher != 0 and ping_request.cipher != 1:
                    print(f"{YELLOW}[NOTICE]{RESET} Received PING with unsupported cipher ({ping_request.cipher}); forcing cipher to 0 in response.")
                    ping_request.cipher = 0
                
                # Store cipher type for future encrypted messages
                self.cipher_type = ping_request.cipher
                    
                # Store session nonce if not already set
                if self.session_nonce is None:
                    self.session_nonce = ping_request.session_nonce
                    print(f"{YELLOW}[NOTICE]{RESET} Stored session nonce from received PING.")
                
                index = len(self.inbound_messages)
                msg = {
                    "type": "PING_REQUEST",
                    "raw": data,
                    "parsed": ping_request,
                    "connection": conn
                }
                self.inbound_messages.append(msg)
                
                # Handle in auto mode (if active)
                self.auto_mode.handle_ping_received(index)
                
                # Legacy auto-responder (for backward compatibility)
                if self.auto_responder and not self.auto_mode.active:
                    timer = threading.Timer(2.0, self._auto_respond_ping, args=[index])
                    timer.daemon = True
                    timer.start()
                return

        # PING RESPONSE (10 bytes)
        elif len(data) == 10:
            ping_response = PingResponse.deserialize(data)
            if ping_response:
                # Store negotiated cipher type
                self.cipher_type = ping_response.cipher
                
                index = len(self.inbound_messages)
                msg = {
                    "type": "PING_RESPONSE",
                    "raw": data,
                    "parsed": ping_response,
                    "connection": conn
                }
                self.inbound_messages.append(msg)
                
                # Notify auto mode (if active)
                self.auto_mode.handle_ping_response_received(ping_response.answer == 1)
                return

        # HANDSHAKE message (168 bytes)
        elif len(data) == 168:
            handshake = Handshake.deserialize(data)
            if handshake:
                self.state["handshake_received"] = True
                index = len(self.inbound_messages)
                msg = {
                    "type": "HANDSHAKE",
                    "raw": data,
                    "parsed": handshake,
                    "connection": conn
                }
                self.inbound_messages.append(msg)
                
                # Notify auto mode (if active)
                self.auto_mode.handle_handshake_received(index)
                
                # Legacy auto-responder
                if self.auto_responder and not self.auto_mode.active:
                    timer = threading.Timer(2.0, self._auto_respond_handshake, args=[index])
                    timer.daemon = True
                    timer.start()
                return

        # Check if the message might be an encrypted message (e.g. header of 18 bytes at start)
        elif len(data) >= 18:
            # Try to parse header
            try:
                header = MessageHeader.unpack(data[:18])
                # If header unpacking is successful and data length matches header expectations
                expected_len = 18 + header.data_len + 16  # Header + payload + tag
                
                # Check if CRC is included
                has_crc = (header.connection_status & 0x01) != 0
                if has_crc:
                    expected_len += 4  # Add CRC32 length
                    
                if len(data) >= expected_len:
                    index = len(self.inbound_messages)
                    msg = {
                        "type": "ENCRYPTED_MESSAGE",
                        "raw": data,
                        "parsed": header,
                        "connection": conn
                    }
                    self.inbound_messages.append(msg)
                    print(f"{YELLOW}[NOTICE]{RESET} Stored inbound ENCRYPTED_MESSAGE at index={index}.")
                    
                    # Notify auto mode
                    self.auto_mode.handle_encrypted_received(index)
                    return
            except Exception as e:
                print(f"{RED}[ERROR]{RESET} Failed to parse message header: {e}")

        # Otherwise, unrecognized/malformed message.
        index = len(self.inbound_messages)
        msg = {
            "type": "UNKNOWN",
            "raw": data,
            "parsed": None,
            "connection": conn
        }
        self.inbound_messages.append(msg)
        print(f"{RED}[WARNING]{RESET} Unrecognized or malformed message stored at index={index}.")


    # -------------------------------------------------------------------------
    # HKDF Derivation
    # -------------------------------------------------------------------------

    def derive_hkdf(self):
        """
        Derives a 256-bit key using HKDF.
        Uses as input keying material (IKM) the shared secret from ECDH.
        The salt is computed as SHA256(session_nonce || pfs_param), where:
          - session_nonce is taken from self.session_nonce (17 bytes, 129 bits) or defaults to zeros.
          - pfs_param is taken from the first inbound HANDSHAKE's pfs_hash field (32 bytes) or zeros.
        """
        if not self.shared_secret:
            print(f"{RED}[ERROR]{RESET} No shared secret available; cannot derive HKDF key.")
            return

        # IKM: shared secret converted from hex to bytes.
        ikm = bytes.fromhex(self.shared_secret)
        # Use stored session_nonce if available; otherwise default to zeros.
        session_nonce = self.session_nonce if self.session_nonce is not None else (b"\x00" * 17)

        # Determine pfs_param from first HANDSHAKE message (if any)
        pfs_param = None
        for msg in self.inbound_messages:
            if msg["type"] == "HANDSHAKE":
                try:
                    handshake = msg["parsed"]
                    pfs_param = handshake.pfs_hash
                except Exception:
                    pfs_param = None
                break
        if pfs_param is None:
            print(f"{RED}[WARNING]{RESET} No HANDSHAKE found; using 32 zero bytes for pfs_param.")
            pfs_param = b"\x00" * 32  # 256-bit zeros

        # Ensure both are bytes
        if isinstance(session_nonce, str):
            session_nonce = session_nonce.encode()
        if isinstance(pfs_param, str):
            pfs_param = pfs_param.encode()

        from cryptography.hazmat.primitives import hashes
        from cryptography.hazmat.primitives.kdf.hkdf import HKDF
        hasher = hashes.Hash(hashes.SHA256())
        hasher.update(session_nonce + pfs_param)
        salt_value = hasher.finalize()

        hkdf = HKDF(
            algorithm=hashes.SHA256(),
            length=32,  # 256 bits
            salt=salt_value,
            info=b"",
        )
        derived_key = hkdf.derive(ikm)
        self.hkdf_key = derived_key.hex()
        self.state["key_exchange_complete"] = True
        print(f"{GREEN}[HKDF]{RESET} Derived HKDF key: {self.hkdf_key}")
        return True

    # -------------------------------------------------------------------------
    # Legacy Auto-responder helpers (kept for backward compatibility)
    # -------------------------------------------------------------------------

    def _auto_respond_ping(self, index: int):
        """
        Called by a Timer to respond automatically to a PING_REQUEST after 2s.
        """
        print(f"{BLUE}[AUTO]{RESET} Delayed responding to PING at index={index}")
        self.respond_to_ping(index, answer=1)  # Accept by default
        self.show_state()

    def _auto_respond_handshake(self, index: int):
        """
        Called by a Timer to handle inbound HANDSHAKE automatically.
        1) Generate ephemeral keys if not already set
        2) Compute ECDH with the inbound ephemeral pub (generate_ecdhe)
        3) Send our handshake back
        4) Show state
        """
        print(f"{BLUE}[AUTO]{RESET} Delayed ECDH process for HANDSHAKE at index={index}")

        # 1) Generate ephemeral keys if we don't have them
        if not self.ephemeral_privkey or not self.ephemeral_pubkey:
            self.generate_ephemeral_keys()

        # 2) Compute ECDH from inbound ephemeral pub
        self.generate_ecdhe(index)

        # 3) Send our handshake to the peer
        self.send_handshake()

        # 4) Show final state
        self.show_state()

    # -------------------------------------------------------------------------
    # Public Methods for Auto Mode Management
    # -------------------------------------------------------------------------
    
    def start_auto_mode(self):
        """Start the automatic protocol operation mode."""
        self.auto_mode.start()
        
    def stop_auto_mode(self):
        """Stop the automatic protocol operation mode."""
        self.auto_mode.stop()
        
    def configure_auto_mode(self, **kwargs):
        """
        Configure the automatic mode parameters.
        
        Args:
            **kwargs: Configuration parameters to set. Supported parameters:
                - ping_response_accept: bool, whether to accept incoming pings
                - ping_auto_initiate: bool, whether to initiate pings on connection
                - ping_retry_count: int, number of ping retries
                - ping_retry_delay: float, seconds between ping retries
                - ping_timeout: float, seconds to wait for ping response
                - preferred_cipher: int, preferred cipher (0=AES-GCM, 1=ChaCha20-Poly1305)
                - handshake_retry_count: int, number of handshake retries
                - handshake_retry_delay: float, seconds between handshake retries
                - handshake_timeout: float, seconds to wait for handshake
                - auto_message_enabled: bool, whether to auto-send messages
                - message_interval: float, seconds between auto messages
                - message_content: str, default message content
                - active_mode: bool, whether to actively initiate protocol
        """
        for key, value in kwargs.items():
            if hasattr(self.auto_mode.config, key):
                setattr(self.auto_mode.config, key, value)
                print(f"{BLUE}[CONFIG]{RESET} Set auto mode {key} = {value}")
            else:
                print(f"{RED}[ERROR]{RESET} Unknown auto mode configuration parameter: {key}")
                
    def get_auto_mode_config(self):
        """Return the current auto mode configuration."""
        return self.auto_mode.config
        
    def queue_auto_message(self, message: str):
        """
        Add a message to the auto-send queue.
        
        Args:
            message: Message text to send
        """
        self.auto_mode.queue_message(message)
        
    def toggle_auto_mode_logging(self):
        """
        Toggle detailed logging for auto mode.
        When enabled, will show more information about state transitions and decision making.
        """
        if not hasattr(self.auto_mode, 'verbose_logging'):
            self.auto_mode.verbose_logging = True
        else:
            self.auto_mode.verbose_logging = not self.auto_mode.verbose_logging
            
        status = "enabled" if self.auto_mode.verbose_logging else "disabled"
        print(f"{BLUE}[AUTO-LOG]{RESET} Detailed logging {status}")
        
    def debug_message(self, index: int):
        """
        Debug a message in the inbound message queue.
        Prints detailed information about the message.
        
        Args:
            index: The index of the message in the inbound_messages queue
        """
        if index < 0 or index >= len(self.inbound_messages):
            print(f"{RED}[ERROR]{RESET} Invalid message index {index}")
            return
            
        msg = self.inbound_messages[index]
        print(f"\n{YELLOW}=== Message Debug [{index}] ==={RESET}")
        print(f"Type: {msg['type']}")
        print(f"Length: {len(msg['raw'])} bytes = {len(msg['raw'])*8} bits")
        print(f"Raw data: {msg['raw'].hex()}")
        
        if msg['parsed'] is not None:
            print(f"\n{YELLOW}--- Parsed Data ---{RESET}")
            if msg['type'] == 'PING_REQUEST':
                ping = msg['parsed']
                print(f"Version: {ping.version}")
                print(f"Cipher: {ping.cipher} ({'AES-256-GCM' if ping.cipher == 0 else 'ChaCha20-Poly1305' if ping.cipher == 1 else 'Unknown'})")
                print(f"Session nonce: {ping.session_nonce.hex()}")
                print(f"CRC32: {ping.crc32:08x}")
                
            elif msg['type'] == 'PING_RESPONSE':
                resp = msg['parsed']
                print(f"Version: {resp.version}")
                print(f"Cipher: {resp.cipher} ({'AES-256-GCM' if resp.cipher == 0 else 'ChaCha20-Poly1305' if resp.cipher == 1 else 'Unknown'})")
                print(f"Answer: {resp.answer} ({'Accept' if resp.answer == 1 else 'Reject'})")
                print(f"CRC32: {resp.crc32:08x}")
                
            elif msg['type'] == 'HANDSHAKE':
                hs = msg['parsed']
                print(f"Ephemeral pubkey: {hs.ephemeral_pubkey.hex()[:16]}...")
                print(f"Ephemeral signature: {hs.ephemeral_signature.hex()[:16]}...")
                print(f"PFS hash: {hs.pfs_hash.hex()[:16]}...")
                print(f"Timestamp: {hs.timestamp}")
                print(f"CRC32: {hs.crc32:08x}")
                
            elif msg['type'] == 'ENCRYPTED_MESSAGE':
                header = msg['parsed']
                print(f"Flag: 0x{header.flag:04x}")
                print(f"Data length: {header.data_len} bytes")
                print(f"Retry: {header.retry}")
                print(f"Connection status: {header.connection_status} ({'CRC included' if header.connection_status & 0x01 else 'No CRC'})")
                print(f"IV: {header.iv.hex()}")
                
                # Calculate expected message size
                expected_len = 18 + header.data_len + 16  # Header + payload + tag
                if header.connection_status & 0x01:
                    expected_len += 4  # Add CRC
                    
                print(f"Expected total length: {expected_len} bytes")
                print(f"Actual length: {len(msg['raw'])} bytes")
                
                # If we have a key, try to decrypt
                if self.hkdf_key:
                    print("\nAttempting decryption...")
                    try:
                        key = bytes.fromhex(self.hkdf_key)
                        plaintext = decrypt_message(msg['raw'], key, self.cipher_type)
                        print(f"Decrypted: {plaintext.decode('utf-8')}")
                    except Exception as e:
                        print(f"Decryption failed: {e}")
        print()

    # -------------------------------------------------------------------------
    # Public Methods
    # -------------------------------------------------------------------------

    def connect_to_peer(self, port: int):
        conn = transmission.connect_to_peer("127.0.0.1", port, self.on_data_received)
        self.connections.append(conn)
        print(f"{GREEN}[IcingProtocol]{RESET} Outgoing connection to port {port} established.")
        
        # Notify auto mode
        self.auto_mode.handle_connection_established()

    def set_peer_identity(self, peer_pubkey_hex: str):
        pubkey_bytes = bytes.fromhex(peer_pubkey_hex)
        self.peer_identity_pubkey_obj = load_peer_identity_key(pubkey_bytes)
        self.peer_identity_pubkey_bytes = pubkey_bytes
        print(f"{GREEN}[IcingProtocol]{RESET} Stored peer identity pubkey (hex={peer_pubkey_hex[:16]}...).")

    def generate_ephemeral_keys(self):
        self.ephemeral_privkey, self.ephemeral_pubkey = get_ephemeral_keypair()
        print(f"{GREEN}[IcingProtocol]{RESET} Generated ephemeral key pair: pubkey={self.ephemeral_pubkey.hex()[:16]}...")

    # Send PING (session discovery and cipher negotiation)
    def send_ping_request(self, cipher_type=0):
        """
        Send a ping request to the first connected peer.
        
        Args:
            cipher_type: Preferred cipher type (0 = AES-256-GCM, 1 = ChaCha20-Poly1305)
        """
        if not self.connections:
            print(f"{RED}[ERROR]{RESET} No active connections.")
            return False
        
        # Validate cipher type
        if cipher_type not in (0, 1):
            print(f"{YELLOW}[WARNING]{RESET} Invalid cipher type {cipher_type}, defaulting to AES-256-GCM (0)")
            cipher_type = 0
            
        # Create ping request with specified cipher
        ping_request = PingRequest(version=0, cipher=cipher_type)
        
        # Store session nonce if not already set
        if self.session_nonce is None:
            self.session_nonce = ping_request.session_nonce
            print(f"{YELLOW}[NOTICE]{RESET} Stored session nonce from sent PING.")
            
        # Serialize and send
        pkt = ping_request.serialize()
        self._send_packet(self.connections[0], pkt, "PING_REQUEST")
        self.state["ping_sent"] = True
        return True

    def send_handshake(self):
        """
        Build and send handshake:
          - ephemeral_pubkey (64 bytes, raw x||y)
          - ephemeral_signature (64 bytes, raw r||s)
          - pfs_hash (32 bytes)
          - timestamp (32 bits)
          - CRC (32 bits)
        """
        if not self.connections:
            print(f"{RED}[ERROR]{RESET} No active connections.")
            return False
        if not self.ephemeral_privkey or not self.ephemeral_pubkey:
            print(f"{RED}[ERROR]{RESET} Ephemeral keys not generated.")
            return False
        if self.peer_identity_pubkey_bytes is None:
            print(f"{RED}[ERROR]{RESET} Peer identity not set; needed for PFS tracking.")
            return False

        # 1) Sign ephemeral_pubkey using identity key
        sig_der = sign_data(self.identity_privkey, self.ephemeral_pubkey)
        # Convert DER signature to raw r||s format (64 bytes)
        raw_signature = der_to_raw(sig_der)

        # 2) Compute PFS hash
        session_number, last_secret_hex = self.pfs_history.get(self.peer_identity_pubkey_bytes, (-1, ""))
        pfs = compute_pfs_hash(session_number, last_secret_hex)

        # 3) Create handshake object
        handshake = Handshake(
            ephemeral_pubkey=self.ephemeral_pubkey,
            ephemeral_signature=raw_signature,
            pfs_hash=pfs
        )

        # 4) Serialize and send
        pkt = handshake.serialize()
        self._send_packet(self.connections[0], pkt, "HANDSHAKE")
        self.state["handshake_sent"] = True
        return True

    def enable_auto_responder(self, enable: bool):
        """
        Legacy method for enabling/disabling auto responder.
        For new code, use start_auto_mode() and stop_auto_mode() instead.
        """
        self.auto_responder = enable
        print(f"{YELLOW}[LEGACY]{RESET} Auto responder set to {enable}. Consider using auto_mode instead.")

    # -------------------------------------------------------------------------
    # Manual Responses
    # -------------------------------------------------------------------------

    def respond_to_ping(self, index: int, answer: int):
        """
        Respond to a ping request with the specified answer (0 = no, 1 = yes).
        If answer is 1, we accept the connection and use the cipher specified in the request.
        """
        if index < 0 or index >= len(self.inbound_messages):
            print(f"{RED}[ERROR]{RESET} Invalid index {index}.")
            return False
        msg = self.inbound_messages[index]
        if msg["type"] != "PING_REQUEST":
            print(f"{RED}[ERROR]{RESET} inbound_messages[{index}] is not a PING_REQUEST.")
            return False

        ping_request = msg["parsed"]
        version = ping_request.version
        cipher = ping_request.cipher
        
        # Force cipher to 0 or 1 (only AES-256-GCM and ChaCha20-Poly1305 are supported)
        if cipher != 0 and cipher != 1:
            print(f"{YELLOW}[NOTICE]{RESET} Received PING with unsupported cipher ({cipher}); forcing cipher to 0 in response.")
            cipher = 0
            
        # Store the negotiated cipher type if we're accepting
        if answer == 1:
            self.cipher_type = cipher

        conn = msg["connection"]
        # Create ping response
        ping_response = PingResponse(version, cipher, answer)
        resp = ping_response.serialize()
        self._send_packet(conn, resp, "PING_RESPONSE")
        print(f"{BLUE}[MANUAL]{RESET} Responded to ping with answer={answer}.")
        return True

    def generate_ecdhe(self, index: int):
        """
        Process a handshake message:
        1. Verify the ephemeral signature
        2. Compute the ECDH shared secret
        3. Update PFS history
        """
        if index < 0 or index >= len(self.inbound_messages):
            print(f"{RED}[ERROR]{RESET} Invalid index {index}.")
            return False
        msg = self.inbound_messages[index]
        if msg["type"] != "HANDSHAKE":
            print(f"{RED}[ERROR]{RESET} inbound_messages[{index}] is not a HANDSHAKE.")
            return False

        handshake = msg["parsed"]
        
        # Convert raw signature to DER for verification
        raw_sig = handshake.ephemeral_signature
        sig_der = raw_signature_to_der(raw_sig)

        # Verify signature
        ok = verify_signature(self.peer_identity_pubkey_obj, sig_der, handshake.ephemeral_pubkey)
        if not ok:
            print(f"{RED}[ERROR]{RESET} Ephemeral signature invalid.")
            return False
        print(f"{GREEN}[OK]{RESET} Ephemeral signature verified.")

        # Check if we have ephemeral keys
        if not self.ephemeral_privkey:
            print(f"{YELLOW}[WARN]{RESET} No ephemeral_privkey available, cannot compute shared secret.")
            return False
            
        # Compute ECDH shared secret
        shared = compute_ecdh_shared_key(self.ephemeral_privkey, handshake.ephemeral_pubkey)
        self.shared_secret = shared.hex()
        print(f"{GREEN}[OK]{RESET} Computed ECDH shared key = {self.shared_secret}")

        # Update PFS history
        old_session, _ = self.pfs_history.get(self.peer_identity_pubkey_bytes, (-1, ""))
        new_session = 1 if old_session < 0 else old_session + 1
        self.pfs_history[self.peer_identity_pubkey_bytes] = (new_session, self.shared_secret)
        return True

    # -------------------------------------------------------------------------
    # Utility
    # -------------------------------------------------------------------------

    def _send_packet(self, conn: transmission.PeerConnection, data: bytes, label: str):
        bits_count = len(data) * 8
        print(f"{BLUE}[SEND]{RESET} {label} -> {bits_count} bits: {data.hex()[:60]}{'...' if len(data.hex())>60 else ''}")
        conn.send(data)

    def show_state(self):
        print(f"\n{YELLOW}=== Global State ==={RESET}")
        print(f"Listening Port: {self.local_port}")
        print(f"Identity PubKey: 512 bits => {self.identity_pubkey.hex()[:16]}...")

        if self.peer_identity_pubkey_bytes:
            print(f"Peer Identity PubKey: 512 bits => {self.peer_identity_pubkey_bytes.hex()[:16]}...")
        else:
            print("Peer Identity PubKey: [None]")

        print("\nEphemeral Keys:")
        if self.ephemeral_pubkey:
            print(f"  ephemeral_pubkey: 512 bits => {self.ephemeral_pubkey.hex()[:16]}...")
        else:
            print("  ephemeral_pubkey: [None]")

        print(f"\nShared Secret: {self.shared_secret if self.shared_secret else '[None]'}")

        if self.hkdf_key:
            print(f"HKDF Derived Key: {self.hkdf_key} (size: {len(self.hkdf_key)*8} bits)")
        else:
            print("HKDF Derived Key: [None]")
            
        print(f"Negotiated Cipher: {'AES-256-GCM' if self.cipher_type == 0 else 'ChaCha20-Poly1305'} (code: {self.cipher_type})")
        
        if self.session_nonce:
            print(f"Session Nonce: {self.session_nonce.hex()} (129 bits)")
        else:
            print("Session Nonce: [None]")
            
        if self.last_iv:
            print(f"Last IV: {self.last_iv.hex()} (96 bits)")
        else:
            print("Last IV: [None]")
            
        print("\nProtocol Flags:")
        for k, v in self.state.items():
            print(f"  {k}: {v}")

        print("\nAuto Mode Active:", self.auto_mode.active)
        print("Auto Mode State:", self.auto_mode.state)
        print("Legacy Auto Responder:", self.auto_responder)

        print("\nActive Connections:")
        for i, c in enumerate(self.connections):
            print(f"  [{i}] Alive={c.alive}")

        print("\nInbound Message Queue:")
        for i, m in enumerate(self.inbound_messages):
            print(f"  [{i}] type={m['type']} len={len(m['raw'])} bytes => {len(m['raw']) * 8} bits")
        print()

    def stop(self):
        """Stop the protocol and clean up resources."""
        # Stop auto mode first
        self.auto_mode.stop()
        
        # Stop server listener
        self.server_listener.stop()
        
        # Close all connections
        for c in self.connections:
            c.close()
        self.connections.clear()
        self.inbound_messages.clear()
        print(f"{RED}[STOP]{RESET} Protocol stopped.")

    # -------------------------------------------------------------------------
    # Encrypted Messaging
    # -------------------------------------------------------------------------

    def send_encrypted_message(self, plaintext: str):
        """
        Encrypts and sends a message using the derived HKDF key and negotiated cipher.
        The message format is:
        - Header (18 bytes): flag, data_len, retry, connection_status, IV
        - Payload: variable length encrypted data
        - Footer: Authentication tag (16 bytes) + optional CRC32 (4 bytes)
        """
        if not self.connections:
            print(f"{RED}[ERROR]{RESET} No active connections.")
            return False
        if not self.hkdf_key:
            print(f"{RED}[ERROR]{RESET} No HKDF key derived. Cannot encrypt message.")
            return False
            
        # Get the encryption key
        key = bytes.fromhex(self.hkdf_key)
        
        # Convert plaintext to bytes
        plaintext_bytes = plaintext.encode('utf-8')
        
        # Generate or increment the IV
        if self.last_iv is None:
            # First message, generate random IV
            iv = generate_iv(initial=True)
        else:
            # Subsequent message, increment previous IV
            iv = generate_iv(initial=False, previous_iv=self.last_iv)
            
        # Store the new IV
        self.last_iv = iv
        
        # Create encrypted message (connection_status 0 = no CRC)
        encrypted = encrypt_message(
            plaintext=plaintext_bytes,
            key=key,
            flag=0xBEEF,  # Default flag
            retry=0,
            connection_status=0,  # No CRC
            iv=iv,
            cipher_type=self.cipher_type
        )
        
        # Send the encrypted message
        self._send_packet(self.connections[0], encrypted, "ENCRYPTED_MESSAGE")
        print(f"{GREEN}[SEND_ENCRYPTED]{RESET} Encrypted message sent.")
        return True

    def decrypt_received_message(self, index: int):
        """
        Decrypt a received encrypted message using the HKDF key and negotiated cipher.
        """
        if index < 0 or index >= len(self.inbound_messages):
            print(f"{RED}[ERROR]{RESET} Invalid message index.")
            return None
            
        msg = self.inbound_messages[index]
        if msg["type"] != "ENCRYPTED_MESSAGE":
            print(f"{RED}[ERROR]{RESET} Message at index {index} is not an ENCRYPTED_MESSAGE.")
            return None
            
        # Get the encrypted message
        encrypted = msg["raw"]
        
        if not self.hkdf_key:
            print(f"{RED}[ERROR]{RESET} No HKDF key derived. Cannot decrypt message.")
            return None
            
        # Get the encryption key
        key = bytes.fromhex(self.hkdf_key)
        
        try:
            # Decrypt the message
            plaintext = decrypt_message(encrypted, key, self.cipher_type)
            
            # Convert to string
            plaintext_str = plaintext.decode('utf-8')
            
            # Update last IV from the header
            header = MessageHeader.unpack(encrypted[:18])
            self.last_iv = header.iv
            
            print(f"{GREEN}[DECRYPTED]{RESET} Decrypted message: {plaintext_str}")
            return plaintext_str
        except Exception as e:
            print(f"{RED}[ERROR]{RESET} Decryption failed: {e}")
            return None
