newton-core 0.4.16

//! testing utils for newton prover rs

use crate::common::{
    address,
    chain::{is_source_chain, is_supported_chain},
};

use crate::{
    challenge_verifier::ChallengeVerifier, newton_policy_client::NewtonPolicyClient,
    newton_prover_task_manager::NewtonProverTaskManager,
};
use alloy::{
    primitives::{Address, ChainId, B256},
    providers::Provider,
};
use newton_common::get_provider;
use tracing::info;

/// Get the newton prover registry coordinator address.
///
/// Newton's `OperatorRegistry` is deployed on both source and destination chains.
pub async fn get_newton_operator_registry(chain_id: ChainId) -> eyre::Result<Address> {
    eyre::ensure!(is_supported_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_newton_operator_registry(chain_id).await
}

/// Get the prover service manager address.
///
/// EigenLayer's service manager lives only on source chains.
pub async fn get_newton_prover_service_manager(chain_id: ChainId) -> eyre::Result<Address> {
    eyre::ensure!(is_source_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_newton_prover_service_manager(chain_id).await
}

/// Get the task manager address.
///
/// `NewtonProverTaskManager` (source) and `NewtonProverDestTaskManager` (destination)
/// share the same deployment key and are accessible on any supported chain.
pub async fn get_newton_prover_task_manager(chain_id: ChainId) -> eyre::Result<Address> {
    eyre::ensure!(is_supported_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_newton_prover_task_manager(chain_id).await
}

/// Get the newton prover operator state retriever address.
///
/// EigenLayer middleware contract, source-only.
pub async fn get_newton_operator_state_retriever(chain_id: ChainId) -> eyre::Result<Address> {
    eyre::ensure!(is_source_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_newton_operator_state_retriever(chain_id).await
}

/// Get the challenge verifier address.
///
/// `ChallengeVerifier` is deployed on both source and destination chains.
pub async fn get_challenge_verifier_address(chain_id: ChainId) -> eyre::Result<Address> {
    eyre::ensure!(is_supported_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_challenge_verifier_address(chain_id).await
}

/// Get the attestation validator address.
///
/// `AttestationValidator` is deployed on both source and destination chains.
pub async fn get_attestation_validator_address(chain_id: ChainId) -> eyre::Result<Address> {
    eyre::ensure!(is_supported_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_attestation_validator_address(chain_id).await
}

/// Get the newton prover policy address.
///
/// Policy contracts are deployed on both source and destination chains.
pub async fn get_test_newton_policy_address(chain_id: ChainId) -> eyre::Result<Address> {
    eyre::ensure!(is_supported_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_test_newton_policy_address(chain_id).await
}

/// Get the enclave version registry address.
pub async fn get_enclave_version_registry_address(chain_id: ChainId) -> eyre::Result<Address> {
    eyre::ensure!(is_supported_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_enclave_version_registry_address(chain_id).await
}

/// Get the test newton policy client address.
///
/// `PolicyClient` contracts are deployed on both source and destination chains.
pub async fn get_mock_newton_policy_client_address(chain_id: ChainId) -> eyre::Result<Address> {
    eyre::ensure!(is_supported_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_mock_newton_policy_client_address(chain_id).await
}

/// Get the policy id for the given policy client address
pub async fn get_policy_id_for_policy_client(rpc_url: &str, policy_client_address: Address) -> eyre::Result<B256> {
    let policy = NewtonPolicyClient::new(policy_client_address, get_provider(rpc_url));
    let policy_id = policy.getPolicyId().call().await.map_err(|e| {
        eyre::eyre!(
            "failed to get policy id for policy client {}: {}",
            policy_client_address,
            e
        )
    })?;
    Ok(policy_id)
}

/// Get the policy address for the given policy client address
pub async fn get_policy_address_for_policy_client(
    rpc_url: &str,
    policy_client_address: Address,
) -> eyre::Result<Address> {
    let policy = NewtonPolicyClient::new(policy_client_address, get_provider(rpc_url));
    let policy_address = policy.getPolicyAddress().call().await.map_err(|e| {
        eyre::eyre!(
            "failed to get policy address for policy client {}: {}",
            policy_client_address,
            e
        )
    })?;
    Ok(policy_address)
}

/// Get the test newton policy data address.
///
/// `PolicyData` contracts are deployed on both source and destination chains.
pub async fn get_test_newton_policy_data_address(chain_id: ChainId) -> eyre::Result<Address> {
    eyre::ensure!(is_supported_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_test_newton_policy_data_address(chain_id).await
}

/// Get the policy factory address.
///
/// `PolicyFactory` is deployed on both source and destination chains.
pub async fn get_policy_factory_address(chain_id: ChainId) -> eyre::Result<Address> {
    eyre::ensure!(is_supported_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_policy_factory_address(chain_id).await
}

/// Get the policy data factory address.
///
/// `PolicyDataFactory` is deployed on both source and destination chains.
pub async fn get_policy_data_factory_address(chain_id: ChainId) -> eyre::Result<Address> {
    eyre::ensure!(is_supported_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_policy_data_factory_address(chain_id).await
}

/// Get NEWT token address.
///
/// The NEWT mock token is deployed on source chains only.
pub async fn get_mock_token_address(chain_id: ChainId) -> eyre::Result<Address> {
    eyre::ensure!(is_source_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_mock_token_address(chain_id).await
}

// Eigenlayer Core Contracts — source-only (EigenLayer is deployed on source chains)

/// Delegation Manager contract address
pub async fn get_delegation_manager_address(rpc_url: &str) -> eyre::Result<Address> {
    let chain_id = get_chain_id(rpc_url).await;
    eyre::ensure!(is_source_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_delegation_manager_address(chain_id).await
}

/// Strategy Mananger contract address
pub async fn get_strategy_manager_address(rpc_url: &str) -> eyre::Result<Address> {
    let chain_id = get_chain_id(rpc_url).await;
    eyre::ensure!(is_source_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_strategy_manager_address(chain_id).await
}

/// Avs Directory contract address
pub async fn get_avs_directory_address(rpc_url: &str) -> eyre::Result<Address> {
    let chain_id = get_chain_id(rpc_url).await;
    eyre::ensure!(is_source_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_avs_directory_address(chain_id).await
}

/// erc20 mock strategy contract address
pub async fn get_mock_strategy(rpc_url: &str) -> eyre::Result<Address> {
    let chain_id = get_chain_id(rpc_url).await;
    eyre::ensure!(is_source_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_mock_strategy(chain_id).await
}

/// proxy admin contract address
pub async fn get_proxy_admin(rpc_url: &str) -> eyre::Result<Address> {
    let chain_id = get_chain_id(rpc_url).await;
    eyre::ensure!(is_source_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_proxy_admin(chain_id).await
}

/// Rewards contract address
pub async fn get_rewards_coordinator_address(rpc_url: &str) -> eyre::Result<Address> {
    let chain_id = get_chain_id(rpc_url).await;
    eyre::ensure!(is_source_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_rewards_coordinator_address(chain_id).await
}

/// Allocation Manager contract address
pub async fn get_allocation_manager_address(rpc_url: &str) -> eyre::Result<Address> {
    let chain_id = get_chain_id(rpc_url).await;
    eyre::ensure!(is_source_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_allocation_manager_address(chain_id).await
}

/// Permission Controller contract address
pub async fn get_permission_controller_address(rpc_url: &str) -> eyre::Result<Address> {
    let chain_id = get_chain_id(rpc_url).await;
    eyre::ensure!(is_source_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_permission_controller_address(chain_id).await
}

/// Bls Apk Registry contract address
pub async fn get_bls_apk_registry_address(rpc_url: &str) -> eyre::Result<Address> {
    let chain_id = get_chain_id(rpc_url).await;
    eyre::ensure!(is_source_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_bls_apk_registry_address(chain_id).await
}

/// Socket registry contract address
pub async fn get_socket_registry_address(rpc_url: &str) -> eyre::Result<Address> {
    let chain_id = get_chain_id(rpc_url).await;
    eyre::ensure!(is_source_chain(chain_id), "Unsupported chain id: {}", chain_id);
    address::get_socket_registry_address(chain_id).await
}

// Re-export functions from common::chain to avoid duplication
pub use crate::common::chain::{
    get_block_duration_ms, get_block_gas_limit, get_block_time_ms, get_chain_id, get_epoch_blocks,
    get_task_challenge_window_blocks, get_task_response_window_blocks,
};