newton-core 0.7.5

use std::str::FromStr;

use crate::{
    common::chain::is_destination_chain,
    config::contracts::{get_deployment_env, load_deployment_json},
};
use alloy::primitives::{Address, ChainId};
use serde_json::Value;

fn load_newton_prover_data(chain_id: ChainId) -> eyre::Result<Value> {
    let env = get_deployment_env();
    let category = if is_destination_chain(chain_id) {
        "newton-cross-chain"
    } else {
        "newton-prover"
    };
    let raw = load_deployment_json(category, chain_id, &env)
        .map_err(|e| eyre::eyre!("Failed to load Newton prover deployment data: {}", e))?;
    serde_json::from_str(&raw).map_err(|e| eyre::eyre!("Failed to parse Newton prover deployment data: {}", e))
}

fn load_newton_policy_data(chain_id: ChainId, env: &str) -> eyre::Result<Value> {
    let raw = load_deployment_json("policy", chain_id, env)
        .map_err(|e| eyre::eyre!("Failed to load Newton policy deployment data: {}", e))?;
    serde_json::from_str(&raw).map_err(|e| eyre::eyre!("Failed to parse Newton policy deployment data: {}", e))
}

fn load_eigenlayer_data(chain_id: ChainId) -> eyre::Result<Value> {
    let env = get_deployment_env();
    let raw = load_deployment_json("core", chain_id, &env)
        .map_err(|e| eyre::eyre!("Failed to load Eigenlayer deployment data: {}", e))?;
    serde_json::from_str(&raw).map_err(|e| eyre::eyre!("Failed to parse Eigenlayer deployment data: {}", e))
}

fn address_from(data: &Value, key: &str, description: &str) -> eyre::Result<Address> {
    let addr = data["addresses"][key]
        .as_str()
        .ok_or_else(|| eyre::eyre!("{} address missing in deployment data", description))?;
    Address::from_str(addr).map_err(|e| eyre::eyre!("Invalid {} address: {}", description, e))
}

/// Get the newton prover registry coordinator address for anvil
pub async fn get_newton_operator_registry(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_newton_prover_data(chain_id)?;
    address_from(&data, "operatorRegistry", "operator registry")
}

/// Get the prover service manager address for anvil
pub async fn get_newton_prover_service_manager(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_newton_prover_data(chain_id)?;
    address_from(&data, "newtonProverServiceManager", "newton prover service manager")
}

/// Get the newton prover operator state retriever address for anvil
pub async fn get_newton_operator_state_retriever(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_newton_prover_data(chain_id)?;
    address_from(&data, "operatorStateRetriever", "operator state retriever")
}

/// Get the newton prover task manager address for anvil
pub async fn get_newton_prover_task_manager(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_newton_prover_data(chain_id)?;
    address_from(&data, "newtonProverTaskManager", "newton prover task manager")
}

/// Get the challenge verifier address for anvil
pub async fn get_challenge_verifier_address(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_newton_prover_data(chain_id)?;
    address_from(&data, "challengeVerifier", "challenge verifier")
}

/// Get the attestation validator address for anvil
pub async fn get_attestation_validator_address(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_newton_prover_data(chain_id)?;
    address_from(&data, "attestationValidator", "attestation validator")
}

/// Get the enclave version registry address
pub async fn get_enclave_version_registry_address(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_newton_prover_data(chain_id)?;
    address_from(&data, "enclaveVersionRegistry", "enclave version registry")
}

/// Get the newton prover policy address for anvil
pub async fn get_test_newton_policy_address(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_newton_policy_data(chain_id, &get_deployment_env())?;
    address_from(&data, "policy", "newton policy")
}

/// Get the test newton policy client address
pub async fn get_mock_newton_policy_client_address(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_newton_policy_data(chain_id, &get_deployment_env())?;
    address_from(&data, "policyClient", "mock newton policy client")
}

/// Get the test newton policy data address
pub async fn get_test_newton_policy_data_address(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_newton_policy_data(chain_id, &get_deployment_env())?;
    address_from(&data, "policyData", "test newton policy data")
}

/// Get the policy factory address for the `DEPLOYMENT_ENV` environment.
pub async fn get_policy_factory_address(chain_id: ChainId) -> eyre::Result<Address> {
    get_policy_factory_address_for_env(chain_id, &get_deployment_env()).await
}

/// Get the policy factory address for an explicit deployment environment
/// (`prod`, `stagef`), without reading `DEPLOYMENT_ENV`.
pub async fn get_policy_factory_address_for_env(chain_id: ChainId, env: &str) -> eyre::Result<Address> {
    let data = load_newton_policy_data(chain_id, env)?;
    address_from(&data, "policyFactory", "policy factory")
}

/// Get the policy data factory address from anvil
pub async fn get_policy_data_factory_address(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_newton_policy_data(chain_id, &get_deployment_env())?;
    address_from(&data, "policyDataFactory", "policy data factory")
}

/// Get NEWT token address
pub async fn get_mock_token_address(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_newton_prover_data(chain_id)?;
    address_from(&data, "token", "mock token")
}

/// erc20 mock strategy contract address
pub async fn get_mock_strategy(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_newton_prover_data(chain_id)?;
    address_from(&data, "strategy", "mock strategy")
}

/// Bls Apk Registry contract address
pub async fn get_bls_apk_registry_address(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_newton_prover_data(chain_id)?;
    address_from(&data, "blsApkRegistry", "bls apk registry")
}

/// Socket registry contract address
pub async fn get_socket_registry_address(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_newton_prover_data(chain_id)?;
    address_from(&data, "socketRegistry", "socket registry")
}

// Eigenlayer Core Contracts

/// Delegation Manager contract address
pub async fn get_delegation_manager_address(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_eigenlayer_data(chain_id)?;
    address_from(&data, "delegation", "delegation manager")
}

/// Strategy Mananger contract address
pub async fn get_strategy_manager_address(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_eigenlayer_data(chain_id)?;
    address_from(&data, "strategyManager", "strategy manager")
}

/// Avs Directory contract address
pub async fn get_avs_directory_address(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_eigenlayer_data(chain_id)?;
    address_from(&data, "avsDirectory", "avs directory")
}

/// proxy admin contract address
pub async fn get_proxy_admin(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_eigenlayer_data(chain_id)?;
    address_from(&data, "proxyAdmin", "proxy admin")
}

/// Rewards contract address
pub async fn get_rewards_coordinator_address(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_eigenlayer_data(chain_id)?;
    address_from(&data, "rewardsCoordinator", "rewards coordinator")
}

/// Allocation Manager contract address
pub async fn get_allocation_manager_address(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_eigenlayer_data(chain_id)?;
    address_from(&data, "allocationManager", "allocation manager")
}

/// Permission Controller contract address
pub async fn get_permission_controller_address(chain_id: ChainId) -> eyre::Result<Address> {
    let data = load_eigenlayer_data(chain_id)?;
    address_from(&data, "permissionController", "permission controller")
}

#[cfg(test)]
mod tests {
    use super::*;

    fn embedded_policy_factory(chain_id: ChainId, env: &str) -> Address {
        let data = load_newton_policy_data(chain_id, env).expect("embedded policy deployment");
        address_from(&data, "policyFactory", "policy factory").expect("policyFactory")
    }

    #[tokio::test]
    async fn policy_factory_for_env_uses_the_requested_env() {
        let prod = get_policy_factory_address_for_env(11155111, "prod").await.unwrap();
        let stagef = get_policy_factory_address_for_env(11155111, "stagef").await.unwrap();
        assert_eq!(prod, embedded_policy_factory(11155111, "prod"));
        assert_eq!(stagef, embedded_policy_factory(11155111, "stagef"));
        assert_ne!(prod, stagef, "Sepolia prod and stagef policy factories differ");
    }

    #[tokio::test]
    async fn policy_factory_for_env_errors_on_missing_deployment() {
        assert!(get_policy_factory_address_for_env(999_999, "prod").await.is_err());
    }
}