//! Transport-independent MTP handshake engine.
//!
//! This module contains the shared state machine used by both native `MTPHost`
//! and the web server's `MTPWebServer` to perform the MTP opening handshake,
//! version negotiation, authentication, and guest assignment.

use crate::config::HostConfig;
use crate::error::AcceptError;
use mtp_codec::{
    CommunicationType, CommunicationValue, DataType, DataValue, TypeMap, Version,
    registry::{Registry, VersionedCodec},
};
use mtp_common::{CommunicationError, RejectionReason};
#[cfg(feature = "crypto")]
use std::collections::HashSet;
use std::sync::Arc;
#[cfg(feature = "crypto")]
use std::sync::Mutex;

/// Trait for sending handshake messages during the opening exchange.
///
/// Implemented by both the concrete `Sender` and `GenericSender<C>`.
pub trait HandshakeSender: Send + Sync {
    fn send(
        &self,
        msg: &CommunicationValue,
    ) -> impl std::future::Future<Output = Result<(), CommunicationError>> + Send;
    fn finish_stream(
        &self,
    ) -> impl std::future::Future<Output = Result<(), CommunicationError>> + Send;
    fn set_type_map(&self, _type_map: &TypeMap) -> impl std::future::Future<Output = ()> + Send {
        async {}
    }
    fn close(&self);
}

/// Trait for receiving handshake messages during the opening exchange.
///
/// Implemented by both the concrete `Receiver` and `GenericReceiver<C>`.
pub trait HandshakeReceiver: Send + Sync {
    fn receive(
        &self,
    ) -> impl std::future::Future<Output = Result<CommunicationValue, CommunicationError>> + Send;

    /// Bind subsequently decoded frames to the negotiated type map.
    ///
    /// The opening frame must be decoded with the transport's bootstrap map so
    /// that it can reveal the version. Once negotiation succeeds, all later
    /// frames—including the remainder of the authentication exchange—must use
    /// the negotiated map rather than whichever map happens to be latest at
    /// compile time.
    fn set_type_map(&self, _type_map: &TypeMap) -> impl std::future::Future<Output = ()> + Send {
        async {}
    }
}

/// A non-zero guest ID reserved for the lifetime of a connected session.
///
/// The lease is moved into the resulting `MTPConnection`, so dropping that
/// connection releases the ID for a later guest session.
#[cfg(feature = "crypto")]
#[derive(Debug)]
pub struct GuestIdLease {
    active_ids: Arc<Mutex<HashSet<u64>>>,
    id: u64,
}

#[cfg(feature = "crypto")]
impl Drop for GuestIdLease {
    fn drop(&mut self) {
        if let Ok(mut active_ids) = self.active_ids.lock() {
            active_ids.remove(&self.id);
        }
    }
}

/// The result of a successful handshake, containing everything needed to
/// construct the final `MTPConnection`.
#[derive(Debug)]
pub struct HandshakeResult {
    pub negotiated_version: Version,
    pub codec: VersionedCodec,
    pub description: Option<String>,
    #[cfg(feature = "crypto")]
    pub auth_state: crate::error::AuthState,
    #[cfg(feature = "crypto")]
    pub client_id: u64,
    #[cfg(feature = "crypto")]
    pub client_public_key: Option<mtp_crypto::PublicKeyBundle>,
    #[cfg(feature = "crypto")]
    pub guest_id_lease: Option<GuestIdLease>,
}

/// Transport-independent handshake state machine.
///
/// Both `MTPHost` and `MTPWebServer` create a `HandshakeEngine` with the
/// shared `HostConfig` and delegate the full opening handshake to it.
pub struct HandshakeEngine {
    registry: Registry,
    #[cfg(feature = "crypto")]
    config: Arc<HostConfig>,
}

impl HandshakeEngine {
    #[cfg(feature = "crypto")]
    pub fn new(registry: Registry, config: Arc<HostConfig>) -> Self {
        Self { registry, config }
    }

    #[cfg(not(feature = "crypto"))]
    pub fn new(registry: Registry, _config: Arc<HostConfig>) -> Self {
        Self { registry }
    }

    /// Run the complete opening handshake with the given transport pair.
    ///
    /// This handles:
    /// - Opening-frame timeout (when crypto is enabled)
    /// - Opening-type classification (Identification, Register, or other)
    /// - Version negotiation
    /// - Authentication-policy selection (Unauthenticated, AllowAuthentication, ForceAuthentication)
    /// - Guest allocation and collision avoidance
    /// - Full challenge/response authentication when required
    /// - PQ preflight checks and dual-signature verification
    /// - Rejection response construction on failure
    pub async fn accept<S: HandshakeSender, R: HandshakeReceiver>(
        &self,
        sender: &S,
        receiver: &R,
    ) -> Result<HandshakeResult, AcceptError> {
        #[cfg(feature = "crypto")]
        {
            self.accept_until(
                sender,
                receiver,
                tokio::time::Instant::now() + self.config.auth_timeout,
            )
            .await
        }
        #[cfg(not(feature = "crypto"))]
        {
            let result = self.accept_inner(sender, receiver).await;
            if result.is_err() {
                sender.close();
            }
            result
        }
    }

    /// Run the crypto handshake until an absolute deadline.
    ///
    /// WebTransport authentication may wait for a shared semaphore before it
    /// reaches this engine. Passing the deadline through keeps that queueing
    /// time from silently granting the handshake another full timeout.
    #[cfg(feature = "crypto")]
    pub async fn accept_until<S: HandshakeSender, R: HandshakeReceiver>(
        &self,
        sender: &S,
        receiver: &R,
        deadline: tokio::time::Instant,
    ) -> Result<HandshakeResult, AcceptError> {
        match tokio::time::timeout_at(deadline, self.accept_inner(sender, receiver)).await {
            Ok(result) => {
                if result.is_err() {
                    sender.close();
                }
                result
            }
            Err(_) => {
                let error = AcceptError::AuthenticationTimedOut;
                send_rejection_generic(
                    sender,
                    RejectionReason::AuthenticationFailed {
                        detail: error.to_string(),
                    },
                    None,
                )
                .await;
                sender.close();
                Err(error)
            }
        }
    }

    async fn accept_inner<S: HandshakeSender, R: HandshakeReceiver>(
        &self,
        sender: &S,
        receiver: &R,
    ) -> Result<HandshakeResult, AcceptError> {
        let mut first_msg = receiver.receive().await.map_err(AcceptError::Receive)?;

        let version_str = match first_msg.get_data(DataType::Version) {
            Some(DataValue::Str(s)) => s.clone(),
            _ => {
                send_rejection_generic(
                    sender,
                    RejectionReason::AuthenticationFailed {
                        detail: "opening message omitted a valid protocol version".into(),
                    },
                    None,
                )
                .await;
                sender.close();
                return Err(AcceptError::MissingVersion);
            }
        };
        let client_version = match Version::parse(&version_str) {
            Some(v) => v,
            _ => {
                send_rejection_generic(
                    sender,
                    RejectionReason::AuthenticationFailed {
                        detail: "opening message omitted a valid protocol version".into(),
                    },
                    None,
                )
                .await;
                sender.close();
                return Err(AcceptError::MissingVersion);
            }
        };

        let negotiated = match self
            .registry
            .negotiate(std::slice::from_ref(&client_version))
        {
            Some(v) => v,
            None => {
                send_rejection_generic(
                    sender,
                    RejectionReason::BadVersion {
                        supported_versions: self
                            .registry
                            .versions()
                            .map(|v| v.to_string())
                            .collect(),
                    },
                    None,
                )
                .await;
                sender.close();
                return Err(AcceptError::UnsupportedVersion(client_version));
            }
        };

        let codec = VersionedCodec::for_version(self.registry.clone(), negotiated.clone())
            .ok_or_else(|| AcceptError::UnsupportedVersion(negotiated.clone()))?;

        sender.set_type_map(codec.type_map()).await;
        receiver.set_type_map(codec.type_map()).await;
        first_msg.set_type_map(codec.type_map());

        let description = match first_msg.get_data(DataType::Description) {
            Some(DataValue::Str(s)) => Some(s.clone()),
            _ => None,
        };

        #[cfg(feature = "crypto")]
        {
            match self.config.authentication_policy {
                crate::config::AuthenticationPolicy::ForceAuthentication => {
                    self.force_auth_handshake(
                        sender,
                        receiver,
                        first_msg,
                        negotiated,
                        codec,
                        description,
                        &version_str,
                        client_version,
                    )
                    .await
                }
                crate::config::AuthenticationPolicy::AllowAuthentication => {
                    self.allow_auth_handshake(
                        sender,
                        receiver,
                        first_msg,
                        negotiated,
                        codec,
                        description,
                        &version_str,
                        client_version,
                    )
                    .await
                }
                crate::config::AuthenticationPolicy::Unauthenticated => {
                    self.unauthenticated_handshake(
                        sender,
                        first_msg,
                        negotiated,
                        codec,
                        description,
                    )
                    .await
                }
            }
        }

        #[cfg(not(feature = "crypto"))]
        {
            let authentication_requested = Some(first_msg.get_type())
                == CommunicationType::Register.try_to_id(codec.type_map())
                || first_msg.get_data(DataType::PublicKeys).is_some();
            if authentication_requested {
                let error = AcceptError::AuthenticationFailed(
                    "authentication is unavailable on this non-crypto host".into(),
                );
                send_rejection_generic(
                    sender,
                    RejectionReason::AuthenticationFailed {
                        detail: error.to_string(),
                    },
                    Some(codec.type_map()),
                )
                .await;
                sender.close();
                return Err(error);
            }
            let _ = receiver;
            send_accepted_generic(sender, &negotiated, codec.type_map(), Some(0))
                .await
                .map_err(AcceptError::Send)?;
            Ok(HandshakeResult {
                negotiated_version: negotiated,
                codec,
                description,
            })
        }
    }

    #[cfg(feature = "crypto")]
    async fn unauthenticated_handshake<S: HandshakeSender>(
        &self,
        sender: &S,
        first_msg: CommunicationValue,
        negotiated: Version,
        codec: VersionedCodec,
        description: Option<String>,
    ) -> Result<HandshakeResult, AcceptError> {
        let tm = codec.type_map();

        // Reject explicit authentication attempts on unauthenticated hosts.
        // Authenticated clients include PublicKeys in Identification as an
        // intent marker; this avoids acknowledging the opening as a guest
        // connection and leaving the client waiting for a Challenge.
        if Some(first_msg.get_type()) == CommunicationType::Register.try_to_id(&tm)
            || first_msg.get_data(DataType::PublicKeys).is_some()
        {
            send_rejection_generic(
                sender,
                RejectionReason::AuthenticationFailed {
                    detail: "authentication not allowed on this host".into(),
                },
                Some(tm),
            )
            .await;
            sender.close();
            return Err(AcceptError::AuthenticationFailed(
                "authentication not allowed on this host".into(),
            ));
        }

        let guest_id_lease = match self.assign_guest_id().await {
            Ok(lease) => lease,
            Err(error) => {
                reject_error_generic(sender, &error, tm).await;
                return Err(error);
            }
        };
        let guest_id = guest_id_lease.id;
        send_accepted_generic(sender, &negotiated, tm, Some(guest_id))
            .await
            .map_err(AcceptError::Send)?;

        Ok(HandshakeResult {
            negotiated_version: negotiated,
            codec,
            description,
            auth_state: crate::error::AuthState::Unauthenticated,
            client_id: guest_id,
            client_public_key: None,
            guest_id_lease: Some(guest_id_lease),
        })
    }

    #[cfg(feature = "crypto")]
    #[allow(clippy::too_many_arguments)]
    async fn allow_auth_handshake<S: HandshakeSender, R: HandshakeReceiver>(
        &self,
        sender: &S,
        receiver: &R,
        first_msg: CommunicationValue,
        negotiated: Version,
        codec: VersionedCodec,
        description: Option<String>,
        version_str: &str,
        client_version: Version,
    ) -> Result<HandshakeResult, AcceptError> {
        let tm = codec.type_map();

        // Register frames always go through full authentication
        if Some(first_msg.get_type()) == CommunicationType::Register.try_to_id(&tm) {
            let bundle = match extract_register_bundle(&first_msg) {
                Ok(bundle) => bundle,
                Err(error) => {
                    reject_error_generic(sender, &error, tm).await;
                    return Err(error);
                }
            };
            let pk_bytes = bundle.as_bytes();
            return self
                .complete_auth_handshake(
                    sender,
                    receiver,
                    Flow::Register { bundle, pk_bytes },
                    CommunicationType::RegisterResponse,
                    &negotiated,
                    &codec,
                    description,
                    version_str,
                    client_version,
                )
                .await;
        }

        // Identification: try lookup, fall back to guest
        if Some(first_msg.get_type()) == CommunicationType::Identification.try_to_id(&tm) {
            let cid = match first_msg.get_data(DataType::Id) {
                Some(DataValue::UnsignedNumber(n)) => u64::try_from(*n).unwrap_or(0),
                _ => 0,
            };

            if cid > 0
                && let Some(bundle) =
                    (self.config.get_existing_client)(cid, description.clone()).await
            {
                return self
                    .complete_auth_handshake(
                        sender,
                        receiver,
                        Flow::Login { id: cid, bundle },
                        CommunicationType::IdentificationResponse,
                        &negotiated,
                        &codec,
                        description,
                        version_str,
                        client_version,
                    )
                    .await;
            }

            // Unknown or zero ID: an Identification carrying PublicKeys is an
            // explicit authentication attempt, not a guest connection.
            if first_msg.get_data(DataType::PublicKeys).is_some() {
                let error = AcceptError::AuthenticationFailed(
                    "unknown authenticated client identity".into(),
                );
                reject_error_generic(sender, &error, tm).await;
                return Err(error);
            }

            // Unknown or zero ID: fall back to guest
            let guest_id_lease = match self.assign_guest_id().await {
                Ok(lease) => lease,
                Err(error) => {
                    reject_error_generic(sender, &error, tm).await;
                    return Err(error);
                }
            };
            let guest_id = guest_id_lease.id;
            send_accepted_generic(sender, &negotiated, tm, Some(guest_id))
                .await
                .map_err(AcceptError::Send)?;
            return Ok(HandshakeResult {
                negotiated_version: negotiated,
                codec,
                description,
                auth_state: crate::error::AuthState::Unauthenticated,
                client_id: guest_id,
                client_public_key: None,
                guest_id_lease: Some(guest_id_lease),
            });
        }

        let error = AcceptError::AuthenticationFailed("unexpected message type".into());
        reject_error_generic(sender, &error, tm).await;
        Err(error)
    }

    #[cfg(feature = "crypto")]
    #[allow(clippy::too_many_arguments)]
    async fn force_auth_handshake<S: HandshakeSender, R: HandshakeReceiver>(
        &self,
        sender: &S,
        receiver: &R,
        first_msg: CommunicationValue,
        negotiated: Version,
        codec: VersionedCodec,
        description: Option<String>,
        version_str: &str,
        client_version: Version,
    ) -> Result<HandshakeResult, AcceptError> {
        let tm = codec.type_map();

        let (flow, response_type) = if Some(first_msg.get_type())
            == CommunicationType::Identification.try_to_id(&tm)
        {
            let cid = match first_msg.get_data(DataType::Id) {
                Some(DataValue::UnsignedNumber(n)) => match u64::try_from(*n) {
                    Ok(id) => id,
                    Err(_) => {
                        let error =
                            AcceptError::AuthenticationFailed("client id is out of range".into());
                        reject_error_generic(sender, &error, tm).await;
                        return Err(error);
                    }
                },
                _ => {
                    let error = AcceptError::AuthenticationFailed("missing client id".into());
                    reject_error_generic(sender, &error, tm).await;
                    return Err(error);
                }
            };
            let bundle = match (self.config.get_existing_client)(cid, description.clone()).await {
                Some(b) => b,
                None => {
                    let rejection = CommunicationValue::new_with_type_map(
                        CommunicationType::IdentificationResponse,
                        tm,
                    )
                    .add_typed_default(DataType::Connected, DataValue::BoolFalse)
                    .add_typed_default(
                        DataType::ErrorMessage,
                        DataValue::Str("unknown client id".into()),
                    );
                    let _ = sender.send(&rejection).await;
                    sender.close();
                    return Err(AcceptError::AuthenticationFailed(
                        "unknown client id".into(),
                    ));
                }
            };
            (
                Flow::Login { id: cid, bundle },
                CommunicationType::IdentificationResponse,
            )
        } else if Some(first_msg.get_type()) == CommunicationType::Register.try_to_id(&tm) {
            let bundle = match extract_register_bundle(&first_msg) {
                Ok(bundle) => bundle,
                Err(error) => {
                    reject_error_generic(sender, &error, tm).await;
                    return Err(error);
                }
            };
            let pk_bytes = bundle.as_bytes();
            (
                Flow::Register { bundle, pk_bytes },
                CommunicationType::RegisterResponse,
            )
        } else {
            let error =
                AcceptError::AuthenticationFailed("unexpected authentication message".into());
            reject_error_generic(sender, &error, tm).await;
            return Err(error);
        };

        self.complete_auth_handshake(
            sender,
            receiver,
            flow,
            response_type,
            &negotiated,
            &codec,
            description,
            version_str,
            client_version,
        )
        .await
    }

    #[cfg(feature = "crypto")]
    #[allow(clippy::too_many_arguments)]
    async fn complete_auth_handshake<S: HandshakeSender, R: HandshakeReceiver>(
        &self,
        sender: &S,
        receiver: &R,
        flow: Flow,
        response_type: CommunicationType,
        negotiated: &Version,
        codec: &VersionedCodec,
        description: Option<String>,
        version_str: &str,
        _client_version: Version,
    ) -> Result<HandshakeResult, AcceptError> {
        use mtp_crypto::{Ed25519Signer, MlDsaSigner, SignatureScheme, auth, verify_ed25519};

        let tm = codec.type_map();

        // PQ preflight: host requiring PQ must have a PQ key
        let pq_enabled = !self
            .config
            .host_keyring
            .sig_pq_secret_key
            .as_bytes()
            .is_empty();
        if self.config.require_pq
            && (!pq_enabled
                || self
                    .config
                    .host_keyring
                    .sig_pq_public_key
                    .as_bytes()
                    .is_empty())
        {
            send_rejection_generic(
                sender,
                RejectionReason::AuthenticationFailed {
                    detail: "host requires PQ authentication but has no PQ signing key".into(),
                },
                Some(tm),
            )
            .await;
            sender.close();
            return Err(AcceptError::AuthenticationFailed(
                "PQ authentication is required but the host PQ key is absent".into(),
            ));
        }

        // Initialize host signers
        let host_pq_signer = if pq_enabled {
            Some(Arc::new(
                match MlDsaSigner::new(
                    &self.config.host_keyring.sig_pq_secret_key,
                    &self.config.host_keyring.sig_pq_public_key,
                ) {
                    Ok(signer) => signer,
                    Err(error) => {
                        let error = AcceptError::AuthenticationFailed(error.to_string());
                        reject_error_generic(sender, &error, tm).await;
                        return Err(error);
                    }
                },
            ))
        } else {
            None
        };

        let host_sign = |payload: Vec<u8>| async {
            let signer = Ed25519Signer::new(&self.config.host_keyring.sig_cl_secret_key)
                .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?;
            if let Some(pq_signer) = host_pq_signer.as_ref() {
                mtp_crypto::sign_parallel::sign_dual_parallel_shared_pq(
                    signer,
                    Arc::clone(pq_signer),
                    payload,
                )
                .await
                .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))
            } else {
                let sig = signer
                    .sign(&payload)
                    .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?;
                Ok((sig, Vec::new()))
            }
        };

        // Sign and send challenge
        let challenge_id = match &flow {
            Flow::Login { id, .. } => *id,
            Flow::Register { .. } => 0,
        };

        let server_challenge: u128 = rand::random();
        let (chal_sig, chal_pq_sig) =
            match host_sign(auth::challenge_payload(challenge_id, server_challenge)).await {
                Ok(signatures) => signatures,
                Err(error) => {
                    reject_error_generic(sender, &error, tm).await;
                    return Err(error);
                }
            };

        let mut challenge_msg =
            CommunicationValue::new_with_type_map(CommunicationType::Challenge, tm)
                .add_typed_default(
                    DataType::ServerNonce,
                    DataValue::UnsignedNumber(server_challenge),
                )
                .add_typed_default(DataType::Signature, DataValue::Bytes(chal_sig));
        challenge_msg = challenge_msg.add_typed_default(
            DataType::RequirePq,
            if self.config.require_pq {
                DataValue::BoolTrue
            } else {
                DataValue::BoolFalse
            },
        );
        if pq_enabled {
            challenge_msg = challenge_msg
                .add_typed_default(DataType::PqSignature, DataValue::Bytes(chal_pq_sig));
        }
        if let Err(e) = sender.send(&challenge_msg).await {
            sender.close();
            return Err(AcceptError::Send(e));
        }

        // Receive and verify client proof
        let proof = receiver.receive().await.map_err(|e| {
            sender.close();
            AcceptError::Receive(e)
        })?;
        if Some(proof.get_type()) != CommunicationType::ChallengeResponse.try_to_id(&tm) {
            let error = AcceptError::AuthenticationFailed("missing challenge response".into());
            reject_error_generic(sender, &error, tm).await;
            return Err(error);
        }
        let client_nonce = match proof.get_data(DataType::ClientNonce) {
            Some(DataValue::UnsignedNumber(n)) => *n,
            _ => {
                let error = AcceptError::AuthenticationFailed("missing client nonce".into());
                reject_error_generic(sender, &error, tm).await;
                return Err(error);
            }
        };
        let sig_bytes = match proof.get_data(DataType::Signature) {
            Some(DataValue::Bytes(b)) => b.clone(),
            _ => {
                let error = AcceptError::AuthenticationFailed("missing challenge signature".into());
                reject_error_generic(sender, &error, tm).await;
                return Err(error);
            }
        };
        let pq_sig_bytes: Vec<u8> = match proof.get_data(DataType::PqSignature) {
            Some(DataValue::Bytes(b)) => b.clone(),
            _ => vec![],
        };

        let (proof_payload, bundle) = match &flow {
            Flow::Login { id, bundle } => (
                auth::login_proof_payload(version_str, *id, server_challenge, client_nonce),
                bundle,
            ),
            Flow::Register {
                bundle, pk_bytes, ..
            } => (
                auth::register_proof_payload(version_str, pk_bytes, server_challenge, client_nonce),
                bundle,
            ),
        };

        let has_client_pq_key = !bundle.sig_pq_public_key.as_bytes().is_empty();
        let proof_ok = if pq_sig_bytes.is_empty() {
            !self.config.require_pq
                && verify_ed25519(&bundle.sig_cl_public_key, &proof_payload, &sig_bytes).is_ok()
        } else if has_client_pq_key {
            mtp_crypto::sign_parallel::verify_dual_parallel(
                bundle.sig_cl_public_key.clone(),
                bundle.sig_pq_public_key.clone(),
                proof_payload,
                sig_bytes,
                pq_sig_bytes,
            )
            .await
            .is_ok()
        } else {
            false
        };

        if !proof_ok {
            send_rejection_generic(
                sender,
                RejectionReason::AuthenticationFailed {
                    detail: "client proof signature invalid".into(),
                },
                Some(tm),
            )
            .await;
            sender.close();
            return Err(AcceptError::AuthenticationFailed(
                "client proof signature invalid".into(),
            ));
        }

        // Register or login
        let (assigned_id, client_bundle) = match flow {
            Flow::Login { id, bundle } => (id, bundle),
            Flow::Register { bundle, .. } => {
                let _registration_guard = self.config.registration_lock.lock().await;
                let identity = bundle.as_bytes();
                let cached_id = self
                    .config
                    .registration_ids
                    .lock()
                    .ok()
                    .and_then(|registrations| registrations.get(&identity).copied());
                let new_id = if let Some(id) = cached_id {
                    id
                } else if let Some(lookup) = &self.config.find_registered_client {
                    match lookup(bundle.clone(), description.clone()).await {
                        Some(id) => id,
                        None => {
                            (self.config.complete_register)(bundle.clone(), description.clone())
                                .await
                        }
                    }
                } else {
                    (self.config.complete_register)(bundle.clone(), description.clone()).await
                };
                if new_id != 0
                    && let Ok(mut registrations) = self.config.registration_ids.lock()
                {
                    registrations.insert(identity, new_id);
                }
                if new_id == 0 {
                    let error = AcceptError::AuthenticationFailed(
                        "registration callback returned reserved client id 0".into(),
                    );
                    reject_error_generic(sender, &error, tm).await;
                    return Err(error);
                }
                (new_id, bundle)
            }
        };
        if assigned_id == 0 {
            let error = AcceptError::AuthenticationFailed(
                "client id 0 is reserved for no authenticated identity".into(),
            );
            reject_error_generic(sender, &error, tm).await;
            return Err(error);
        }

        // Sign and send final response
        let (host_sig, host_pq_sig) = match host_sign(auth::host_final_payload(
            assigned_id,
            client_nonce,
            server_challenge,
        ))
        .await
        {
            Ok(signatures) => signatures,
            Err(error) => {
                reject_error_generic(sender, &error, tm).await;
                return Err(error);
            }
        };

        let mut response = CommunicationValue::new_with_type_map(response_type, tm)
            .add_typed_default(DataType::Connected, DataValue::BoolTrue)
            .add_typed_default(DataType::Id, DataValue::UnsignedNumber(assigned_id as u128))
            .add_typed_default(
                DataType::ClientNonce,
                DataValue::UnsignedNumber(client_nonce),
            )
            .add_typed_default(DataType::Signature, DataValue::Bytes(host_sig));
        response =
            response.add_typed_default(DataType::Version, DataValue::Str(negotiated.to_string()));
        if pq_enabled {
            response =
                response.add_typed_default(DataType::PqSignature, DataValue::Bytes(host_pq_sig));
        }

        if let Err(e) = sender.send(&response).await {
            sender.close();
            return Err(AcceptError::Send(e));
        }
        if let Err(e) = sender.finish_stream().await {
            sender.close();
            return Err(AcceptError::Send(e));
        }

        Ok(HandshakeResult {
            negotiated_version: negotiated.clone(),
            codec: codec.clone(),
            description,
            auth_state: crate::error::AuthState::Authenticated,
            client_id: assigned_id,
            client_public_key: Some(client_bundle),
            guest_id_lease: None,
        })
    }
}

#[cfg(feature = "crypto")]
enum Flow {
    Login {
        id: u64,
        bundle: mtp_crypto::PublicKeyBundle,
    },
    Register {
        bundle: mtp_crypto::PublicKeyBundle,
        pk_bytes: Vec<u8>,
    },
}

// ---------------------------------------------------------------------------
// Guest ID allocation
// ---------------------------------------------------------------------------

#[cfg(feature = "crypto")]
impl HandshakeEngine {
    const GUEST_ID_MAX_RETRIES: u32 = 100;

    async fn assign_guest_id(&self) -> Result<GuestIdLease, AcceptError> {
        if let Some(ref generator) = self.config.guest_id_generator {
            let id = generator().await.ok_or_else(|| {
                AcceptError::AuthenticationFailed(
                    "guest id generator rejected the connection".into(),
                )
            })?;
            if id == 0 {
                return Err(AcceptError::AuthenticationFailed(
                    "guest id 0 is reserved for no authenticated identity".into(),
                ));
            }
            if let Some(lease) = self.try_reserve_guest_id(id).await? {
                return Ok(lease);
            }
        }
        self.random_guest_id().await
    }

    async fn random_guest_id(&self) -> Result<GuestIdLease, AcceptError> {
        for _ in 0..Self::GUEST_ID_MAX_RETRIES {
            let id = rand::random::<u64>();
            if let Some(lease) = self.try_reserve_guest_id(id).await? {
                return Ok(lease);
            }
        }
        Err(AcceptError::AuthenticationFailed(
            "failed to allocate a unique guest id after retries".into(),
        ))
    }

    async fn try_reserve_guest_id(&self, id: u64) -> Result<Option<GuestIdLease>, AcceptError> {
        if id == 0 || (self.config.get_existing_client)(id, None).await.is_some() {
            return Ok(None);
        }
        let mut active_ids = self.config.active_guest_ids.lock().map_err(|_| {
            AcceptError::AuthenticationFailed("guest ID registry is poisoned".into())
        })?;
        if !active_ids.insert(id) {
            return Ok(None);
        }
        Ok(Some(GuestIdLease {
            active_ids: Arc::clone(&self.config.active_guest_ids),
            id,
        }))
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

#[cfg(feature = "crypto")]
fn extract_register_bundle(
    msg: &CommunicationValue,
) -> Result<mtp_crypto::PublicKeyBundle, AcceptError> {
    match msg.get_data(DataType::PublicKeys) {
        Some(DataValue::Bytes(b)) => mtp_crypto::PublicKeyBundle::from_bytes(b)
            .map_err(|_| AcceptError::AuthenticationFailed("invalid public key bundle".into())),
        _ => Err(AcceptError::AuthenticationFailed(
            "missing public keys".into(),
        )),
    }
}

async fn send_rejection_generic<S: HandshakeSender>(
    sender: &S,
    reason: RejectionReason,
    type_map: Option<&TypeMap>,
) {
    let type_map = type_map.cloned().unwrap_or_else(TypeMap::latest);
    let response = match &reason {
        RejectionReason::BadVersion { supported_versions } => {
            CommunicationValue::new_with_type_map(CommunicationType::ErrorBadVersion, &type_map)
                .add_typed_default(
                    DataType::Version,
                    DataValue::Str(supported_versions.join(",")),
                )
                .add_typed_default(DataType::ErrorMessage, DataValue::Str(reason.to_string()))
        }
        _ => CommunicationValue::new_with_type_map(
            CommunicationType::IdentificationResponse,
            &type_map,
        )
        .add_typed_default(DataType::Connected, DataValue::BoolFalse)
        .add_typed_default(DataType::ErrorMessage, DataValue::Str(reason.to_string())),
    };
    let _ = sender.send(&response).await;
}

#[cfg(feature = "crypto")]
async fn reject_error_generic<S: HandshakeSender>(
    sender: &S,
    error: &AcceptError,
    type_map: &TypeMap,
) {
    send_rejection_generic(
        sender,
        RejectionReason::AuthenticationFailed {
            detail: error.to_string(),
        },
        Some(type_map),
    )
    .await;
    sender.close();
}

async fn send_accepted_generic<S: HandshakeSender>(
    sender: &S,
    version: &Version,
    type_map: &TypeMap,
    assigned_id: Option<u64>,
) -> Result<(), CommunicationError> {
    let mut response =
        CommunicationValue::new_with_type_map(CommunicationType::IdentificationResponse, type_map)
            .add_typed_default(DataType::Connected, DataValue::BoolTrue)
            .add_typed_default(DataType::Version, DataValue::Str(version.to_string()));
    if let Some(id) = assigned_id {
        response = response.add_typed_default(DataType::Id, DataValue::UnsignedNumber(id as u128));
    }
    sender.send(&response).await?;
    sender.finish_stream().await
}

// ---------------------------------------------------------------------------
// Trait implementations for concrete transport types
// ---------------------------------------------------------------------------

impl HandshakeSender for mtp_transport::Sender {
    fn send(
        &self,
        msg: &CommunicationValue,
    ) -> impl std::future::Future<Output = Result<(), CommunicationError>> + Send {
        mtp_transport::Sender::send(self, msg)
    }
    fn finish_stream(
        &self,
    ) -> impl std::future::Future<Output = Result<(), CommunicationError>> + Send {
        mtp_transport::Sender::finish_stream(self)
    }
    fn set_type_map(&self, type_map: &TypeMap) -> impl std::future::Future<Output = ()> + Send {
        async move { self.set_type_map(type_map).await }
    }
    fn close(&self) {
        let sender = self.clone();
        tokio::spawn(async move { sender.close().await });
    }
}

impl HandshakeReceiver for mtp_transport::Receiver {
    fn receive(
        &self,
    ) -> impl std::future::Future<Output = Result<CommunicationValue, CommunicationError>> + Send
    {
        mtp_transport::Receiver::receive(self)
    }

    fn set_type_map(&self, type_map: &TypeMap) -> impl std::future::Future<Output = ()> + Send {
        async move { self.set_type_map(type_map).await }
    }
}

impl<C: mtp_transport::TransportConnection> HandshakeSender for mtp_transport::GenericSender<C> {
    fn send(
        &self,
        msg: &CommunicationValue,
    ) -> impl std::future::Future<Output = Result<(), CommunicationError>> + Send {
        mtp_transport::GenericSender::send(self, msg)
    }
    fn finish_stream(
        &self,
    ) -> impl std::future::Future<Output = Result<(), CommunicationError>> + Send {
        mtp_transport::GenericSender::finish_stream(self)
    }
    fn set_type_map(&self, type_map: &TypeMap) -> impl std::future::Future<Output = ()> + Send {
        async move { self.set_type_map(type_map).await }
    }
    fn close(&self) {
        mtp_transport::GenericSender::close(self);
    }
}

impl<C: mtp_transport::TransportConnection> HandshakeReceiver
    for mtp_transport::GenericReceiver<C>
{
    fn receive(
        &self,
    ) -> impl std::future::Future<Output = Result<CommunicationValue, CommunicationError>> + Send
    {
        mtp_transport::GenericReceiver::receive(self)
    }

    fn set_type_map(&self, type_map: &TypeMap) -> impl std::future::Future<Output = ()> + Send {
        async move { self.set_type_map(type_map).await }
    }
}

#[cfg(all(test, feature = "crypto"))]
mod tests {
    use super::*;
    use std::sync::Mutex;

    #[tokio::test]
    async fn guest_generator_accepts_full_width_id_after_collision_check()
    -> Result<(), Box<dyn std::error::Error>> {
        let lookups = Arc::new(Mutex::new(Vec::new()));
        let recorded_lookups = Arc::clone(&lookups);

        let mut config = HostConfig::new("127.0.0.1".parse()?, 4433, Vec::new(), Vec::new())
            .with_guest_id_generator(Box::new(|| Box::pin(async { Some(u64::MAX) })));
        config.get_existing_client = Box::new(move |id, description| {
            let recorded_lookups = Arc::clone(&recorded_lookups);
            Box::pin(async move {
                recorded_lookups
                    .lock()
                    .expect("guest ID lookup mutex should not be poisoned")
                    .push((id, description));
                None
            })
        });

        let engine = HandshakeEngine::new(Registry::builtin(), Arc::new(config));

        assert_eq!(engine.assign_guest_id().await?.id, u64::MAX);
        assert_eq!(
            *lookups
                .lock()
                .expect("guest ID lookup mutex should not be poisoned"),
            vec![(u64::MAX, None)]
        );
        Ok(())
    }

    #[tokio::test]
    async fn active_guest_ids_are_unique_and_released() -> Result<(), Box<dyn std::error::Error>> {
        let config = HostConfig::new("127.0.0.1".parse()?, 4433, Vec::new(), Vec::new());
        let engine = HandshakeEngine::new(Registry::builtin(), Arc::new(config));

        let first = engine.try_reserve_guest_id(1).await?.unwrap();
        assert!(engine.try_reserve_guest_id(1).await?.is_none());

        drop(first);
        assert!(engine.try_reserve_guest_id(1).await?.is_some());
        Ok(())
    }

    #[tokio::test]
    async fn zero_is_rejected_as_a_guest_id() -> Result<(), Box<dyn std::error::Error>> {
        let config = HostConfig::new("127.0.0.1".parse()?, 4433, Vec::new(), Vec::new())
            .with_guest_id_generator(Box::new(|| Box::pin(async { Some(0) })));
        let engine = HandshakeEngine::new(Registry::builtin(), Arc::new(config));

        assert!(engine.assign_guest_id().await.is_err());
        Ok(())
    }
}