predict_sdk/

onchain.rs

//! On-chain operations for Predict.fun
//!
//! This module handles direct contract interactions for split/merge/redeem operations.

use crate::constants::Addresses;
use crate::errors::{Error, Result};
use crate::types::ChainId;
use alloy::primitives::{Address, Bytes, FixedBytes, U256};
use alloy::providers::{Provider, ProviderBuilder};
use alloy::signers::local::PrivateKeySigner;
use alloy::sol;
use alloy::sol_types::SolCall;
use tracing::{debug, info};

// Define the contract ABIs using alloy's sol! macro

sol! {
    /// Conditional Tokens contract for non-neg-risk markets
    #[sol(rpc)]
    interface IConditionalTokens {
        function splitPosition(
            address collateralToken,
            bytes32 parentCollectionId,
            bytes32 conditionId,
            uint256[] calldata partition,
            uint256 amount
        ) external;

        function mergePositions(
            address collateralToken,
            bytes32 parentCollectionId,
            bytes32 conditionId,
            uint256[] calldata partition,
            uint256 amount
        ) external;
    }
}

sol! {
    /// Neg Risk Adapter for neg-risk markets
    #[sol(rpc)]
    interface INegRiskAdapter {
        #[allow(non_snake_case)]
        function splitPosition(bytes32 conditionId, uint256 amount) external;

        #[allow(non_snake_case)]
        function mergePositions(bytes32 conditionId, uint256 amount) external;
    }
}

sol! {
    /// Kernel smart wallet contract
    #[sol(rpc)]
    interface IKernel {
        function execute(bytes32 mode, bytes calldata executionCalldata) external payable returns (bytes memory);
    }
}

sol! {
    /// ERC20 interface for approvals
    #[sol(rpc)]
    interface IERC20 {
        function approve(address spender, uint256 amount) external returns (bool);
        function allowance(address owner, address spender) external view returns (uint256);
        function balanceOf(address account) external view returns (uint256);
    }
}

sol! {
    /// ERC1155 interface for operator approvals and transfers (Conditional Tokens)
    #[sol(rpc)]
    interface IERC1155 {
        function setApprovalForAll(address operator, bool approved) external;
        function isApprovedForAll(address account, address operator) external view returns (bool);
        function safeTransferFrom(address from, address to, uint256 id, uint256 amount, bytes calldata data) external;
        function balanceOf(address account, uint256 id) external view returns (uint256);
    }
}

/// Execution mode for Kernel smart wallet (single call)
const KERNEL_EXEC_MODE: [u8; 32] = [0u8; 32];

/// Options for split operation
#[derive(Debug, Clone)]
pub struct SplitOptions {
    /// Condition ID of the market
    pub condition_id: String,
    /// Amount to split (in USDT, will be converted to wei - 18 decimals)
    pub amount: f64,
    /// Whether this is a neg-risk market
    pub is_neg_risk: bool,
    /// Whether this is a yield-bearing market
    pub is_yield_bearing: bool,
}

/// On-chain client for Predict.fun operations
pub struct OnchainClient {
    chain_id: ChainId,
    signer: PrivateKeySigner,
    addresses: Addresses,
    rpc_url: String,
    /// Predict Account address for smart wallet operations
    predict_account: Option<Address>,
}

impl OnchainClient {
    /// Create a new OnchainClient for direct EOA operations
    pub fn new(chain_id: ChainId, signer: PrivateKeySigner) -> Self {
        let addresses = Addresses::for_chain(chain_id);
        let rpc_url = match chain_id {
            ChainId::BnbMainnet => "https://bsc-dataseed.bnbchain.org/".to_string(),
            ChainId::BnbTestnet => "https://bsc-testnet-dataseed.bnbchain.org/".to_string(),
        };

        Self {
            chain_id,
            signer,
            addresses,
            rpc_url,
            predict_account: None,
        }
    }

    /// Create a new OnchainClient for Predict Account (smart wallet) operations
    pub fn with_predict_account(
        chain_id: ChainId,
        signer: PrivateKeySigner,
        predict_account: &str,
    ) -> Result<Self> {
        let mut client = Self::new(chain_id, signer);
        client.predict_account = Some(
            predict_account
                .parse()
                .map_err(|e| Error::Other(format!("Invalid predict account address: {}", e)))?,
        );
        Ok(client)
    }

    /// Get the signer address
    pub fn signer_address(&self) -> Address {
        self.signer.address()
    }

    /// Get the trading address (predict_account if set, otherwise signer)
    pub fn trading_address(&self) -> Address {
        self.predict_account.unwrap_or_else(|| self.signer.address())
    }

    /// Check if using smart wallet
    pub fn is_smart_wallet(&self) -> bool {
        self.predict_account.is_some()
    }

    /// Get the addresses configuration
    pub fn addresses(&self) -> &Addresses {
        &self.addresses
    }

    /// Set all necessary approvals for trading on the CTF Exchange.
    ///
    /// This sets:
    /// 1. ERC-1155 `setApprovalForAll` on Conditional Tokens → CTF Exchange
    ///    (allows exchange to move outcome tokens on behalf of the trader)
    /// 2. ERC-20 `approve` on USDT → CTF Exchange
    ///    (allows exchange to take USDT collateral)
    ///
    /// For neg-risk markets, also approves the Neg Risk Adapter.
    pub async fn set_approvals(
        &self,
        is_neg_risk: bool,
        is_yield_bearing: bool,
    ) -> Result<()> {
        let provider = ProviderBuilder::new()
            .wallet(alloy::network::EthereumWallet::from(self.signer.clone()))
            .connect_http(self.rpc_url.parse().unwrap());

        let owner = self.trading_address();

        // 1. ERC-1155 approval: Conditional Tokens → CTF Exchange
        let ct_address: Address = self.addresses.get_conditional_tokens(is_yield_bearing, is_neg_risk)
            .parse().unwrap();
        let exchange_address: Address = self.addresses.get_ctf_exchange(is_yield_bearing, is_neg_risk)
            .parse().unwrap();

        let ct = IERC1155::new(ct_address, provider.clone());
        let is_approved = ct
            .isApprovedForAll(owner, exchange_address)
            .call()
            .await
            .map_err(|e| Error::Other(format!("Failed to check ERC-1155 approval: {}", e)))?;

        if !is_approved {
            info!("Setting ERC-1155 approval: {} → {}", ct_address, exchange_address);
            let tx = ct
                .setApprovalForAll(exchange_address, true)
                .send()
                .await
                .map_err(|e| Error::Other(format!("Failed to send setApprovalForAll: {}", e)))?;
            let receipt = tx
                .get_receipt()
                .await
                .map_err(|e| Error::Other(format!("Failed to get approval receipt: {}", e)))?;
            if !receipt.status() {
                return Err(Error::Other(format!(
                    "setApprovalForAll reverted: {:?}", receipt.transaction_hash
                )));
            }
            info!("ERC-1155 approval set: {:?}", receipt.transaction_hash);
        } else {
            debug!("ERC-1155 already approved: {} → {}", ct_address, exchange_address);
        }

        // 2. ERC-20 approval: USDT → CTF Exchange
        let usdt_address: Address = self.addresses.usdt.parse().unwrap();
        let usdt = IERC20::new(usdt_address, provider.clone());
        let allowance = usdt
            .allowance(owner, exchange_address)
            .call()
            .await
            .map_err(|e| Error::Other(format!("Failed to check USDT allowance: {}", e)))?;

        if allowance < U256::from(1_000_000_000_000_000_000_000_u128) {
            // Less than 1000 USDT allowance, approve max
            info!("Approving USDT for CTF Exchange: {}", exchange_address);
            let tx = usdt
                .approve(exchange_address, U256::MAX)
                .send()
                .await
                .map_err(|e| Error::Other(format!("Failed to send USDT approval: {}", e)))?;
            let receipt = tx
                .get_receipt()
                .await
                .map_err(|e| Error::Other(format!("Failed to get USDT approval receipt: {}", e)))?;
            if !receipt.status() {
                return Err(Error::Other(format!(
                    "USDT approval reverted: {:?}", receipt.transaction_hash
                )));
            }
            info!("USDT approval set: {:?}", receipt.transaction_hash);
        } else {
            debug!("USDT already approved for CTF Exchange");
        }

        // 3. For neg-risk markets, also approve the Neg Risk Adapter
        if is_neg_risk {
            let adapter_address: Address = if is_yield_bearing {
                self.addresses.yield_bearing_neg_risk_adapter
            } else {
                self.addresses.neg_risk_adapter
            }.parse().unwrap();

            let is_adapter_approved = ct
                .isApprovedForAll(owner, adapter_address)
                .call()
                .await
                .map_err(|e| Error::Other(format!("Failed to check adapter approval: {}", e)))?;

            if !is_adapter_approved {
                info!("Setting ERC-1155 approval for Neg Risk Adapter: {}", adapter_address);
                let tx = ct
                    .setApprovalForAll(adapter_address, true)
                    .send()
                    .await
                    .map_err(|e| Error::Other(format!("Failed to approve adapter: {}", e)))?;
                let receipt = tx
                    .get_receipt()
                    .await
                    .map_err(|e| Error::Other(format!("Failed to get adapter approval receipt: {}", e)))?;
                if !receipt.status() {
                    return Err(Error::Other(format!(
                        "Adapter approval reverted: {:?}", receipt.transaction_hash
                    )));
                }
                info!("Neg Risk Adapter approval set: {:?}", receipt.transaction_hash);
            }
        }

        Ok(())
    }

    /// Split USDT into UP/DOWN outcome tokens
    ///
    /// # Arguments
    /// * `options` - Split options including condition_id, amount, and market type
    ///
    /// # Returns
    /// Transaction hash on success
    pub async fn split_positions(&self, options: SplitOptions) -> Result<String> {
        info!(
            "Splitting {} USDT for condition {} (neg_risk={}, yield_bearing={})",
            options.amount, options.condition_id, options.is_neg_risk, options.is_yield_bearing
        );

        // Convert amount to wei (USDT has 18 decimals on BNB Chain)
        let amount_wei = U256::from((options.amount * 1e18) as u128);

        // Parse condition ID
        let condition_id: FixedBytes<32> = options
            .condition_id
            .parse()
            .map_err(|e| Error::Other(format!("Invalid condition ID: {}", e)))?;

        // Create provider with signer
        let provider = ProviderBuilder::new()
            .wallet(alloy::network::EthereumWallet::from(self.signer.clone()))
            .connect_http(self.rpc_url.parse().unwrap());

        // First, ensure USDT approval
        self.ensure_usdt_approval(&provider, amount_wei, options.is_neg_risk, options.is_yield_bearing)
            .await?;

        // Execute split based on wallet type and market type
        let tx_hash = if self.is_smart_wallet() {
            self.split_via_kernel(&provider, condition_id, amount_wei, &options)
                .await?
        } else {
            self.split_direct(&provider, condition_id, amount_wei, &options)
                .await?
        };

        info!("Split transaction submitted: {}", tx_hash);
        Ok(tx_hash)
    }

    /// Ensure USDT is approved for the target contract
    async fn ensure_usdt_approval<P: Provider + Clone>(
        &self,
        provider: &P,
        amount: U256,
        is_neg_risk: bool,
        is_yield_bearing: bool,
    ) -> Result<()> {
        let usdt_address: Address = self.addresses.usdt.parse().unwrap();
        let spender = self.get_target_contract(is_neg_risk, is_yield_bearing);
        let owner = self.trading_address();

        let usdt = IERC20::new(usdt_address, provider.clone());

        // Check current allowance
        let allowance = usdt
            .allowance(owner, spender)
            .call()
            .await
            .map_err(|e| Error::Other(format!("Failed to check allowance: {}", e)))?;

        if allowance < amount {
            info!("Approving USDT spend for {:?}", spender);

            // Approve max amount
            let approve_call = usdt.approve(spender, U256::MAX);

            if self.is_smart_wallet() {
                // Execute approval through Kernel
                let encoded = approve_call.calldata().clone();
                self.execute_via_kernel(provider, usdt_address, encoded)
                    .await?;
            } else {
                // Direct approval
                let tx = approve_call
                    .send()
                    .await
                    .map_err(|e| Error::Other(format!("Failed to send approval: {}", e)))?;
                let receipt = tx
                    .get_receipt()
                    .await
                    .map_err(|e| Error::Other(format!("Failed to get approval receipt: {}", e)))?;

                // Check transaction status
                if !receipt.status() {
                    return Err(Error::Other(format!(
                        "Approval transaction reverted: {:?}",
                        receipt.transaction_hash
                    )));
                }

                debug!("Approval tx: {:?}", receipt.transaction_hash);
            }
        }

        Ok(())
    }

    /// Get the target contract address for split operations
    fn get_target_contract(&self, is_neg_risk: bool, is_yield_bearing: bool) -> Address {
        let addr_str = if is_neg_risk {
            if is_yield_bearing {
                self.addresses.yield_bearing_neg_risk_adapter
            } else {
                self.addresses.neg_risk_adapter
            }
        } else if is_yield_bearing {
            self.addresses.yield_bearing_conditional_tokens
        } else {
            self.addresses.conditional_tokens
        };
        addr_str.parse().unwrap()
    }

    /// Split directly (EOA wallet)
    async fn split_direct<P: Provider + Clone>(
        &self,
        provider: &P,
        condition_id: FixedBytes<32>,
        amount: U256,
        options: &SplitOptions,
    ) -> Result<String> {
        let target = self.get_target_contract(options.is_neg_risk, options.is_yield_bearing);

        if options.is_neg_risk {
            // Neg risk: splitPosition(bytes32, uint256)
            let contract = INegRiskAdapter::new(target, provider.clone());
            let tx = contract
                .splitPosition(condition_id, amount)
                .send()
                .await
                .map_err(|e| Error::Other(format!("Failed to send split tx: {}", e)))?;

            let receipt = tx
                .get_receipt()
                .await
                .map_err(|e| Error::Other(format!("Failed to get receipt: {}", e)))?;

            // Check transaction status (0 = reverted, 1 = success)
            if !receipt.status() {
                return Err(Error::Other(format!(
                    "Transaction reverted: {:?}",
                    receipt.transaction_hash
                )));
            }

            Ok(format!("{:?}", receipt.transaction_hash))
        } else {
            // Regular: splitPosition(address, bytes32, bytes32, uint256[], uint256)
            let contract = IConditionalTokens::new(target, provider.clone());
            let usdt: Address = self.addresses.usdt.parse().unwrap();
            let parent_collection = FixedBytes::<32>::ZERO;
            let partition = vec![U256::from(1), U256::from(2)];

            let tx = contract
                .splitPosition(usdt, parent_collection, condition_id, partition, amount)
                .send()
                .await
                .map_err(|e| Error::Other(format!("Failed to send split tx: {}", e)))?;

            let receipt = tx
                .get_receipt()
                .await
                .map_err(|e| Error::Other(format!("Failed to get receipt: {}", e)))?;

            // Check transaction status (0 = reverted, 1 = success)
            if !receipt.status() {
                return Err(Error::Other(format!(
                    "Transaction reverted: {:?}",
                    receipt.transaction_hash
                )));
            }

            Ok(format!("{:?}", receipt.transaction_hash))
        }
    }

    /// Split via Kernel smart wallet
    async fn split_via_kernel<P: Provider + Clone>(
        &self,
        provider: &P,
        condition_id: FixedBytes<32>,
        amount: U256,
        options: &SplitOptions,
    ) -> Result<String> {
        let target = self.get_target_contract(options.is_neg_risk, options.is_yield_bearing);

        // Encode the split call
        let calldata = if options.is_neg_risk {
            // Neg risk: splitPosition(bytes32, uint256)
            let call = INegRiskAdapter::splitPositionCall {
                conditionId: condition_id,
                amount,
            };
            Bytes::from(call.abi_encode())
        } else {
            // Regular: splitPosition(address, bytes32, bytes32, uint256[], uint256)
            let usdt: Address = self.addresses.usdt.parse().unwrap();
            let parent_collection = FixedBytes::<32>::ZERO;
            let partition = vec![U256::from(1), U256::from(2)];

            let call = IConditionalTokens::splitPositionCall {
                collateralToken: usdt,
                parentCollectionId: parent_collection,
                conditionId: condition_id,
                partition,
                amount,
            };
            Bytes::from(call.abi_encode())
        };

        self.execute_via_kernel(provider, target, calldata).await
    }

    /// Execute a call through the Kernel smart wallet
    async fn execute_via_kernel<P: Provider + Clone>(
        &self,
        provider: &P,
        target: Address,
        calldata: Bytes,
    ) -> Result<String> {
        let predict_account = self
            .predict_account
            .ok_or_else(|| Error::Other("No predict account configured".to_string()))?;

        // Encode execution calldata: target (20 bytes) + value (32 bytes) + calldata
        let mut execution_calldata = Vec::new();
        execution_calldata.extend_from_slice(target.as_slice());
        execution_calldata.extend_from_slice(&U256::ZERO.to_be_bytes::<32>());
        execution_calldata.extend_from_slice(&calldata);

        let kernel = IKernel::new(predict_account, provider.clone());
        let mode = FixedBytes::<32>::from(KERNEL_EXEC_MODE);

        let tx = kernel
            .execute(mode, Bytes::from(execution_calldata))
            .send()
            .await
            .map_err(|e| Error::Other(format!("Failed to send kernel execute: {}", e)))?;

        let receipt = tx
            .get_receipt()
            .await
            .map_err(|e| Error::Other(format!("Failed to get receipt: {}", e)))?;

        // Check transaction status (0 = reverted, 1 = success)
        if !receipt.status() {
            return Err(Error::Other(format!(
                "Kernel execute reverted: {:?}",
                receipt.transaction_hash
            )));
        }

        Ok(format!("{:?}", receipt.transaction_hash))
    }
}

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

    #[test]
    fn test_create_onchain_client() {
        let signer = PrivateKeySigner::random();
        let client = OnchainClient::new(ChainId::BnbTestnet, signer);
        assert!(!client.is_smart_wallet());
    }

    #[test]
    fn test_create_smart_wallet_client() {
        let signer = PrivateKeySigner::random();
        let client = OnchainClient::with_predict_account(
            ChainId::BnbTestnet,
            signer,
            "0x1234567890123456789012345678901234567890",
        )
        .unwrap();
        assert!(client.is_smart_wallet());
    }
}