sal/src/rhai/vault.rs

//! Rhai bindings for SAL crypto functionality

use rhai::{Engine, Dynamic, EvalAltResult};
use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
use std::fs;
use std::sync::Mutex;
use once_cell::sync::{Lazy, OnceCell};
use tokio::runtime::Runtime;
use ethers::types::{Address, U256};
use std::str::FromStr;
use cfg_if::cfg_if;

use crate::hero_vault::{keypair, symmetric, ethereum, kvs};
use crate::hero_vault::kvs::{KVStore, DefaultStore};
use crate::hero_vault::ethereum::prepare_function_arguments;

// Global Tokio runtime for blocking async operations
static RUNTIME: Lazy<Mutex<Runtime>> = Lazy::new(|| {
    Mutex::new(Runtime::new().expect("Failed to create Tokio runtime"))
});

// Helper function to run async operations and handle errors consistently
fn run_async<F, T, E>(future: F) -> Result<T, String> 
where 
    F: std::future::Future<Output = Result<T, E>>,
    E: std::fmt::Display,
{
    let rt = RUNTIME.lock().map_err(|e| format!("Failed to acquire runtime lock: {}", e))?;
    rt.block_on(async { future.await.map_err(|e| e.to_string()) })
}

// Get a platform-specific DefaultStore implementation for Rhai bindings
// This function is implemented differently based on the target platform
cfg_if! {
    if #[cfg(target_arch = "wasm32")] {
        fn get_key_store() -> DefaultStore {
            use wasm_bindgen_futures::JsFuture;
            use once_cell::sync::OnceCell;
            use std::future::Future;
            
            // Static store instance
            static KEY_STORE: OnceCell<DefaultStore> = OnceCell::new();
            
            // Initialize if not already done
            KEY_STORE.get_or_init(|| {
                // In WebAssembly, we need to use a blocking approach for Rhai
                let store_future = async {
                    match kvs::open_default_store("rhai-vault", None).await {
                        Ok(store) => store,
                        Err(_) => {
                            // Try to create the store if opening failed
                            kvs::create_default_store("rhai-vault", false, None).await
                                .expect("Failed to create key store")
                        }
                    }
                };
                
                // Block on the async operation
                let rt = RUNTIME.lock().unwrap();
                rt.block_on(store_future)
            }).clone()
        }
    } else {
        fn get_key_store() -> DefaultStore {
            use once_cell::sync::OnceCell;
            
            // Static store instance
            static KEY_STORE: OnceCell<DefaultStore> = OnceCell::new();
            
            // Initialize if not already done
            KEY_STORE.get_or_init(|| {
                // For native platforms, the operations are synchronous
                match kvs::open_default_store("rhai-vault", None) {
                    Ok(store) => store,
                    Err(_) => {
                        // Try to create the store if opening failed
                        kvs::create_default_store("rhai-vault", false, None)
                            .expect("Failed to create key store")
                    }
                }
            }).clone()
        }
    }
}

// Key space management functions
fn load_key_space(name: &str, password: &str) -> bool {
    let store = get_key_store();
    kvs::load_key_space(&store, name, password)
}

fn create_key_space(name: &str, password: &str) -> bool {
    let store = get_key_store();
    kvs::create_key_space(&store, name, password)
}

// Auto-save function for internal use
fn auto_save_key_space(password: &str) -> bool {
    let store = get_key_store();
    kvs::save_key_space(&store, password)
}

// Export the current key space to a JSON string
fn encrypt_key_space(password: &str) -> String {
    match keypair::get_current_space()
        .and_then(|space| symmetric::encrypt_key_space(&space, password))
        .and_then(|encrypted_space| symmetric::serialize_encrypted_space(&encrypted_space)) 
    {
        Ok(json) => json,
        Err(e) => {
            log::error!("Error encrypting key space: {}", e);
            String::new()
        }
    }
}

// Import a key space from a JSON string
fn decrypt_key_space(encrypted: &str, password: &str) -> bool {
    match symmetric::deserialize_encrypted_space(encrypted)
        .and_then(|encrypted_space| symmetric::decrypt_key_space(&encrypted_space, password))
        .and_then(|space| keypair::set_current_space(space))
    {
        Ok(_) => true,
        Err(e) => {
            log::error!("Error decrypting key space: {}", e);
            false
        }
    }
}

// Keypair management functions
fn create_keypair(name: &str, password: &str) -> bool {
    match keypair::create_keypair(name) {
        Ok(_) => {
            // Auto-save the key space after creating a keypair
            auto_save_key_space(password)
        },
        Err(e) => {
            log::error!("Error creating keypair: {}", e);
            false
        }
    }
}

fn select_keypair(name: &str) -> bool {
    match keypair::select_keypair(name) {
        Ok(_) => true,
        Err(e) => {
            log::error!("Error selecting keypair: {}", e);
            false
        }
    }
}

fn list_keypairs() -> Vec<String> {
    match keypair::list_keypairs() {
        Ok(keypairs) => keypairs,
        Err(e) => {
            log::error!("Error listing keypairs: {}", e);
            Vec::new()
        }
    }
}

// Cryptographic operations
fn sign(message: &str) -> String {
    let message_bytes = message.as_bytes();
    match keypair::keypair_sign(message_bytes) {
        Ok(signature) => BASE64.encode(signature),
        Err(e) => {
            log::error!("Error signing message: {}", e);
            String::new()
        }
    }
}

fn verify(message: &str, signature: &str) -> bool {
    let message_bytes = message.as_bytes();
    match BASE64.decode(signature)
        .map_err(|e| e.to_string())
        .and_then(|sig_bytes| keypair::keypair_verify(message_bytes, &sig_bytes)
            .map_err(|e| e.to_string()))
    {
        Ok(is_valid) => is_valid,
        Err(e) => {
            log::error!("Error verifying signature: {}", e);
            false
        }
    }
}

// Symmetric encryption
fn generate_key() -> String {
    BASE64.encode(symmetric::generate_symmetric_key())
}

fn encrypt(key: &str, message: &str) -> String {
    match BASE64.decode(key)
        .map_err(|e| format!("Error decoding key: {}", e))
        .and_then(|key_bytes| {
            symmetric::encrypt_symmetric(&key_bytes, message.as_bytes())
                .map_err(|e| format!("Error encrypting message: {}", e))
        })
    {
        Ok(ciphertext) => BASE64.encode(ciphertext),
        Err(e) => {
            log::error!("{}", e);
            String::new()
        }
    }
}

fn decrypt(key: &str, ciphertext: &str) -> String {
    match BASE64.decode(key)
        .map_err(|e| format!("Error decoding key: {}", e))
        .and_then(|key_bytes| {
            BASE64.decode(ciphertext)
                .map_err(|e| format!("Error decoding ciphertext: {}", e))
                .and_then(|cipher_bytes| {
                    symmetric::decrypt_symmetric(&key_bytes, &cipher_bytes)
                        .map_err(|e| format!("Error decrypting ciphertext: {}", e))
                })
        })
        .and_then(|plaintext| {
            String::from_utf8(plaintext)
                .map_err(|e| format!("Error converting plaintext to string: {}", e))
        })
    {
        Ok(text) => text,
        Err(e) => {
            log::error!("{}", e);
            String::new()
        }
    }
}

// Ethereum operations

// Create a network-independent Ethereum wallet
fn create_ethereum_wallet() -> bool {
    match ethereum::create_ethereum_wallet() {
        Ok(_) => true,
        Err(e) => {
            log::error!("Error creating Ethereum wallet: {}", e);
            false
        }
    }
}

// Get the Ethereum wallet address (same for all networks)
fn get_ethereum_address() -> String {
    match ethereum::get_current_ethereum_wallet() {
        Ok(wallet) => wallet.address_string(),
        Err(e) => {
            log::error!("Error getting Ethereum address: {}", e);
            String::new()
        }
    }
}

// Create a wallet from a name
fn create_ethereum_wallet_from_name(name: &str) -> bool {
    match ethereum::create_ethereum_wallet_from_name(name) {
        Ok(_) => true,
        Err(e) => {
            log::error!("Error creating Ethereum wallet from name: {}", e);
            false
        }
    }
}

// Create a wallet from a private key
fn create_ethereum_wallet_from_private_key(private_key: &str) -> bool {
    match ethereum::create_ethereum_wallet_from_private_key(private_key) {
        Ok(_) => true,
        Err(e) => {
            log::error!("Error creating Ethereum wallet from private key: {}", e);
            false
        }
    }
}

// Clear all Ethereum wallets
fn clear_ethereum_wallets() -> bool {
    ethereum::clear_ethereum_wallets();
    true // Always return true since the operation doesn't have a failure mode
}

// Network registry functions

// Register a new network
fn register_network(name: &str, chain_id: i64, rpc_url: &str, explorer_url: &str, token_symbol: &str, decimals: i64) -> bool {
    ethereum::register_network(
        name,
        chain_id as u64,
        rpc_url,
        explorer_url,
        token_symbol,
        decimals as u8
    )
}

// Remove a network
fn remove_network(name: &str) -> bool {
    ethereum::remove_network(name)
}

// List supported networks
fn list_supported_networks() -> rhai::Array {
    let mut arr = rhai::Array::new();
    for name in ethereum::list_network_names() {
        arr.push(Dynamic::from(name.to_lowercase()));
    }
    arr
}

// Create a provider for a specific network
fn create_provider(network_name: &str) -> String {
    match ethereum::create_provider(network_name) {
        Ok(_) => network_name.to_string(),
        Err(e) => {
            log::error!("Error creating provider: {}", e);
            String::new()
        }
    }
}

// Get network token symbol
fn get_network_token_symbol(network_name: &str) -> String {
    match ethereum::get_network_by_name(network_name) {
        Some(network) => network.token_symbol,
        None => {
            log::error!("Unknown network: {}", network_name);
            String::new()
        }
    }
}

// Get network explorer URL
fn get_network_explorer_url(network_name: &str) -> String {
    match ethereum::get_network_by_name(network_name) {
        Some(network) => network.explorer_url,
        None => {
            log::error!("Unknown network: {}", network_name);
            String::new()
        }
    }
}

// Get the balance of an address on a specific network
fn get_balance(network_name: &str, address: &str) -> String {
    // Parse the address
    let addr = match Address::from_str(address) {
        Ok(addr) => addr,
        Err(e) => {
            log::error!("Invalid address format: {}", e);
            return String::new();
        }
    };

    // Execute the balance query in a blocking manner
    match run_async(ethereum::get_balance(network_name, addr)) {
        Ok(balance) => {
            // Format the balance with the network's token symbol
            match ethereum::get_network_by_name(network_name) {
                Some(network) => ethereum::format_balance(balance, &network),
                None => balance.to_string()
            }
        },
        Err(e) => {
            log::error!("Failed to get balance: {}", e);
            String::new()
        }
    }
}

// Send ETH from one address to another
fn send_eth(network_name: &str, to_address: &str, amount_str: &str) -> String {
    // Parse the address
    let to_addr = match Address::from_str(to_address) {
        Ok(addr) => addr,
        Err(e) => {
            log::error!("Invalid address format: {}", e);
            return String::new();
        }
    };

    // Parse the amount (using string to handle large numbers)
    let amount = match U256::from_dec_str(amount_str) {
        Ok(amt) => amt,
        Err(e) => {
            log::error!("Invalid amount format: {}", e);
            return String::new();
        }
    };

    // Get the wallet
    let wallet = match ethereum::get_current_ethereum_wallet() {
        Ok(w) => w,
        Err(e) => {
            log::error!("Failed to get wallet: {}", e);
            return String::new();
        }
    };

    // Execute the transaction in a blocking manner
    match run_async(ethereum::send_eth(&wallet, network_name, to_addr, amount)) {
        Ok(tx_hash) => format!("{:?}", tx_hash),
        Err(e) => {
            log::error!("Transaction failed: {}", e);
            String::new()
        }
    }
}

// Smart contract operations

// Load a contract ABI from a JSON string and create a contract instance
fn load_contract_abi(network_name: &str, address: &str, abi_json: &str) -> String {
    // Get the network
    let network = match ethereum::get_network_by_name(network_name) {
        Some(network) => network,
        None => {
            log::error!("Unknown network: {}", network_name);
            return String::new();
        }
    };

    // Parse the ABI
    let abi = match ethereum::load_abi_from_json(abi_json) {
        Ok(abi) => abi,
        Err(e) => {
            log::error!("Error parsing ABI: {}", e);
            return String::new();
        }
    };

    // Create the contract
    match ethereum::Contract::from_address_string(address, abi, network) {
        Ok(contract) => {
            // Serialize the contract to JSON for storage
            match serde_json::to_string(&contract) {
                Ok(json) => json,
                Err(e) => {
                    log::error!("Error serializing contract: {}", e);
                    String::new()
                }
            }
        },
        Err(e) => {
            log::error!("Error creating contract: {}", e);
            String::new()
        }
    }
}

// Load a contract ABI from a file
fn load_contract_abi_from_file(network_name: &str, address: &str, file_path: &str) -> String {
    // Read the ABI file
    match fs::read_to_string(file_path) {
        Ok(abi_json) => load_contract_abi(network_name, address, &abi_json),
        Err(e) => {
            log::error!("Error reading ABI file: {}", e);
            String::new()
        }
    }
}

// Call a read-only function on a contract (no arguments version)
fn call_contract_read_no_args(contract_json: &str, function_name: &str) -> Dynamic {
    call_contract_read(contract_json, function_name, rhai::Array::new())
}

// Call a read-only function on a contract with arguments
fn call_contract_read(contract_json: &str, function_name: &str, args: rhai::Array) -> Dynamic {
    // Deserialize the contract
    let contract: ethereum::Contract = match serde_json::from_str(contract_json) {
        Ok(contract) => contract,
        Err(e) => {
            log::error!("Error deserializing contract: {}", e);
            return Dynamic::UNIT;
        }
    };

    // Prepare the arguments
    let tokens = match prepare_function_arguments(&contract.abi, function_name, &args) {
        Ok(tokens) => tokens,
        Err(e) => {
            log::error!("Error preparing arguments: {}", e);
            return Dynamic::UNIT;
        }
    };

    // Create a provider
    let provider = match ethereum::create_provider(&contract.network.name) {
        Ok(p) => p,
        Err(e) => {
            log::error!("Failed to create provider: {}", e);
            return Dynamic::UNIT;
        }
    };

    // Execute the call in a blocking manner
    match run_async(ethereum::call_read_function(&contract, &provider, function_name, tokens)) {
        Ok(result) => ethereum::token_to_dynamic(&result),
        Err(e) => {
            log::error!("Failed to call contract function: {}", e);
            Dynamic::UNIT
        }
    }
}

// Call a state-changing function on a contract (no arguments version)
fn call_contract_write_no_args(contract_json: &str, function_name: &str) -> String {
    call_contract_write(contract_json, function_name, rhai::Array::new())
}

// Call a state-changing function on a contract with arguments
fn call_contract_write(contract_json: &str, function_name: &str, args: rhai::Array) -> String {
    // Deserialize the contract
    let contract: ethereum::Contract = match serde_json::from_str(contract_json) {
        Ok(contract) => contract,
        Err(e) => {
            log::error!("Error deserializing contract: {}", e);
            return String::new();
        }
    };

    // Prepare the arguments
    let tokens = match prepare_function_arguments(&contract.abi, function_name, &args) {
        Ok(tokens) => tokens,
        Err(e) => {
            log::error!("Error preparing arguments: {}", e);
            return String::new();
        }
    };

    // Get the wallet
    let wallet = match ethereum::get_current_ethereum_wallet() {
        Ok(w) => w,
        Err(e) => {
            log::error!("Failed to get wallet: {}", e);
            return String::new();
        }
    };

    // Create a provider
    let provider = match ethereum::create_provider(&contract.network.name) {
        Ok(p) => p,
        Err(e) => {
            log::error!("Failed to create provider: {}", e);
            return String::new();
        }
    };

    // Execute the transaction in a blocking manner
    match run_async(ethereum::call_write_function(&contract, &wallet, &provider, function_name, tokens)) {
        Ok(tx_hash) => format!("{:?}", tx_hash),
        Err(e) => {
            // Log the error details for debugging
            log::debug!("\nERROR DETAILS: Transaction failed: {}", e);
            log::debug!("Contract address: {}", contract.address);
            log::debug!("Function: {}", function_name);
            log::debug!("Arguments: {:?}", args);
            log::debug!("Wallet address: {}", wallet.address);
            log::debug!("Network: {}", contract.network.name);
            log::error!("Transaction failed: {}", e);
            String::new()
        }
    }
}

/// Register crypto functions with the Rhai engine
pub fn register_crypto_module(engine: &mut Engine) -> Result<(), Box<EvalAltResult>> {
    // Register key space functions
    engine.register_fn("load_key_space", load_key_space);
    engine.register_fn("create_key_space", create_key_space);
    engine.register_fn("encrypt_key_space", encrypt_key_space);
    engine.register_fn("decrypt_key_space", decrypt_key_space);

    // Register keypair functions
    engine.register_fn("create_keypair", create_keypair);
    engine.register_fn("select_keypair", select_keypair);
    engine.register_fn("list_keypairs", list_keypairs);

    // Register signing/verification functions
    engine.register_fn("sign", sign);
    engine.register_fn("verify", verify);

    // Register symmetric encryption functions
    engine.register_fn("generate_key", generate_key);
    engine.register_fn("encrypt", encrypt);
    engine.register_fn("decrypt", decrypt);

    // Register Ethereum wallet functions
    engine.register_fn("create_ethereum_wallet", create_ethereum_wallet);
    engine.register_fn("get_ethereum_address", get_ethereum_address);
    engine.register_fn("create_ethereum_wallet_from_name", create_ethereum_wallet_from_name);
    engine.register_fn("create_ethereum_wallet_from_private_key", create_ethereum_wallet_from_private_key);
    engine.register_fn("clear_ethereum_wallets", clear_ethereum_wallets);

    // Register network registry functions
    engine.register_fn("register_network", register_network);
    engine.register_fn("remove_network", remove_network);
    engine.register_fn("list_supported_networks", list_supported_networks);
    engine.register_fn("get_network_token_symbol", get_network_token_symbol);
    engine.register_fn("get_network_explorer_url", get_network_explorer_url);

    // Register provider functions
    engine.register_fn("create_provider", create_provider);

    // Register transaction functions
    engine.register_fn("send_eth", send_eth);
    engine.register_fn("get_balance", get_balance);

    // Register smart contract functions
    engine.register_fn("load_contract_abi", load_contract_abi);
    engine.register_fn("load_contract_abi_from_file", load_contract_abi_from_file);
    engine.register_fn("call_contract_read", call_contract_read_no_args);
    engine.register_fn("call_contract_read", call_contract_read);
    engine.register_fn("call_contract_write", call_contract_write_no_args);
    engine.register_fn("call_contract_write", call_contract_write);

    Ok(())
}