predict-sdk 0.1.0

use crate::{
    constants::{self, Addresses, ZERO_ADDRESS},
    errors::{Error, Result},
    internal::{amounts, signing},
    types::*,
};
use alloy::primitives::Address;
use alloy::signers::local::PrivateKeySigner;
use std::time::{SystemTime, UNIX_EPOCH};

/// Options for creating an OrderBuilder
#[derive(Debug, Clone)]
pub struct OrderBuilderOptions {
    /// Custom addresses (defaults to chain-specific addresses)
    pub addresses: Option<Addresses>,

    /// Function to generate salt for orders
    pub generate_salt: Option<fn() -> String>,
}

impl Default for OrderBuilderOptions {
    fn default() -> Self {
        Self {
            addresses: None,
            generate_salt: Some(generate_order_salt),
        }
    }
}

/// Default function to generate a random salt for orders
pub fn generate_order_salt() -> String {
    use rand::Rng;
    let mut rng = rand::rng();
    let salt: u64 = rng.random_range(0..constants::MAX_SALT);
    salt.to_string()
}

/// Main OrderBuilder struct for creating and signing orders
///
/// This is the primary interface for interact with predict.fun's CTF Exchange
pub struct OrderBuilder {
    chain_id: ChainId,
    signer: Option<PrivateKeySigner>,
    addresses: Addresses,
    generate_salt: fn() -> String,
    /// Predict Account address for smart wallet signing (Kernel-based)
    /// When set, orders use this address as `maker` and signatures use Kernel wrapping
    predict_account: Option<Address>,
}

impl OrderBuilder {
    /// Create a new OrderBuilder
    ///
    /// # Arguments
    ///
    /// * `chain_id` - The chain ID (BNB Mainnet or Testnet)
    /// * `signer` - Optional signer for signing orders
    /// * `options` - Optional configuration options
    ///
    /// # Returns
    ///
    /// A new OrderBuilder instance
    pub fn new(
        chain_id: ChainId,
        signer: Option<PrivateKeySigner>,
        options: Option<OrderBuilderOptions>,
    ) -> Result<Self> {
        let opts = options.unwrap_or_default();
        let addresses = opts.addresses.unwrap_or_else(|| Addresses::for_chain(chain_id));
        let generate_salt = opts.generate_salt.unwrap_or(generate_order_salt);

        Ok(Self {
            chain_id,
            signer,
            addresses,
            generate_salt,
            predict_account: None,
        })
    }

    /// Create a new OrderBuilder with Predict Account support
    ///
    /// Use this constructor when trading via a Predict Smart Wallet (Kernel-based).
    /// The `predict_account` address will be used as the order `maker`, and
    /// signatures will use Kernel EIP-712 wrapping.
    ///
    /// # Arguments
    ///
    /// * `chain_id` - The chain ID (BNB Mainnet or Testnet)
    /// * `signer` - The Privy private key signer
    /// * `predict_account` - The Predict Account (smart wallet) address
    /// * `options` - Optional configuration options
    ///
    /// # Returns
    ///
    /// A new OrderBuilder instance configured for Predict Account signing
    pub fn with_predict_account(
        chain_id: ChainId,
        signer: PrivateKeySigner,
        predict_account: &str,
        options: Option<OrderBuilderOptions>,
    ) -> Result<Self> {
        let opts = options.unwrap_or_default();
        let addresses = opts.addresses.unwrap_or_else(|| Addresses::for_chain(chain_id));
        let generate_salt = opts.generate_salt.unwrap_or(generate_order_salt);

        let predict_account_addr = predict_account.parse::<Address>()
            .map_err(|e| Error::Other(format!("Invalid predict account address: {}", e)))?;

        Ok(Self {
            chain_id,
            signer: Some(signer),
            addresses,
            generate_salt,
            predict_account: Some(predict_account_addr),
        })
    }

    /// Check if this OrderBuilder uses Predict Account signing
    pub fn uses_predict_account(&self) -> bool {
        self.predict_account.is_some()
    }

    /// Get the Predict Account address if set
    pub fn predict_account(&self) -> Option<Address> {
        self.predict_account
    }

    /// Get the signer address
    ///
    /// # Returns
    ///
    /// The signer address, or an error if no signer is configured
    pub fn signer_address(&self) -> Result<Address> {
        self.signer
            .as_ref()
            .map(|s| s.address())
            .ok_or_else(|| Error::Other("No signer configured".to_string()))
    }

    /// Get a clone of the signer (for on-chain operations)
    ///
    /// # Returns
    ///
    /// A clone of the signer, or None if no signer is configured
    pub fn signer(&self) -> Option<PrivateKeySigner> {
        self.signer.clone()
    }

    /// Helper function to calculate the amounts for a LIMIT strategy order
    ///
    /// # Arguments
    ///
    /// * `data` - The limit order data (side, price, quantity)
    ///
    /// # Returns
    ///
    /// Order amounts including maker/taker amounts and price per share
    ///
    /// # Errors
    ///
    /// Returns an error if the quantity is too small (< 1e16)
    pub fn get_limit_order_amounts(&self, data: LimitOrderData) -> Result<LimitOrderAmounts> {
        amounts::get_limit_order_amounts(data)
    }

    /// Build an order struct
    ///
    /// # Arguments
    ///
    /// * `strategy` - The order strategy (MARKET or LIMIT)
    /// * `input` - The order input data
    ///
    /// # Returns
    ///
    /// A constructed Order ready for signing
    ///
    /// # Errors
    ///
    /// Returns an error if the input data is invalid or expiration is in the past
    pub fn build_order(&self, strategy: OrderStrategy, input: BuildOrderInput) -> Result<Order> {
        // Get signer address if available
        let signer_address = self.signer_address()
            .unwrap_or_else(|_| {
                input.signer.as_ref()
                    .and_then(|s| s.parse::<Address>().ok())
                    .unwrap_or(ZERO_ADDRESS.parse().unwrap())
            });

        let signer_str = format!("{}", signer_address);

        // When using a Predict Account, maker and signer must be the predict_account address.
        // The Predict API verifies signatures via EIP-1271 (Kernel smart wallet).
        // When using EOA directly, maker and signer are the EOA address.
        let (maker_str, order_signer_str) = if let Some(predict_account) = self.predict_account {
            let pa = format!("{}", predict_account);
            (pa.clone(), pa)
        } else {
            (signer_str.clone(), signer_str.clone())
        };

        // Calculate expiration
        let expiration = if let Some(expires_at) = input.expires_at {
            let timestamp = expires_at.timestamp();
            let now = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_secs() as i64;

            if timestamp <= now {
                return Err(Error::InvalidOrderData("Expiration must be in the future".to_string()));
            }
            timestamp.to_string()
        } else {
            // Default expiration based on strategy
            let now = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_secs();

            let expiration_secs = match strategy {
                OrderStrategy::Market => now + constants::FIVE_MINUTES_SECONDS,
                OrderStrategy::Limit => now + (365 * 24 * 60 * 60), // 1 year
            };

            expiration_secs.to_string()
        };

        Ok(Order {
            salt: input.salt.unwrap_or_else(|| (self.generate_salt)()),
            maker: input.maker.unwrap_or_else(|| maker_str.clone()),
            signer: input.signer.unwrap_or_else(|| order_signer_str.clone()),
            taker: input.taker.unwrap_or_else(|| ZERO_ADDRESS.to_string()),
            token_id: input.token_id,
            maker_amount: input.maker_amount,
            taker_amount: input.taker_amount,
            expiration,
            nonce: input.nonce.unwrap_or_else(|| "0".to_string()),
            fee_rate_bps: input.fee_rate_bps.to_string(),
            side: input.side,
            signature_type: input.signature_type.unwrap_or(SignatureType::Eoa),
        })
    }

    /// Build EIP-712 typed data for an order
    ///
    /// # Arguments
    ///
    /// * `order` - The order to build typed data for
    /// * `is_neg_risk` - Whether this is a neg risk market (winner takes all)
    /// * `is_yield_bearing` - Whether this market has yield enabled
    ///
    /// # Returns
    ///
    /// The verifying contract address
    pub fn get_verifying_contract(&self, is_neg_risk: bool, is_yield_bearing: bool) -> Address {
        let address_str = self.addresses.get_ctf_exchange(is_yield_bearing, is_neg_risk);
        address_str.parse().unwrap()
    }

    /// Build the EIP-712 typed data hash for an order
    ///
    /// # Arguments
    ///
    /// * `order` - The order to hash
    /// * `is_neg_risk` - Whether this is a neg risk market
    /// * `is_yield_bearing` - Whether this market has yield enabled
    ///
    /// # Returns
    ///
    /// The hash to be signed
    ///
    /// # Errors
    ///
    /// Returns an error if the order data is invalid
    pub fn build_typed_data_hash(
        &self,
        order: &Order,
        is_neg_risk: bool,
        is_yield_bearing: bool,
    ) -> Result<String> {
        let verifying_contract = self.get_verifying_contract(is_neg_risk, is_yield_bearing);
        let domain = signing::get_domain(self.chain_id, verifying_contract);
        let hash = signing::build_typed_data_hash(order, &domain)?;
        Ok(format!("0x{}", hex::encode(hash.as_slice())))
    }

    /// Sign an order using EIP-712
    ///
    /// This method automatically uses the appropriate signing method:
    /// - For regular EOA wallets: Standard EIP-712 signing
    /// - For Predict Account (Kernel): Kernel-wrapped EIP-712 signing
    ///
    /// # Arguments
    ///
    /// * `order` - The order to sign
    /// * `is_neg_risk` - Whether this is a neg risk market
    /// * `is_yield_bearing` - Whether this market has yield enabled
    ///
    /// # Returns
    ///
    /// A signed order with signature
    ///
    /// # Errors
    ///
    /// Returns an error if no signer is configured or signing fails
    pub async fn sign_typed_data_order(
        &self,
        order: Order,
        is_neg_risk: bool,
        is_yield_bearing: bool,
    ) -> Result<SignedOrder> {
        let signer = self.signer.as_ref()
            .ok_or_else(|| Error::Other("No signer configured".to_string()))?;

        let verifying_contract = self.get_verifying_contract(is_neg_risk, is_yield_bearing);
        let hash = self.build_typed_data_hash(&order, is_neg_risk, is_yield_bearing)?;

        let signature = if let Some(predict_account) = self.predict_account {
            // Use Kernel-wrapped signing for Predict Account
            // verifyingContract in Kernel domain = predict_account (user's smart wallet),
            // NOT the global Kernel contract address
            let ecdsa_validator = self.addresses.ecdsa_validator.parse::<Address>()
                .map_err(|e| Error::Other(format!("Invalid ECDSA validator address: {}", e)))?;

            signing::sign_order_for_predict_account(
                &order,
                self.chain_id,
                verifying_contract,
                predict_account,
                ecdsa_validator,
                signer,
            ).await?
        } else {
            // Standard EOA signing
            signing::sign_order(&order, self.chain_id, verifying_contract, signer).await?
        };

        Ok(SignedOrder {
            order,
            hash: Some(hash),
            signature,
        })
    }

    /// Sign an order using standard EOA EIP-712 (never Kernel-wrapped)
    ///
    /// This is used for REST API order placement, where the server does plain ecrecover
    /// to verify the signature. Kernel-wrapped signatures are only needed for on-chain
    /// settlement, which the platform handles internally.
    pub async fn sign_typed_data_order_eoa(
        &self,
        order: Order,
        is_neg_risk: bool,
        is_yield_bearing: bool,
    ) -> Result<SignedOrder> {
        let signer = self.signer.as_ref()
            .ok_or_else(|| Error::Other("No signer configured".to_string()))?;

        let verifying_contract = self.get_verifying_contract(is_neg_risk, is_yield_bearing);
        let hash = self.build_typed_data_hash(&order, is_neg_risk, is_yield_bearing)?;

        // Always use standard EOA signing, regardless of predict_account
        let signature = signing::sign_order(&order, self.chain_id, verifying_contract, signer).await?;

        Ok(SignedOrder {
            order,
            hash: Some(hash),
            signature,
        })
    }

    /// Get the EOA signer address as a formatted string
    pub fn signer_address_string(&self) -> Result<String> {
        self.signer_address().map(|addr| format!("{}", addr))
    }

    /// Get the chain ID
    pub fn chain_id(&self) -> ChainId {
        self.chain_id
    }

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

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

    #[test]
    fn test_new_order_builder() {
        let builder = OrderBuilder::new(ChainId::BnbTestnet, None, None).unwrap();
        assert_eq!(builder.chain_id(), ChainId::BnbTestnet);
    }

    #[test]
    fn test_get_limit_order_amounts() {
        let builder = OrderBuilder::new(ChainId::BnbTestnet, None, None).unwrap();

        let data = LimitOrderData {
            side: Side::Buy,
            price_per_share_wei: dec!(500000000000000000),
            quantity_wei: dec!(10000000000000000000),
        };

        let amounts = builder.get_limit_order_amounts(data).unwrap();
        assert!(amounts.maker_amount > rust_decimal::Decimal::ZERO);
    }

    #[test]
    fn test_build_order() {
        let signer = PrivateKeySigner::random();
        let builder = OrderBuilder::new(ChainId::BnbTestnet, Some(signer), None).unwrap();

        let input = BuildOrderInput {
            side: Side::Buy,
            token_id: "12345".to_string(),
            maker_amount: "1000000000000000000".to_string(),
            taker_amount: "2000000000000000000".to_string(),
            fee_rate_bps: 100,
            signer: None,
            nonce: None,
            salt: None,
            maker: None,
            taker: None,
            signature_type: None,
            expires_at: None,
        };

        let order = builder.build_order(OrderStrategy::Limit, input).unwrap();
        assert_eq!(order.side, Side::Buy);
        assert_eq!(order.token_id, "12345");
    }

    #[tokio::test]
    async fn test_sign_order() {
        let signer = PrivateKeySigner::random();
        let builder = OrderBuilder::new(ChainId::BnbTestnet, Some(signer), None).unwrap();

        let input = BuildOrderInput {
            side: Side::Buy,
            token_id: "12345".to_string(),
            maker_amount: "1000000000000000000".to_string(),
            taker_amount: "2000000000000000000".to_string(),
            fee_rate_bps: 100,
            signer: None,
            nonce: None,
            salt: None,
            maker: None,
            taker: None,
            signature_type: None,
            expires_at: None,
        };

        let order = builder.build_order(OrderStrategy::Limit, input).unwrap();
        let signed_order = builder.sign_typed_data_order(order, false, false).await.unwrap();

        assert!(signed_order.signature.starts_with("0x"));
        assert!(signed_order.hash.is_some());
    }
}