from __future__ import annotations

import binascii
from typing import Callable

from .._hazmat import (
    CryptoContext as RustCryptoContext,
)
from .._hazmat import (
    CryptoError,
)
from ..tls import CipherSuite, cipher_suite_hash, hkdf_expand_label, hkdf_extract
from .packet import QuicProtocolVersion

CIPHER_SUITES = {
    CipherSuite.AES_128_GCM_SHA256: (b"aes-128-ecb", b"aes-128-gcm"),
    CipherSuite.CHACHA20_POLY1305_SHA256: (b"chacha20", b"chacha20-poly1305"),
    CipherSuite.AES_256_GCM_SHA384: (b"aes-256-ecb", b"aes-256-gcm"),
}
INITIAL_CIPHER_SUITE = CipherSuite.AES_128_GCM_SHA256
INITIAL_SALT_VERSION_1 = binascii.unhexlify("38762cf7f55934b34d179ae6a4c80cadccbb7f0a")
INITIAL_SALT_VERSION_2 = binascii.unhexlify("0dede3def700a6db819381be6e269dcbf9bd2ed9")
SAMPLE_SIZE = 16


Callback = Callable[[str], None]


def NoCallback(trigger: str) -> None:
    pass


class KeyUnavailableError(CryptoError):
    pass


def derive_key_iv_hp(
    *, cipher_suite: CipherSuite, secret: bytes, version: int
) -> tuple[bytes, bytes, bytes]:
    algorithm = cipher_suite_hash(cipher_suite)

    if cipher_suite in [
        CipherSuite.AES_256_GCM_SHA384,
        CipherSuite.CHACHA20_POLY1305_SHA256,
    ]:
        key_size = 32
    else:
        key_size = 16

    if version == QuicProtocolVersion.VERSION_2:
        return (
            hkdf_expand_label(algorithm, secret, b"quicv2 key", b"", key_size),
            hkdf_expand_label(algorithm, secret, b"quicv2 iv", b"", 12),
            hkdf_expand_label(algorithm, secret, b"quicv2 hp", b"", key_size),
        )
    else:
        return (
            hkdf_expand_label(algorithm, secret, b"quic key", b"", key_size),
            hkdf_expand_label(algorithm, secret, b"quic iv", b"", 12),
            hkdf_expand_label(algorithm, secret, b"quic hp", b"", key_size),
        )


class CryptoContext:
    __slots__ = (
        "_inner",
        "cipher_suite",
        "key_phase",
        "secret",
        "version",
        "_setup_cb",
        "_teardown_cb",
    )

    def __init__(
        self,
        key_phase: int = 0,
        setup_cb: Callback = NoCallback,
        teardown_cb: Callback = NoCallback,
    ) -> None:
        self._inner: RustCryptoContext | None = None
        self.cipher_suite: CipherSuite | None = None
        self.key_phase = key_phase
        self.secret: bytes | None = None
        self.version: int | None = None
        self._setup_cb = setup_cb
        self._teardown_cb = teardown_cb

    def decrypt_packet(
        self, packet: bytes, encrypted_offset: int, expected_packet_number: int
    ) -> tuple[bytes, bytes, int, bool]:
        if self._inner is None:
            raise KeyUnavailableError("Decryption key is not available")

        # HP removal + PN decode + AEAD decrypt
        plain_header, payload, packet_number, key_phase_changed = (
            self._inner.decrypt_packet(packet, encrypted_offset, expected_packet_number)
        )

        if key_phase_changed:
            # Key phase changed — create next-phase context and decrypt with it
            crypto = next_key_phase(self)
            payload = crypto._inner.decrypt_payload(
                packet[len(plain_header) :], plain_header, packet_number
            )
            return plain_header, payload, packet_number, True

        return plain_header, payload, packet_number, False

    def encrypt_packet(
        self, plain_header: bytes, plain_payload: bytes, packet_number: int
    ) -> bytes:
        assert self.is_valid(), "Encryption key is not available"

        # AEAD encrypt + HP apply
        return self._inner.encrypt_packet(plain_header, plain_payload, packet_number)

    def is_valid(self) -> bool:
        return self._inner is not None

    def setup(self, *, cipher_suite: CipherSuite, secret: bytes, version: int) -> None:
        hp_cipher_name, aead_cipher_name = CIPHER_SUITES[cipher_suite]

        key, iv, hp = derive_key_iv_hp(
            cipher_suite=cipher_suite,
            secret=secret,
            version=version,
        )

        self._inner = RustCryptoContext(
            aead_cipher_name.decode(),
            hp_cipher_name.decode(),
            key,
            iv,
            hp,
            self.key_phase,
        )

        self.cipher_suite = cipher_suite
        self.secret = secret
        self.version = version

        # trigger callback
        self._setup_cb("tls")

    def teardown(self) -> None:
        self._inner = None
        self.cipher_suite = None
        self.secret = None

        # trigger callback
        self._teardown_cb("tls")


def apply_key_phase(self: CryptoContext, crypto: CryptoContext, trigger: str) -> None:
    # Update the AEAD key material in the Rust CryptoContext without changing HP
    key, iv, _ = derive_key_iv_hp(
        cipher_suite=crypto.cipher_suite,
        secret=crypto.secret,
        version=crypto.version,
    )
    self._inner.update_aead(key, iv, crypto.key_phase)
    self.key_phase = crypto.key_phase
    self.secret = crypto.secret

    # trigger callback
    self._setup_cb(trigger)


def next_key_phase(self: CryptoContext) -> CryptoContext:
    algorithm = cipher_suite_hash(self.cipher_suite)

    crypto = CryptoContext(key_phase=int(not self.key_phase))
    crypto.setup(
        cipher_suite=self.cipher_suite,
        secret=hkdf_expand_label(
            algorithm, self.secret, b"quic ku", b"", int(algorithm / 8)
        ),
        version=self.version,
    )
    return crypto


class CryptoPair:
    __slots__ = (
        "aead_tag_size",
        "recv",
        "send",
        "_update_key_requested",
        "_previous_recv",
        "_previous_recv_expires_at",
    )

    def __init__(
        self,
        recv_setup_cb: Callback = NoCallback,
        recv_teardown_cb: Callback = NoCallback,
        send_setup_cb: Callback = NoCallback,
        send_teardown_cb: Callback = NoCallback,
    ) -> None:
        self.aead_tag_size = 16
        self.recv = CryptoContext(setup_cb=recv_setup_cb, teardown_cb=recv_teardown_cb)
        self.send = CryptoContext(setup_cb=send_setup_cb, teardown_cb=send_teardown_cb)
        self._update_key_requested = False
        # RFC 9001 6.5: keep the previous receive context for 3*PTO so
        # reordered packets sent under the previous key can be decrypted.
        self._previous_recv: CryptoContext | None = None
        self._previous_recv_expires_at: float | None = None

    def expire_previous_keys(self, now: float) -> None:
        """Drop the retained previous receive key once 3*PTO has elapsed."""
        if (
            self._previous_recv is not None
            and self._previous_recv_expires_at is not None
            and now >= self._previous_recv_expires_at
        ):
            self._previous_recv.teardown()
            self._previous_recv = None
            self._previous_recv_expires_at = None

    def decrypt_packet(
        self, packet: bytes, encrypted_offset: int, expected_packet_number: int
    ) -> tuple[bytes, bytes, int]:
        try:
            plain_header, payload, packet_number, update_key = self.recv.decrypt_packet(
                packet, encrypted_offset, expected_packet_number
            )
        except CryptoError:
            # AEAD failed — possibly a reordered packet sent under the
            # previous key (RFC 9001 6.5). Try the retained snapshot
            # before giving up; on success we MUST NOT rotate.
            if self._previous_recv is not None:
                plain_header, payload, packet_number, _ = (
                    self._previous_recv.decrypt_packet(
                        packet, encrypted_offset, expected_packet_number
                    )
                )
                return plain_header, payload, packet_number
            raise
        if update_key:
            # The packet's key phase differs from our current one. It
            # could be either (a) the peer initiating the next phase, or
            # (b) a delayed packet sent before a previous local-initiated
            # update. Try the retained previous key first to avoid
            # spuriously rotating when (b) applies.
            if self._previous_recv is not None:
                try:
                    plain_header2, payload2, pn2, _ = (
                        self._previous_recv.decrypt_packet(
                            packet, encrypted_offset, expected_packet_number
                        )
                    )
                    return plain_header2, payload2, pn2
                except CryptoError:
                    pass
            self._update_key("remote_update")
        return plain_header, payload, packet_number

    def encrypt_packet(
        self, plain_header: bytes, plain_payload: bytes, packet_number: int
    ) -> bytes:
        if self._update_key_requested:
            self._update_key("local_update")
        return self.send.encrypt_packet(plain_header, plain_payload, packet_number)

    def finalize_packet(
        self,
        buffer,
        packet_start: int,
        packet_size: int,
        padding_size: int,
        header_size: int,
        is_long_header: bool,
        version: int,
        packet_type: int,
        peer_cid: bytes,
        host_cid: bytes,
        peer_token: bytes,
        spin_bit: int,
        packet_number: int,
    ) -> int:
        """Finalize a packet: write header, pad, encrypt, apply HP in one call."""
        if self._update_key_requested:
            self._update_key("local_update")
        return self.send._inner.finalize_packet(
            buffer,
            packet_start,
            packet_size,
            padding_size,
            header_size,
            is_long_header,
            version,
            packet_type,
            peer_cid,
            host_cid,
            peer_token,
            spin_bit,
            packet_number,
        )

    def setup_initial(self, cid: bytes, is_client: bool, version: int) -> None:
        if is_client:
            recv_label, send_label = b"server in", b"client in"
        else:
            recv_label, send_label = b"client in", b"server in"

        if version == QuicProtocolVersion.VERSION_2:
            initial_salt = INITIAL_SALT_VERSION_2
        else:
            initial_salt = INITIAL_SALT_VERSION_1

        algorithm = cipher_suite_hash(INITIAL_CIPHER_SUITE)
        digest_size = int(algorithm / 8)
        initial_secret = hkdf_extract(algorithm, initial_salt, cid)
        self.recv.setup(
            cipher_suite=INITIAL_CIPHER_SUITE,
            secret=hkdf_expand_label(
                algorithm, initial_secret, recv_label, b"", digest_size
            ),
            version=version,
        )
        self.send.setup(
            cipher_suite=INITIAL_CIPHER_SUITE,
            secret=hkdf_expand_label(
                algorithm, initial_secret, send_label, b"", digest_size
            ),
            version=version,
        )

    def teardown(self) -> None:
        self.recv.teardown()
        self.send.teardown()

    def update_key(self) -> None:
        self._update_key_requested = True

    @property
    def key_phase(self) -> int:
        if self._update_key_requested:
            return int(not self.recv.key_phase)
        else:
            return self.recv.key_phase

    def _update_key(self, trigger: str) -> None:
        # Snapshot the current receive context so we can keep decrypting
        # reordered packets under the previous key for 3*PTO
        # (RFC 9001 6.5). The caller is responsible for setting
        # _previous_recv_expires_at via retain_previous_keys().
        if self.recv.is_valid():
            snapshot = CryptoContext(key_phase=self.recv.key_phase)
            snapshot.setup(
                cipher_suite=self.recv.cipher_suite,
                secret=self.recv.secret,
                version=self.recv.version,
            )
            # Replace any existing snapshot.
            if self._previous_recv is not None:
                self._previous_recv.teardown()
            self._previous_recv = snapshot

        apply_key_phase(self.recv, next_key_phase(self.recv), trigger=trigger)
        apply_key_phase(self.send, next_key_phase(self.send), trigger=trigger)
        self._update_key_requested = False

    def retain_previous_keys(self, expires_at: float) -> None:
        """
        Schedule the retained previous receive key to be discarded at
        ``expires_at``. Called by the connection after a key update.
        """
        if self._previous_recv is not None:
            self._previous_recv_expires_at = expires_at