newton-chainio 0.5.2

//! Policy Client

use crate::error::ChainIoError;
use alloy::{
    primitives::{Address, Bytes, B256, U256},
    rpc::types::TransactionReceipt,
    sol_types::SolValue,
};
use eigensdk::common::{get_provider, get_signer};
use newton_core::{
    mock_newton_policy_client::{
        INewtonProverTaskManager::{Task as MockTask, TaskResponse as MockTaskResponse},
        MockNewtonPolicyClient, NewtonMessage,
    },
    newton_policy_client::{INewtonPolicy, NewtonPolicyClient},
    newton_prover_task_manager::{
        INewtonPolicy as IINewtonPolicy,
        INewtonProverTaskManager::{Task, TaskResponse as ContractTaskResponse},
        NewtonProverTaskManager,
    },
};

use tracing::info;

/// PolicyClientController struct
#[derive(Debug)]
pub struct PolicyClientController {
    rpc_url: String,
    signer: String,
    client_address: Address,
}

impl PolicyClientController {
    /// new instance
    pub fn new(signer: String, rpc_url: String, client_address: Address) -> Self {
        PolicyClientController {
            signer,
            rpc_url,
            client_address,
        }
    }

    /// Get the policy address associated with this client
    pub async fn get_policy_address(&self) -> Result<Address, ChainIoError> {
        let provider = get_provider(&self.rpc_url);
        let policy_client = NewtonPolicyClient::new(self.client_address, provider);

        info!("Getting policy address for client: {}", self.client_address);

        let get_policy_address_call = policy_client.getPolicyAddress();

        let get_policy_address_result = get_policy_address_call.call().await;

        match get_policy_address_result {
            Ok(result) => {
                let policy_address = result;
                info!("Policy address for client {}: {}", self.client_address, policy_address);
                Ok(policy_address)
            }
            Err(e) => Err(ChainIoError::ContractError(e)),
        }
    }

    /// Get the policy configuration for this client
    pub async fn get_policy_config(&self) -> Result<INewtonPolicy::PolicyConfig, ChainIoError> {
        get_policy_config_for_client(self.client_address, self.rpc_url.clone())
            .await
            .map(|config| INewtonPolicy::PolicyConfig {
                policyParams: config.policyParams,
                expireAfter: config.expireAfter,
            })
    }

    /// Get the policy ID for this client
    pub async fn get_policy_id(&self) -> Result<B256, ChainIoError> {
        get_policy_id_for_client(self.client_address, self.rpc_url.clone()).await
    }

    /// Get the Newton Policy Task Manager address
    pub async fn get_newton_policy_task_manager(&self) -> Result<Address, ChainIoError> {
        let provider = get_provider(&self.rpc_url);
        let policy_client = NewtonPolicyClient::new(self.client_address, provider);

        info!(
            "Getting Newton Policy Task Manager address for client: {}",
            self.client_address
        );

        let get_task_manager_call = policy_client.getNewtonPolicyTaskManager();

        let get_task_manager_result = get_task_manager_call.call().await;

        match get_task_manager_result {
            Ok(result) => {
                let task_manager_address = result;
                info!(
                    "Newton Policy Task Manager address for client {}: {}",
                    self.client_address, task_manager_address
                );
                Ok(task_manager_address)
            }
            Err(e) => Err(ChainIoError::ContractError(e)),
        }
    }

    /// Set the policy configuration for this client
    pub async fn set_policy(
        &self,
        policy_params: Bytes,
        expire_after: u32,
    ) -> Result<(TransactionReceipt, B256), ChainIoError> {
        let wallet = get_signer(&self.signer, &self.rpc_url);
        let policy_client = NewtonPolicyClient::new(self.client_address, wallet);

        let policy_config = INewtonPolicy::PolicyConfig {
            policyParams: policy_params.clone(),
            expireAfter: expire_after,
        };

        info!(
            "Setting policy for client: {} with params: {}, expire_after: {}",
            self.client_address,
            hex!(policy_params),
            expire_after
        );

        let set_policy_call = policy_client.setPolicy(policy_config);

        let set_policy_result = set_policy_call.send().await;

        match set_policy_result {
            Ok(set_policy) => {
                let receipt_result = set_policy.get_receipt().await;

                match receipt_result {
                    Ok(receipt) => {
                        // Get the policy ID after setting
                        let policy_id = self.get_policy_id().await?;

                        info!(
                            "Policy set successfully for client: {}, policy_id: {}",
                            self.client_address, policy_id
                        );
                        Ok((receipt, policy_id))
                    }
                    Err(e) => Err(ChainIoError::AlloyProviderError(e)),
                }
            }
            Err(e) => Err(ChainIoError::ContractError(e)),
        }
    }

    /// Set the policy client owner
    pub async fn set_policy_client_owner(&self, new_owner: Address) -> Result<TransactionReceipt, ChainIoError> {
        let wallet = get_signer(&self.signer, &self.rpc_url);
        let policy_client = NewtonPolicyClient::new(self.client_address, wallet);

        info!(
            "Setting policy client owner for client: {} to: {}",
            self.client_address, new_owner
        );

        let set_owner_call = policy_client.setPolicyClientOwner(new_owner);

        let set_owner_result = set_owner_call.send().await;

        match set_owner_result {
            Ok(set_owner) => {
                let receipt_result = set_owner.get_receipt().await;
                match receipt_result {
                    Ok(receipt) => {
                        info!(
                            "Policy client owner set successfully for client: {} to: {}",
                            self.client_address, new_owner
                        );
                        Ok(receipt)
                    }
                    Err(e) => Err(ChainIoError::AlloyProviderError(e)),
                }
            }
            Err(e) => Err(ChainIoError::ContractError(e)),
        }
    }
}

/// MockPolicyClientController struct for mock policy client operations
#[derive(Debug)]
pub struct MockPolicyClientController {
    rpc_url: String,
    signer: String,
    client_address: Address,
}

impl MockPolicyClientController {
    /// new instance
    pub fn new(signer: String, rpc_url: String, client_address: Address) -> Self {
        MockPolicyClientController {
            signer,
            rpc_url,
            client_address,
        }
    }

    /// Get the balance of a token for this client
    pub async fn balance_of(&self, token: Address) -> Result<U256, ChainIoError> {
        let provider = get_provider(&self.rpc_url);
        let mock_policy_client = MockNewtonPolicyClient::new(self.client_address, provider);

        info!(
            "Getting balance for token: {} for client: {}",
            token, self.client_address
        );

        let balance_of_call = mock_policy_client.balanceOf(token);

        let balance_of_result = balance_of_call.call().await;

        match balance_of_result {
            Ok(result) => {
                let balance = result;
                info!(
                    "Balance for token {} for client {}: {}",
                    token, self.client_address, balance
                );
                Ok(balance)
            }
            Err(e) => Err(ChainIoError::ContractError(e)),
        }
    }

    /// Deposit tokens to this client
    pub async fn deposit(&self, token: Address, token_amount: U256) -> Result<TransactionReceipt, ChainIoError> {
        let wallet = get_signer(&self.signer, &self.rpc_url);
        let mock_policy_client = MockNewtonPolicyClient::new(self.client_address, wallet);

        info!(
            "Depositing {} tokens of {} for client: {}",
            token_amount, token, self.client_address
        );

        let deposit_call = mock_policy_client.deposit(token, token_amount);

        let deposit_result = deposit_call.send().await;

        match deposit_result {
            Ok(deposit) => {
                let receipt_result = deposit.get_receipt().await;
                match receipt_result {
                    Ok(receipt) => {
                        info!(
                            "Deposit successful for client: {} - {} tokens of {}",
                            self.client_address, token_amount, token
                        );
                        Ok(receipt)
                    }
                    Err(e) => Err(ChainIoError::AlloyProviderError(e)),
                }
            }
            Err(e) => Err(ChainIoError::ContractError(e)),
        }
    }

    /// Withdraw tokens from this client
    pub async fn withdraw(&self, token: Address, token_amount: U256) -> Result<TransactionReceipt, ChainIoError> {
        let wallet = get_signer(&self.signer, &self.rpc_url);
        let mock_policy_client = MockNewtonPolicyClient::new(self.client_address, wallet);

        info!(
            "Withdrawing {} tokens of {} for client: {}",
            token_amount, token, self.client_address
        );

        let withdraw_call = mock_policy_client.withdraw(token, token_amount);

        let withdraw_result = withdraw_call.send().await;

        match withdraw_result {
            Ok(withdraw) => {
                let receipt_result = withdraw.get_receipt().await;
                match receipt_result {
                    Ok(receipt) => {
                        info!(
                            "Withdrawal successful for client: {} - {} tokens of {}",
                            self.client_address, token_amount, token
                        );
                        Ok(receipt)
                    }
                    Err(e) => Err(ChainIoError::AlloyProviderError(e)),
                }
            }
            Err(e) => Err(ChainIoError::ContractError(e)),
        }
    }

    /// Execute an intent using an attestation
    pub async fn execute_intent(
        &self,
        attestation: NewtonMessage::Attestation,
    ) -> Result<TransactionReceipt, ChainIoError> {
        let wallet = get_signer(&self.signer, &self.rpc_url);
        let mock_policy_client = MockNewtonPolicyClient::new(self.client_address, wallet);

        info!(
            "Executing intent for client: {} with task_id: {}, policy_id: {}",
            self.client_address, attestation.taskId, attestation.policyId
        );

        let execute_intent_call = mock_policy_client.executeIntent(attestation);

        let execute_intent_result = execute_intent_call.send().await;

        match execute_intent_result {
            Ok(execute_intent) => {
                let receipt_result = execute_intent.get_receipt().await;
                match receipt_result {
                    Ok(receipt) => {
                        info!("Intent executed successfully for client: {}", self.client_address);
                        Ok(receipt)
                    }
                    Err(e) => Err(ChainIoError::AlloyProviderError(e)),
                }
            }
            Err(e) => Err(ChainIoError::ContractError(e)),
        }
    }

    /// Get the owner of this mock policy client
    pub async fn owner(&self) -> Result<Address, ChainIoError> {
        let provider = get_provider(&self.rpc_url);
        let mock_policy_client = MockNewtonPolicyClient::new(self.client_address, provider);

        info!("Getting owner for mock client: {}", self.client_address);

        let owner_call = mock_policy_client.owner();

        let owner_result = owner_call.call().await;

        match owner_result {
            Ok(result) => {
                let owner = result;
                info!("Owner for mock client {}: {}", self.client_address, owner);
                Ok(owner)
            }
            Err(e) => Err(ChainIoError::ContractError(e)),
        }
    }

    /// Set the owner of this mock policy client
    pub async fn set_owner(&self, new_owner: Address) -> Result<TransactionReceipt, ChainIoError> {
        let wallet = get_signer(&self.signer, &self.rpc_url);
        let mock_policy_client = MockNewtonPolicyClient::new(self.client_address, wallet);

        info!(
            "Setting owner for mock client: {} to: {}",
            self.client_address, new_owner
        );

        let set_owner_call = mock_policy_client.setOwner(new_owner);

        let set_owner_result = set_owner_call.send().await;

        match set_owner_result {
            Ok(set_owner) => {
                let receipt_result = set_owner.get_receipt().await;
                match receipt_result {
                    Ok(receipt) => {
                        info!(
                            "Owner set successfully for mock client: {} to: {}",
                            self.client_address, new_owner
                        );
                        Ok(receipt)
                    }
                    Err(e) => Err(ChainIoError::AlloyProviderError(e)),
                }
            }
            Err(e) => Err(ChainIoError::ContractError(e)),
        }
    }

    /// Execute intent using direct attestation validation (on-chain signature verification).
    ///
    /// Calls `validateAttestationDirect` through the policy client contract, which
    /// is required because the TaskManager enforces `msg.sender == task.policyClient`.
    ///
    /// # Arguments
    /// * `task_manager_address` - Address of the NewtonProverTaskManager contract (unused, kept for API compatibility)
    /// * `task` - The Task struct
    /// * `task_response` - The TaskResponse struct (contract format)
    /// * `signature_data` - ABI-encoded signature data (NonSignerStakesAndSignature or BN254Certificate)
    pub async fn execute_intent_direct(
        &self,
        _task_manager_address: Address,
        task: Task,
        task_response: ContractTaskResponse,
        signature_data: Bytes,
    ) -> Result<bool, ChainIoError> {
        let wallet = get_signer(&self.signer, &self.rpc_url);
        let policy_client = MockNewtonPolicyClient::new(self.client_address, wallet);

        info!(
            "Executing validateAttestationDirect via policy client {} for task_id: {}",
            self.client_address, task_response.taskId
        );

        // ABI roundtrip to convert between structurally identical but type-incompatible
        // Task/TaskResponse from different generated modules
        let mock_task = MockTask::abi_decode(&task.abi_encode()).expect("Task ABI roundtrip");
        let mock_task_response =
            MockTaskResponse::abi_decode(&task_response.abi_encode()).expect("TaskResponse ABI roundtrip");

        let validate_call = policy_client.validateAttestationDirect(mock_task, mock_task_response, signature_data);

        let validate_result = validate_call.send().await;

        match validate_result {
            Ok(pending_tx) => {
                let receipt = pending_tx
                    .get_receipt()
                    .await
                    .map_err(ChainIoError::AlloyProviderError)?;

                if !receipt.status() {
                    tracing::error!(
                        tx_hash = %receipt.transaction_hash,
                        "validateAttestationDirect transaction reverted on-chain"
                    );
                    return Err(ChainIoError::TransactionReverted(receipt.transaction_hash));
                }

                info!(
                    "validateAttestationDirect executed successfully, tx_hash: {}",
                    receipt.transaction_hash
                );
                Ok(true)
            }
            Err(e) => {
                info!("validateAttestationDirect failed: {}", e);
                Err(ChainIoError::ContractError(e))
            }
        }
    }
}

// Standalone functions for policy client operations

/// Get the policy address for a given policy client address
pub async fn get_policy_address_for_client(client_address: Address, rpc_url: String) -> Result<Address, ChainIoError> {
    let provider = get_provider(&rpc_url);
    let policy_client = NewtonPolicyClient::new(client_address, provider);

    info!("Getting policy address for client: {}", client_address);

    let get_policy_address_call = policy_client.getPolicyAddress();

    let get_policy_address_result = get_policy_address_call.call().await;

    match get_policy_address_result {
        Ok(result) => {
            let policy_address = result;
            info!("Policy address for client {}: {}", client_address, policy_address);
            Ok(policy_address)
        }
        Err(e) => Err(ChainIoError::ContractError(e)),
    }
}

/// Get the policy ID for a given policy client address
pub async fn get_policy_id_for_client(client_address: Address, rpc_url: String) -> Result<B256, ChainIoError> {
    let provider = get_provider(&rpc_url);
    let policy_client = NewtonPolicyClient::new(client_address, provider);

    info!("Getting policy ID for client: {}", client_address);

    let get_policy_id_call = policy_client.getPolicyId();

    let get_policy_id_result = get_policy_id_call.call().await;

    match get_policy_id_result {
        Ok(result) => {
            let policy_id = result;
            info!("Policy ID for client {}: {}", client_address, policy_id);
            Ok(policy_id)
        }
        Err(e) => Err(ChainIoError::ContractError(e)),
    }
}

/// Get the policy configuration for a given policy client address
pub async fn get_policy_config_for_client(
    client_address: Address,
    rpc_url: String,
) -> Result<IINewtonPolicy::PolicyConfig, ChainIoError> {
    let provider = get_provider(&rpc_url);
    let policy_client = NewtonPolicyClient::new(client_address, provider);

    info!("Getting policy config for client: {}", client_address);

    let get_policy_config_call = policy_client.getPolicyConfig();

    let get_policy_config_result = get_policy_config_call.call().await;

    match get_policy_config_result {
        Ok(result) => {
            let policy_config = result;
            info!(
                "Policy config for client {}: params: {}, expire_after: {}",
                client_address,
                hex!(policy_config.policyParams.clone()),
                policy_config.expireAfter
            );
            Ok(IINewtonPolicy::PolicyConfig {
                policyParams: policy_config.policyParams,
                expireAfter: policy_config.expireAfter,
            })
        }
        Err(e) => Err(ChainIoError::ContractError(e)),
    }
}

/// Get the PolicyConfig for a given policyId on a specific policy contract.
///
/// This calls `NewtonPolicy.getPolicyConfig(policyId)` directly.
/// The policyId pins the config to a specific generation — the returned
/// config is consensus-safe for caching keyed by policyId.
pub async fn get_policy_config_by_id(
    policy_address: Address,
    policy_id: B256,
    rpc_url: &str,
) -> Result<newton_core::newton_prover_task_manager::INewtonPolicy::PolicyConfig, ChainIoError> {
    use newton_core::newton_policy::NewtonPolicy;

    let provider = get_provider(rpc_url);
    let policy = NewtonPolicy::new(policy_address, provider);

    let config = policy
        .getPolicyConfig(policy_id)
        .call()
        .await
        .map_err(ChainIoError::ContractError)?;

    Ok(newton_core::newton_prover_task_manager::INewtonPolicy::PolicyConfig {
        policyParams: config.policyParams,
        expireAfter: config.expireAfter,
    })
}