bolt-cw-sdk 1.0.0

use super::error::CosmosClientError;
use super::gas_info::ClientGasInfo;
use std::future::Future;
use tendermint_rpc::event::Event;
use tendermint_rpc::query::{Condition, EventType, Query};
use tendermint_rpc::{HttpClient, Subscription, SubscriptionClient, WebSocketClient};
use tokio::sync::mpsc;
use tokio::sync::mpsc::Receiver;
use tokio::task::JoinHandle;
use tonic::codegen::tokio_stream::StreamExt;
use tonic::Status;
use tracing::info;

#[derive(Debug)]
pub struct CosmosClient<ClientT = HttpClient> {
    pub(crate) client: ClientT,
    pub(crate) driver_handle: Option<JoinHandle<Result<(), CosmosClientError>>>,
    pub(crate) gas_info: Option<ClientGasInfo>,
}

impl CosmosClient<WebSocketClient> {
    pub async fn new(websocket_url: &str) -> Result<Self, CosmosClientError> {
        info!(
            "Creating a new CosmosClient with WebSocket URL: {}",
            websocket_url
        );
        let (client, driver) = WebSocketClient::new(websocket_url).await?;

        // Run the driver in a separate thread
        let driver_handle = Some(tokio::spawn(async move {
            driver.run().await.map_err(CosmosClientError::from)
        }));

        Ok(Self {
            client,
            driver_handle,
            gas_info: None,
        })
    }

    /// Subscribe to a specific wasm action.
    /// When specifying a contract, it'll only listen to the specific event from a specific contract address.
    pub async fn subscribe_to_action(
        &self,
        action: &str,
        from_contract: Option<&str>,
    ) -> Result<Subscription, CosmosClientError> {
        let conditions = vec![from_contract
            .map(|addr| Condition::eq(format!("{action}._contract_address"), addr.into()))
            .unwrap_or(Condition::exists(action.to_string()))];

        self.client
            .subscribe(Query {
                event_type: Some(EventType::Tx),
                conditions,
            })
            .await
            .map_err(CosmosClientError::from)
    }

    /// Shorthand for quickly creating a mpsc stream, the function acts as a map
    pub async fn map_subscription<F, T, Fut>(
        &self,
        action: &str,
        from_contract: Option<&str>,
        process_event: F,
    ) -> Result<Receiver<Result<T, Status>>, CosmosClientError>
    where
        T: Send + 'static,
        F: Fn(Result<Event, tendermint_rpc::Error>) -> Fut + Send + 'static,
        Fut: Future<Output = Result<T, Status>> + Send,
    {
        info!(
            "Subscribing to action: {} from contract: {:?}",
            action, from_contract
        );
        let mut subscription = self.subscribe_to_action(action, from_contract).await?;
        let (tx, rx) = mpsc::channel(5);

        tokio::spawn(async move {
            while let Some(next) = subscription.next().await {
                // The processed event can return errors, but this won't explicitly close the connection
                let res = process_event(next).await;

                // This only errors if the connection is dropped, so we use it to drop our side too
                if tx.send(res).await.is_err() {
                    break;
                }
            }
        });

        Ok(rx)
    }

    /// A `map_subscription` that can filter out events
    /// This means that instead of returning a value,
    /// it expects an optional value, if `None` nothing happens.
    pub async fn filtered_subscription<F, T, Fut>(
        &self,
        action: &str,
        from_contract: Option<&str>,
        process_event: F,
    ) -> Result<Receiver<Result<T, Status>>, CosmosClientError>
    where
        T: Send + 'static,
        F: Fn(Result<Event, tendermint_rpc::Error>) -> Fut + Send + 'static,
        Fut: Future<Output = Result<Option<T>, Status>> + Send,
    {
        info!(
            "Subscribing to action: {} from contract: {:?}",
            action, from_contract
        );
        let mut subscription = self.subscribe_to_action(action, from_contract).await?;
        let (tx, rx) = mpsc::channel(5);

        tokio::spawn(async move {
            while let Some(next) = subscription.next().await {
                if let Some(res) = process_event(next).await.transpose() {
                    if tx.send(res).await.is_err() {
                        break;
                    }
                }
            }
        });

        Ok(rx)
    }

    pub async fn close(self) -> Result<(), CosmosClientError> {
        info!("Closing the WebSocketClient connection");
        let Self {
            client,
            driver_handle,
            ..
        } = self;
        client.close()?;
        driver_handle.unwrap().await??;

        Ok(())
    }
}