"""Noise_XX handshake and encrypted framing.

Implements a simplified Noise_XX pattern using:
- X25519 for Diffie-Hellman key exchange
- HKDF-SHA256 for key derivation
- AESGCM for authenticated encryption

The XX pattern provides mutual authentication and forward secrecy:
  -> e                   (initiator sends ephemeral public key)
  <- e, ee, s, es       (responder sends ephemeral + encrypted static)
  -> s, se              (initiator sends encrypted static)

After the 3-message handshake, both sides have:
- Verified each other's static (long-term) public keys
- Derived shared symmetric keys with forward secrecy
- An encrypted channel for all subsequent communication
"""

from __future__ import annotations

import hashlib
import struct
from dataclasses import dataclass, field

from cryptography.hazmat.primitives.asymmetric.x25519 import (
    X25519PrivateKey,
    X25519PublicKey,
)
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes

from a2a_p2p.errors import CryptoError


_PROTOCOL_NAME = b"Noise_XX_25519_AESGCM_SHA256"
_EMPTY_KEY = b"\x00" * 32


def _dh(private: X25519PrivateKey, public: X25519PublicKey) -> bytes:
    return private.exchange(public)


def _pub_bytes(key: X25519PublicKey) -> bytes:
    from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat
    return key.public_bytes(Encoding.Raw, PublicFormat.Raw)


def _pub_from_bytes(data: bytes) -> X25519PublicKey:
    return X25519PublicKey.from_public_bytes(data)


def _hkdf2(ck: bytes, ikm: bytes) -> tuple[bytes, bytes]:
    """HKDF-SHA256 extract+expand → (new_ck, key)."""
    hkdf = HKDF(algorithm=hashes.SHA256(), length=64, salt=ck, info=b"")
    out = hkdf.derive(ikm)
    return out[:32], out[32:]


@dataclass
class _SymState:
    """Noise SymmetricState — tracks ck, h, and an optional cipher key."""
    ck: bytes = field(default_factory=lambda: hashlib.sha256(_PROTOCOL_NAME).digest())
    h: bytes = field(default_factory=lambda: hashlib.sha256(_PROTOCOL_NAME).digest())
    cipher_key: bytes = _EMPTY_KEY
    cipher_nonce: int = 0

    def mix_hash(self, data: bytes) -> None:
        self.h = hashlib.sha256(self.h + data).digest()

    def mix_key(self, ikm: bytes) -> None:
        """HKDF: derive new ck and cipher key from input key material."""
        self.ck, self.cipher_key = _hkdf2(self.ck, ikm)
        self.cipher_nonce = 0  # reset nonce on re-key

    def has_key(self) -> bool:
        return self.cipher_key != _EMPTY_KEY

    def _aead_encrypt(self, plaintext: bytes, ad: bytes) -> bytes:
        nonce = struct.pack(">4sQ", b"\x00" * 4, self.cipher_nonce)
        self.cipher_nonce += 1
        return AESGCM(self.cipher_key).encrypt(nonce, plaintext, ad)

    def _aead_decrypt(self, ciphertext: bytes, ad: bytes) -> bytes:
        nonce = struct.pack(">4sQ", b"\x00" * 4, self.cipher_nonce)
        self.cipher_nonce += 1
        try:
            return AESGCM(self.cipher_key).decrypt(nonce, ciphertext, ad)
        except Exception as e:
            raise CryptoError(f"Decryption failed: {e}") from e

    def encrypt_and_hash(self, plaintext: bytes) -> bytes:
        """Encrypt plaintext (if we have a key), mix ciphertext into h."""
        if self.has_key():
            ct = self._aead_encrypt(plaintext, ad=self.h)
        else:
            ct = plaintext  # no encryption yet (before first mix_key)
        self.mix_hash(ct)
        return ct

    def decrypt_and_hash(self, ciphertext: bytes) -> bytes:
        """Decrypt ciphertext (if we have a key), mix ciphertext into h."""
        if self.has_key():
            pt = self._aead_decrypt(ciphertext, ad=self.h)
        else:
            pt = ciphertext
        self.mix_hash(ciphertext)
        return pt

    def split(self) -> tuple[bytes, bytes]:
        """Derive two transport keys from the final chaining key."""
        k1, k2 = _hkdf2(self.ck, b"")
        return k1, k2


@dataclass
class _CipherState:
    """Transport cipher for one direction."""
    key: bytes
    nonce: int = 0

    def encrypt(self, plaintext: bytes) -> bytes:
        nonce = struct.pack(">4sQ", b"\x00" * 4, self.nonce)
        self.nonce += 1
        return AESGCM(self.key).encrypt(nonce, plaintext, b"")

    def decrypt(self, ciphertext: bytes) -> bytes:
        nonce = struct.pack(">4sQ", b"\x00" * 4, self.nonce)
        self.nonce += 1
        try:
            return AESGCM(self.key).decrypt(nonce, ciphertext, b"")
        except Exception as e:
            raise CryptoError(f"Decryption failed: {e}") from e


class NoiseSession:
    """An established Noise session with send/receive cipher states."""

    def __init__(
        self,
        send_cipher: _CipherState,
        recv_cipher: _CipherState,
        remote_static_pub: bytes,
        handshake_hash: bytes,
    ) -> None:
        self._send = send_cipher
        self._recv = recv_cipher
        self.remote_static_pub = remote_static_pub
        self.handshake_hash = handshake_hash

    def encrypt(self, plaintext: bytes) -> bytes:
        return self._send.encrypt(plaintext)

    def decrypt(self, ciphertext: bytes) -> bytes:
        return self._recv.decrypt(ciphertext)


class NoiseInitiator:
    """Initiator side of the Noise_XX handshake.

    Usage::

        initiator = NoiseInitiator(static_private)
        msg1 = initiator.write_message_1()
        msg3 = initiator.read_message_2(msg2)
        session = initiator.finalize()
    """

    def __init__(self, static_private: X25519PrivateKey) -> None:
        self._s = static_private
        self._e = X25519PrivateKey.generate()
        self._ss = _SymState()
        self._remote_s_pub: X25519PublicKey | None = None
        self._re_pub: X25519PublicKey | None = None

    def write_message_1(self) -> bytes:
        """-> e"""
        e_pub = _pub_bytes(self._e.public_key())
        self._ss.mix_hash(e_pub)
        return e_pub

    def read_message_2(self, msg: bytes) -> bytes:
        """<- e, ee, s, es  then  -> s, se.  Returns message 3."""
        if len(msg) < 80:
            raise CryptoError(f"Message 2 too short: {len(msg)}")

        # Parse: re_pub(32) + encrypted_s(48)
        re_pub_bytes = msg[:32]
        encrypted_rs = msg[32:80]

        self._re_pub = _pub_from_bytes(re_pub_bytes)
        self._ss.mix_hash(re_pub_bytes)

        # ee
        self._ss.mix_key(_dh(self._e, self._re_pub))
        # decrypt responder static
        rs_bytes = self._ss.decrypt_and_hash(encrypted_rs)
        self._remote_s_pub = _pub_from_bytes(rs_bytes)
        # es
        self._ss.mix_key(_dh(self._e, self._remote_s_pub))

        # Build message 3: -> s, se
        s_pub = _pub_bytes(self._s.public_key())
        encrypted_s = self._ss.encrypt_and_hash(s_pub)
        # se
        self._ss.mix_key(_dh(self._s, self._re_pub))

        return encrypted_s

    def finalize(self) -> NoiseSession:
        if self._remote_s_pub is None:
            raise CryptoError("Handshake not complete")
        k1, k2 = self._ss.split()
        return NoiseSession(
            send_cipher=_CipherState(key=k1),
            recv_cipher=_CipherState(key=k2),
            remote_static_pub=_pub_bytes(self._remote_s_pub),
            handshake_hash=self._ss.h,
        )


class NoiseResponder:
    """Responder side of the Noise_XX handshake.

    Usage::

        responder = NoiseResponder(static_private)
        msg2 = responder.read_message_1(msg1)
        responder.read_message_3(msg3)
        session = responder.finalize()
    """

    def __init__(self, static_private: X25519PrivateKey) -> None:
        self._s = static_private
        self._e = X25519PrivateKey.generate()
        self._ss = _SymState()
        self._remote_s_pub: X25519PublicKey | None = None
        self._re_pub: X25519PublicKey | None = None

    def read_message_1(self, msg: bytes) -> bytes:
        """<- e  then  -> e, ee, s, es.  Returns message 2."""
        if len(msg) < 32:
            raise CryptoError(f"Message 1 too short: {len(msg)}")

        re_pub_bytes = msg[:32]
        self._re_pub = _pub_from_bytes(re_pub_bytes)
        self._ss.mix_hash(re_pub_bytes)

        # Build message 2: -> e, ee, s, es
        e_pub = _pub_bytes(self._e.public_key())
        self._ss.mix_hash(e_pub)
        # ee
        self._ss.mix_key(_dh(self._e, self._re_pub))
        # encrypt our static
        s_pub = _pub_bytes(self._s.public_key())
        encrypted_s = self._ss.encrypt_and_hash(s_pub)
        # es
        self._ss.mix_key(_dh(self._s, self._re_pub))

        return e_pub + encrypted_s

    def read_message_3(self, msg: bytes) -> None:
        """<- s, se"""
        if len(msg) < 48:
            raise CryptoError(f"Message 3 too short: {len(msg)}")

        encrypted_is = msg[:48]
        is_bytes = self._ss.decrypt_and_hash(encrypted_is)
        self._remote_s_pub = _pub_from_bytes(is_bytes)
        # se
        self._ss.mix_key(_dh(self._e, self._remote_s_pub))

    def finalize(self) -> NoiseSession:
        if self._remote_s_pub is None:
            raise CryptoError("Handshake not complete")
        k1, k2 = self._ss.split()
        # Responder's directions are flipped
        return NoiseSession(
            send_cipher=_CipherState(key=k2),
            recv_cipher=_CipherState(key=k1),
            remote_static_pub=_pub_bytes(self._remote_s_pub),
            handshake_hash=self._ss.h,
        )


def generate_x25519_from_ed25519_seed(ed25519_private_key_bytes: bytes) -> X25519PrivateKey:
    """Derive an X25519 private key from Ed25519 private key seed."""
    hkdf = HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=b"a2a-p2p-x25519-derivation",
        info=b"ed25519-to-x25519",
    )
    return X25519PrivateKey.from_private_bytes(hkdf.derive(ed25519_private_key_bytes))


# ── Framing ─────────────────────────────────────────────────────────


def frame_encrypt(session: NoiseSession, plaintext: bytes) -> bytes:
    ct = session.encrypt(plaintext)
    return struct.pack(">I", len(ct)) + ct


def frame_decrypt_length(header: bytes) -> int:
    if len(header) < 4:
        raise CryptoError("Frame header too short")
    return struct.unpack(">I", header[:4])[0]


def frame_decrypt(session: NoiseSession, ciphertext: bytes) -> bytes:
    return session.decrypt(ciphertext)
