"""Wrapper for Noise XK handshake over GSM simulator"""

import struct
from dissononce.processing.impl.handshakestate import HandshakeState
from dissononce.processing.impl.symmetricstate import SymmetricState
from dissononce.processing.impl.cipherstate import CipherState
from dissononce.processing.handshakepatterns.interactive.XK import XKHandshakePattern
from dissononce.cipher.chachapoly import ChaChaPolyCipher
from dissononce.dh.x25519.x25519 import X25519DH
from dissononce.dh.keypair import KeyPair
from dissononce.dh.x25519.public import PublicKey
from dissononce.hash.sha256 import SHA256Hash

class NoiseXKWrapper:
    """Wrapper for Noise XK that works over message-passing instead of direct sockets"""
    
    def __init__(self, keypair, peer_pubkey, debug_callback=None):
        self.keypair = keypair
        self.peer_pubkey = peer_pubkey
        self.debug = debug_callback or print
        
        # Build handshake state
        cipher = ChaChaPolyCipher()
        dh = X25519DH()
        hshash = SHA256Hash()
        symmetric = SymmetricState(CipherState(cipher), hshash)
        self._hs = HandshakeState(symmetric, dh)
        
        self._send_cs = None
        self._recv_cs = None
        self.handshake_complete = False
        self.is_initiator = None  # Track initiator status
        
        # Message buffers
        self.outgoing_messages = []
        self.incoming_messages = []
        
    def start_handshake(self, initiator):
        """Start the handshake process"""
        self.debug(f"Starting Noise XK handshake as {'initiator' if initiator else 'responder'}")
        self.is_initiator = initiator  # Store initiator status
        
        if initiator:
            # Initiator knows peer's static out-of-band
            self._hs.initialize(
                XKHandshakePattern(),
                True,
                b'',
                s=self.keypair,
                rs=self.peer_pubkey
            )
            # Generate first message
            buf = bytearray()
            self._hs.write_message(b'', buf)
            self.outgoing_messages.append(bytes(buf))
            self.debug(f"Generated handshake message 1: {len(buf)} bytes")
        else:
            # Responder doesn't know peer's static yet
            self._hs.initialize(
                XKHandshakePattern(),
                False,
                b'',
                s=self.keypair
            )
            self.debug("Responder initialized, waiting for first message")
    
    def process_handshake_message(self, data):
        """Process incoming handshake message and generate response if needed"""
        self.debug(f"Processing handshake message: {len(data)} bytes")
        
        try:
            # Read the message
            payload = bytearray()
            cs_pair = self._hs.read_message(data, payload)
            
            # Check if we need to send a response
            if not cs_pair:
                # More messages needed
                buf = bytearray()
                cs_pair = self._hs.write_message(b'', buf)
                self.outgoing_messages.append(bytes(buf))
                self.debug(f"Generated handshake response: {len(buf)} bytes")
                
                # Check if handshake completed after writing (for initiator)
                if cs_pair:
                    self._complete_handshake(cs_pair)
            else:
                # Handshake complete after reading (for responder)
                self._complete_handshake(cs_pair)
                
        except Exception as e:
            self.debug(f"Handshake error: {e}")
            raise
    
    def get_next_handshake_message(self):
        """Get next outgoing handshake message"""
        if self.outgoing_messages:
            return self.outgoing_messages.pop(0)
        return None
    
    def encrypt(self, plaintext):
        """Encrypt a message"""
        if not self.handshake_complete:
            raise RuntimeError("Handshake not complete")
        return self._send_cs.encrypt_with_ad(b'', plaintext)
    
    def decrypt(self, ciphertext):
        """Decrypt a message"""
        if not self.handshake_complete:
            raise RuntimeError("Handshake not complete")
        return self._recv_cs.decrypt_with_ad(b'', ciphertext)
    
    def _complete_handshake(self, cs_pair):
        """Complete the handshake with the given cipher states"""
        self.debug("Handshake complete, setting up cipher states")
        cs0, cs1 = cs_pair
        
        # Use stored initiator status
        if self.is_initiator:
            self._send_cs, self._recv_cs = cs0, cs1
            self.debug("Set up cipher states as initiator")
        else:
            self._send_cs, self._recv_cs = cs1, cs0
            self.debug("Set up cipher states as responder")
        
        self.handshake_complete = True
        self.debug("Cipher states established")