volt-client-rs 0.1.12

use crate::websocket_rpc::{Event, EventType, WebsocketRpc};
use futures::{stream::StreamExt, SinkExt};
use serde_json;
use std::collections::HashMap;
use std::sync::{Arc, Mutex as StdMutex};
use tokio::sync::{
    mpsc::{self, Receiver, Sender},
    Mutex,
};
use tokio_tungstenite::{connect_async, tungstenite::protocol::Message};
use volt_ws_protocol;

/// The VoltClient struct is the main entry point for interacting with the Volt API.
/// It provides methods for making unary, server streaming, and client streaming RPC calls.
/// The client maintains a connection to the Volt server and handles the WebSocket communication,
/// it manages the RPC lifecycle, including sending requests and receiving responses.
/// The client also provides a set of convenience methods for making calls to the Volt API.
pub struct VoltClient {
    send_channel: mpsc::Sender<Vec<u8>>,
    protocol_manager: Arc<Mutex<volt_ws_protocol::rpc_manager::RpcManager>>,
    active_rpc: Arc<Mutex<HashMap<u64, Arc<Mutex<WebsocketRpc>>>>>,
}

impl VoltClient {
    /// Creates a new API client instance and connects to the WebSocket server.
    /// The client is configured with the provided JSON configuration.
    /// # Arguments
    /// * `config_json` - A JSON string containing the configuration settings.
    /// # Returns
    /// A Result containing the API client instance or an error message.
    pub async fn create(config_json: &str) -> Result<VoltClient, String> {
        // Create the protocol manager.
        let protocol_manager =
            Arc::new(Mutex::new(volt_ws_protocol::rpc_manager::RpcManager::new()));

        let url: String;
        {
            let mut protocol_manager = protocol_manager.lock().await;

            // Initialise the protocol manager and get the WebSocket server URL.
            url = match protocol_manager.set_configuration(config_json) {
                Ok(url) => url,
                Err(e) => return Err(e),
            };

            println!("Connecting to WebSocket server at {}", url);
        }

        // Connect to the WebSocket server.
        let ws_stream = match connect_async(url).await {
            Ok((ws_stream, _)) => ws_stream,
            Err(e) => return Err(format!("Failed to connect to the WebSocket server: {}", e)),
        };

        // Split the WebSocket stream into a write and read stream.
        let (write, read) = ws_stream.split();

        // Create a channel for sending requests
        let (send_channel, mut receive_channel) = mpsc::channel(100);

        // Spawn a task to handle sending requests
        tokio::spawn(async move {
            let mut write = write;
            while let Some(req) = receive_channel.recv().await {
                if write.send(Message::Binary(req)).await.is_err() {
                    break;
                }
            }
        });

        // Create a map to store active RPCs.
        let active_rpc = Arc::new(Mutex::new(HashMap::<u64, Arc<Mutex<WebsocketRpc>>>::new()));

        // Clone the active RPC map and the send channel for use in the read task.
        let active_rpc_clone = Arc::clone(&active_rpc);
        let send_channel_clone = Arc::<mpsc::Sender<Vec<u8>>>::new(send_channel.clone());
        let protocol_manager_clone = Arc::clone(&protocol_manager);

        // Spawn a task to handle reading responses
        tokio::spawn(async move {
            // Capture the read stream in the closure.
            let mut read = read;

            // Loop to read messages from the WebSocket server.
            while let Some(msg) = read.next().await {
                match msg {
                    Ok(payload) => {
                        // Decode the payload using the wasm protocol manager.
                        let protocol = protocol_manager_clone.lock().await;
                        let decoded = match protocol.decode_payload(payload.into_data()) {
                            Ok(decoded) => decoded,
                            Err(e) => {
                                println!("Failed to decode payload: {:?}", e);
                                continue;
                            }
                        };

                        // Parse the decoded payload as JSON.
                        let response: serde_json::Value = match serde_json::from_str(&decoded) {
                            Ok(response) => response,
                            Err(e) => {
                                println!("Failed to parse JSON: {:?}", e);
                                continue;
                            }
                        };

                        // Extract the target rpc method_id from the response.
                        let method_id = match response["method_id"].as_u64() {
                            Some(method_id) => method_id,
                            None => {
                                println!("Received response with no method_id");
                                continue;
                            }
                        };

                        // Check if the response contains a 'key_exchanged' field.
                        if !response["key_exchanged"].is_null() {
                            // Handle key exchange.
                            let mut pending_payload = match protocol.pending_payload(&method_id) {
                                Ok(pending_payload) => pending_payload,
                                Err(e) => {
                                    println!(
                                        "Failure fetching pending payload for method_id: {} {}",
                                        method_id, e
                                    );
                                    continue;
                                }
                            };

                            while pending_payload.len() > 0 {
                                // Send the pending request.
                                let send_result = send_channel_clone.send(pending_payload).await;

                                match send_result {
                                    Ok(_) => println!("Sent pending payload"),
                                    Err(e) => {
                                        println!("Failed to send pending payload: {:?}", e);
                                        break;
                                    }
                                }

                                // Fetch the next pending payload.
                                pending_payload = match protocol.pending_payload(&method_id) {
                                    Ok(pending_payload) => pending_payload,
                                    Err(e) => {
                                        println!(
                                            "Failure fetching pending payload for method_id: {} {}",
                                            method_id, e
                                        );
                                        break;
                                    }
                                };
                            }
                        } else {
                            // Not a key exchange response, handle as a normal response.
                            let mut active_rpc = active_rpc_clone.lock().await;

                            if !active_rpc.contains_key(&method_id) {
                                println!("Received response for unknown method_id: {}", method_id);
                            } else {
                                // Let the RPC handle the response.
                                let mut rpc = match active_rpc.get_mut(&method_id) {
                                    Some(rpc) => rpc.lock().await,
                                    None => {
                                        println!("Failed to get RPC for method_id: {}", method_id);
                                        continue;
                                    }
                                };

                                rpc.handle_response(&response);
                            }
                        }
                    }
                    _ => {
                        println!("Failure in received message");
                    }
                }
            }
        });

        // Create the API client instance.
        let client = VoltClient {
            send_channel,
            protocol_manager,
            active_rpc,
        };

        Ok(client)
    }

    async fn call_internal(
        &mut self,
        method: &str,
        service: &str,
    ) -> Result<Arc<Mutex<WebsocketRpc>>, String> {
        if method.is_empty() {
            return Err("Method cannot be empty".to_string());
        }

        let rpc_id: u64;

        {
            let mut rpc_manager = self.protocol_manager.lock().await;
            rpc_id = match rpc_manager.create_rpc(method, service) {
                Ok(rpc_id) => rpc_id,
                Err(e) => return Err(e),
            };
        }

        let rpc = WebsocketRpc::new(
            rpc_id,
            Arc::clone(&self.protocol_manager),
            self.send_channel.clone(),
        );

        // Register the end event to remove the RPC from the active RPC map.
        rpc.on(EventType::End, {
            let rpc_id = rpc_id;
            let active_rpc = Arc::clone(&self.active_rpc);
            move |_response| {
                let active_rpc = Arc::clone(&active_rpc);
                tokio::spawn(async move {
                    println!("api_client: removing rpc: {}", rpc_id);
                    let mut active_rpc = active_rpc.lock().await;
                    active_rpc.remove(&rpc_id);
                });
            }
        });

        // Cache the RPC in the active RPC map.
        let rpc_cached = Arc::new(Mutex::new(rpc));
        self.active_rpc
            .lock()
            .await
            .insert(rpc_id, rpc_cached.clone());

        Ok(rpc_cached)
    }

    /// Generic method for unary RPC calls.
    /// Sends a request and awaits the response.
    /// # Arguments
    /// * `method` - The method to call.
    /// * `request` - The request payload.
    /// * `service` - The service to call.
    /// # Returns
    /// A Result containing the response or an error message.
    pub async fn unary_rpc(
        &mut self,
        method: &str,
        request: &serde_json::Value,
        service: &str,
    ) -> Result<serde_json::Value, String> {
        if request.is_null() {
            return Err("Request cannot be null".to_string());
        }

        // Start the call.
        let rpc = match self.call_internal(method, service).await {
            Ok(rpc) => rpc,
            Err(e) => return Err(e),
        };

        // Create a channel to asynchronously receive the response.
        let (tx, mut rx): (Sender<()>, Receiver<()>) = mpsc::channel(1);

        // Clone the channel for use in the event handlers.
        let tx = Arc::new(tx);

        // Create a mutex to store the unary response or error.
        let unary_response = Arc::new(StdMutex::new(serde_json::Value::Null));
        let unary_error = Arc::new(StdMutex::new(String::new()));

        {
            // Register to receive rpc events
            let rpc = rpc.lock().await;

            rpc.on(EventType::Data, {
                let unary_response = Arc::clone(&unary_response);
                move |response| {
                    println!("api_client: received response: {:?}", response);
                    if let Event::Data(response) = response {
                        match unary_response.lock() {
                            Ok(mut unary_response) => {
                                // Copy the response into the unary_response so we can emit it from the `end` event.
                                *unary_response = response;
                            }
                            Err(e) => {
                                println!("failed to lock unary_response: {:?}", e);
                                return;
                            }
                        };
                    }
                }
            });

            rpc.on(EventType::Error, {
                let unary_error = Arc::clone(&unary_error);
                move |response| {
                    println!("api_client: received error");
                    if let Event::Error(response) = response {
                        match unary_error.lock() {
                            Ok(mut unary_error) => {
                                // Copy the response into the unary_response so we can emit it from the `end` event.
                                *unary_error = response;
                            }
                            Err(e) => {
                                println!("failed to lock unary_response: {:?}", e);
                                return;
                            }
                        };
                    }
                }
            });

            rpc.on(EventType::End, {
                let tx = Arc::clone(&tx);
                move |_response| {
                    println!("api_client: received end");
                    let tx = Arc::clone(&tx);
                    tokio::spawn(async move {
                        match tx.send(()).await {
                            Ok(_) => println!("unary response sent successfully"),
                            Err(e) => println!("failed to send unary response: {:?}", e),
                        }
                    });
                }
            });

            // Send the payload.
            match rpc.send(request).await {
                Ok(_) => println!("sent request"),
                Err(e) => {
                    return Err(e);
                }
            }

            let _ = rpc.end().await;
        }

        // Await the response.
        rx.recv().await;

        // Determine if we received an error or a response.
        let result = match unary_error.lock() {
            Ok(unary_error) => {
                if unary_error.is_empty() {
                    // There is no error so return the response.
                    match unary_response.lock() {
                        Ok(unary_response) => Ok(unary_response.clone()),
                        Err(e) => Err(format!("failed to lock unary_response: {:?}", e)),
                    }
                } else {
                    // Propagate the error.
                    Err(unary_error.clone())
                }
            }
            Err(e) => Err(format!("failed to lock unary_error: {:?}", e)),
        };

        result
    }

    /// Generic method for server streaming RPC calls.
    /// Sends a request and returns an RPC instance.
    /// # Arguments
    /// * `method` - The method to call.
    /// * `request` - The request payload.
    /// * `service` - The service to call.
    /// # Returns
    /// A Result containing the RPC event emitter or an error message.
    pub async fn server_streaming_call(
        &mut self,
        method: &str,
        request: &serde_json::Value,
        service: &str,
    ) -> Result<Arc<Mutex<WebsocketRpc>>, String> {
        if request.is_null() {
            return Err("Request cannot be null".to_string());
        }

        let rpc = match self.call_internal(method, service).await {
            Ok(rpc) => rpc,
            Err(e) => return Err(e),
        };

        {
            let rpc = rpc.lock().await;

            // Send the payload.
            match rpc.send(request).await {
                Ok(_) => println!("sent request"),
                Err(e) => {
                    return Err(e);
                }
            }

            let _ = rpc.end().await;
        }

        Ok(rpc)
    }

    /// Generic method for client streaming RPC calls.
    /// Sends a request and returns an RPC instance.
    /// # Arguments
    /// * `method` - The method to call.
    /// * `request` - The request payload.
    /// * `service` - The service to call.
    /// # Returns
    /// A Result containing the RPC event emitter or an error message.
    pub async fn streaming_call(
        &mut self,
        method: &str,
        request: &serde_json::Value,
        service: &str,
    ) -> Result<Arc<Mutex<WebsocketRpc>>, String> {
        let rpc = match self.call_internal(method, service).await {
            Ok(rpc) => rpc,
            Err(e) => return Err(e),
        };

        if !request.is_null() {
            let rpc = rpc.lock().await;

            // Send the payload.
            match rpc.send(request).await {
                Ok(_) => println!("sent request"),
                Err(e) => {
                    println!("failed to send request: {:?}", e);
                    return Err(e);
                }
            }
        } else {
            println!("no initial request");
        }

        Ok(rpc)
    }

    ///
    /// RESOURCE API
    ///

    /// See https://docs.tdxvolt.com/en/api/volt_api#CanAccessResource
    pub async fn can_access_resource(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc(
            "/tdx.volt_api.volt.v1.VoltAPI/CanAccessResource",
            request,
            "",
        )
        .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#Connect
    pub async fn connect(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<Arc<Mutex<WebsocketRpc>>, String> {
        self.streaming_call("/tdx.volt_api.volt.v1.VoltAPI/Connect", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#DeleteResource
    pub async fn delete_resource(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/DeleteResource", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#DiscoverServices
    pub async fn discover_services(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc(
            "/tdx.volt_api.volt.v1.VoltAPI/DiscoverServices",
            request,
            "",
        )
        .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#GetResource
    pub async fn get_resource(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/GetResource", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#GetResources
    pub async fn get_resources(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/GetResources", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#GetResourceAncestors
    pub async fn get_resource_ancestors(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc(
            "/tdx.volt_api.volt.v1.VoltAPI/GetResourceAncestors",
            request,
            "",
        )
        .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#GetResourceDescendants
    pub async fn get_resource_descendants(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc(
            "/tdx.volt_api.volt.v1.VoltAPI/GetResourceDescendants",
            request,
            "",
        )
        .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#RequestAccess
    pub async fn request_access(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/RequestAccess", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#SaveResource
    pub async fn save_resource(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/SaveResource", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#SaveResourceAttribute
    pub async fn save_resource_attribute(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc(
            "/tdx.volt_api.volt.v1.VoltAPI/SaveResourceAttribute",
            request,
            "",
        )
        .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#SetServiceStatus
    pub async fn set_service_status(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc(
            "/tdx.volt_api.volt.v1.VoltAPI/SetServiceStatus",
            request,
            "",
        )
        .await
    }

    ///
    /// FILE API
    ///

    /// See https://docs.tdxvolt.com/en/api/file_api#DownloadFile
    pub async fn download_file(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<Arc<Mutex<WebsocketRpc>>, String> {
        self.server_streaming_call("/tdx.volt_api.volt.v1.FileAPI/DownloadFile", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/file_api#GetFile
    pub async fn get_file(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.FileAPI/GetFile", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/file_api#GetFileContent
    pub async fn get_file_content(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<Arc<Mutex<WebsocketRpc>>, String> {
        self.server_streaming_call("/tdx.volt_api.volt.v1.FileAPI/GetFileContent", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/file_api#GetFileDescendants
    pub async fn get_file_descendants(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc(
            "/tdx.volt_api.volt.v1.FileAPI/GetFileDescendants",
            request,
            "",
        )
        .await
    }

    /// See https://docs.tdxvolt.com/en/api/file_api#SetFileContent
    pub async fn set_file_content(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<Arc<Mutex<WebsocketRpc>>, String> {
        self.streaming_call("/tdx.volt_api.volt.v1.FileAPI/SetFileContent", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/file_api#UploadFile
    pub async fn upload_file(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<Arc<Mutex<WebsocketRpc>>, String> {
        self.streaming_call("/tdx.volt_api.volt.v1.FileAPI/UploadFile", request, "")
            .await
    }

    ///
    /// VOLT MANAGEMENT API
    ///

    /// See https://docs.tdxvolt.com/en/api/volt_api#Authenticate
    pub async fn authenticate(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/Authenticate", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#DeleteAccess
    pub async fn delete_access(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/DeleteAccess", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#DeleteVolt
    pub async fn delete_volt(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/DeleteVolt", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#GetAccess
    pub async fn get_access(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/GetAccess", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#GetIdentities
    pub async fn get_identities(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/GetIdentities", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#GetIdentity
    pub async fn get_identity(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/GetIdentity", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#GetOneTimeToken
    pub async fn get_one_time_token(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/GetOneTimeToken", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#GetPolicy
    pub async fn get_policy(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/GetPolicy", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#GetSettings
    pub async fn get_settings(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/GetSettings", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#Invoke
    pub async fn invoke(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/Invoke", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#SaveAccess
    pub async fn save_access(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/SaveAccess", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#SaveIdentity
    pub async fn save_identity(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/SaveIdentity", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#SaveSettings
    pub async fn save_settings(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/SaveSettings", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#SetAccessRequestDecision
    pub async fn set_access_request_decision(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc(
            "/tdx.volt_api.volt.v1.VoltAPI/SetAccessRequestDecision",
            request,
            "",
        )
        .await
    }

    /// See https://docs.tdxvolt.com/en/api/volt_api#SignVerify
    pub async fn sign_verify(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc("/tdx.volt_api.volt.v1.VoltAPI/SignVerify", request, "")
            .await
    }

    ///
    /// WIRE API
    ///

    /// See https://docs.tdxvolt.com/en/api/wire_api#PublishWire
    pub async fn publish_wire(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<Arc<Mutex<WebsocketRpc>>, String> {
        self.streaming_call("/tdx.volt_api.volt.v1.WireAPI/PublishWire", request, "")
            .await
    }

    /// See https://docs.tdxvolt.com/en/api/wire_api#SubscribeWire
    pub async fn subscribe_wire(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<Arc<Mutex<WebsocketRpc>>, String> {
        self.streaming_call("/tdx.volt_api.volt.v1.WireAPI/SubscribeWire", request, "")
            .await
    }

    ///
    /// DATABASE API
    ///

    /// See https://docs.tdxvolt.com/en/api/sqlite_database_api#BulkUpdate
    pub async fn bulk_update(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<Arc<Mutex<WebsocketRpc>>, String> {
        self.streaming_call(
            "/tdx.volt_api.data.v1.SqliteDatabaseAPI/BulkUpdate",
            request,
            "",
        )
        .await
    }

    /// https://docs.tdxvolt.com/en/api/sqlite_server_api#CreateDatabase
    pub async fn create_database(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.unary_rpc(
            "/tdx.volt_api.data.v1.SqliteServerAPI/CreateDatabase",
            request,
            "",
        )
        .await
    }

    /// See https://docs.tdxvolt.com/en/api/sqlite_database_api#Execute
    pub async fn sql_execute(
        &mut self,
        request: &serde_json::Value,
        service: &str,
    ) -> Result<Arc<Mutex<WebsocketRpc>>, String> {
        self.streaming_call(
            "/tdx.volt_api.data.v1.SqliteDatabaseAPI/Execute",
            request,
            service,
        )
        .await
    }

    /// See https://docs.tdxvolt.com/en/api/sqlite_database_api#Execute
    pub async fn sql_execute_json(
        &mut self,
        request: &serde_json::Value,
        service: &str,
    ) -> Result<serde_json::Value, String> {
        let rpc = match self.sql_execute(request, service).await {
            Ok(call) => call,
            Err(e) => return Err(e),
        };

        // Create a channel to asynchronously receive the response.
        let (tx, mut rx): (Sender<()>, Receiver<()>) = mpsc::channel(1);

        // Clone the channel for use in the event handlers.
        let tx = Arc::new(tx);

        // Create a mutex to store the unary response or error.
        let header = Arc::new(StdMutex::new(Vec::<serde_json::Value>::new()));
        let rows = Arc::new(StdMutex::new(Vec::<serde_json::Value>::new()));
        let json_error = Arc::new(StdMutex::new(String::new()));

        {
            // Register to receive rpc events
            let rpc = rpc.lock().await;

            rpc.on(EventType::Data, {
                let header = Arc::clone(&header);
                let rows = Arc::clone(&rows);
                move |response| {
                    println!("sql_execute_json: received response: {:?}", response);
                    if let Event::Data(response) = response {
                        match response["header"].as_object() {
                            Some(header_obj) => {
                                match header.lock() {
                                    Ok(mut header) => {
                                        // Copy the response into the unary_response so we can emit it from the `end` event.
                                        *header = header_obj["column"].as_array().unwrap().clone();
                                    }
                                    Err(e) => {
                                        println!("failed to lock unary_response: {:?}", e);
                                        return;
                                    }
                                };
                            }
                            None => {
                                let row = response["row"].as_object().unwrap()["column"]
                                    .as_array()
                                    .unwrap();

                                let columns = header.lock().unwrap();

                                let mut row_map = serde_json::Map::new();
                                for (i, column) in columns.iter().enumerate() {
                                    let name = column["name"].as_str().unwrap();

                                    let cell = row[i].as_object().unwrap();

                                    if cell.contains_key("null") {
                                        row_map.insert(name.to_string(), serde_json::Value::Null);
                                    } else if cell.contains_key("text") {
                                        row_map.insert(name.to_string(), cell["text"].clone());
                                    } else if cell.contains_key("integer") {
                                        row_map.insert(
                                            name.to_string(),
                                            cell["integer"]
                                                .as_str()
                                                .unwrap()
                                                .parse::<i64>()
                                                .unwrap()
                                                .into(),
                                        );
                                    } else if cell.contains_key("real") {
                                        // Store as real.
                                        row_map.insert(
                                            name.to_string(),
                                            cell["real"]
                                                .as_str()
                                                .unwrap()
                                                .parse::<f64>()
                                                .unwrap()
                                                .into(),
                                        );
                                    } else if cell.contains_key("blob") {
                                        row_map.insert(name.to_string(), cell["blob"].clone());
                                    }
                                }

                                match rows.lock() {
                                    Ok(mut rows) => {
                                        // Copy the response into the unary_response so we can emit it from the `end` event.
                                        (*rows).push(row_map.into());
                                    }
                                    Err(e) => {
                                        println!("failed to lock unary_response: {:?}", e);
                                        return;
                                    }
                                };
                            }
                        }
                    }
                }
            });

            rpc.on(EventType::Error, {
                let json_error = Arc::clone(&json_error);
                move |response| {
                    println!("api_client: received error");
                    if let Event::Error(response) = response {
                        match json_error.lock() {
                            Ok(mut json_error) => {
                                // Copy the response into the unary_response so we can emit it from the `end` event.
                                *json_error = response;
                            }
                            Err(e) => {
                                println!("failed to lock unary_response: {:?}", e);
                                return;
                            }
                        };
                    }
                }
            });

            rpc.on(EventType::End, {
                let tx = Arc::clone(&tx);
                move |_response| {
                    println!("api_client: received end");
                    let tx = Arc::clone(&tx);
                    tokio::spawn(async move {
                        match tx.send(()).await {
                            Ok(_) => println!("unary response sent successfully"),
                            Err(e) => println!("failed to send unary response: {:?}", e),
                        }
                    });
                }
            });

            // Send the payload.
            match rpc.send(request).await {
                Ok(_) => println!("sent request"),
                Err(e) => {
                    return Err(e);
                }
            }
        }

        // Await the response.
        rx.recv().await;

        // Determine if we received an error or a response.
        let result = match json_error.lock() {
            Ok(json_error) => {
                if json_error.is_empty() {
                    // There is no error so return the response.
                    match rows.lock() {
                        Ok(rows) => Ok(rows.clone().into()),
                        Err(e) => Err(format!("failed to lock unary_response: {:?}", e)),
                    }
                } else {
                    // Propagate the error.
                    Err(json_error.clone())
                }
            }
            Err(e) => Err(format!("failed to lock unary_error: {:?}", e)),
        };

        result
    }

    /// See https://docs.tdxvolt.com/en/api/sqlite_database_api#ImportCSV
    pub async fn import_csv(
        &mut self,
        request: &serde_json::Value,
    ) -> Result<Arc<Mutex<WebsocketRpc>>, String> {
        self.streaming_call(
            "/tdx.volt_api.data.v1.SqliteDatabaseAPI/ImportCSV",
            request,
            "",
        )
        .await
    }
}