# noise_xk/session.py

import socket
import logging
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

# Configure root logger for debug output
logging.basicConfig(level=logging.DEBUG, format="%(message)s")

class NoiseXKSession:
    @staticmethod
    def generate_keypair() -> KeyPair:
        """
        Generate a static X25519 KeyPair.
        Returns:
            KeyPair object with .private and .public attributes.
        """
        return X25519DH().generate_keypair()

    def __init__(self, local_kp: KeyPair, peer_pubkey: PublicKey):
        """
        Initialize with our KeyPair and the peer's PublicKey.
        """
        self.local_kp: KeyPair = local_kp
        self.peer_pubkey: PublicKey = peer_pubkey

        # Build the Noise handshake state (X25519 DH, ChaChaPoly cipher, SHA256 hash)
        cipher = ChaChaPolyCipher()
        dh = X25519DH()
        hshash = SHA256Hash()
        symmetric = SymmetricState(CipherState(cipher), hshash)
        self._hs = HandshakeState(symmetric, dh)

        self._send_cs = None  # type: CipherState
        self._recv_cs = None

    def handshake(self, sock: socket.socket, initiator: bool) -> None:
        """
        Perform the XK handshake over the socket. Branches on initiator/responder
        so that each side reads or writes in the correct message order.
        On completion, self._send_cs and self._recv_cs hold the two CipherStates.
        """

        logging.debug(f"[handshake] start (initiator={initiator})")
        # initialize with our KeyPair and their PublicKey
        if initiator:
            # initiator knows peer’s static out-of-band
            self._hs.initialize(
                XKHandshakePattern(),
                True,
                b'',
                s=self.local_kp,
                rs=self.peer_pubkey
            )
        else:
            logging.debug("[handshake] responder initializing without rs")
            # responder must NOT supply rs here
            self._hs.initialize(
                XKHandshakePattern(),
                False,
                b'',
                s=self.local_kp
            )

        cs_pair = None
        if initiator:
            # 1) -> e
            buf1 = bytearray()
            cs_pair = self._hs.write_message(b'', buf1)
            logging.debug(f"[-> e] {buf1.hex()}")
            self._send_all(sock, buf1)

            # 2) <- e, es, s, ss
            msg2 = self._recv_all(sock)
            logging.debug(f"[<- msg2] {msg2.hex()}")
            self._hs.read_message(msg2, bytearray())

            # 3) -> se (final)
            buf3 = bytearray()
            cs_pair = self._hs.write_message(b'', buf3)
            logging.debug(f"[-> se] {buf3.hex()}")
            self._send_all(sock, buf3)
        else:
            # 1) <- e
            msg1 = self._recv_all(sock)
            logging.debug(f"[<- e] {msg1.hex()}")
            self._hs.read_message(msg1, bytearray())

            # 2) -> e, es, s, ss
            buf2 = bytearray()
            cs_pair = self._hs.write_message(b'', buf2)
            logging.debug(f"[-> msg2] {buf2.hex()}")
            self._send_all(sock, buf2)

            # 3) <- se (final)
            msg3 = self._recv_all(sock)
            logging.debug(f"[<- se] {msg3.hex()}")
            cs_pair = self._hs.read_message(msg3, bytearray())

        # on the final step, we must get exactly two CipherStates
        if not cs_pair or len(cs_pair) != 2:
            raise RuntimeError("Handshake did not complete properly")
        cs0, cs1 = cs_pair
        # the library returns (cs_encrypt_for_initiator, cs_decrypt_for_initiator)
        if initiator:
            # initiator: cs0 encrypts, cs1 decrypts
            self._send_cs, self._recv_cs = cs0, cs1
        else:
            # responder must swap
            self._send_cs, self._recv_cs = cs1, cs0

        # dump the raw symmetric keys & nonces (if available)
        self._dump_cipherstate("HANDSHAKE→ SEND", self._send_cs)
        self._dump_cipherstate("HANDSHAKE→ RECV", self._recv_cs)

    def send(self, sock: socket.socket, plaintext: bytes) -> None:
        """
        Encrypt and send a message.
        """
        if self._send_cs is None:
            raise RuntimeError("Handshake not complete")
        ct = self._send_cs.encrypt_with_ad(b'', plaintext)
        logging.debug(f"[ENCRYPT] {ct.hex()}")
        self._dump_cipherstate("SEND→ after encrypt", self._send_cs)
        self._send_all(sock, ct)

    def receive(self, sock: socket.socket) -> bytes:
        """
        Receive and decrypt a message.
        """
        if self._recv_cs is None:
            raise RuntimeError("Handshake not complete")
        ct = self._recv_all(sock)
        logging.debug(f"[CIPHERTEXT] {ct.hex()}")
        self._dump_cipherstate("RECV→ before decrypt", self._recv_cs)
        pt = self._recv_cs.decrypt_with_ad(b'', ct)
        logging.debug(f"[DECRYPT] {pt!r}")
        return pt

    def _send_all(self, sock: socket.socket, data: bytes) -> None:
        # Length-prefix (2 bytes big-endian) + data
        length = len(data).to_bytes(2, 'big')
        sock.sendall(length + data)

    def _recv_all(self, sock: socket.socket) -> bytes:
        # Read 2-byte length prefix, then the payload
        hdr = self._read_exact(sock, 2)
        length = int.from_bytes(hdr, 'big')
        return self._read_exact(sock, length)

    @staticmethod
    def _read_exact(sock: socket.socket, n: int) -> bytes:
        buf = bytearray()
        while len(buf) < n:
            chunk = sock.recv(n - len(buf))
            if not chunk:
                raise ConnectionError("Socket closed during read")
            buf.extend(chunk)
        return bytes(buf)

    def _dump_cipherstate(self, label: str, cs: CipherState) -> None:
        """
        Print the symmetric key (cs._k) and nonce counter (cs._n) for inspection.
        """
        key = cs._key
        nonce = getattr(cs, "_n", None)
        if isinstance(key, (bytes, bytearray)):
            key_hex = key.hex()
        else:
            key_hex = repr(key)
        logging.debug(f"[{label}] key={key_hex}")
