"""Noise_XX handshake and encrypted framing.

X25519 DH + HKDF-SHA256 + AESGCM. Provides mutual authentication and
forward secrecy in a 3-message handshake — no CA required.

    -> e                   (initiator ephemeral)
    <- e, ee, s, es       (responder ephemeral + encrypted static)
    -> s, se              (initiator encrypted static)
"""

from __future__ import annotations

import hashlib
import struct
from dataclasses import dataclass, field

from cryptography.hazmat.primitives import hashes
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 p2p.errors import CryptoError

_PROTO = b"Noise_XX_25519_AESGCM_SHA256"
_ZERO = b"\x00" * 32


def _dh(priv: X25519PrivateKey, pub: X25519PublicKey) -> bytes:
    return priv.exchange(pub)


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


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


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


# ── Symmetric state ────────────────────────────────────────────────


@dataclass
class _Sym:
    ck: bytes = field(default_factory=lambda: hashlib.sha256(_PROTO).digest())
    h: bytes = field(default_factory=lambda: hashlib.sha256(_PROTO).digest())
    key: bytes = _ZERO
    n: int = 0

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

    def mix_key(self, ikm: bytes) -> None:
        self.ck, self.key = _hkdf2(self.ck, ikm)
        self.n = 0

    def _aead(self, encrypt: bool, data: bytes) -> bytes:
        if self.n >= (2**63 - 1):
            raise CryptoError("Nonce exhausted — renegotiate session")
        nonce = struct.pack(">4sQ", b"\x00" * 4, self.n)
        self.n += 1
        aes = AESGCM(self.key)
        if encrypt:
            return aes.encrypt(nonce, data, self.h)
        try:
            return aes.decrypt(nonce, data, self.h)
        except Exception as e:
            raise CryptoError(f"Decryption failed: {e}") from e

    def encrypt_and_hash(self, pt: bytes) -> bytes:
        ct = self._aead(True, pt) if self.key != _ZERO else pt
        self.mix_hash(ct)
        return ct

    def decrypt_and_hash(self, ct: bytes) -> bytes:
        pt = self._aead(False, ct) if self.key != _ZERO else ct
        self.mix_hash(ct)
        return pt

    def split(self) -> tuple[bytes, bytes]:
        # Per Noise spec section 5.2: Split() derives from chaining key alone.
        # Empty IKM is intentional — not a bug.
        return _hkdf2(self.ck, b"")


# ── Transport cipher ───────────────────────────────────────────────


@dataclass
class _Cipher:
    key: bytes
    n: int = 0

    def encrypt(self, pt: bytes) -> bytes:
        if self.n >= (2**63 - 1):
            raise CryptoError("Nonce exhausted — renegotiate session")
        nonce = struct.pack(">4sQ", b"\x00" * 4, self.n)
        self.n += 1
        return AESGCM(self.key).encrypt(nonce, pt, b"")

    def decrypt(self, ct: bytes) -> bytes:
        if self.n >= (2**63 - 1):
            raise CryptoError("Nonce exhausted — renegotiate session")
        nonce = struct.pack(">4sQ", b"\x00" * 4, self.n)
        self.n += 1
        try:
            return AESGCM(self.key).decrypt(nonce, ct, b"")
        except Exception as e:
            raise CryptoError(f"Decryption failed: {e}") from e


# ── Session ─────────────────────────────────────────────────────────


class NoiseSession:
    """Established encrypted session after Noise_XX handshake."""

    def __init__(self, send: _Cipher, recv: _Cipher,
                 remote_static: bytes, h: bytes) -> None:
        self._send = send
        self._recv = recv
        self.remote_static_pub = remote_static
        self.handshake_hash = h

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

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


# ── Initiator ───────────────────────────────────────────────────────


class NoiseInitiator:
    """Initiator side of Noise_XX.

    ::
        init = NoiseInitiator(static_priv)
        msg1 = init.write_message_1()
        msg3 = init.read_message_2(msg2)
        session = init.finalize()
    """

    def __init__(self, static: X25519PrivateKey) -> None:
        self._s = static
        self._e = X25519PrivateKey.generate()
        self._ss = _Sym()
        self._rs: X25519PublicKey | None = None
        self._re: X25519PublicKey | None = None

    def write_message_1(self) -> bytes:
        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:
        if len(msg) < 80:
            raise CryptoError(f"msg2 too short: {len(msg)}")
        self._re = _pub_from(msg[:32])
        self._ss.mix_hash(msg[:32])
        self._ss.mix_key(_dh(self._e, self._re))                   # ee
        rs = self._ss.decrypt_and_hash(msg[32:80])
        self._rs = _pub_from(rs)
        self._ss.mix_key(_dh(self._e, self._rs))                   # es
        enc_s = self._ss.encrypt_and_hash(_pub_bytes(self._s.public_key()))
        self._ss.mix_key(_dh(self._s, self._re))                   # se
        return enc_s

    def finalize(self) -> NoiseSession:
        if not self._rs:
            raise CryptoError("Handshake not complete")
        k1, k2 = self._ss.split()
        return NoiseSession(_Cipher(k1), _Cipher(k2),
                            _pub_bytes(self._rs), self._ss.h)


# ── Responder ───────────────────────────────────────────────────────


class NoiseResponder:
    """Responder side of Noise_XX.

    ::
        resp = NoiseResponder(static_priv)
        msg2 = resp.read_message_1(msg1)
        resp.read_message_3(msg3)
        session = resp.finalize()
    """

    def __init__(self, static: X25519PrivateKey) -> None:
        self._s = static
        self._e = X25519PrivateKey.generate()
        self._ss = _Sym()
        self._rs: X25519PublicKey | None = None
        self._re: X25519PublicKey | None = None

    def read_message_1(self, msg: bytes) -> bytes:
        if len(msg) < 32:
            raise CryptoError(f"msg1 too short: {len(msg)}")
        self._re = _pub_from(msg[:32])
        self._ss.mix_hash(msg[:32])
        e_pub = _pub_bytes(self._e.public_key())
        self._ss.mix_hash(e_pub)
        self._ss.mix_key(_dh(self._e, self._re))                   # ee
        enc_s = self._ss.encrypt_and_hash(_pub_bytes(self._s.public_key()))
        self._ss.mix_key(_dh(self._s, self._re))                   # es
        return e_pub + enc_s

    def read_message_3(self, msg: bytes) -> None:
        if len(msg) < 48:
            raise CryptoError(f"msg3 too short: {len(msg)}")
        is_bytes = self._ss.decrypt_and_hash(msg[:48])
        self._rs = _pub_from(is_bytes)
        self._ss.mix_key(_dh(self._e, self._rs))                   # se

    def finalize(self) -> NoiseSession:
        if not self._rs:
            raise CryptoError("Handshake not complete")
        k1, k2 = self._ss.split()
        return NoiseSession(_Cipher(k2), _Cipher(k1),
                            _pub_bytes(self._rs), self._ss.h)


# ── Helpers ─────────────────────────────────────────────────────────


def derive_x25519_from_ed25519(ed_seed: bytes) -> X25519PrivateKey:
    """Derive X25519 transport key from Ed25519 identity seed."""
    seed = HKDF(
        algorithm=hashes.SHA256(), length=32,
        salt=b"openclaw-p2p-x25519", info=b"ed25519-to-x25519",
    ).derive(ed_seed)
    return X25519PrivateKey.from_private_bytes(seed)


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


def frame_decrypt_length(header: bytes) -> int:
    return struct.unpack(">I", header[:4])[0]


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