nucleus-sdk 0.1.0

use std::sync::Arc;
use ethers::{
    abi::{Token, parse_abi},
    types::{Bytes, U256, Address, TransactionRequest, H256, H256 as Bytes32},
    signers::Signer,
    providers::{Middleware, Provider, Http},
    contract::Contract,
};
use serde::{Deserialize, Serialize};
use crate::{
    client::Client,
    error::{Result, InvalidInputsError, Transaction},
    utils::encode_with_signature,
};
use ethers_middleware::SignerMiddleware;

#[derive(Debug)]
pub struct CalldataQueue {
    pub client: Arc<Client>,
    pub manager_address: Address,
    pub chain_id: u64,
    pub root: String,
    pub calls: Vec<Call>,
}

#[derive(Debug, Serialize, Deserialize)]
struct Call {
    target_address: Address,
    data: Bytes,
    value: U256,
    args: Vec<Token>,  // Using ethers Token for type-safe args
    function_signature: String,
    proof_data: Vec<Bytes>,
    decoder_and_sanitizer: Address,
}

impl CalldataQueue {
    pub async fn new(
        chain_id: u64,
        strategist_address: String,
        rpc_url: String,
        symbol: String,
        client: Arc<Client>,
    ) -> Result<Self> {
        // Convert chain_id to string for address book lookup
        let network_string = chain_id.to_string();

        // get the manager address from the address book using chain_id
        let manager_address = client
            .address_book
            .get(&network_string)
            .and_then(|network| network.get("nucleus"))
            .and_then(|nucleus| nucleus.get(&symbol))
            .and_then(|symbol_data| symbol_data.get("manager"))
            .and_then(|addr| addr.as_str())
            .ok_or_else(|| {
                InvalidInputsError(format!(
                    "Could not find manager address for network '{}' and symbol '{}'. Please check the network and symbol are valid.",
                    network_string, symbol
                ))
            })?;

        // parse the manager address
        let manager_address = manager_address.parse()
            .map_err(|_| InvalidInputsError("Invalid manager address format".to_string()))?;

        // Create a provider to interact with the blockchain
        let provider = Arc::new(
            Provider::<Http>::try_from(&rpc_url)
                .map_err(|_| InvalidInputsError(format!("Invalid RPC URL: {}", rpc_url)))?
        );

        // Test the connection
        provider.get_block_number().await
            .map_err(|_| InvalidInputsError(format!("Could not connect to RPC URL '{}'. Please check the RPC URL is valid and accessible.", rpc_url)))?;

        // Parse the strategist address
        let strategist = strategist_address.parse::<Address>()
            .map_err(|_| InvalidInputsError("Invalid strategist address format".to_string()))?;

        // Create contract instance
        let abi = parse_abi(&[
            "function manageRoot(address strategist) view returns (bytes32)"
        ]).expect("Failed to parse ABI");

        let contract = Contract::new(
            manager_address,
            abi,
            Arc::clone(&provider),
        );

        // Call manageRoot function
        let root: Bytes32 = contract
            .method::<_, Bytes32>("manageRoot", strategist)
            .map_err(|e| InvalidInputsError(format!("Failed to create manageRoot method call: {}", e)))?
            .call()
            .await
            .map_err(|e| InvalidInputsError(format!("Failed to call manageRoot: {}", e)))?;

        // Convert root to hex string
        let root = format!("{:?}", root);

        // Check if root is zero
        if root == "0x0000000000000000000000000000000000000000000000000000000000000000" {
            return Err(InvalidInputsError(format!(
                "Could not find root for strategist '{}'. Please check the strategist address is valid.",
                strategist_address
            )).into());
        }


        Ok(Self {
            client,
            manager_address,
            chain_id,
            root,
            calls: Vec::new(),
        })
    }

    // add a call to the manager call
    pub fn add_call(
        &mut self,
        target_address: Address,
        function_signature: String,
        args: Vec<Token>,
        value: U256,
    ) -> Result<()> {
        // encode the function call
        let data = encode_with_signature(&function_signature, &args)
            .map_err(|e| InvalidInputsError(format!("Failed to encode function data: {}", e)))?;

        // push the call to the calls vector
        self.calls.push(Call {
            target_address,
            data: data.into(),
            value,
            args,
            function_signature,
            proof_data: vec![],
            decoder_and_sanitizer: Address::zero(),
        });

        Ok(())
    }

    pub async fn get_calldata(&mut self) -> Result<Bytes> {
        // Collect all calls that need proof data
        let proof_requests: Vec<ProofRequest> = self.calls.iter()
            .filter(|call| call.proof_data.is_empty())
            .map(|call| ProofRequest {
                target: format!("{:?}", call.target_address),
                calldata: format!("0x{}", hex::encode(&call.data)),
                value: call.value.as_u64(),
            })
            .collect();

        // If we have any calls needing proofs, get them in batch
        if !proof_requests.is_empty() {
            let batch_response = self._get_batch_proofs_and_decoders(proof_requests).await?;
            
            // Update the calls with their proof data
            let mut proof_index = 0;
            for call in &mut self.calls {
                if call.proof_data.is_empty() {
                    call.proof_data = batch_response.proofs[proof_index].clone();
                    call.decoder_and_sanitizer = batch_response.decoder_and_sanitizer_addresses[proof_index];
                    proof_index += 1;
                }
            }
        }

        // Then build the vectors for the final call
        let mut proofs = Vec::new();
        let mut decoders = Vec::new();
        let mut targets = Vec::new();
        let mut data = Vec::new();
        let mut values = Vec::new();

        for call in &self.calls {
            proofs.push(Token::Array(
                call.proof_data.iter()
                    .map(|p| Token::FixedBytes(p.to_vec()))
                    .collect()
            ));
            decoders.push(Token::Address(call.decoder_and_sanitizer));
            targets.push(Token::Address(call.target_address));
            data.push(Token::Bytes(call.data.to_vec()));
            values.push(Token::Uint(call.value));
        }

        // Create the manager function
        let function = "manageVaultWithMerkleVerification(bytes32[][],address[],address[],bytes[],uint256[])";

        let args = vec![
            Token::Array(proofs),
            Token::Array(decoders),
            Token::Array(targets),
            Token::Array(data),
            Token::Array(values),
        ];

        let encoded = encode_with_signature(function, &args)
            .map_err(|e| InvalidInputsError(format!("Failed to encode manager function: {}", e)))?;

        Ok(encoded.into())
    }

    /// Execute the queued calls using the provided signer.
    ///
    /// The signer contains both the provider connection and the account information.
    pub async fn execute<M: Middleware, S: Signer>(
        &mut self, 
        signer: &SignerMiddleware<M, S>
    ) -> Result<H256> {
        if self.calls.is_empty() {
            return Err(InvalidInputsError("No calls to execute".to_string()).into());
        }

        let calldata = self.get_calldata().await?;
        
        // Get the sender's address from the signer
        let from_address = signer.address();

        let tx = TransactionRequest {
            from: Some(from_address),
            to: Some(self.manager_address.into()),
            data: Some(calldata),
            value: Some(U256::zero()),
            ..Default::default()
        };

        // Send the transaction using the signer and map the error
        let pending_tx = signer.send_transaction(tx, None)
            .await
            .map_err(|e| Transaction(format!("Transaction failed: {}", e)))?;
        
        // Return the transaction hash immediately (like the Python version)
        Ok(pending_tx.tx_hash())
    }

    async fn _get_batch_proofs_and_decoders(
        &self,
        leaves: Vec<ProofRequest>
    ) -> Result<BatchProofResponse> {
        // Create the batch request with chain_id at the top level
        let request = BatchProofRequest {
            chain: self.chain_id,
            calls: leaves,
        };

        // Convert request to JSON value
        let json_value = serde_json::to_value(&request)
            .map_err(|e| InvalidInputsError(format!("Failed to serialize batch proof request: {}", e)))?;
        
        // Make the API call
        let response: BatchProofResponse = self.client
            .post(&format!("multiproofs/{}", self.root), Some(&json_value))
            .await?;

        // Verify the response arrays are the same length
        if response.proofs.len() != response.decoder_and_sanitizer_addresses.len() {
            return Err(InvalidInputsError(
                "Mismatched lengths in proof response arrays".to_string()
            ).into());
        }

        Ok(response)
    }
}

#[derive(Serialize)]
struct BatchProofRequest {
    chain: u64,  // Chain ID at the top level
    calls: Vec<ProofRequest>,
}

#[derive(Deserialize, Serialize)]
struct BatchProofResponse {
    proofs: Vec<Vec<Bytes>>,
    #[serde(rename = "decoderAndSanitizerAddress")]
    decoder_and_sanitizer_addresses: Vec<Address>,
}

#[derive(Serialize)]
struct ProofRequest {
    target: String,
    calldata: String,
    value: u64,
}

#[derive(Deserialize)]
struct ProofResponse {
    proof: Vec<Bytes>,
    decoder_and_sanitizer: Address,
}