polymarket 0.1.0

//! Signer abstractions for different wallet types

use alloy::network::EthereumWallet;
use alloy::primitives::Address;
use alloy::signers::local::PrivateKeySigner;
use alloy::signers::Signer;

use crate::onchain::{OnchainError, Result, WalletType};

/// Unified signer interface supporting multiple wallet types
#[derive(Clone)]
pub enum OnchainSigner {
    /// Direct private key signer (EOA - Externally Owned Account)
    EOA(PrivateKeySigner),

    /// Proxy wallet signer (Email/Magic accounts via ProxyFactory)
    Proxy {
        /// The underlying EOA that controls the proxy
        signer: PrivateKeySigner,
        /// The proxy wallet address
        proxy_address: Address,
    },

    /// Safe (Gnosis Safe) wallet signer
    Safe {
        /// The EOA that will propose transactions
        signer: PrivateKeySigner,
        /// The Safe wallet address
        safe_address: Address,
        /// Threshold for multi-sig (default: 1)
        threshold: u64,
    },
}

impl OnchainSigner {
    /// Create a new private key signer from a hex string
    ///
    /// # Example
    /// ```no_run
    /// # use polymarket_sdk::onchain::OnchainSigner;
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let signer = OnchainSigner::from_private_key("0x1234...")?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn from_private_key(private_key: impl AsRef<str>) -> Result<Self> {
        let key_str = private_key.as_ref();
        let key_str = key_str.strip_prefix("0x").unwrap_or(key_str);

        let signer = PrivateKeySigner::from_slice(
            &hex::decode(key_str)
                .map_err(|e| OnchainError::InvalidPrivateKey(format!("Invalid hex: {}", e)))?,
        )
        .map_err(|e| OnchainError::InvalidPrivateKey(format!("Invalid key: {}", e)))?;

        Ok(Self::EOA(signer))
    }

    /// Create a proxy wallet signer
    ///
    /// # Arguments
    /// * `private_key` - The EOA private key that controls the proxy
    /// * `proxy_address` - The proxy wallet address (derived from salt)
    pub fn from_proxy(private_key: impl AsRef<str>, proxy_address: Address) -> Result<Self> {
        let key_str = private_key.as_ref();
        let key_str = key_str.strip_prefix("0x").unwrap_or(key_str);

        let signer = PrivateKeySigner::from_slice(
            &hex::decode(key_str)
                .map_err(|e| OnchainError::InvalidPrivateKey(format!("Invalid hex: {}", e)))?,
        )
        .map_err(|e| OnchainError::InvalidPrivateKey(format!("Invalid key: {}", e)))?;

        Ok(Self::Proxy {
            signer,
            proxy_address,
        })
    }

    /// Create a Safe wallet signer
    ///
    /// # Arguments
    /// * `private_key` - The EOA private key for proposing transactions
    /// * `safe_address` - The Safe wallet address
    /// * `threshold` - Number of required signatures (default: 1)
    pub fn from_safe(
        private_key: impl AsRef<str>,
        safe_address: Address,
        threshold: Option<u64>,
    ) -> Result<Self> {
        let key_str = private_key.as_ref();
        let key_str = key_str.strip_prefix("0x").unwrap_or(key_str);

        let signer = PrivateKeySigner::from_slice(
            &hex::decode(key_str)
                .map_err(|e| OnchainError::InvalidPrivateKey(format!("Invalid hex: {}", e)))?,
        )
        .map_err(|e| OnchainError::InvalidPrivateKey(format!("Invalid key: {}", e)))?;

        Ok(Self::Safe {
            signer,
            safe_address,
            threshold: threshold.unwrap_or(1),
        })
    }

    /// Get the wallet type
    pub fn wallet_type(&self) -> WalletType {
        match self {
            Self::EOA(_) => WalletType::PrivateKey,
            Self::Proxy { .. } => WalletType::Proxy,
            Self::Safe { .. } => WalletType::Safe,
        }
    }

    /// Get the address that will send transactions
    ///
    /// - For EOA: returns the EOA address
    /// - For Proxy: returns the proxy wallet address
    /// - For Safe: returns the Safe wallet address
    pub fn address(&self) -> Address {
        match self {
            Self::EOA(signer) => signer.address(),
            Self::Proxy { proxy_address, .. } => *proxy_address,
            Self::Safe { safe_address, .. } => *safe_address,
        }
    }

    /// Get the underlying EOA signer address
    ///
    /// For all wallet types, this returns the EOA that controls the wallet
    pub fn signer_address(&self) -> Address {
        match self {
            Self::EOA(signer) => signer.address(),
            Self::Proxy { signer, .. } => signer.address(),
            Self::Safe { signer, .. } => signer.address(),
        }
    }

    /// Get the underlying PrivateKeySigner reference
    pub(crate) fn private_key_signer(&self) -> &PrivateKeySigner {
        match self {
            Self::EOA(signer) => signer,
            Self::Proxy { signer, .. } => signer,
            Self::Safe { signer, .. } => signer,
        }
    }

    /// Create an EthereumWallet from this signer (for contract calls)
    ///
    /// # Arguments
    /// * `chain_id` - The chain ID to use for signing (e.g., 137 for Polygon, 1 for Ethereum)
    pub fn to_ethereum_wallet(&self, chain_id: u64) -> EthereumWallet {
        let mut signer = self.private_key_signer().clone();
        // Set the chain ID on the signer
        signer = signer.with_chain_id(Some(chain_id));
        EthereumWallet::from(signer)
    }

    /// Sign a message with the underlying EOA
    pub async fn sign_message(&self, message: &[u8]) -> Result<Vec<u8>> {
        let signer = self.private_key_signer();
        let signature = signer
            .sign_message(message)
            .await
            .map_err(|e| OnchainError::SignerError(e.to_string()))?;

        Ok(signature.as_bytes().to_vec())
    }

    /// Check if this is a smart contract wallet (Proxy or Safe)
    pub fn is_smart_wallet(&self) -> bool {
        matches!(self, Self::Proxy { .. } | Self::Safe { .. })
    }

    /// Get proxy address if this is a proxy wallet
    pub fn proxy_address(&self) -> Option<Address> {
        match self {
            Self::Proxy { proxy_address, .. } => Some(*proxy_address),
            _ => None,
        }
    }

    /// Get Safe address if this is a Safe wallet
    pub fn safe_address(&self) -> Option<Address> {
        match self {
            Self::Safe { safe_address, .. } => Some(*safe_address),
            _ => None,
        }
    }

    /// Get Safe threshold if this is a Safe wallet
    pub fn safe_threshold(&self) -> Option<u64> {
        match self {
            Self::Safe { threshold, .. } => Some(*threshold),
            _ => None,
        }
    }
}

impl std::fmt::Debug for OnchainSigner {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::EOA(signer) => f
                .debug_struct("EOA")
                .field("address", &signer.address())
                .finish(),
            Self::Proxy {
                signer,
                proxy_address,
            } => f
                .debug_struct("Proxy")
                .field("signer_address", &signer.address())
                .field("proxy_address", proxy_address)
                .finish(),
            Self::Safe {
                signer,
                safe_address,
                threshold,
            } => f
                .debug_struct("Safe")
                .field("signer_address", &signer.address())
                .field("safe_address", safe_address)
                .field("threshold", threshold)
                .finish(),
        }
    }
}

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

    const TEST_PRIVATE_KEY: &str =
        "0x0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";

    #[test]
    fn test_from_private_key() {
        let signer = OnchainSigner::from_private_key(TEST_PRIVATE_KEY).unwrap();
        assert_eq!(signer.wallet_type(), WalletType::PrivateKey);
        assert!(!signer.is_smart_wallet());
    }

    #[test]
    fn test_from_private_key_without_0x() {
        let key_without_prefix = TEST_PRIVATE_KEY.strip_prefix("0x").unwrap();
        let signer = OnchainSigner::from_private_key(key_without_prefix).unwrap();
        assert_eq!(signer.wallet_type(), WalletType::PrivateKey);
    }

    #[test]
    fn test_from_proxy() {
        let proxy_address = Address::repeat_byte(0x42);
        let signer = OnchainSigner::from_proxy(TEST_PRIVATE_KEY, proxy_address).unwrap();

        assert_eq!(signer.wallet_type(), WalletType::Proxy);
        assert!(signer.is_smart_wallet());
        assert_eq!(signer.address(), proxy_address);
        assert_eq!(signer.proxy_address(), Some(proxy_address));
        assert_ne!(signer.signer_address(), proxy_address);
    }

    #[test]
    fn test_from_safe() {
        let safe_address = Address::repeat_byte(0x99);
        let signer = OnchainSigner::from_safe(TEST_PRIVATE_KEY, safe_address, Some(2)).unwrap();

        assert_eq!(signer.wallet_type(), WalletType::Safe);
        assert!(signer.is_smart_wallet());
        assert_eq!(signer.address(), safe_address);
        assert_eq!(signer.safe_address(), Some(safe_address));
        assert_eq!(signer.safe_threshold(), Some(2));
        assert_ne!(signer.signer_address(), safe_address);
    }

    #[test]
    fn test_invalid_private_key() {
        let result = OnchainSigner::from_private_key("invalid_hex");
        assert!(result.is_err());

        let result = OnchainSigner::from_private_key("0x123"); // too short
        assert!(result.is_err());
    }

    #[test]
    fn test_debug_format() {
        let signer = OnchainSigner::from_private_key(TEST_PRIVATE_KEY).unwrap();
        let debug_str = format!("{:?}", signer);
        assert!(debug_str.contains("EOA"));
        assert!(debug_str.contains("address"));
    }
}