playtron-core 1.0.0

use std::{fmt, time::Duration};

use futures_util::StreamExt;
use uuid::Uuid;
use zbus::{proxy::PropertyStream, Connection};

use crate::dbus::wallet_interface::WalletInterfaceProxy;

#[derive(Debug, Clone, PartialEq)]
pub enum SuiError {
    DbusConnectionFailed(zbus::Error),
    DbusMessageFailed(zbus::fdo::Error),
    DbusInvalidArgument(zbus::Error),
    GetAddressError(zbus::Error),
    UserCancelled,
    Error(String),
    Timeout,
}

impl fmt::Display for SuiError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            SuiError::DbusConnectionFailed(e) => write!(f, "D-Bus connection failed: {e}"),
            SuiError::DbusMessageFailed(e) => write!(f, "D-Bus message failed: {e}"),
            SuiError::DbusInvalidArgument(e) => write!(f, "Invalid D-Bus argument: {e}"),
            SuiError::GetAddressError(e) => write!(f, "Failed to get wallet address: {e}"),
            SuiError::UserCancelled => write!(f, "User cancelled the operation"),
            SuiError::Error(msg) => write!(f, "{}", msg),
            SuiError::Timeout => write!(f, "Timeout"),
        }
    }
}

#[derive(Clone, Debug)]
pub struct TransactionSuccessPayload {
    pub result: String,
    pub completed_at: u64,
}

pub struct Sui {
    connection: Connection,
}

impl Sui {
    pub async fn new() -> Result<Self, SuiError> {
        let connection = Connection::session()
            .await
            .map_err(SuiError::DbusConnectionFailed)?;

        Ok(Sui { connection })
    }

    pub async fn get_address(&self) -> Result<String, SuiError> {
        let proxy = self.get_proxy().await?;
        proxy.address().await.map_err(SuiError::GetAddressError)
    }

    pub async fn receive_address_changed(&self) -> Result<PropertyStream<String>, SuiError> {
        let proxy = self.get_proxy().await?;
        Ok(proxy.receive_address_changed().await)
    }

    pub async fn sign_message(
        &self,
        provider_app_id: String,
        msg: String,
    ) -> Result<String, SuiError> {
        let proxy = self.get_proxy().await?;
        let msg_id = Uuid::new_v4().to_string();

        // Set up both signal streams before calling
        let mut signed_stream = proxy
            .receive_on_message_sign_success()
            .await
            .map_err(SuiError::DbusConnectionFailed)?;

        let mut cancelled_stream = proxy
            .receive_on_message_sign_cancelled()
            .await
            .map_err(SuiError::DbusConnectionFailed)?;

        // Trigger the signing
        proxy
            .sign_message(provider_app_id, msg_id.clone(), msg)
            .await
            .map_err(SuiError::DbusMessageFailed)?;

        let result = tokio::time::timeout(Duration::from_secs(180), async {
            loop {
                tokio::select! {
                    Some(success) = signed_stream.next() => {
                        let args = success.args().map_err(SuiError::DbusInvalidArgument)?;
                        if args.msg_id == msg_id {
                            return Ok(args.signature);
                        }
                    }
                    Some(cancelled) = cancelled_stream.next() => {
                        let args = cancelled.args().map_err(SuiError::DbusInvalidArgument)?;
                        if args.msg_id == msg_id {
                            let err = if args.error.is_empty() {
                                SuiError::UserCancelled
                            } else {
                                SuiError::Error(args.error)
                            };

                            return Err(err);
                        }
                    }
                    else => break,
                }
            }

            Err(SuiError::Timeout)
        })
        .await;

        match result {
            Ok(Ok(payload)) => Ok(payload),
            Ok(Err(e)) => Err(e),
            Err(_) => Err(SuiError::Timeout),
        }
    }

    pub async fn sign_and_execute_transaction(
        &self,
        provider_app_id: String,
        tx_bytes: String,
    ) -> Result<TransactionSuccessPayload, SuiError> {
        let proxy = self.get_proxy().await?;
        let tx_id = Uuid::new_v4().to_string();

        // Set up signal streams before triggering the call
        let mut success_stream = proxy
            .receive_on_transaction_success()
            .await
            .map_err(SuiError::DbusConnectionFailed)?;

        let mut cancelled_stream = proxy
            .receive_on_transaction_cancelled()
            .await
            .map_err(SuiError::DbusConnectionFailed)?;

        // Send the D-Bus call
        proxy
            .sign_and_execute_transaction(provider_app_id, tx_id.clone(), tx_bytes)
            .await
            .map_err(SuiError::DbusMessageFailed)?;

        let result = tokio::time::timeout(Duration::from_secs(180), async {
            loop {
                tokio::select! {
                    Some(success) = success_stream.next() => {
                        let args = success.args().map_err(SuiError::DbusInvalidArgument)?;
                        if args.tx_id == tx_id {
                            return Ok(TransactionSuccessPayload {
                                result: args.result,
                                completed_at: args.completed_at,
                            });
                        }
                    }
                    Some(cancelled) = cancelled_stream.next() => {
                        let args = cancelled.args().map_err(SuiError::DbusInvalidArgument)?;
                        if args.tx_id == tx_id {
                            let err = if args.error.is_empty() {
                                SuiError::UserCancelled
                            } else {
                                SuiError::Error(args.error)
                            };

                            return Err(err);
                        }
                    }
                    else => break,
                }
            }

            Err(SuiError::Timeout)
        })
        .await;

        match result {
            Ok(Ok(payload)) => Ok(payload),
            Ok(Err(e)) => Err(e),
            Err(_) => Err(SuiError::Timeout),
        }
    }

    pub async fn sign_transaction(
        &self,
        provider_app_id: String,
        tx_bytes: String,
    ) -> Result<String, SuiError> {
        let proxy = self.get_proxy().await?;
        let tx_id = Uuid::new_v4().to_string();

        // Set up signal streams before triggering the call
        let mut success_stream = proxy
            .receive_on_transaction_sign_success()
            .await
            .map_err(SuiError::DbusConnectionFailed)?;

        let mut cancelled_stream = proxy
            .receive_on_transaction_sign_cancelled()
            .await
            .map_err(SuiError::DbusConnectionFailed)?;

        // Send the D-Bus call
        proxy
            .sign_transaction(provider_app_id, tx_id.clone(), tx_bytes)
            .await
            .map_err(SuiError::DbusMessageFailed)?;

        let result = tokio::time::timeout(Duration::from_secs(180), async {
            loop {
                tokio::select! {
                    Some(success) = success_stream.next() => {
                        let args = success.args().map_err(SuiError::DbusInvalidArgument)?;
                        if args.tx_id == tx_id {
                            return Ok(args.signature);
                        }
                    }
                    Some(cancelled) = cancelled_stream.next() => {
                        let args = cancelled.args().map_err(SuiError::DbusInvalidArgument)?;
                        if args.tx_id == tx_id {
                            let err = if args.error.is_empty() {
                                SuiError::UserCancelled
                            } else {
                                SuiError::Error(args.error)
                            };

                            return Err(err);
                        }
                    }
                    else => break,
                }
            }

            Err(SuiError::Timeout)
        })
        .await;

        match result {
            Ok(Ok(payload)) => Ok(payload),
            Ok(Err(e)) => Err(e),
            Err(_) => Err(SuiError::Timeout),
        }
    }

    async fn get_proxy(&self) -> Result<WalletInterfaceProxy<'_>, SuiError> {
        WalletInterfaceProxy::builder(&self.connection)
            .build()
            .await
            .map_err(SuiError::DbusConnectionFailed)
    }
}