ostium-rust-sdk 0.1.0

use alloy::network::Network;
use alloy::primitives::{aliases::U192, Address, U256};
use alloy::providers::Provider;
use alloy::sol_types::sol;

use crate::error::{OstiumError, Result};
use crate::types::{OpenOrderType, Trade, UnsignedTransaction, UnsignedTransactionParams};

// Define the Trading contract interface
sol! {
    #[sol(rpc)]
    contract OstiumTrading {
        // Structs
        struct Trade {
            uint256 collateral;
            uint192 openPrice;
            uint192 tp;
            uint192 sl;
            address trader;
            uint32 leverage;
            uint16 pairIndex;
            uint8 index;
            bool buy;
        }

        // Enums
        enum OpenOrderType {
            /// Market order - executed immediately at current market price
            MARKET,
            /// Limit order - executed when price reaches specified level
            LIMIT,
            /// Stop order - executed when price crosses stop level
            STOP
        }

        enum AutomationOrderStatus {
            /// Order executed successfully
            SUCCESS,
            /// Order execution failed
            FAILED
        }

        // Functions
        function openTrade(Trade calldata t, OpenOrderType orderType, uint256 slippageP) external;
        function closeTradeMarket(uint16 pairIndex, uint8 index, uint16 closePercentage) external;
        function cancelOpenLimitOrder(uint16 pairIndex, uint8 index) external;
        function updateTp(uint16 pairIndex, uint8 index, uint192 newTp) external;
        function updateSl(uint16 pairIndex, uint8 index, uint192 newSl) external;
        function updateOpenLimitOrder(uint16 pairIndex, uint8 index, uint192 price, uint192 tp, uint192 sl) external;
        function topUpCollateral(uint16 pairIndex, uint8 index, uint256 topUpAmount) external;
        function removeCollateral(uint16 pairIndex, uint8 index, uint256 removeAmount) external;
        function setDelegate(address delegate) external;
        function removeDelegate() external;
        function delegatedAction(address trader, bytes calldata call_data) external returns (bytes memory);
        function openTradeMarketTimeout(uint256 _order) external;
        function closeTradeMarketTimeout(uint256 _order, bool retry) external;
    }
}

/// Trading contract wrapper
pub struct TradingContract<P: Provider<N>, N: Network = alloy::network::Ethereum> {
    contract: OstiumTrading::OstiumTradingInstance<P, N>,
}

impl<P: Provider<N>, N: Network> TradingContract<P, N> {
    /// Create a new Trading contract instance
    pub fn new(address: Address, provider: P) -> Self {
        let contract = OstiumTrading::new(address, provider);
        Self { contract }
    }

    /// Open a new trade
    pub async fn open_trade(
        &self,
        trade: Trade,
        order_type: OpenOrderType,
        slippage_p: U256,
    ) -> Result<()> {
        let sol_trade = OstiumTrading::Trade {
            collateral: trade.collateral,
            openPrice: U192::from(trade.open_price),
            tp: U192::from(trade.tp),
            sl: U192::from(trade.sl),
            trader: trade.trader,
            leverage: trade.leverage,
            pairIndex: trade.pair_index,
            index: trade.index,
            buy: trade.buy,
        };

        let sol_order_type = match order_type {
            OpenOrderType::Market => OstiumTrading::OpenOrderType::MARKET,
            OpenOrderType::Limit => OstiumTrading::OpenOrderType::LIMIT,
            OpenOrderType::Stop => OstiumTrading::OpenOrderType::STOP,
        };

        let _receipt = self
            .contract
            .openTrade(sol_trade, sol_order_type, slippage_p)
            .send()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to open trade: {}", e)))?
            .get_receipt()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to get receipt: {}", e)))?;

        Ok(())
    }

    /// Close a trade at market price
    pub async fn close_trade_market(
        &self,
        pair_index: u16,
        index: u8,
        close_percentage: u16,
    ) -> Result<()> {
        let _receipt = self
            .contract
            .closeTradeMarket(pair_index, index, close_percentage)
            .send()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to close trade: {}", e)))?
            .get_receipt()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to get receipt: {}", e)))?;

        Ok(())
    }

    /// Cancel an open limit order
    pub async fn cancel_open_limit_order(&self, pair_index: u16, index: u8) -> Result<()> {
        let _receipt = self
            .contract
            .cancelOpenLimitOrder(pair_index, index)
            .send()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to cancel order: {}", e)))?
            .get_receipt()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to get receipt: {}", e)))?;

        Ok(())
    }

    /// Update take profit price
    pub async fn update_tp(&self, pair_index: u16, index: u8, new_tp: U192) -> Result<()> {
        let _receipt = self
            .contract
            .updateTp(pair_index, index, new_tp)
            .send()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to update TP: {}", e)))?
            .get_receipt()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to get receipt: {}", e)))?;

        Ok(())
    }

    /// Update stop loss price
    pub async fn update_sl(&self, pair_index: u16, index: u8, new_sl: U192) -> Result<()> {
        let _receipt = self
            .contract
            .updateSl(pair_index, index, new_sl)
            .send()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to update SL: {}", e)))?
            .get_receipt()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to get receipt: {}", e)))?;

        Ok(())
    }

    /// Update an open limit order
    pub async fn update_open_limit_order(
        &self,
        pair_index: u16,
        index: u8,
        price: U192,
        tp: U192,
        sl: U192,
    ) -> Result<()> {
        let _receipt = self
            .contract
            .updateOpenLimitOrder(pair_index, index, price, tp, sl)
            .send()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to update order: {}", e)))?
            .get_receipt()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to get receipt: {}", e)))?;

        Ok(())
    }

    /// Add collateral to an existing position
    pub async fn top_up_collateral(&self, pair_index: u16, index: u8, amount: U256) -> Result<()> {
        let _receipt = self
            .contract
            .topUpCollateral(pair_index, index, amount)
            .send()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to add collateral: {}", e)))?
            .get_receipt()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to get receipt: {}", e)))?;

        Ok(())
    }

    /// Remove collateral from an existing position
    pub async fn remove_collateral(&self, pair_index: u16, index: u8, amount: U256) -> Result<()> {
        let _receipt = self
            .contract
            .removeCollateral(pair_index, index, amount)
            .send()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to remove collateral: {}", e)))?
            .get_receipt()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to get receipt: {}", e)))?;

        Ok(())
    }

    /// Set a delegate for trading operations
    pub async fn set_delegate(&self, delegate: Address) -> Result<()> {
        let _receipt = self
            .contract
            .setDelegate(delegate)
            .send()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to set delegate: {}", e)))?
            .get_receipt()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to get receipt: {}", e)))?;

        Ok(())
    }

    /// Remove the current delegate
    pub async fn remove_delegate(&self) -> Result<()> {
        let _receipt = self
            .contract
            .removeDelegate()
            .send()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to remove delegate: {}", e)))?
            .get_receipt()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to get receipt: {}", e)))?;

        Ok(())
    }

    /// Build unsigned transaction for opening a trade
    pub async fn open_trade_unsigned(
        &self,
        trade: Trade,
        order_type: OpenOrderType,
        slippage_p: U256,
        tx_params: UnsignedTransactionParams,
    ) -> Result<UnsignedTransaction> {
        let sol_trade = OstiumTrading::Trade {
            collateral: trade.collateral,
            openPrice: U192::from(trade.open_price),
            tp: U192::from(trade.tp),
            sl: U192::from(trade.sl),
            trader: trade.trader,
            leverage: trade.leverage,
            pairIndex: trade.pair_index,
            index: trade.index,
            buy: trade.buy,
        };

        let sol_order_type = match order_type {
            OpenOrderType::Market => OstiumTrading::OpenOrderType::MARKET,
            OpenOrderType::Limit => OstiumTrading::OpenOrderType::LIMIT,
            OpenOrderType::Stop => OstiumTrading::OpenOrderType::STOP,
        };

        // Create the function call
        let call = OstiumTrading::openTradeCall {
            t: sol_trade,
            orderType: sol_order_type,
            slippageP: slippage_p,
        };

        self.build_unsigned_transaction_from_call(call, tx_params)
            .await
    }

    /// Build unsigned transaction for closing a trade
    pub async fn close_trade_market_unsigned(
        &self,
        pair_index: u16,
        index: u8,
        close_percentage: u16,
        tx_params: UnsignedTransactionParams,
    ) -> Result<UnsignedTransaction> {
        let call = OstiumTrading::closeTradeMarketCall {
            pairIndex: pair_index,
            index,
            closePercentage: close_percentage,
        };

        self.build_unsigned_transaction_from_call(call, tx_params)
            .await
    }

    /// Build unsigned transaction for canceling a limit order
    pub async fn cancel_open_limit_order_unsigned(
        &self,
        pair_index: u16,
        index: u8,
        tx_params: UnsignedTransactionParams,
    ) -> Result<UnsignedTransaction> {
        let call = OstiumTrading::cancelOpenLimitOrderCall {
            pairIndex: pair_index,
            index,
        };

        self.build_unsigned_transaction_from_call(call, tx_params)
            .await
    }

    /// Build unsigned transaction for updating take profit
    pub async fn update_tp_unsigned(
        &self,
        pair_index: u16,
        index: u8,
        new_tp: U192,
        tx_params: UnsignedTransactionParams,
    ) -> Result<UnsignedTransaction> {
        let call = OstiumTrading::updateTpCall {
            pairIndex: pair_index,
            index,
            newTp: new_tp,
        };

        self.build_unsigned_transaction_from_call(call, tx_params)
            .await
    }

    /// Build unsigned transaction for updating stop loss
    pub async fn update_sl_unsigned(
        &self,
        pair_index: u16,
        index: u8,
        new_sl: U192,
        tx_params: UnsignedTransactionParams,
    ) -> Result<UnsignedTransaction> {
        let call = OstiumTrading::updateSlCall {
            pairIndex: pair_index,
            index,
            newSl: new_sl,
        };

        self.build_unsigned_transaction_from_call(call, tx_params)
            .await
    }

    /// Build unsigned transaction for updating a limit order
    pub async fn update_open_limit_order_unsigned(
        &self,
        pair_index: u16,
        index: u8,
        price: U192,
        tp: U192,
        sl: U192,
        tx_params: UnsignedTransactionParams,
    ) -> Result<UnsignedTransaction> {
        let call = OstiumTrading::updateOpenLimitOrderCall {
            pairIndex: pair_index,
            index,
            price,
            tp,
            sl,
        };

        self.build_unsigned_transaction_from_call(call, tx_params)
            .await
    }

    /// Generic method to build unsigned transactions from sol calls
    async fn build_unsigned_transaction_from_call<C>(
        &self,
        call: C,
        tx_params: UnsignedTransactionParams,
    ) -> Result<UnsignedTransaction>
    where
        C: alloy::sol_types::SolCall,
    {
        // Get the contract address
        let to = *self.contract.address();

        // Get the encoded call data
        let data = call.abi_encode();

        // Get chain ID from provider
        let chain_id = self
            .contract
            .provider()
            .get_chain_id()
            .await
            .map_err(|e| OstiumError::contract(format!("Failed to get chain ID: {}", e)))?;

        // Estimate gas if requested
        let gas_limit = if tx_params.include_gas_estimates {
            // For gas estimation, we can use a simple eth_estimateGas call
            // This is a simplified approach - in production you might want more sophisticated gas estimation
            Some(U256::from(200_000)) // Default gas limit
        } else {
            None
        };

        // Get gas price estimate if requested
        let gas_price = if tx_params.include_gas_estimates {
            match self.contract.provider().get_gas_price().await {
                Ok(price) => Some(U256::from(price)),
                Err(e) => {
                    tracing::warn!("Failed to get gas price: {}", e);
                    None
                }
            }
        } else {
            None
        };

        // Get nonce if requested
        let nonce = if tx_params.include_nonce {
            match self
                .contract
                .provider()
                .get_transaction_count(tx_params.from)
                .await
            {
                Ok(count) => Some(count),
                Err(e) => {
                    tracing::warn!("Failed to get nonce: {}", e);
                    None
                }
            }
        } else {
            None
        };

        Ok(UnsignedTransaction {
            to,
            data,
            value: U256::ZERO, // Contract calls typically don't send ETH
            gas_limit,
            gas_price,
            chain_id,
            nonce,
        })
    }
}