polynode 0.6.0

//! Wallet signer abstraction for EIP-712 typed data signing.

use alloy_primitives::Address;
use alloy_signer::Signer;
use alloy_signer_local::PrivateKeySigner as AlloySigner;
use async_trait::async_trait;

use crate::error::{Error, Result};
use super::types::Eip712Payload;

/// Trait for signing EIP-712 typed data and raw messages.
///
/// Implement this trait for custom signer backends (e.g. HSM, remote KMS, Privy).
/// The SDK provides [`PrivateKeySigner`] for raw private keys.
#[async_trait]
pub trait TradingSigner: Send + Sync {
    /// The EOA address (signer address, not the Safe/Proxy).
    fn address(&self) -> Address;

    /// Sign EIP-712 typed data, returning the 65-byte signature.
    async fn sign_typed_data(&self, payload: &Eip712Payload) -> Result<Vec<u8>>;

    /// Sign a raw message (personal_sign), returning the 65-byte signature.
    async fn sign_message(&self, message: &[u8]) -> Result<Vec<u8>>;

    /// Sign a raw 32-byte hash, returning the 65-byte signature (r, s, v).
    /// Used for on-chain transaction signing (wrap/unwrap, approvals).
    /// Default implementation returns an error — override in signers that support raw hash signing.
    async fn sign_hash(&self, _hash: &[u8; 32]) -> Result<Vec<u8>> {
        Err(Error::Signing("sign_hash not supported by this signer".into()))
    }
}

/// Signer backed by a raw secp256k1 private key.
pub struct PrivateKeySigner {
    inner: AlloySigner,
}

impl PrivateKeySigner {
    /// Create from a hex-encoded private key (with or without 0x prefix).
    pub fn from_hex(key: &str) -> Result<Self> {
        let key = key.strip_prefix("0x").unwrap_or(key);
        let signer: AlloySigner = key
            .parse()
            .map_err(|e| Error::Signing(format!("Invalid private key: {}", e)))?;
        Ok(Self { inner: signer })
    }

    /// Generate a new random private key. Returns (signer, hex_private_key).
    pub fn generate() -> (Self, String) {
        let signer = AlloySigner::random();
        let key_hex = hex::encode(signer.credential().to_bytes());
        (Self { inner: signer }, format!("0x{}", key_hex))
    }

    /// Get the underlying alloy signer for direct use.
    pub fn inner(&self) -> &AlloySigner {
        &self.inner
    }
}

#[async_trait]
impl TradingSigner for PrivateKeySigner {
    fn address(&self) -> Address {
        self.inner.address()
    }

    async fn sign_typed_data(&self, payload: &Eip712Payload) -> Result<Vec<u8>> {
        // Build the EIP-712 hash manually:
        // 1. Hash the domain separator
        // 2. Hash the struct data
        // 3. Combine: keccak256(0x1901 || domainSeparator || structHash)
        // 4. Sign the result
        let signing_hash = compute_eip712_hash(payload)?;
        let sig = self.inner
            .sign_hash(&signing_hash)
            .await
            .map_err(|e| Error::Signing(format!("sign_typed_data failed: {}", e)))?;
        let mut bytes = sig.as_bytes().to_vec();
        // Normalize v to raw recovery id (0/1): CLOB expects v=0/1, not Ethereum's 27/28
        if bytes.len() == 65 && bytes[64] >= 27 {
            bytes[64] -= 27;
        }
        Ok(bytes)
    }

    async fn sign_message(&self, message: &[u8]) -> Result<Vec<u8>> {
        let sig = self.inner
            .sign_message(message)
            .await
            .map_err(|e| Error::Signing(format!("sign_message failed: {}", e)))?;
        Ok(sig.as_bytes().to_vec())
    }

    async fn sign_hash(&self, hash: &[u8; 32]) -> Result<Vec<u8>> {
        let hash = alloy_primitives::B256::from_slice(hash);
        let sig = self.inner
            .sign_hash(&hash)
            .await
            .map_err(|e| Error::Signing(format!("sign_hash failed: {}", e)))?;
        Ok(sig.as_bytes().to_vec())
    }
}

/// Compute the EIP-712 signing hash from a payload.
///
/// This implements the EIP-712 hashing algorithm:
/// `keccak256(0x1901 || domainSeparator || structHash)`
pub fn compute_eip712_hash(payload: &Eip712Payload) -> Result<alloy_primitives::B256> {
    use alloy_primitives::keccak256;

    // Build domain separator from the JSON domain object
    let domain = &payload.domain;
    let domain_sep = hash_eip712_domain(domain)?;

    // Filter EIP712Domain from types before hashing (matches TS SDK behavior)
    let mut filtered_types = payload.types.clone();
    if let Some(obj) = filtered_types.as_object_mut() {
        obj.remove("EIP712Domain");
    }

    // Hash the struct
    let struct_hash = hash_eip712_struct(&payload.primary_type, &filtered_types, &payload.message)?;

    // Final hash: keccak256(0x1901 || domainSeparator || structHash)
    let mut buf = Vec::with_capacity(2 + 32 + 32);
    buf.extend_from_slice(&[0x19, 0x01]);
    buf.extend_from_slice(domain_sep.as_slice());
    buf.extend_from_slice(struct_hash.as_slice());
    Ok(keccak256(&buf))
}

fn hash_eip712_domain(domain: &serde_json::Value) -> Result<alloy_primitives::B256> {
    use alloy_primitives::keccak256;

    // Build the domain type string dynamically based on which fields are present.
    // EIP-712 spec: the type hash must match exactly the fields in the domain object.
    let mut type_parts = Vec::new();
    if domain.get("name").is_some() { type_parts.push("string name"); }
    if domain.get("version").is_some() { type_parts.push("string version"); }
    if domain.get("chainId").is_some() { type_parts.push("uint256 chainId"); }
    if domain.get("verifyingContract").is_some() { type_parts.push("address verifyingContract"); }
    if domain.get("salt").is_some() { type_parts.push("bytes32 salt"); }
    let type_string = format!("EIP712Domain({})", type_parts.join(","));
    let type_hash = keccak256(type_string.as_bytes());

    let mut encoded = Vec::new();
    encoded.extend_from_slice(type_hash.as_slice());

    // name
    if let Some(name) = domain.get("name").and_then(|v| v.as_str()) {
        encoded.extend_from_slice(keccak256(name.as_bytes()).as_slice());
    }
    // version
    if let Some(version) = domain.get("version").and_then(|v| v.as_str()) {
        encoded.extend_from_slice(keccak256(version.as_bytes()).as_slice());
    }
    // chainId
    if let Some(chain_id) = domain.get("chainId") {
        let id = chain_id.as_u64().unwrap_or(0);
        let mut word = [0u8; 32];
        word[24..32].copy_from_slice(&id.to_be_bytes());
        encoded.extend_from_slice(&word);
    }
    // verifyingContract
    if let Some(addr) = domain.get("verifyingContract").and_then(|v| v.as_str()) {
        let addr_bytes = parse_address(addr)?;
        let mut word = [0u8; 32];
        word[12..32].copy_from_slice(&addr_bytes);
        encoded.extend_from_slice(&word);
    }

    Ok(keccak256(&encoded))
}

fn hash_eip712_struct(
    primary_type: &str,
    types: &serde_json::Value,
    message: &serde_json::Value,
) -> Result<alloy_primitives::B256> {
    use alloy_primitives::keccak256;

    // Build the type string for the primary type
    let type_string = build_type_string(primary_type, types)?;
    let type_hash = keccak256(type_string.as_bytes());

    let mut encoded = Vec::new();
    encoded.extend_from_slice(type_hash.as_slice());

    // Encode each field
    if let Some(fields) = types.get(primary_type).and_then(|v| v.as_array()) {
        for field in fields {
            let name = field.get("name").and_then(|v| v.as_str()).unwrap_or("");
            let typ = field.get("type").and_then(|v| v.as_str()).unwrap_or("");
            let value = &message[name];
            encode_field(typ, value, types, &mut encoded)?;
        }
    }

    Ok(keccak256(&encoded))
}

fn build_type_string(primary_type: &str, types: &serde_json::Value) -> Result<String> {
    let fields = types
        .get(primary_type)
        .and_then(|v| v.as_array())
        .ok_or_else(|| Error::Signing(format!("Missing type definition for {}", primary_type)))?;

    let mut result = format!("{}(", primary_type);
    for (i, field) in fields.iter().enumerate() {
        if i > 0 {
            result.push(',');
        }
        let typ = field.get("type").and_then(|v| v.as_str()).unwrap_or("");
        let name = field.get("name").and_then(|v| v.as_str()).unwrap_or("");
        result.push_str(typ);
        result.push(' ');
        result.push_str(name);
    }
    result.push(')');
    Ok(result)
}

fn encode_field(
    typ: &str,
    value: &serde_json::Value,
    types: &serde_json::Value,
    out: &mut Vec<u8>,
) -> Result<()> {
    use alloy_primitives::keccak256;

    match typ {
        "string" => {
            let s = value.as_str().unwrap_or("");
            out.extend_from_slice(keccak256(s.as_bytes()).as_slice());
        }
        "bytes" => {
            let s = value.as_str().unwrap_or("0x");
            let bytes = hex::decode(s.strip_prefix("0x").unwrap_or(s))
                .map_err(|e| Error::Signing(format!("Invalid bytes: {}", e)))?;
            out.extend_from_slice(keccak256(&bytes).as_slice());
        }
        "address" => {
            let addr = value.as_str().unwrap_or("0x0000000000000000000000000000000000000000");
            let addr_bytes = parse_address(addr)?;
            let mut word = [0u8; 32];
            word[12..32].copy_from_slice(&addr_bytes);
            out.extend_from_slice(&word);
        }
        "bool" => {
            let b = value.as_bool().unwrap_or(false);
            let mut word = [0u8; 32];
            if b {
                word[31] = 1;
            }
            out.extend_from_slice(&word);
        }
        t if t.starts_with("uint") || t.starts_with("int") => {
            // Handle uint256, int256, uint8, etc. — encode as 32-byte big-endian word.
            // Values can be JSON numbers or decimal/hex strings.
            let word = encode_uint_value(value)?;
            out.extend_from_slice(&word);
        }
        t if t.starts_with("bytes") && t.len() > 5 => {
            // Fixed-size bytes (bytes1..bytes32)
            let s = value.as_str().unwrap_or("0x");
            let bytes = hex::decode(s.strip_prefix("0x").unwrap_or(s))
                .map_err(|e| Error::Signing(format!("Invalid {}: {}", t, e)))?;
            let mut word = [0u8; 32];
            let len = bytes.len().min(32);
            word[..len].copy_from_slice(&bytes[..len]);
            out.extend_from_slice(&word);
        }
        _ => {
            // Struct type — hash recursively
            if types.get(typ).is_some() {
                let hash = hash_eip712_struct(typ, types, value)?;
                out.extend_from_slice(hash.as_slice());
            } else {
                // Unknown type, treat as bytes32(0)
                out.extend_from_slice(&[0u8; 32]);
            }
        }
    }
    Ok(())
}

fn parse_address(addr: &str) -> Result<[u8; 20]> {
    let hex_str = addr.strip_prefix("0x").unwrap_or(addr);
    let bytes = hex::decode(hex_str)
        .map_err(|e| Error::Signing(format!("Invalid address {}: {}", addr, e)))?;
    if bytes.len() != 20 {
        return Err(Error::Signing(format!("Address wrong length: {} bytes", bytes.len())));
    }
    let mut arr = [0u8; 20];
    arr.copy_from_slice(&bytes);
    Ok(arr)
}

/// Encode a uint/int value to a 32-byte big-endian word.
/// Handles JSON numbers, decimal strings (any size up to uint256), and hex strings.
fn encode_uint_value(value: &serde_json::Value) -> Result<[u8; 32]> {
    use alloy_primitives::U256;

    let mut word = [0u8; 32];

    if let Some(n) = value.as_u64() {
        word[24..32].copy_from_slice(&n.to_be_bytes());
    } else if let Some(n) = value.as_i64() {
        if n >= 0 {
            word[24..32].copy_from_slice(&n.to_be_bytes());
        } else {
            // Signed negative: two's complement
            let bytes = n.to_be_bytes();
            word[..24].fill(0xff);
            word[24..32].copy_from_slice(&bytes);
        }
    } else if let Some(s) = value.as_str() {
        if s.starts_with("0x") || s.starts_with("0X") {
            // Hex string
            let v = U256::from_str_radix(&s[2..], 16)
                .map_err(|e| Error::Signing(format!("Invalid hex uint: {} ({})", s, e)))?;
            word = v.to_be_bytes::<32>();
        } else {
            // Decimal string — use U256 for arbitrary precision
            let v = U256::from_str_radix(s, 10)
                .map_err(|e| Error::Signing(format!("Invalid decimal uint: {} ({})", s, e)))?;
            word = v.to_be_bytes::<32>();
        }
    }

    Ok(word)
}