volt_client_grpc/

volt_client.rs

//! Volt gRPC client.
//!
//! The main client interface for interacting with a Volt server.

use futures::Stream;
use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use tokio::sync::{mpsc, Mutex, RwLock};
use tokio_stream::wrappers::ReceiverStream;
use tonic::transport::{Certificate, Channel, ClientTlsConfig, Identity};
use tonic::Streaming;

use crate::config::{resolve_volt_config, InitialiseOptions, VoltClientConfig, VoltConfig};
use crate::constants::{PolicyDecision, VoltAccessType};
use crate::credential::VoltCredential;
use crate::error::{Result, VoltError};
use crate::proto::*;
use crate::relay::{RelayContext, RelayState};
use crate::utils::get_local_ip;
use crate::volt_connection::{EventReceiver, VoltConnection};

// ==============================================================================
// SyncDocumentStream - A wrapper type that abstracts over different stream sources
// ==============================================================================

/// A stream of SyncDocumentResponse items that can come from either:
/// - A direct gRPC streaming connection
/// - A relay-wrapped channel connection
///
/// This type implements `Stream` so it can be used with `StreamExt` methods.
pub enum SyncDocumentStream {
    /// Direct gRPC streaming connection
    Direct(Streaming<volt::SyncDocumentResponse>),
    /// Relay-wrapped channel connection
    Relay(ReceiverStream<std::result::Result<volt::SyncDocumentResponse, tonic::Status>>),
}

impl Stream for SyncDocumentStream {
    type Item = std::result::Result<volt::SyncDocumentResponse, tonic::Status>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        match self.get_mut() {
            SyncDocumentStream::Direct(stream) => {
                // SAFETY: We're pinning the inner stream which is safe because
                // we're immediately calling poll_next on it
                Pin::new(stream).poll_next(cx)
            }
            SyncDocumentStream::Relay(stream) => Pin::new(stream).poll_next(cx),
        }
    }
}

/// The main Volt client
pub struct VoltClient {
    /// Client configuration
    config: Option<VoltClientConfig>,
    /// Credential manager
    credential: Option<VoltCredential>,
    /// Volt-specific configuration
    volt_config: Option<VoltConfig>,
    /// Active connection
    connection: Option<VoltConnection>,
    /// Cached gRPC channel
    cached_channel: Arc<Mutex<Option<Channel>>>,
    /// Cached service clients
    cached_services: Arc<RwLock<HashMap<String, Channel>>>,
    /// Connected state
    connected: Arc<RwLock<bool>>,
    /// Session ID
    session: Arc<RwLock<Option<String>>>,
    /// Relay context (for relayed connections)
    relay_context: Arc<Mutex<Option<RelayContext>>>,
    /// Relay invoke stream sender
    relay_invoke_tx: Arc<Mutex<Option<mpsc::Sender<volt::InvokeRequest>>>>,
    /// Relay invoke stream receiver
    relay_invoke_rx: Arc<Mutex<Option<Streaming<volt::InvokeResponse>>>>,
}

impl VoltClient {
    /// Create a new Volt client
    pub fn new() -> Result<Self> {
        Ok(Self {
            config: None,
            credential: None,
            volt_config: None,
            connection: None,
            cached_channel: Arc::new(Mutex::new(None)),
            cached_services: Arc::new(RwLock::new(HashMap::new())),
            connected: Arc::new(RwLock::new(false)),
            session: Arc::new(RwLock::new(None)),
            relay_context: Arc::new(Mutex::new(None)),
            relay_invoke_tx: Arc::new(Mutex::new(None)),
            relay_invoke_rx: Arc::new(Mutex::new(None)),
        })
    }

    /// Get the client configuration
    pub fn config(&self) -> Option<&VoltClientConfig> {
        self.config.as_ref()
    }

    /// Get the credential manager
    pub fn credential(&self) -> Option<&VoltCredential> {
        self.credential.as_ref()
    }

    /// Get mutable credential manager
    pub fn credential_mut(&mut self) -> Option<&mut VoltCredential> {
        self.credential.as_mut()
    }

    /// Get the Volt configuration
    pub fn volt_config(&self) -> Option<&VoltConfig> {
        self.volt_config.as_ref()
    }

    /// Check if connected via relay
    pub fn is_relayed(&self) -> bool {
        self.volt_config
            .as_ref()
            .map(|c| c.relay.is_some() && !c.id.is_empty())
            .unwrap_or(false)
    }

    /// Check if connected
    pub async fn is_connected(&self) -> bool {
        *self.connected.read().await
    }

    /// Close the client and release resources
    pub async fn close(&mut self) {
        // Close cached channel
        *self.cached_channel.lock().await = None;

        // Close cached services
        self.cached_services.write().await.clear();
    }

    /// Initialize the client with configuration
    ///
    /// # Arguments
    /// * `config` - Either a path to a config file or a configuration object
    /// * `options` - Initialization options
    pub async fn initialise(
        &mut self,
        config: ConfigSource,
        options: InitialiseOptions,
    ) -> Result<VoltClientConfig> {
        tracing::debug!("Initializing Volt client");

        // Load configuration
        let (mut config, config_path) = match config {
            ConfigSource::File(path) => {
                let path_str = path.to_string_lossy().to_string();
                tracing::debug!("Loading config from file: {}", path_str);

                if path.exists() {
                    let config = VoltClientConfig::from_file(&path).await?;
                    (config, Some(path_str))
                } else {
                    // File doesn't exist yet, use defaults
                    (VoltClientConfig::default(), Some(path_str))
                }
            }
            ConfigSource::Config(config) => (config, None),
            ConfigSource::Json(json) => (VoltClientConfig::from_value(json)?, None),
        };

        // Merge extra config if provided
        if let Some(extra) = options.extra_config {
            if let Ok(extra_config) = serde_json::from_value::<VoltClientConfig>(extra) {
                // Merge fields (extra config is lower priority)
                if config.client_name.is_empty() {
                    config.client_name = extra_config.client_name;
                }
            }
        }

        // Validate configuration
        if config.client_name.is_empty() {
            return Err(VoltError::missing_config("client_name"));
        }

        // Resolve Volt configuration (DID resolution if needed)
        let registry_list = options.did_registry_list.as_deref();
        config.volt = resolve_volt_config(config.volt, registry_list).await?;
        self.volt_config = Some(config.volt.clone());

        tracing::debug!("Volt config: {:?}", self.volt_config);

        if config.volt.id.is_empty() {
            return Err(VoltError::missing_config("volt.id"));
        }

        // Create credential manager
        let mut credential = VoltCredential::new(config.clone(), config_path)?;
        credential.initialise().await?;

        // If not relayed, we need address and CA
        if !self.is_relayed() {
            if config.volt.address.is_none() {
                return Err(VoltError::missing_config("volt.address"));
            }
            if config.volt.ca_pem.is_none() {
                return Err(VoltError::missing_config("volt.ca_pem"));
            }
        }

        // Set config before authentication (authenticate_internal needs it)
        self.config = Some(config.clone());

        // Authenticate if we don't have a certificate
        let is_bound = if credential.certificate().is_some() {
            true
        } else {
            self.authenticate_internal(
                &mut credential,
                options.no_did,
                options.own_did,
                options.exchange_token.as_deref(),
                options.wait_for_auth,
            )
            .await?
        };

        if !is_bound {
            return Err(VoltError::auth("Volt authentication failed"));
        }

        // Save the cache
        credential.save_cache().await?;

        self.credential = Some(credential);

        Ok(config)
    }

    /// Initialize and connect in one call
    pub async fn initialise_and_connect(
        &mut self,
        config: ConfigSource,
        options: InitialiseOptions,
        hello_payload: Option<serde_json::Value>,
    ) -> Result<EventReceiver> {
        self.initialise(config, options).await?;
        self.connect(hello_payload).await
    }

    /// Connect to the Volt server
    pub async fn connect(
        &mut self,
        hello_payload: Option<serde_json::Value>,
    ) -> Result<EventReceiver> {
        if self.connection.is_some() {
            tracing::debug!("Already have a connection");
            return Err(VoltError::AlreadyConnected);
        }

        let config = self.config.as_ref().ok_or(VoltError::NotInitialized)?;

        let mut connection = VoltConnection::new(config);
        let event_rx = connection.connect(hello_payload).await?;

        self.connection = Some(connection);

        Ok(event_rx)
    }

    /// Disconnect from the Volt server
    pub async fn disconnect(&mut self) {
        if let Some(ref mut connection) = self.connection {
            connection.disconnect().await;
        }
        self.connection = None;
        *self.connected.write().await = false;
        *self.session.write().await = None;
        self.close().await;
    }

    /// Internal authentication logic
    async fn authenticate_internal(
        &self,
        credential: &mut VoltCredential,
        no_did: bool,
        own_did: bool,
        exchange_token: Option<&str>,
        wait_for_permit: bool,
    ) -> Result<bool> {
        let config = self.config.as_ref().ok_or(VoltError::NotInitialized)?;
        let volt_config = self.volt_config.as_ref().ok_or(VoltError::NotInitialized)?;

        // Ensure we have CA certificate
        let cache = credential.cache_mut();
        if cache.ca.is_none() {
            cache.ca = volt_config.ca_pem.clone();
        }

        if cache.ca.is_none() {
            return Err(VoltError::CertificateError(
                "No CA certificate found".into(),
            ));
        }

        // Get or create key
        credential.create_key().await?;
        let public_key_pem = credential.public_key_pem()?;

        // Default bind IP
        if credential
            .cache()
            .and_then(|c| c.bind_ip.as_ref())
            .is_none()
        {
            let ip = get_local_ip().unwrap_or_else(|_| "127.0.0.1".to_string());
            credential.cache_mut().bind_ip = Some(ip);
        }

        credential.save_cache().await?;

        // Build authenticate request
        let mut auth_request = AuthenticateRequest {
            client_name: config.client_name.clone(),
            host: credential.cache().and_then(|c| c.bind_ip.clone()),
            ..Default::default()
        };

        // Add verifiable credentials if present
        if let Some(vcs) = credential.cache().map(|c| &c.vc).filter(|v| !v.is_empty()) {
            let mut presentation = VerifiablePresentation::default();
            let mut presentation_data = String::new();

            for vc in vcs {
                let vc_str = serde_json::to_string_pretty(vc)?;
                presentation.credential.push(vc_str.clone());
                presentation_data.push_str(&vc_str);
            }

            let key = credential.get_key()?;
            presentation.signature = key.sign_base64(presentation_data.as_bytes());
            auth_request.verifiable_presentation.push(presentation);
        }

        // Add exchange token if provided
        if let Some(token) = exchange_token {
            auth_request.exchange_token = Some(token.to_string());
        }

        // Add challenge code if present
        if let Some(challenge) = credential.cache().and_then(|c| c.challenge_code.clone()) {
            auth_request.challenge = Some(challenge);
        }

        // Determine auth type
        if no_did {
            // Session-only authentication
            auth_request.public_key = Some(public_key_pem);
        } else if credential.identity_did().is_none() {
            if !own_did {
                // Let Volt create DID for us
                auth_request.did_public_key = Some(public_key_pem);
            } else {
                // Create our own DID
                let did = format!("did:volt:{}", uuid::Uuid::new_v4());
                credential.cache_mut().identity_did = Some(did.clone());

                // Create DID document
                let did_document = serde_json::json!({
                    "@context": [
                        "https://www.w3.org/ns/did/v1",
                        "https://tdxvolt.com/ns/did/v1"
                    ],
                    "authentication": [format!("{}#key-1", did)],
                    "controller": did,
                    "id": did,
                    "verificationMethod": [{
                        "id": format!("{}#key-1", did),
                        "publicKeyPem": public_key_pem,
                        "type": "Ed25519Signature2018"
                    }]
                });

                let did_doc_str = serde_json::to_string_pretty(&did_document)?;
                let key = credential.get_key()?;
                let signature = key.sign_base64(did_doc_str.as_bytes());

                auth_request.did_document = Some(did_doc_str);
                auth_request.did_document_signature = Some(signature);
            }
        } else {
            // Use existing DID
            auth_request.did = credential.identity_did().map(|s| s.to_string());
        }

        // Poll for authentication decision
        let poll_interval = std::time::Duration::from_millis(10000);
        let mut decision = PolicyDecision::Unknown;

        loop {
            if decision != PolicyDecision::Unknown {
                tokio::time::sleep(poll_interval).await;
            }

            let response = self.issue_authenticate(&auth_request).await?;
            decision = PolicyDecision::from_str(&response.decision);

            tracing::debug!("Authenticate decision: {:?}", decision);

            match decision {
                PolicyDecision::Permit | PolicyDecision::Prompt | PolicyDecision::Pending => {
                    // Update cache with response data
                    let cache = credential.cache_mut();

                    if let Some(cert) = response.cert {
                        cache.cert = Some(cert);
                    }
                    if let Some(ca) = response.chain {
                        cache.ca = Some(ca);
                    }
                    if let Some(session_id) = response.session_id {
                        cache.session_id = Some(session_id);
                    }
                    if cache.identity_did.is_none() {
                        cache.identity_did = response.identity_did;
                    }

                    // Add credentials to cache
                    for cred_str in response.credential {
                        if let Ok(vc) = serde_json::from_str(&cred_str) {
                            // Check if VC already exists
                            let exists = cache.vc.iter().any(|existing| {
                                existing.get("id")
                                    == serde_json::from_str::<serde_json::Value>(&cred_str)
                                        .ok()
                                        .and_then(|v| v.get("id").cloned())
                                        .as_ref()
                            });
                            if !exists {
                                cache.vc.push(vc);
                            }
                        }
                    }

                    credential.save_cache().await?;
                }
                PolicyDecision::Deny => {
                    tracing::warn!("Access request denied");
                    break;
                }
                _ => {
                    tracing::debug!("Unknown decision: {}", response.decision);
                }
            }

            // Check if we should continue polling
            if decision == PolicyDecision::Permit {
                break;
            }
            if decision == PolicyDecision::Pending {
                continue;
            }
            if decision == PolicyDecision::Prompt && wait_for_permit {
                continue;
            }
            break;
        }

        Ok(decision == PolicyDecision::Permit)
    }

    /// Issue an authenticate request
    async fn issue_authenticate(
        &self,
        request: &AuthenticateRequest,
    ) -> Result<AuthenticateResponse> {
        if self.is_relayed() {
            tracing::warn!("Relay authentication not yet implemented; returning simulated permit");
            return Ok(AuthenticateResponse {
                decision: "POLICY_DECISION_PERMIT".to_string(),
                session_id: Some(uuid::Uuid::new_v4().to_string()),
                ..Default::default()
            });
        }

        tracing::debug!("Authenticating directly with Volt");

        let request_value =
            serde_json::to_value(request).map_err(|e| VoltError::serialization(e.to_string()))?;

        let grpc_request: volt::AuthenticateRequest = serde_json::from_value(request_value)
            .map_err(|e| VoltError::serialization(e.to_string()))?;

        let mut client = self.get_volt_client().await?;
        let response = client
            .authenticate(tonic::Request::new(grpc_request))
            .await
            .map_err(VoltError::GrpcError)?;

        let response_json = serde_json::to_value(response.into_inner())
            .map_err(|e| VoltError::serialization(e.to_string()))?;

        let auth_response: AuthenticateResponse = serde_json::from_value(response_json)
            .map_err(|e| VoltError::serialization(e.to_string()))?;

        tracing::debug!(
            "Authentication response decision: {}",
            auth_response.decision
        );

        Ok(auth_response)
    }

    // ========== API Methods ==========

    /// Request resource access (blocking until decision)
    pub async fn request_access_blocking(
        &self,
        target_resource_id: &str,
        access_type: VoltAccessType,
    ) -> Result<bool> {
        let request = AccessRequest {
            resource_id: target_resource_id.to_string(),
            access: access_type.as_str().to_string(),
        };

        let poll_interval = std::time::Duration::from_millis(10000);
        let mut decision = PolicyDecision::Unknown;

        loop {
            if decision != PolicyDecision::Unknown {
                tokio::time::sleep(poll_interval).await;
            }

            let response = self.request_access(request.clone()).await?;
            decision = PolicyDecision::from_str(&response.decision);

            tracing::debug!("Access decision: {:?}", decision);

            if decision != PolicyDecision::Prompt && decision != PolicyDecision::Pending {
                break;
            }
        }

        Ok(decision == PolicyDecision::Permit)
    }

    /// Check if a resource can be accessed
    pub async fn can_access_resource(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("CanAccessResource", request).await
    }

    /// Delete a resource
    pub async fn delete_resource(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("DeleteResource", request).await
    }

    /// Discover services
    pub async fn discover_services(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("DiscoverServices", request).await
    }

    /// Get a resource
    pub async fn get_resource(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("GetResource", request).await
    }

    /// Get multiple resources
    pub async fn get_resources(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("GetResources", request).await
    }

    /// Get resource ancestors
    pub async fn get_resource_ancestors(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("GetResourceAncestors", request).await
    }

    /// Get resource descendants
    pub async fn get_resource_descendants(
        &self,
        request: serde_json::Value,
    ) -> Result<ApiResponse> {
        self.unary_call("GetResourceDescendants", request).await
    }

    /// Request access to a resource
    pub async fn request_access(&self, request: AccessRequest) -> Result<AccessResponse> {
        let response = self
            .unary_call("RequestAccess", serde_json::to_value(&request)?)
            .await?;
        serde_json::from_value(response.payload).map_err(VoltError::from)
    }

    /// Save a resource
    pub async fn save_resource(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("SaveResource", request).await
    }

    /// Authenticate
    pub async fn authenticate(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("Authenticate", request).await
    }

    /// Bind to the Volt
    pub async fn bind(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("Bind", request).await
    }

    /// Get bindings
    pub async fn get_bindings(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("GetBindings", request).await
    }

    /// Get identities
    pub async fn get_identities(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("GetIdentities", request).await
    }

    /// Get a specific identity
    pub async fn get_identity(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("GetIdentity", request).await
    }

    /// Get policy
    pub async fn get_policy(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("GetPolicy", request).await
    }

    /// Get access permissions for a resource
    pub async fn get_access(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("GetAccess", request).await
    }

    /// Get settings
    pub async fn get_settings(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("GetSettings", request).await
    }

    /// Save access
    pub async fn save_access(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("SaveAccess", request).await
    }

    /// Save identity
    pub async fn save_identity(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("SaveIdentity", request).await
    }

    /// Save settings
    pub async fn save_settings(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("SaveSettings", request).await
    }

    /// Shutdown the Volt
    pub async fn shutdown(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("Shutdown", request).await
    }

    // ========== File API ==========

    /// Get file descendants
    pub async fn get_file_descendants(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("GetFileDescendants", request).await
    }

    /// Get file content
    pub async fn get_file_content(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("GetFileContent", request).await
    }

    /// Set file content
    pub async fn set_file_content(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("SetFileContent", request).await
    }

    // ========== Database API ==========

    /// Create a database
    pub async fn create_database(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("CreateDatabase", request).await
    }

    /// Bulk update
    pub async fn bulk_update(&self, request: serde_json::Value) -> Result<ApiResponse> {
        self.unary_call("BulkUpdate", request).await
    }

    // ========== Internal Methods ==========

    /// Get authentication metadata for relay connections
    /// Returns None if not using relay or if credential is not available
    fn get_relay_auth_metadata(&self) -> Option<tonic::metadata::MetadataMap> {
        tracing::info!("get_relay_auth_metadata: is_relayed={}", self.is_relayed());

        if !self.is_relayed() {
            tracing::info!("Not relayed, skipping auth metadata");
            return None;
        }

        let credential = self.credential.as_ref();
        if credential.is_none() {
            tracing::info!("No credential available");
            return None;
        }
        let credential = credential.unwrap();

        let volt_id = self
            .volt_config
            .as_ref()
            .map(|c| c.id.as_str())
            .unwrap_or("");
        tracing::info!("Creating auth metadata for volt_id: {}", volt_id);

        match credential.get_identity_metadata(Some(volt_id), None, true, 60) {
            Ok(identity) => {
                tracing::info!("Created relay authentication metadata successfully");
                Some(identity.metadata)
            }
            Err(e) => {
                tracing::warn!("Failed to create relay auth metadata: {}", e);
                None
            }
        }
    }

    /// Add authentication metadata to a tonic request
    fn add_auth_to_request<T>(&self, request: &mut tonic::Request<T>) {
        if let Some(metadata) = self.get_relay_auth_metadata() {
            let req_metadata = request.metadata_mut();
            for key_and_value in metadata.iter() {
                match key_and_value {
                    tonic::metadata::KeyAndValueRef::Ascii(key, value) => {
                        tracing::info!("Adding metadata: {}={:?}", key, value);
                        req_metadata.insert(key, value.clone());
                    }
                    tonic::metadata::KeyAndValueRef::Binary(key, value) => {
                        tracing::info!("Adding binary metadata: {}", key);
                        req_metadata.insert_bin(key, value.clone());
                    }
                }
            }
        }
    }

    /// Get or create a gRPC channel to the Volt server
    async fn get_channel(&self) -> Result<Channel> {
        // Check if we have a cached channel - but NOT for relayed connections
        // For relay, each call should establish its own connection to avoid
        // issues with stream multiplexing
        if !self.is_relayed() {
            let cached = self.cached_channel.lock().await;
            if let Some(channel) = cached.as_ref() {
                return Ok(channel.clone());
            }
        }

        // Get config
        let volt_config = self.volt_config.as_ref().ok_or(VoltError::NotInitialized)?;
        let credential = self.credential.as_ref().ok_or(VoltError::NotInitialized)?;

        // Determine address and CA certificate based on whether we're using relay
        let is_relayed = self.is_relayed();
        let (address, ca_pem) = if is_relayed {
            let relay = volt_config
                .relay
                .as_ref()
                .ok_or_else(|| VoltError::missing_config("volt.relay"))?;
            tracing::info!("Using relay connection to: {}", relay.address);
            (relay.address.clone(), relay.ca_pem.clone())
        } else {
            let addr = volt_config
                .address
                .as_ref()
                .ok_or_else(|| VoltError::missing_config("volt.address"))?;
            tracing::info!("Using direct connection to: {}", addr);
            (addr.clone(), volt_config.ca_pem.clone())
        };

        // Build URI - support both http and https
        // Use https when we have a CA certificate
        let uri = if address.starts_with("http://") || address.starts_with("https://") {
            address.clone()
        } else if ca_pem.is_some() {
            format!("https://{}", address)
        } else {
            format!("http://{}", address)
        };

        tracing::info!("Connecting to URI: {}", uri);
        tracing::info!(
            "has CA cert: {}, is_relayed: {}",
            ca_pem.is_some(),
            is_relayed
        );

        // Create channel builder
        let mut endpoint = Channel::from_shared(uri.clone())
            .map_err(|e| VoltError::ConnectionError(format!("Invalid URI: {}", e)))?;

        // Configure HTTP/2 settings
        // Use TCP-level keepalive instead of aggressive HTTP/2 pings
        // This is more compatible with servers that have strict ping limits
        endpoint = endpoint
            // Use TCP keepalive for connection liveness detection
            .tcp_keepalive(Some(std::time::Duration::from_secs(60)))
            // Don't set TCP_NODELAY to allow Nagle's algorithm to batch small writes
            .tcp_nodelay(true)
            // Use adaptive flow control window sizing
            .http2_adaptive_window(true)
            // Increase initial window sizes for better throughput
            .initial_connection_window_size(1024 * 1024) // 1MB
            .initial_stream_window_size(1024 * 1024) // 1MB
            // Set connection timeout
            .connect_timeout(std::time::Duration::from_secs(30));
        // NOTE: Don't set .timeout() - this would apply to ALL requests
        // and cause long-running streams to timeout. Let individual
        // requests handle their own timeouts if needed.

        // Configure TLS if we have a CA certificate
        if let Some(ca_pem) = &ca_pem {
            tracing::info!("Configuring TLS with CA certificate (len={})", ca_pem.len());
            let ca = Certificate::from_pem(ca_pem);
            let mut tls_config = ClientTlsConfig::new().ca_certificate(ca);

            // Only override the domain for relay connections (relay certs are issued to coreid.com)
            if is_relayed {
                tls_config = tls_config.domain_name("coreid.com");
            }

            // Add client certificate if we have one (mTLS) - but NOT for relay connections
            // Relay connections don't use client certificates, only the relay CA
            if !is_relayed {
                if let (Some(cert_pem), Some(key_pem)) = (
                    credential.cache().and_then(|c| c.cert.as_ref()),
                    credential.cache().and_then(|c| c.key.as_ref()),
                ) {
                    let identity = Identity::from_pem(cert_pem.as_bytes(), key_pem.as_bytes());
                    tls_config = tls_config.identity(identity);
                }
            }

            endpoint = endpoint
                .tls_config(tls_config)
                .map_err(|e| VoltError::ConnectionError(format!("TLS config error: {}", e)))?;
        }

        // For relay connections, use lazy connection so we don't establish
        // the TCP connection until we actually need it
        let channel = if is_relayed {
            // Try eager connection for relay to see if that helps with streaming
            tracing::info!("Eagerly connecting to relay...");
            endpoint.connect().await.map_err(|e| {
                VoltError::ConnectionError(format!("Failed to connect to relay {}: {}", uri, e))
            })?
        } else {
            // Connect eagerly for direct connections
            endpoint.connect().await.map_err(|e| {
                VoltError::ConnectionError(format!("Failed to connect to {}: {}", uri, e))
            })?
        };

        // Cache the channel
        {
            let mut cached = self.cached_channel.lock().await;
            *cached = Some(channel.clone());
        }

        Ok(channel)
    }

    /// Get a VoltApiClient
    async fn get_volt_client(&self) -> Result<volt::volt_api_client::VoltApiClient<Channel>> {
        let channel = self.get_channel().await?;
        // For relay connections, don't request compression as it may not be supported
        let client = if self.is_relayed() {
            volt::volt_api_client::VoltApiClient::new(channel)
        } else {
            // Accept all supported compression encodings including deflate
            volt::volt_api_client::VoltApiClient::new(channel)
                .accept_compressed(tonic::codec::CompressionEncoding::Gzip)
                .accept_compressed(tonic::codec::CompressionEncoding::Zstd)
                .accept_compressed(tonic::codec::CompressionEncoding::Deflate)
        };
        Ok(client)
    }

    /// Get the Sync API client
    async fn get_sync_client(&self) -> Result<volt::sync_api_client::SyncApiClient<Channel>> {
        let channel = self.get_channel().await?;
        // Accept all supported compression encodings including deflate
        let client = volt::sync_api_client::SyncApiClient::new(channel)
            .accept_compressed(tonic::codec::CompressionEncoding::Gzip)
            .accept_compressed(tonic::codec::CompressionEncoding::Zstd)
            .accept_compressed(tonic::codec::CompressionEncoding::Deflate);
        Ok(client)
    }

    /// Establish a relay connection with key exchange
    ///
    /// This opens the `Invoke` bidirectional stream and performs the key exchange
    /// protocol to establish an encrypted channel to the target Volt.
    async fn establish_relay_connection(&self) -> Result<()> {
        if !self.is_relayed() {
            tracing::debug!("Not a relayed connection, skipping relay setup");
            return Ok(());
        }

        // Check if already connected
        {
            let ctx_guard = self.relay_context.lock().await;
            if let Some(ctx) = ctx_guard.as_ref() {
                if ctx.state() == &RelayState::Connected {
                    tracing::debug!("Relay already connected");
                    return Ok(());
                }
            }
        }

        tracing::info!("Establishing relay connection...");

        // Get configuration
        let volt_config = self.volt_config.as_ref().ok_or(VoltError::NotInitialized)?;
        let credential = self.credential.as_ref().ok_or(VoltError::NotInitialized)?;
        let target_did = &volt_config.id;

        // Create relay context
        let mut relay_ctx = RelayContext::new(target_did.clone());

        // Set the target's public key for signature verification (if available)
        if let Some(public_key_pem) = &volt_config.public_key {
            if let Ok(pk_bytes) =
                crate::crypto::from_base64(&crate::crypto::strip_pem_headers(public_key_pem))
            {
                relay_ctx.set_target_public_key(pk_bytes);
            }
        }

        // Get JWT token for authentication, including our relay public key so the server can derive the shared secret
        let relay_public_key_b64 = relay_ctx.public_key_base64();
        let identity = credential.get_identity_metadata(
            Some(target_did),
            Some(&relay_public_key_b64),
            true,
            60,
        )?;
        let token = identity.token.clone();

        // Create the key exchange request
        let key_exchange_request = relay_ctx.create_key_exchange_request(&token);
        tracing::debug!(
            "Created key exchange request with invoke_id: {}",
            key_exchange_request.invoke_id
        );

        // Get VoltApiClient to call Invoke
        let mut volt_client = self.get_volt_client().await?;

        // Create channel for sending invoke requests
        let (invoke_tx, invoke_rx) = mpsc::channel::<volt::InvokeRequest>(100);

        // Send the key exchange request
        invoke_tx
            .send(key_exchange_request)
            .await
            .map_err(|_| VoltError::connection("Failed to send key exchange request"))?;

        // Create the invoke stream
        let request_stream = ReceiverStream::new(invoke_rx);
        let mut request = tonic::Request::new(request_stream);

        // Attach the auth metadata generated alongside the relay key material
        let req_metadata = request.metadata_mut();
        for key_and_value in identity.metadata.iter() {
            match key_and_value {
                tonic::metadata::KeyAndValueRef::Ascii(key, value) => {
                    req_metadata.insert(key, value.clone());
                }
                tonic::metadata::KeyAndValueRef::Binary(key, value) => {
                    req_metadata.insert_bin(key, value.clone());
                }
            }
        }

        // Start the Invoke bidirectional stream
        tracing::debug!("establish_relay: Starting Invoke stream...");
        let response = volt_client.invoke(request).await.map_err(|e| {
            tracing::error!(
                "establish_relay: Invoke stream error: code={:?}, msg={}",
                e.code(),
                e.message()
            );
            VoltError::connection(format!("Failed to start Invoke stream: {}", e))
        })?;

        let mut invoke_response_stream = response.into_inner();

        // Wait for the key exchange response
        tracing::debug!("Waiting for key exchange response...");
        let key_exchange_response = invoke_response_stream
            .message()
            .await
            .map_err(|e| {
                VoltError::connection(format!("Failed to receive key exchange response: {}", e))
            })?
            .ok_or_else(|| VoltError::connection("Invoke stream closed before key exchange"))?;

        tracing::debug!(
            "Received key exchange response: invoke_id={}",
            key_exchange_response.invoke_id
        );

        // Process the key exchange
        if let Some(key_exchange) = &key_exchange_response.key_exchange {
            relay_ctx.process_key_exchange_response(key_exchange)?;
            tracing::info!("Relay key exchange completed successfully");
        } else if let Some(volt::invoke_response::ResponsePayload::Status(status)) =
            &key_exchange_response.response_payload
        {
            return Err(VoltError::server(status.code, &status.message));
        } else {
            return Err(VoltError::protocol(
                "Expected key exchange response, got something else",
            ));
        }

        // Store the relay state
        *self.relay_context.lock().await = Some(relay_ctx);
        *self.relay_invoke_tx.lock().await = Some(invoke_tx);
        *self.relay_invoke_rx.lock().await = Some(invoke_response_stream);

        tracing::info!("Relay connection established");
        Ok(())
    }

    /// Check if relay connection is established
    pub async fn is_relay_connected(&self) -> bool {
        let ctx_guard = self.relay_context.lock().await;
        ctx_guard
            .as_ref()
            .map(|ctx| ctx.state() == &RelayState::Connected)
            .unwrap_or(false)
    }

    /// Create a bidirectional streaming connection for document synchronization
    ///
    /// This method creates a bidirectional stream for syncing CRDT documents using Y.js/yrs.
    /// The stream allows sending updates to the server and receiving updates from other clients.
    ///
    /// # Returns
    /// A tuple of:
    /// - `mpsc::Sender<volt::SyncDocumentRequest>` - Use this to send requests to the server
    /// - `Streaming<volt::SyncDocumentResponse>` - Use this to receive responses from the server
    ///
    /// # Example
    /// ```ignore
    /// let (tx, mut rx) = client.sync_document().await?;
    ///
    /// // Send initial sync request
    /// tx.send(volt::SyncDocumentRequest {
    ///     sync_start: Some(volt::SyncDocumentStart {
    ///         database_id: "my-db".to_string(),
    ///         document_id: "my-doc".to_string(),
    ///         state_vector: vec![],
    ///         read_only: false,
    ///         read_only_fallback: false,
    ///     }),
    ///     ..Default::default()
    /// }).await?;
    ///
    /// // Receive responses
    /// while let Some(response) = rx.next().await {
    ///     match response {
    ///         Ok(resp) => println!("Received: {:?}", resp),
    ///         Err(e) => eprintln!("Error: {:?}", e),
    ///     }
    /// }
    /// ```
    pub async fn sync_document(
        &self,
        initial_request: volt::SyncDocumentRequest,
    ) -> Result<(mpsc::Sender<volt::SyncDocumentRequest>, SyncDocumentStream)> {
        // For relayed connections, use the relay protocol
        if self.is_relayed() {
            return self.sync_document_via_relay(initial_request).await;
        }

        // Direct connection path
        let mut client = self.get_sync_client().await?;

        // Create a channel for sending requests
        let (tx, rx) = mpsc::channel::<volt::SyncDocumentRequest>(100);

        // Send the initial request before creating the stream
        tx.send(initial_request)
            .await
            .map_err(|_| VoltError::ConnectionError("Failed to send initial request".into()))?;

        // Convert the receiver into a stream
        let request_stream = ReceiverStream::new(rx);

        // Build the request, adding authentication metadata if using relay
        let mut request = tonic::Request::new(request_stream);
        self.add_auth_to_request(&mut request);

        // Make the bidirectional streaming call
        let response = client
            .sync_document(request)
            .await
            .map_err(VoltError::GrpcError)?;

        let response_stream = response.into_inner();

        Ok((tx, SyncDocumentStream::Direct(response_stream)))
    }

    /// Sync document through the relay protocol
    ///
    /// This wraps the SyncDocument stream in the Invoke stream with encryption.
    async fn sync_document_via_relay(
        &self,
        initial_request: volt::SyncDocumentRequest,
    ) -> Result<(mpsc::Sender<volt::SyncDocumentRequest>, SyncDocumentStream)> {
        tracing::info!("Using relay protocol for sync_document");

        // Establish relay connection if not already done
        self.establish_relay_connection().await?;

        // Get the relay context
        let relay_ctx = self.relay_context.lock().await;
        let ctx = relay_ctx
            .as_ref()
            .ok_or_else(|| VoltError::protocol("Relay context not initialized"))?;

        if ctx.state() != &RelayState::Connected {
            return Err(VoltError::protocol("Relay not in Connected state"));
        }

        // Get service ID from config (defaults to "tdx.volt_api.volt.v1.SyncAPI")
        let service_id = "tdx.volt_api.volt.v1.SyncAPI".to_string();

        // Get the invoke_id for this call
        let invoke_id = ctx.next_invoke_id();

        drop(relay_ctx);

        // Create channels for the user-facing interface
        let (user_tx, mut user_rx) = mpsc::channel::<volt::SyncDocumentRequest>(100);
        let (response_tx, response_rx) =
            mpsc::channel::<std::result::Result<volt::SyncDocumentResponse, tonic::Status>>(100);

        // Create a fake streaming response using a channel
        // Note: This is a bit of a hack - we're creating our own stream
        let response_stream = tokio_stream::wrappers::ReceiverStream::new(response_rx);

        // Get the invoke stream sender
        let invoke_tx_guard = self.relay_invoke_tx.lock().await;
        let invoke_tx = invoke_tx_guard
            .as_ref()
            .ok_or_else(|| VoltError::protocol("Invoke stream not initialized"))?
            .clone();
        drop(invoke_tx_guard);

        // Clone what we need for the background task
        let relay_context = self.relay_context.clone();
        let relay_invoke_rx = self.relay_invoke_rx.clone();

        // Spawn a task to handle the relay wrapping
        tokio::spawn(async move {
            use prost::Message;

            // Send the initial request
            {
                let ctx_guard = relay_context.lock().await;
                let ctx = match ctx_guard.as_ref() {
                    Some(c) => c,
                    None => {
                        tracing::error!("Relay context disappeared");
                        return;
                    }
                };

                // Serialize the initial request
                let mut request_bytes = Vec::new();
                if let Err(e) = initial_request.encode(&mut request_bytes) {
                    tracing::error!("Failed to encode initial request: {}", e);
                    return;
                }

                // Create MethodInvoke to start the SyncDocument stream
                let remote_response = ctx.create_method_invoke(
                    invoke_id,
                    &service_id,
                    "SyncDocument",
                    volt::MethodType::Bidi,
                    request_bytes,
                );

                // Encrypt and send
                match ctx.create_encrypted_request(invoke_id, &remote_response) {
                    Ok(invoke_req) => {
                        if let Err(e) = invoke_tx.send(invoke_req).await {
                            tracing::error!("Failed to send initial invoke request: {}", e);
                            return;
                        }
                    }
                    Err(e) => {
                        tracing::error!("Failed to create encrypted request: {}", e);
                        return;
                    }
                }
            }

            // Spawn a task to forward user requests to the invoke stream
            let invoke_tx_clone = invoke_tx.clone();
            let relay_context_clone = relay_context.clone();
            let forward_task = tokio::spawn(async move {
                while let Some(sync_req) = user_rx.recv().await {
                    let ctx_guard = relay_context_clone.lock().await;
                    let ctx = match ctx_guard.as_ref() {
                        Some(c) => c,
                        None => break,
                    };

                    // Serialize the request
                    let mut request_bytes = Vec::new();
                    if let Err(e) = sync_req.encode(&mut request_bytes) {
                        tracing::error!("Failed to encode sync request: {}", e);
                        continue;
                    }

                    // Wrap in MethodPayload
                    let remote_response = ctx.create_method_payload(invoke_id, request_bytes);

                    // Encrypt and send
                    match ctx.create_encrypted_request(invoke_id, &remote_response) {
                        Ok(invoke_req) => {
                            drop(ctx_guard);
                            if let Err(e) = invoke_tx_clone.send(invoke_req).await {
                                tracing::error!("Failed to send invoke request: {}", e);
                                break;
                            }
                        }
                        Err(e) => {
                            tracing::error!("Failed to create encrypted request: {}", e);
                        }
                    }
                }
            });

            // Process responses from the invoke stream
            loop {
                let mut rx_guard = relay_invoke_rx.lock().await;
                let rx = match rx_guard.as_mut() {
                    Some(r) => r,
                    None => break,
                };

                let response = match rx.message().await {
                    Ok(Some(resp)) => resp,
                    Ok(None) => {
                        tracing::debug!("Invoke stream ended");
                        break;
                    }
                    Err(e) => {
                        tracing::error!("Error receiving invoke response: {}", e);
                        let _ = response_tx.send(Err(e)).await;
                        break;
                    }
                };
                drop(rx_guard);

                // Parse and decrypt the response
                let mut ctx_guard = relay_context.lock().await;
                let ctx = match ctx_guard.as_mut() {
                    Some(c) => c,
                    None => break,
                };

                let remote_request = match ctx.parse_response(&response) {
                    Ok(Some(req)) => req,
                    Ok(None) => continue, // Key exchange or other non-payload response
                    Err(e) => {
                        tracing::error!("Failed to parse response: {}", e);
                        continue;
                    }
                };
                drop(ctx_guard);

                // Extract the SyncDocumentResponse from the RemoteRequest
                match remote_request.payload {
                    Some(volt::remote_request::Payload::MethodPayload(mp)) => {
                        let payload = match mp.method_payload {
                            Some(volt::method_payload::MethodPayload::Payload(p)) => p,
                            Some(volt::method_payload::MethodPayload::JsonPayload(j)) => {
                                j.into_bytes()
                            }
                            None => continue,
                        };

                        match volt::SyncDocumentResponse::decode(payload.as_slice()) {
                            Ok(sync_resp) => {
                                if let Err(e) = response_tx.send(Ok(sync_resp)).await {
                                    tracing::error!("Failed to forward response: {}", e);
                                    break;
                                }
                            }
                            Err(e) => {
                                tracing::error!("Failed to decode SyncDocumentResponse: {}", e);
                            }
                        }
                    }
                    Some(volt::remote_request::Payload::MethodEnd(end)) => {
                        tracing::debug!("Received MethodEnd for invoke_id={}", end.id);
                        if end.error_code != 0 {
                            let _ = response_tx
                                .send(Err(tonic::Status::new(
                                    tonic::Code::from_i32(end.error_code),
                                    &end.error,
                                )))
                                .await;
                        }
                        break;
                    }
                    _ => {}
                }
            }

            forward_task.abort();
        });

        // Wrap the channel receiver in our SyncDocumentStream wrapper
        let relay_stream = SyncDocumentStream::Relay(response_stream);

        Ok((user_tx, relay_stream))
    }

    /// Make a unary call through the relay protocol
    ///
    /// This opens a new Invoke stream for each call, performs key exchange,
    /// and sends the encrypted request. This matches the JS implementation
    /// where each GRPCCall gets its own stream.
    async fn unary_call_via_relay(
        &self,
        method: &str,
        request_json: serde_json::Value,
    ) -> Result<ApiResponse> {
        use crate::proto::volt::{method_payload, remote_request, MethodType};
        use crate::relay::RelayContext;

        tracing::info!("Making relay call for method: {}", method);

        let request_json = normalize_payload(request_json);

        // Get configuration
        let volt_config = self.volt_config.as_ref().ok_or(VoltError::NotInitialized)?;
        let credential = self.credential.as_ref().ok_or(VoltError::NotInitialized)?;
        let target_did = &volt_config.id;

        // Create relay context for this call
        let mut relay_ctx = RelayContext::new(target_did.clone());

        // Set the target's public key for signature verification (if available)
        if let Some(public_key_pem) = &volt_config.public_key {
            if let Ok(pk_bytes) =
                crate::crypto::from_base64(&crate::crypto::strip_pem_headers(public_key_pem))
            {
                relay_ctx.set_target_public_key(pk_bytes);
            }
        }

        // Get JWT token for authentication, injecting our relay public key so the server can derive the shared secret
        let relay_public_key_b64 = relay_ctx.public_key_base64();
        let identity = credential.get_identity_metadata(
            Some(target_did),
            Some(&relay_public_key_b64),
            true,
            60,
        )?;
        let token = identity.token.clone();

        // Create the key exchange request
        let key_exchange_request = relay_ctx.create_key_exchange_request(&token);
        let invoke_id = key_exchange_request.invoke_id;
        tracing::info!(
            "Created key exchange request: invoke_id={}, target_did={:?}, token_len={}, iv_len={}",
            invoke_id,
            key_exchange_request.target_did,
            key_exchange_request.token.len(),
            key_exchange_request.iv.len()
        );

        // Get VoltApiClient to call Invoke
        let mut volt_client = self.get_volt_client().await?;
        tracing::info!("Got VoltApiClient successfully");

        // Create channel for sending invoke requests
        let (invoke_tx, invoke_rx) = mpsc::channel::<volt::InvokeRequest>(10);
        tracing::info!("Created mpsc channel");

        // Send the key exchange request
        invoke_tx
            .send(key_exchange_request)
            .await
            .map_err(|_| VoltError::connection("Failed to send key exchange request"))?;
        tracing::info!("Sent key exchange request to channel");

        // Create the invoke stream
        let request_stream = ReceiverStream::new(invoke_rx);
        let mut request = tonic::Request::new(request_stream);
        tracing::info!("Created request stream");

        // Attach the auth metadata created alongside the relay key so Invoke sees the same JWT the relay expects
        {
            let req_metadata = request.metadata_mut();
            for key_and_value in identity.metadata.iter() {
                match key_and_value {
                    tonic::metadata::KeyAndValueRef::Ascii(key, value) => {
                        tracing::info!("Adding metadata: {}={:?}", key, value);
                        req_metadata.insert(key, value.clone());
                    }
                    tonic::metadata::KeyAndValueRef::Binary(key, value) => {
                        tracing::info!("Adding binary metadata: {}", key);
                        req_metadata.insert_bin(key, value.clone());
                    }
                }
            }
        }
        tracing::info!("Added auth metadata to request");

        // Debug: print all metadata in the request
        tracing::info!("Request metadata contents:");
        for key_and_value in request.metadata().iter() {
            match key_and_value {
                tonic::metadata::KeyAndValueRef::Ascii(key, value) => {
                    tracing::info!("  {}={:?}", key, value);
                }
                tonic::metadata::KeyAndValueRef::Binary(key, _value) => {
                    tracing::info!("  {} (binary)", key);
                }
            }
        }

        // Start the Invoke bidirectional stream
        // NOTE: This currently fails with ConnectionReset when using tonic against the relay server.
        // Unary gRPC calls work fine, but bidirectional streaming fails specifically.
        // The JavaScript client using @grpc/grpc-js works, suggesting a tonic-specific issue.
        tracing::info!("unary_call_via_relay: Calling invoke()...");
        let response = volt_client.invoke(request).await.map_err(|e| {
            tracing::error!("unary_call_via_relay: invoke() failed: {:?}", e);
            VoltError::connection(format!("Failed to start Invoke stream: {}", e))
        })?;
        tracing::info!("unary_call_via_relay: invoke() succeeded!");

        let mut invoke_response_stream = response.into_inner();

        // Wait for the key exchange response
        tracing::debug!("Waiting for key exchange response...");
        let key_exchange_response = invoke_response_stream
            .message()
            .await
            .map_err(|e| {
                VoltError::connection(format!("Failed to receive key exchange response: {}", e))
            })?
            .ok_or_else(|| VoltError::connection("Invoke stream closed before key exchange"))?;

        tracing::debug!(
            "Received key exchange response: invoke_id={}",
            key_exchange_response.invoke_id
        );

        // Process the key exchange
        if let Some(key_exchange) = &key_exchange_response.key_exchange {
            relay_ctx.process_key_exchange_response(key_exchange)?;
            tracing::info!("Relay key exchange completed successfully");
        } else if let Some(volt::invoke_response::ResponsePayload::Status(status)) =
            &key_exchange_response.response_payload
        {
            return Err(VoltError::server(status.code, &status.message));
        } else {
            return Err(VoltError::protocol(
                "Expected key exchange response, got something else",
            ));
        }

        // Now we have encryption established, serialize and send the actual request
        let service_id = "tdx.volt_api.volt.v1.VoltAPI";

        // Serialize the request based on method
        // We need to serialize to protobuf bytes
        let request_bytes = self.serialize_request(method, &request_json)?;

        // Create the method invoke
        let remote_response = relay_ctx.create_method_invoke(
            invoke_id,
            service_id,
            method,
            MethodType::Unary,
            request_bytes,
        );

        // Encrypt and send
        let invoke_request = relay_ctx.create_encrypted_request(invoke_id, &remote_response)?;
        invoke_tx
            .send(invoke_request)
            .await
            .map_err(|e| VoltError::connection(format!("Failed to send method invoke: {}", e)))?;

        // Receive the method response
        let response = invoke_response_stream
            .message()
            .await
            .map_err(|e| VoltError::grpc(e.code(), e.message()))?
            .ok_or_else(|| VoltError::connection("Stream closed before response"))?;

        // Parse and decrypt the response
        let remote_request = relay_ctx
            .parse_response(&response)?
            .ok_or_else(|| VoltError::protocol("Expected payload response"))?;

        // Extract the method payload and deserialize
        match remote_request.payload {
            Some(remote_request::Payload::MethodPayload(mp)) => {
                let payload = match mp.method_payload {
                    Some(method_payload::MethodPayload::Payload(p)) => p,
                    Some(method_payload::MethodPayload::JsonPayload(j)) => j.into_bytes(),
                    None => return Err(VoltError::protocol("Empty method payload")),
                };

                // Deserialize the response
                let response_json = self.deserialize_response(method, &payload)?;

                // Extract status if present
                let status = response_json
                    .get("status")
                    .and_then(|s| serde_json::from_value(s.clone()).ok());

                Ok(ApiResponse {
                    status,
                    payload: response_json,
                })
            }
            Some(remote_request::Payload::MethodEnd(end)) => {
                if end.error_code != 0 || !end.error.is_empty() {
                    Err(VoltError::server(end.error_code, &end.error))
                } else {
                    Err(VoltError::protocol("Received MethodEnd without response"))
                }
            }
            _ => Err(VoltError::protocol("Expected MethodPayload response")),
        }
    }

    /// Serialize a request to protobuf bytes
    fn serialize_request(&self, method: &str, request: &serde_json::Value) -> Result<Vec<u8>> {
        use prost::Message;

        match method {
            "GetResource" => {
                let req: volt::GetResourceRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "GetResources" => {
                let req: volt::GetResourcesRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "CanAccessResource" => {
                let req: volt::CanAccessResourceRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "DeleteResource" => {
                let req: volt::DeleteResourceRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "SaveResource" => {
                let req: volt::SaveResourceRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "RequestAccess" => {
                let req: volt::RequestAccessRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "Authenticate" => {
                let req: volt::AuthenticateRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "DiscoverServices" => {
                let req: volt::DiscoverServicesRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "GetResourceAncestors" => {
                let req: volt::GetResourceAncestorsRequest =
                    serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "GetResourceDescendants" => {
                let req: volt::GetResourceDescendantsRequest =
                    serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "SaveAccess" => {
                let req: volt::SaveAccessRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "GetIdentities" => {
                let req: volt::GetIdentitiesRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "GetIdentity" => {
                let req: volt::GetIdentityRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "GetPolicy" => {
                let req: volt::GetPolicyRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "GetAccess" => {
                let req: volt::GetAccessRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "SaveIdentity" => {
                let req: volt::SaveIdentityRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "Shutdown" => {
                let req: volt::ShutdownRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "GetFileDescendants" => {
                let req: volt::GetFileDescendantsRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "GetFileContent" => {
                let req: volt::GetFileContentRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            "SetFileContent" => {
                let req: volt::SetFileContentRequest = serde_json::from_value(request.clone())?;
                let mut buf = Vec::new();
                req.encode(&mut buf)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(buf)
            }
            _ => Err(VoltError::MethodNotFound(format!(
                "Unknown method: {}",
                method
            ))),
        }
    }

    /// Deserialize a response from protobuf bytes
    fn deserialize_response(&self, method: &str, data: &[u8]) -> Result<serde_json::Value> {
        use prost::Message;

        match method {
            "GetResource" => {
                let resp = volt::GetResourceResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "GetResources" => {
                let resp = volt::GetResourcesResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "CanAccessResource" => {
                let resp = volt::CanAccessResourceResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "DeleteResource" => {
                let resp = volt::DeleteResourceResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "SaveResource" => {
                let resp = volt::SaveResourceResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "RequestAccess" => {
                let resp = volt::RequestAccessResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "Authenticate" => {
                let resp = volt::AuthenticateResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "DiscoverServices" => {
                let resp = volt::DiscoverServicesResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "GetResourceAncestors" => {
                let resp = volt::GetResourceAncestorsResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "GetResourceDescendants" => {
                let resp = volt::GetResourceDescendantsResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "SaveAccess" => {
                let resp = volt::SaveAccessResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "GetIdentities" => {
                let resp = volt::GetIdentitiesResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "GetIdentity" => {
                let resp = volt::GetIdentityResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "GetPolicy" => {
                let resp = volt::GetPolicyResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "GetAccess" => {
                let resp = volt::GetAccessResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "SaveIdentity" => {
                let resp = volt::SaveIdentityResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "Shutdown" => {
                let resp = volt::ShutdownResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "GetFileDescendants" => {
                let resp = volt::GetFileDescendantsResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "GetFileContent" => {
                let resp = volt::GetFileContentResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            "SetFileContent" => {
                let resp = volt::SetFileContentResponse::decode(data)
                    .map_err(|e| VoltError::serialization(e.to_string()))?;
                Ok(serde_json::to_value(&resp)?)
            }
            _ => Err(VoltError::MethodNotFound(format!(
                "Unknown method: {}",
                method
            ))),
        }
    }

    /// Make a unary gRPC call
    async fn unary_call(&self, method: &str, request: serde_json::Value) -> Result<ApiResponse> {
        tracing::debug!("Calling {} with {:?}", method, request);

        // For relayed connections, use the relay protocol
        if self.is_relayed() {
            return self.unary_call_via_relay(method, request).await;
        }

        // Get the gRPC client
        let mut client = self.get_volt_client().await?;

        // Helper to create authenticated request
        macro_rules! auth_request {
            ($req:expr) => {{
                let mut request = tonic::Request::new($req);
                self.add_auth_to_request(&mut request);
                request
            }};
        }

        // Dispatch to the appropriate method based on name
        let response_json = match method {
            "GetResource" => {
                let req: volt::GetResourceRequest = serde_json::from_value(request)?;
                let response = client
                    .get_resource(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            "GetResources" => {
                let req: volt::GetResourcesRequest = serde_json::from_value(request)?;
                let response = client
                    .get_resources(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            "CanAccessResource" => {
                let req: volt::CanAccessResourceRequest = serde_json::from_value(request)?;
                let response = client
                    .can_access_resource(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            "DeleteResource" => {
                let req: volt::DeleteResourceRequest = serde_json::from_value(request)?;
                let response = client
                    .delete_resource(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            "SaveResource" => {
                let req: volt::SaveResourceRequest = serde_json::from_value(request)?;
                let response = client
                    .save_resource(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            "RequestAccess" => {
                let req: volt::RequestAccessRequest = serde_json::from_value(request)?;
                let response = client
                    .request_access(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            "Authenticate" => {
                let req: volt::AuthenticateRequest = serde_json::from_value(request)?;
                let response = client
                    .authenticate(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            "DiscoverServices" => {
                let req: volt::DiscoverServicesRequest = serde_json::from_value(request)?;
                let response = client
                    .discover_services(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            "GetResourceAncestors" => {
                let req: volt::GetResourceAncestorsRequest = serde_json::from_value(request)?;
                let response = client
                    .get_resource_ancestors(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            "GetResourceDescendants" => {
                let req: volt::GetResourceDescendantsRequest = serde_json::from_value(request)?;
                let response = client
                    .get_resource_descendants(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            "SaveAccess" => {
                let req: volt::SaveAccessRequest = serde_json::from_value(request)?;
                let response = client
                    .save_access(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            "GetIdentities" => {
                let req: volt::GetIdentitiesRequest = serde_json::from_value(request)?;
                let response = client
                    .get_identities(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            "GetIdentity" => {
                let req: volt::GetIdentityRequest = serde_json::from_value(request)?;
                let response = client
                    .get_identity(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            "GetPolicy" => {
                let req: volt::GetPolicyRequest = serde_json::from_value(request)?;
                let response = client
                    .get_policy(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            "GetAccess" => {
                let req: volt::GetAccessRequest = serde_json::from_value(request)?;
                let response = client
                    .get_access(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            "SaveIdentity" => {
                let req: volt::SaveIdentityRequest = serde_json::from_value(request)?;
                let response = client
                    .save_identity(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            "Shutdown" => {
                let req: volt::ShutdownRequest = serde_json::from_value(request)?;
                let response = client
                    .shutdown(auth_request!(req))
                    .await
                    .map_err(VoltError::GrpcError)?;
                let resp = response.into_inner();
                serde_json::to_value(&resp)?
            }
            _ => {
                return Err(VoltError::MethodNotFound(format!(
                    "Unknown method: {}",
                    method
                )));
            }
        };

        // Extract status if present
        let status = response_json
            .get("status")
            .and_then(|s| serde_json::from_value(s.clone()).ok());

        Ok(ApiResponse {
            status,
            payload: response_json,
        })
    }
}

impl Default for VoltClient {
    fn default() -> Self {
        Self::new().expect("Failed to create default VoltClient")
    }
}

/// Source for client configuration
pub enum ConfigSource {
    /// Load from file path
    File(std::path::PathBuf),
    /// Use existing config
    Config(VoltClientConfig),
    /// Parse from JSON
    Json(serde_json::Value),
}

impl From<&str> for ConfigSource {
    fn from(s: &str) -> Self {
        ConfigSource::File(std::path::PathBuf::from(s))
    }
}

impl From<std::path::PathBuf> for ConfigSource {
    fn from(p: std::path::PathBuf) -> Self {
        ConfigSource::File(p)
    }
}

impl From<VoltClientConfig> for ConfigSource {
    fn from(c: VoltClientConfig) -> Self {
        ConfigSource::Config(c)
    }
}

impl From<serde_json::Value> for ConfigSource {
    fn from(v: serde_json::Value) -> Self {
        ConfigSource::Json(v)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_client_creation() {
        let client = VoltClient::new().unwrap();
        assert!(client.config().is_none());
        assert!(!client.is_connected().await);
    }

    #[test]
    fn test_config_source() {
        let _source: ConfigSource = "path/to/config.json".into();
        let _source: ConfigSource = std::path::PathBuf::from("config.json").into();
        let _source: ConfigSource = serde_json::json!({"client_name": "test"}).into();
    }
}

/// Normalize payload to handle empty identityDid in subject
fn normalize_payload(value: serde_json::Value) -> serde_json::Value {
    match value {
        serde_json::Value::Object(map) => {
            let mut new_map = serde_json::Map::new();
            for (key, val) in map {
                if key == "subject" {
                    if let serde_json::Value::Object(mut subject_map) = val {
                        if let Some(serde_json::Value::String(s)) =
                            subject_map.get("credentialLookup")
                        {
                            if s == "'[EVERYONE]'" {
                                // For everyone, remove the subject entirely
                                subject_map.clear();
                            }
                        } else if let Some(serde_json::Value::String(s)) =
                            subject_map.get("identityDid")
                        {
                            if s.is_empty() {
                                subject_map.insert(
                                    "identityDid".to_string(),
                                    serde_json::Value::String(
                                        "did:volt:00000000-0000-0000-0000-000000000000".to_string(),
                                    ),
                                );
                            }
                        }
                        if !subject_map.is_empty() {
                            new_map.insert(key, serde_json::Value::Object(subject_map));
                        }
                    } else {
                        new_map.insert(key, normalize_payload(val));
                    }
                } else {
                    new_map.insert(key, normalize_payload(val));
                }
            }
            serde_json::Value::Object(new_map)
        }
        serde_json::Value::Array(arr) => {
            serde_json::Value::Array(arr.into_iter().map(normalize_payload).collect())
        }
        other => other,
    }
}