ap_client/clients/

remote_client.rs

use std::time::Duration;

use ap_noise::{InitiatorHandshake, MultiDeviceTransport, Psk};
use ap_proxy_client::IncomingMessage;
use ap_proxy_protocol::{IdentityFingerprint, RendezvousCode};
use base64::{Engine, engine::general_purpose::STANDARD};
use rand::RngCore;

use crate::compat::{now_millis, timeout};
use crate::proxy::ProxyClient;
use tokio::sync::{mpsc, oneshot};
use tracing::{debug, warn};

use super::notify;
use crate::traits::{ConnectionInfo, ConnectionStore, ConnectionUpdate, IdentityProvider};
use crate::{
    error::ClientError,
    types::{
        CredentialData, CredentialQuery, CredentialRequestPayload, CredentialResponsePayload,
        ProtocolMessage,
    },
};

const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120);

// =============================================================================
// Public types: Notifications (fire-and-forget) and Requests (with reply)
// =============================================================================

/// Fire-and-forget status updates emitted by the remote client.
#[derive(Debug, Clone)]
pub enum RemoteClientNotification {
    /// Connecting to the proxy server
    Connecting,
    /// Successfully connected to the proxy
    Connected {
        /// The device's identity fingerprint (hex-encoded)
        fingerprint: IdentityFingerprint,
    },
    /// Reconnecting to an existing session
    ReconnectingToSession {
        /// The fingerprint being reconnected to
        fingerprint: IdentityFingerprint,
    },
    /// Rendezvous code resolution starting
    RendezvousResolving {
        /// The rendezvous code being resolved
        code: String,
    },
    /// Rendezvous code resolved to fingerprint
    RendezvousResolved {
        /// The resolved identity fingerprint
        fingerprint: IdentityFingerprint,
    },
    /// Using PSK mode for connection
    PskMode {
        /// The fingerprint being connected to
        fingerprint: IdentityFingerprint,
    },
    /// Noise handshake starting
    HandshakeStart,
    /// Noise handshake progress
    HandshakeProgress {
        /// Progress message
        message: String,
    },
    /// Noise handshake complete
    HandshakeComplete,
    /// Handshake fingerprint (informational — for PSK or non-verified connections)
    HandshakeFingerprint {
        /// The 6-character hex fingerprint
        fingerprint: String,
    },
    /// User verified the fingerprint
    FingerprintVerified,
    /// User rejected the fingerprint
    FingerprintRejected {
        /// Reason for rejection
        reason: String,
    },
    /// Client is ready for credential requests
    Ready {
        /// Whether credentials can be requested
        can_request_credentials: bool,
    },
    /// Credential request was sent
    CredentialRequestSent {
        /// The query used for the request
        query: CredentialQuery,
    },
    /// Credential was received
    CredentialReceived {
        /// The credential data
        credential: CredentialData,
    },
    /// An error occurred
    Error {
        /// Error message
        message: String,
        /// Context where error occurred
        context: Option<String>,
    },
    /// Client was disconnected
    Disconnected {
        /// Reason for disconnection
        reason: Option<String>,
    },
}

/// Reply for fingerprint verification requests.
#[derive(Debug)]
pub struct RemoteClientFingerprintReply {
    /// Whether user approved the fingerprint
    pub approved: bool,
}

/// Requests that require a caller response, carrying a oneshot reply channel.
#[derive(Debug)]
pub enum RemoteClientRequest {
    /// Handshake fingerprint requires verification.
    VerifyFingerprint {
        /// The 6-character hex fingerprint for visual verification
        fingerprint: String,
        /// Channel to send the verification reply
        reply: oneshot::Sender<RemoteClientFingerprintReply>,
    },
}

// =============================================================================
// Command channel for RemoteClient handle → event loop communication
// =============================================================================

/// Commands sent from a `RemoteClient` handle to the running event loop.
enum RemoteClientCommand {
    PairWithHandshake {
        rendezvous_code: String,
        verify_fingerprint: bool,
        reply: oneshot::Sender<Result<IdentityFingerprint, ClientError>>,
    },
    PairWithPsk {
        psk: Psk,
        remote_fingerprint: IdentityFingerprint,
        reply: oneshot::Sender<Result<(), ClientError>>,
    },
    LoadCachedConnection {
        remote_fingerprint: IdentityFingerprint,
        reply: oneshot::Sender<Result<(), ClientError>>,
    },
    RequestCredential {
        query: CredentialQuery,
        credential_timeout: Option<Duration>,
        reply: oneshot::Sender<Result<CredentialData, ClientError>>,
    },
    ListConnections {
        reply: oneshot::Sender<Vec<ConnectionInfo>>,
    },
    HasConnection {
        fingerprint: IdentityFingerprint,
        reply: oneshot::Sender<bool>,
    },
}

// =============================================================================
// Handle — cloneable, Send, all methods take &self
// =============================================================================

/// A cloneable handle for controlling the remote client.
///
/// Obtained from [`RemoteClient::connect()`], which authenticates with the proxy,
/// spawns the event loop internally, and returns this handle. All methods
/// communicate with the event loop through an internal command channel.
///
/// `Clone` and `Send` — share freely across tasks and threads.
/// Dropping all handles shuts down the event loop and disconnects from the proxy.
/// Handle returned by [`RemoteClient::connect()`] containing the client and its
/// notification/request channels.
pub struct RemoteClientHandle {
    pub client: RemoteClient,
    pub notifications: mpsc::Receiver<RemoteClientNotification>,
    pub requests: mpsc::Receiver<RemoteClientRequest>,
}

#[derive(Clone)]
pub struct RemoteClient {
    command_tx: mpsc::Sender<RemoteClientCommand>,
}

impl RemoteClient {
    /// Connect to the proxy server, spawn the event loop, and return a handle.
    ///
    /// This is the single entry point. After `connect()` returns, the client is
    /// authenticated with the proxy and ready for pairing. Use one of the pairing
    /// methods to establish a secure channel:
    /// - [`pair_with_handshake()`](Self::pair_with_handshake) for rendezvous-based pairing
    /// - [`pair_with_psk()`](Self::pair_with_psk) for PSK-based pairing
    /// - [`load_cached_connection()`](Self::load_cached_connection) for reconnecting with a cached connection
    pub async fn connect(
        identity_provider: Box<dyn IdentityProvider>,
        connection_store: Box<dyn ConnectionStore>,
        mut proxy_client: Box<dyn ProxyClient>,
    ) -> Result<RemoteClientHandle, ClientError> {
        let identity_keypair = identity_provider.identity().await;
        let own_fingerprint = identity_keypair.identity().fingerprint();

        debug!("Connecting to proxy with identity {:?}", own_fingerprint);

        let (notification_tx, notification_rx) = mpsc::channel(32);
        let (request_tx, request_rx) = mpsc::channel(32);

        notify!(notification_tx, RemoteClientNotification::Connecting);

        let incoming_rx = proxy_client.connect(identity_keypair).await?;

        notify!(
            notification_tx,
            RemoteClientNotification::Connected {
                fingerprint: own_fingerprint,
            }
        );

        debug!("Connected to proxy successfully");

        // Create command channel
        let (command_tx, command_rx) = mpsc::channel(32);

        // Build inner state
        let inner = RemoteClientInner {
            connection_store,
            proxy_client,
            transport: None,
            remote_fingerprint: None,
        };

        // Spawn the event loop — use spawn_local on WASM (no Tokio runtime)
        #[cfg(target_arch = "wasm32")]
        wasm_bindgen_futures::spawn_local(inner.run_event_loop(
            incoming_rx,
            command_rx,
            notification_tx,
            request_tx,
        ));
        #[cfg(not(target_arch = "wasm32"))]
        tokio::spawn(inner.run_event_loop(incoming_rx, command_rx, notification_tx, request_tx));

        Ok(RemoteClientHandle {
            client: Self { command_tx },
            notifications: notification_rx,
            requests: request_rx,
        })
    }

    /// Pair with a remote device using a rendezvous code.
    ///
    /// Resolves the rendezvous code to a fingerprint, performs the Noise handshake,
    /// and optionally waits for user fingerprint verification. If `verify_fingerprint`
    /// is true, a [`RemoteClientRequest::VerifyFingerprint`] will be sent on the
    /// request channel and must be answered before this method returns.
    pub async fn pair_with_handshake(
        &self,
        rendezvous_code: String,
        verify_fingerprint: bool,
    ) -> Result<IdentityFingerprint, ClientError> {
        let (tx, rx) = oneshot::channel();
        self.command_tx
            .send(RemoteClientCommand::PairWithHandshake {
                rendezvous_code,
                verify_fingerprint,
                reply: tx,
            })
            .await
            .map_err(|_| ClientError::ChannelClosed)?;
        rx.await.map_err(|_| ClientError::ChannelClosed)?
    }

    /// Pair with a remote device using a pre-shared key.
    ///
    /// Uses the PSK for authentication, skipping fingerprint verification
    /// since trust is established through the PSK.
    pub async fn pair_with_psk(
        &self,
        psk: Psk,
        remote_fingerprint: IdentityFingerprint,
    ) -> Result<(), ClientError> {
        let (tx, rx) = oneshot::channel();
        self.command_tx
            .send(RemoteClientCommand::PairWithPsk {
                psk,
                remote_fingerprint,
                reply: tx,
            })
            .await
            .map_err(|_| ClientError::ChannelClosed)?;
        rx.await.map_err(|_| ClientError::ChannelClosed)?
    }

    /// Reconnect to a remote device using a cached connection.
    ///
    /// Verifies the connection exists in the connection store and reconnects
    /// without requiring fingerprint verification.
    pub async fn load_cached_connection(
        &self,
        remote_fingerprint: IdentityFingerprint,
    ) -> Result<(), ClientError> {
        let (tx, rx) = oneshot::channel();
        self.command_tx
            .send(RemoteClientCommand::LoadCachedConnection {
                remote_fingerprint,
                reply: tx,
            })
            .await
            .map_err(|_| ClientError::ChannelClosed)?;
        rx.await.map_err(|_| ClientError::ChannelClosed)?
    }

    /// Request a credential over the secure channel.
    ///
    /// An optional `timeout` overrides the default credential response timeout
    /// (120 seconds). Pass `None` to use the default.
    pub async fn request_credential(
        &self,
        query: &CredentialQuery,
        timeout: Option<Duration>,
    ) -> Result<CredentialData, ClientError> {
        let (tx, rx) = oneshot::channel();
        self.command_tx
            .send(RemoteClientCommand::RequestCredential {
                query: query.clone(),
                credential_timeout: timeout,
                reply: tx,
            })
            .await
            .map_err(|_| ClientError::ChannelClosed)?;
        rx.await.map_err(|_| ClientError::ChannelClosed)?
    }

    /// List all cached connections.
    pub async fn list_connections(&self) -> Result<Vec<ConnectionInfo>, ClientError> {
        let (tx, rx) = oneshot::channel();
        self.command_tx
            .send(RemoteClientCommand::ListConnections { reply: tx })
            .await
            .map_err(|_| ClientError::ChannelClosed)?;
        rx.await.map_err(|_| ClientError::ChannelClosed)
    }

    /// Check if a connection exists for a fingerprint.
    pub async fn has_connection(
        &self,
        fingerprint: IdentityFingerprint,
    ) -> Result<bool, ClientError> {
        let (tx, rx) = oneshot::channel();
        self.command_tx
            .send(RemoteClientCommand::HasConnection {
                fingerprint,
                reply: tx,
            })
            .await
            .map_err(|_| ClientError::ChannelClosed)?;
        rx.await.map_err(|_| ClientError::ChannelClosed)
    }
}

// =============================================================================
// Internal state — lives inside the spawned event loop task
// =============================================================================

/// All mutable state for the remote client, owned by the spawned event loop task.
struct RemoteClientInner {
    connection_store: Box<dyn ConnectionStore>,
    proxy_client: Box<dyn ProxyClient>,
    transport: Option<MultiDeviceTransport>,
    remote_fingerprint: Option<IdentityFingerprint>,
}

impl RemoteClientInner {
    /// Run the main event loop (consumes self).
    async fn run_event_loop(
        mut self,
        mut incoming_rx: mpsc::UnboundedReceiver<IncomingMessage>,
        mut command_rx: mpsc::Receiver<RemoteClientCommand>,
        notification_tx: mpsc::Sender<RemoteClientNotification>,
        request_tx: mpsc::Sender<RemoteClientRequest>,
    ) {
        loop {
            tokio::select! {
                msg = incoming_rx.recv() => {
                    match msg {
                        Some(_) => {
                            // Stray message while idle — ignore
                            debug!("Received message while idle");
                        }
                        None => {
                            // Proxy disconnected
                            notify!(notification_tx, RemoteClientNotification::Disconnected {
                                reason: Some("Proxy connection closed".to_string()),
                            });
                            return;
                        }
                    }
                }
                cmd = command_rx.recv() => {
                    match cmd {
                        Some(cmd) => {
                            self.handle_command(
                                cmd,
                                &mut incoming_rx,
                                &notification_tx,
                                &request_tx,
                            ).await;
                        }
                        None => {
                            // All handles dropped — shut down
                            debug!("All RemoteClient handles dropped, shutting down event loop");
                            self.proxy_client.disconnect().await.ok();
                            return;
                        }
                    }
                }
            }
        }
    }

    /// Dispatch a command from the handle.
    async fn handle_command(
        &mut self,
        cmd: RemoteClientCommand,
        incoming_rx: &mut mpsc::UnboundedReceiver<IncomingMessage>,
        notification_tx: &mpsc::Sender<RemoteClientNotification>,
        request_tx: &mpsc::Sender<RemoteClientRequest>,
    ) {
        match cmd {
            RemoteClientCommand::PairWithHandshake {
                rendezvous_code,
                verify_fingerprint,
                reply,
            } => {
                let result = self
                    .do_pair_with_handshake(
                        rendezvous_code,
                        verify_fingerprint,
                        incoming_rx,
                        notification_tx,
                        request_tx,
                    )
                    .await;
                let _ = reply.send(result);
            }
            RemoteClientCommand::PairWithPsk {
                psk,
                remote_fingerprint,
                reply,
            } => {
                let result = self
                    .do_pair_with_psk(psk, remote_fingerprint, incoming_rx, notification_tx)
                    .await;
                let _ = reply.send(result);
            }
            RemoteClientCommand::LoadCachedConnection {
                remote_fingerprint,
                reply,
            } => {
                let result = self
                    .do_load_cached_connection(remote_fingerprint, notification_tx)
                    .await;
                let _ = reply.send(result);
            }
            RemoteClientCommand::RequestCredential {
                query,
                credential_timeout,
                reply,
            } => {
                let result = self
                    .do_request_credential(query, credential_timeout, incoming_rx, notification_tx)
                    .await;
                let _ = reply.send(result);
            }
            RemoteClientCommand::ListConnections { reply } => {
                let connections = self.connection_store.list().await;
                let _ = reply.send(connections);
            }
            RemoteClientCommand::HasConnection { fingerprint, reply } => {
                let has = self.connection_store.get(&fingerprint).await.is_some();
                let _ = reply.send(has);
            }
        }
    }

    // ── Pairing: Rendezvous handshake ────────────────────────────────

    async fn do_pair_with_handshake(
        &mut self,
        rendezvous_code: String,
        verify_fingerprint: bool,
        incoming_rx: &mut mpsc::UnboundedReceiver<IncomingMessage>,
        notification_tx: &mpsc::Sender<RemoteClientNotification>,
        request_tx: &mpsc::Sender<RemoteClientRequest>,
    ) -> Result<IdentityFingerprint, ClientError> {
        // Resolve rendezvous code to fingerprint
        notify!(
            notification_tx,
            RemoteClientNotification::RendezvousResolving {
                code: rendezvous_code.clone(),
            }
        );

        let remote_fingerprint =
            Self::resolve_rendezvous(self.proxy_client.as_ref(), incoming_rx, &rendezvous_code)
                .await?;

        notify!(
            notification_tx,
            RemoteClientNotification::RendezvousResolved {
                fingerprint: remote_fingerprint,
            }
        );

        // Perform Noise handshake (no PSK)
        notify!(notification_tx, RemoteClientNotification::HandshakeStart);

        let (transport, fingerprint_str) = Self::perform_handshake(
            self.proxy_client.as_ref(),
            incoming_rx,
            remote_fingerprint,
            None,
        )
        .await?;

        notify!(notification_tx, RemoteClientNotification::HandshakeComplete);

        // Always emit fingerprint (informational or for verification)
        notify!(
            notification_tx,
            RemoteClientNotification::HandshakeFingerprint {
                fingerprint: fingerprint_str.clone(),
            }
        );

        if verify_fingerprint {
            // Send verification request via request channel
            let (fp_tx, fp_rx) = oneshot::channel();
            if request_tx.capacity() == 0 {
                warn!("Request channel full, waiting for consumer to drain");
            }
            request_tx
                .send(RemoteClientRequest::VerifyFingerprint {
                    fingerprint: fingerprint_str,
                    reply: fp_tx,
                })
                .await
                .map_err(|_| ClientError::ChannelClosed)?;

            // Wait for user verification (60s timeout)
            match timeout(Duration::from_secs(60), fp_rx).await {
                Ok(Ok(RemoteClientFingerprintReply { approved: true })) => {
                    notify!(
                        notification_tx,
                        RemoteClientNotification::FingerprintVerified
                    );
                }
                Ok(Ok(RemoteClientFingerprintReply { approved: false })) => {
                    self.proxy_client.disconnect().await.ok();
                    notify!(
                        notification_tx,
                        RemoteClientNotification::FingerprintRejected {
                            reason: "User rejected fingerprint verification".to_string(),
                        }
                    );
                    return Err(ClientError::FingerprintRejected);
                }
                Ok(Err(_)) => {
                    return Err(ClientError::ChannelClosed);
                }
                Err(_) => {
                    self.proxy_client.disconnect().await.ok();
                    return Err(ClientError::Timeout(
                        "Fingerprint verification timeout".to_string(),
                    ));
                }
            }
        }

        // Finalize connection
        self.finalize_pairing(transport, remote_fingerprint, notification_tx)
            .await?;

        Ok(remote_fingerprint)
    }

    // ── Pairing: PSK ─────────────────────────────────────────────────

    async fn do_pair_with_psk(
        &mut self,
        psk: Psk,
        remote_fingerprint: IdentityFingerprint,
        incoming_rx: &mut mpsc::UnboundedReceiver<IncomingMessage>,
        notification_tx: &mpsc::Sender<RemoteClientNotification>,
    ) -> Result<(), ClientError> {
        notify!(
            notification_tx,
            RemoteClientNotification::PskMode {
                fingerprint: remote_fingerprint,
            }
        );

        // Perform Noise handshake with PSK
        notify!(notification_tx, RemoteClientNotification::HandshakeStart);

        let (transport, _fingerprint_str) = Self::perform_handshake(
            self.proxy_client.as_ref(),
            incoming_rx,
            remote_fingerprint,
            Some(psk),
        )
        .await?;

        notify!(notification_tx, RemoteClientNotification::HandshakeComplete);

        // Skip fingerprint verification (trust via PSK)
        notify!(
            notification_tx,
            RemoteClientNotification::FingerprintVerified
        );

        // Finalize connection
        self.finalize_pairing(transport, remote_fingerprint, notification_tx)
            .await?;

        Ok(())
    }

    // ── Cached session reconnection ──────────────────────────────────

    async fn do_load_cached_connection(
        &mut self,
        remote_fingerprint: IdentityFingerprint,
        notification_tx: &mpsc::Sender<RemoteClientNotification>,
    ) -> Result<(), ClientError> {
        let connection = self
            .connection_store
            .get(&remote_fingerprint)
            .await
            .ok_or(ClientError::ConnectionNotFound)?;

        let transport = connection
            .transport_state
            .ok_or(ClientError::ConnectionNotFound)?;

        notify!(
            notification_tx,
            RemoteClientNotification::ReconnectingToSession {
                fingerprint: remote_fingerprint,
            }
        );

        notify!(notification_tx, RemoteClientNotification::HandshakeComplete);

        // Skip fingerprint verification (already trusted)
        notify!(
            notification_tx,
            RemoteClientNotification::FingerprintVerified
        );

        // Update last_connected_at
        self.connection_store
            .update(ConnectionUpdate {
                fingerprint: remote_fingerprint,
                last_connected_at: crate::compat::now_seconds(),
            })
            .await?;

        self.transport = Some(transport);
        self.remote_fingerprint = Some(remote_fingerprint);

        notify!(
            notification_tx,
            RemoteClientNotification::Ready {
                can_request_credentials: true,
            }
        );

        debug!("Reconnected to cached connection");
        Ok(())
    }

    // ── Shared pairing finalization ──────────────────────────────────

    async fn finalize_pairing(
        &mut self,
        transport: MultiDeviceTransport,
        remote_fingerprint: IdentityFingerprint,
        notification_tx: &mpsc::Sender<RemoteClientNotification>,
    ) -> Result<(), ClientError> {
        let now = crate::compat::now_seconds();
        self.connection_store
            .save(ConnectionInfo {
                fingerprint: remote_fingerprint,
                name: None,
                cached_at: now,
                last_connected_at: now,
                transport_state: Some(transport.clone()),
            })
            .await?;

        // Store transport and remote fingerprint
        self.transport = Some(transport);
        self.remote_fingerprint = Some(remote_fingerprint);

        // Emit Ready event
        notify!(
            notification_tx,
            RemoteClientNotification::Ready {
                can_request_credentials: true,
            }
        );

        debug!("Connection established successfully");
        Ok(())
    }

    // ── Credential request ───────────────────────────────────────────

    async fn do_request_credential(
        &mut self,
        query: CredentialQuery,
        credential_timeout: Option<Duration>,
        incoming_rx: &mut mpsc::UnboundedReceiver<IncomingMessage>,
        notification_tx: &mpsc::Sender<RemoteClientNotification>,
    ) -> Result<CredentialData, ClientError> {
        let remote_fingerprint = self.remote_fingerprint.ok_or(ClientError::NotInitialized)?;

        // Sliced string is a UUID and isn't going to contain wide chars
        #[allow(clippy::string_slice)]
        let request_id = format!("req-{}-{}", now_millis(), &uuid_v4()[..8]);

        debug!("Requesting credential for query: {:?}", query);

        // Create and encrypt request
        let request = CredentialRequestPayload {
            request_type: "credential_request".to_string(),
            query: query.clone(),
            timestamp: now_millis(),
            request_id: request_id.clone(),
        };

        let request_json = serde_json::to_string(&request)?;

        let encrypted_data = {
            let transport = self
                .transport
                .as_mut()
                .ok_or(ClientError::SecureChannelNotEstablished)?;
            let encrypted_packet = transport
                .encrypt(request_json.as_bytes())
                .map_err(|e| ClientError::NoiseProtocol(e.to_string()))?;
            STANDARD.encode(encrypted_packet.encode())
        };

        let msg = ProtocolMessage::CredentialRequest {
            encrypted: encrypted_data,
        };

        // Send via proxy
        let msg_json = serde_json::to_string(&msg)?;
        self.proxy_client
            .send_to(remote_fingerprint, msg_json.into_bytes())
            .await?;

        // Emit event
        notify!(
            notification_tx,
            RemoteClientNotification::CredentialRequestSent {
                query: query.clone(),
            }
        );

        // Wait for matching response inline
        let effective_timeout = credential_timeout.unwrap_or(DEFAULT_REQUEST_TIMEOUT);
        match timeout(
            effective_timeout,
            self.receive_credential_response(&request_id, incoming_rx, notification_tx),
        )
        .await
        {
            Ok(result) => result,
            Err(_) => Err(ClientError::Timeout(format!(
                "Timeout waiting for credential response for query: {query:?}"
            ))),
        }
    }

    /// Wait for a credential response matching the given request_id.
    ///
    /// Stale responses from previous requests (e.g. duplicate multi-device
    /// responses) are decrypted, logged, and silently discarded.
    async fn receive_credential_response(
        &mut self,
        request_id: &str,
        incoming_rx: &mut mpsc::UnboundedReceiver<IncomingMessage>,
        notification_tx: &mpsc::Sender<RemoteClientNotification>,
    ) -> Result<CredentialData, ClientError> {
        loop {
            match incoming_rx.recv().await {
                Some(IncomingMessage::Send { payload, .. }) => {
                    if let Ok(text) = String::from_utf8(payload)
                        && let Ok(ProtocolMessage::CredentialResponse { encrypted }) =
                            serde_json::from_str::<ProtocolMessage>(&text)
                    {
                        match self
                            .decrypt_credential_response(&encrypted, request_id, notification_tx)
                            .await
                        {
                            Ok(credential) => return Ok(credential),
                            Err(ClientError::CredentialRequestFailed(ref msg))
                                if msg.contains("request_id mismatch") =>
                            {
                                // Stale response from a previous request — skip it
                                debug!("Skipping stale credential response: {msg}");
                                continue;
                            }
                            Err(e) => return Err(e),
                        }
                    }
                }
                Some(_) => {
                    // Non-Send messages (RendezvousInfo, IdentityInfo) — ignore
                }
                None => {
                    return Err(ClientError::ChannelClosed);
                }
            }
        }
    }

    /// Decrypt and validate a credential response.
    async fn decrypt_credential_response(
        &mut self,
        encrypted: &str,
        request_id: &str,
        notification_tx: &mpsc::Sender<RemoteClientNotification>,
    ) -> Result<CredentialData, ClientError> {
        let encrypted_bytes = STANDARD
            .decode(encrypted)
            .map_err(|e| ClientError::Serialization(format!("Invalid base64: {e}")))?;

        let packet = ap_noise::TransportPacket::decode(&encrypted_bytes)
            .map_err(|e| ClientError::NoiseProtocol(format!("Invalid packet: {e}")))?;

        let transport = self
            .transport
            .as_mut()
            .ok_or(ClientError::SecureChannelNotEstablished)?;

        let decrypted = transport
            .decrypt(&packet)
            .map_err(|e| ClientError::NoiseProtocol(e.to_string()))?;

        let response: CredentialResponsePayload = serde_json::from_slice(&decrypted)?;

        // Verify request_id matches
        if response.request_id.as_deref() != Some(request_id) {
            warn!(
                "Ignoring response with mismatched request_id: {:?}",
                response.request_id
            );
            return Err(ClientError::CredentialRequestFailed(
                "Response request_id mismatch".to_string(),
            ));
        }

        if let Some(error) = response.error {
            return Err(ClientError::CredentialRequestFailed(error));
        }

        if let Some(credential) = response.credential {
            notify!(
                notification_tx,
                RemoteClientNotification::CredentialReceived {
                    credential: credential.clone(),
                }
            );
            Ok(credential)
        } else {
            Err(ClientError::CredentialRequestFailed(
                "Response contains neither credential nor error".to_string(),
            ))
        }
    }

    // ── Handshake helpers (associated functions) ─────────────────────

    /// Resolve rendezvous code to identity fingerprint.
    async fn resolve_rendezvous(
        proxy_client: &dyn ProxyClient,
        incoming_rx: &mut mpsc::UnboundedReceiver<IncomingMessage>,
        rendezvous_code: &str,
    ) -> Result<IdentityFingerprint, ClientError> {
        // Send GetIdentity request
        proxy_client
            .request_identity(RendezvousCode::from_string(rendezvous_code.to_string()))
            .await
            .map_err(|e| ClientError::RendezvousResolutionFailed(e.to_string()))?;

        // Wait for IdentityInfo response with timeout
        let timeout_duration = Duration::from_secs(10);
        match timeout(timeout_duration, async {
            while let Some(msg) = incoming_rx.recv().await {
                if let IncomingMessage::IdentityInfo { fingerprint, .. } = msg {
                    return Some(fingerprint);
                }
            }
            None
        })
        .await
        {
            Ok(Some(fingerprint)) => Ok(fingerprint),
            Ok(None) => Err(ClientError::RendezvousResolutionFailed(
                "Connection closed while waiting for identity response".to_string(),
            )),
            Err(_) => Err(ClientError::RendezvousResolutionFailed(
                "Timeout waiting for identity response. The rendezvous code may be invalid, expired, or the target client may be disconnected.".to_string(),
            )),
        }
    }

    /// Perform Noise handshake as initiator.
    async fn perform_handshake(
        proxy_client: &dyn ProxyClient,
        incoming_rx: &mut mpsc::UnboundedReceiver<IncomingMessage>,
        remote_fingerprint: IdentityFingerprint,
        psk: Option<Psk>,
    ) -> Result<(MultiDeviceTransport, String), ClientError> {
        // Compute PSK ID before moving the PSK into the handshake
        let psk_id = psk.as_ref().map(|p| p.id());

        // Create initiator handshake (with or without PSK)
        let mut handshake = if let Some(psk) = psk {
            InitiatorHandshake::with_psk(psk)
        } else {
            InitiatorHandshake::new()
        };

        // Generate handshake init
        let init_packet = handshake.send_start()?;

        // Send HandshakeInit message
        let msg = ProtocolMessage::HandshakeInit {
            data: STANDARD.encode(init_packet.encode()?),
            ciphersuite: format!("{:?}", handshake.ciphersuite()),
            psk_id,
        };

        let msg_json = serde_json::to_string(&msg)?;
        proxy_client
            .send_to(remote_fingerprint, msg_json.into_bytes())
            .await?;

        debug!("Sent handshake init");

        // Wait for HandshakeResponse
        let response_timeout = Duration::from_secs(10);
        let response: String = timeout(response_timeout, async {
            loop {
                if let Some(incoming) = incoming_rx.recv().await {
                    match incoming {
                        IncomingMessage::Send { payload, .. } => {
                            // Try to parse as ProtocolMessage
                            if let Ok(text) = String::from_utf8(payload)
                                && let Ok(ProtocolMessage::HandshakeResponse { data, .. }) =
                                    serde_json::from_str::<ProtocolMessage>(&text)
                            {
                                return Ok::<String, ClientError>(data);
                            }
                        }
                        _ => continue,
                    }
                }
            }
        })
        .await
        .map_err(|_| ClientError::Timeout("Waiting for handshake response".to_string()))??;

        // Decode and process response
        let response_bytes = STANDARD
            .decode(&response)
            .map_err(|e| ClientError::Serialization(format!("Invalid base64: {e}")))?;

        let response_packet = ap_noise::HandshakePacket::decode(&response_bytes)
            .map_err(|e| ClientError::NoiseProtocol(format!("Invalid packet: {e}")))?;

        // Complete handshake
        handshake.receive_finish(&response_packet)?;
        let (transport, fingerprint) = handshake.finalize()?;

        debug!("Handshake complete");
        Ok((transport, fingerprint.to_string()))
    }
}

fn uuid_v4() -> String {
    // Simple UUID v4 generation without external dependency
    let mut bytes = [0u8; 16];
    let mut rng = rand::thread_rng();
    rng.fill_bytes(&mut bytes);

    // Set version (4) and variant bits
    bytes[6] = (bytes[6] & 0x0f) | 0x40;
    bytes[8] = (bytes[8] & 0x3f) | 0x80;

    format!(
        "{:02x}{:02x}{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}",
        bytes[0],
        bytes[1],
        bytes[2],
        bytes[3],
        bytes[4],
        bytes[5],
        bytes[6],
        bytes[7],
        bytes[8],
        bytes[9],
        bytes[10],
        bytes[11],
        bytes[12],
        bytes[13],
        bytes[14],
        bytes[15]
    )
}