polymarket 0.1.0

use crate::clob::{
    get_contract_config,
    error::{ClobError, Result},
    types::{SignedOrderArgs, PartialOrderArgs},
    Signer,
};
use alloy::primitives::{Address, U256};
use alloy::sol;
use alloy_sol_types::SolStruct;

// Define the EIP-712 Order struct matching the CTF Exchange contract
sol! {
    struct Order {
        uint256 salt;
        address maker;
        address signer;
        address taker;
        uint256 tokenId;
        uint256 makerAmount;
        uint256 takerAmount;
        uint256 expiration;
        uint256 nonce;
        uint256 feeRateBps;
        uint8 side;
        uint8 signatureType;
    }
}

/// Domain for CTF Exchange EIP-712 signing
const DOMAIN_NAME: &str = "Polymarket CTF Exchange";
const DOMAIN_VERSION: &str = "1";

/// Sign an order using EIP-712
///
/// # Arguments
/// * `partial` - The partial order to sign
/// * `signer` - The signer to use for signing
/// * `neg_risk` - Whether this is a neg-risk order (uses different exchange contract)
///
/// # Returns
/// The fully signed SignedOrderArgs
pub async fn sign_order(
    partial: PartialOrderArgs,
    signer: &Signer,
    neg_risk: bool,
) -> Result<SignedOrderArgs> {
    // Parse amounts as integers (assuming they're in raw units)
    let maker_amount = parse_amount(&partial.maker_amount)?;
    let taker_amount = parse_amount(&partial.taker_amount)?;
    let token_id = parse_token_id(&partial.token_id)?;

    // Create the Order struct for EIP-712 signing
    let order = Order {
        salt: U256::from(partial.salt),
        maker: parse_address(&partial.maker)?,
        signer: parse_address(&partial.signer)?,
        taker: parse_address(&partial.taker)?,
        tokenId: token_id,
        makerAmount: maker_amount,
        takerAmount: taker_amount,
        expiration: U256::from(partial.expiration),
        nonce: U256::from(partial.nonce),
        feeRateBps: U256::from(partial.fee_rate_bps),
        side: partial.side as u8,
        signatureType: partial.signature_type as u8,
    };

    // Get the correct exchange address based on chain ID and neg_risk flag
    let contract_config = get_contract_config(signer.chain_id(), neg_risk)
        .ok_or_else(|| {
            ClobError::InvalidOrder(format!("Unsupported chain ID: {}", signer.chain_id()))
        })?;

    // Create EIP-712 domain with DYNAMIC exchange address
    let domain = alloy::sol_types::eip712_domain! {
        name: DOMAIN_NAME,
        version: DOMAIN_VERSION,
        chain_id: signer.chain_id(),
        verifying_contract: parse_address(&contract_config.exchange)?,
    };

    // Get the signing hash
    let signing_hash = order.eip712_signing_hash(&domain);

    // Sign the hash
    let signature = signer.sign_hash(signing_hash.as_ref()).await?;

    Ok(SignedOrderArgs {
        salt: partial.salt,
        maker: partial.maker,
        signer: partial.signer,
        taker: partial.taker,
        token_id: partial.token_id,
        maker_amount: partial.maker_amount,
        taker_amount: partial.taker_amount,
        side: partial.side,
        fee_rate_bps: partial.fee_rate_bps,
        nonce: partial.nonce,
        signature_type: partial.signature_type,
        expiration: partial.expiration,
        signature,
    })
}

/// Parse an amount string into U256
fn parse_amount(amount: &str) -> Result<U256> {
    // Remove any decimal points and convert to raw units
    // For now, assume amounts are already in the correct format
    U256::from_str_radix(amount, 10).map_err(|e| {
        ClobError::InvalidOrder(format!("Invalid amount '{}': {}", amount, e))
    })
}

/// Parse a token ID string into U256
fn parse_token_id(token_id: &str) -> Result<U256> {
    // Handle both decimal and hex token IDs
    if let Some(hex) = token_id.strip_prefix("0x") {
        U256::from_str_radix(hex, 16)
    } else {
        U256::from_str_radix(token_id, 10)
    }
    .map_err(|e| ClobError::InvalidOrder(format!("Invalid token ID '{}': {}", token_id, e)))
}

/// Parse an address string into Address
fn parse_address(addr: &str) -> Result<Address> {
    addr.parse()
        .map_err(|e| ClobError::InvalidOrder(format!("Invalid address '{}': {}", addr, e)))
}

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

    #[test]
    fn test_parse_amount() {
        let amount = parse_amount("1000000").unwrap();
        assert_eq!(amount, U256::from(1000000));
    }

    #[test]
    fn test_parse_token_id_decimal() {
        let token_id = parse_token_id("123456").unwrap();
        assert_eq!(token_id, U256::from(123456));
    }

    #[test]
    fn test_parse_token_id_hex() {
        let token_id = parse_token_id("0x1e240").unwrap();
        assert_eq!(token_id, U256::from(123456));
    }

    #[test]
    fn test_parse_address() {
        let addr = parse_address("0x0000000000000000000000000000000000000000").unwrap();
        assert_eq!(
            addr,
            "0x0000000000000000000000000000000000000000"
                .parse::<Address>()
                .unwrap()
        );
    }

    #[test]
    fn test_invalid_amount() {
        assert!(parse_amount("not_a_number").is_err());
    }

    #[test]
    fn test_invalid_address() {
        assert!(parse_address("not_an_address").is_err());
    }
}