newton-core 0.4.16

use std::str::FromStr;

use alloy_primitives::Address;
use tracing::{error, info};

// Generated by build.rs — maps (category, chain_id, env) → embedded JSON string.
include!(concat!(env!("OUT_DIR"), "/embedded_deployments.rs"));

/// Load deployment JSON for the given category, chain_id, and env.
///
/// Resolution order:
/// 1. `DEPLOYMENT_DIR` env var → `{DEPLOYMENT_DIR}/{category}/{chain_id}-{env}.json`
/// 2. Embedded data compiled into the binary at build time
///
/// This makes distributed binaries self-contained while allowing operators on
/// private testnets to override addresses via `DEPLOYMENT_DIR`.
pub fn load_deployment_json(category: &str, chain_id: u64, env: &str) -> Result<String, std::io::Error> {
    if let Ok(dir) = std::env::var("DEPLOYMENT_DIR") {
        let path = format!("{}/{}/{}-{}.json", dir, category, chain_id, env);
        info!(path = %path, "loading deployment from DEPLOYMENT_DIR override");
        return std::fs::read_to_string(&path).map_err(|e| {
            error!(path = %path, err = %e, "DEPLOYMENT_DIR set but file unreadable");
            std::io::Error::new(e.kind(), format!("DEPLOYMENT_DIR={dir} set but {path} unreadable: {e}"))
        });
    }

    get_embedded_deployment(category, chain_id, env)
        .map(|s| s.to_string())
        .ok_or_else(|| {
            error!(
                category,
                chain_id, env, "no embedded deployment found — set DEPLOYMENT_DIR for custom chains"
            );
            std::io::Error::other(format!(
                "No embedded deployment found for {}/{}-{}.json. \
                 Set DEPLOYMENT_DIR to a directory containing deployment JSONs for custom chains.",
                category, chain_id, env
            ))
        })
}

/// Grouped configuration for all contracts used by the Newton AVS on source chains.
#[derive(Debug, Clone, Default)]
pub struct ContractsConfig {
    /// Newton AVS contracts deployed for the prover.
    pub avs: NewtonAvsContractsConfig,
    /// EigenLayer protocol contracts the AVS integrates with.
    pub eigenlayer: EigenlayerContractsConfig,
    /// Newton policy-related contracts.
    pub policy: NewtonPolicyContractsConfig,
    /// Destination chain multichain contracts (operator table updater, certificate verifiers).
    /// Only populated when running in multichain mode with `with_destination_chain_id()`.
    pub destination_multichain: Option<DestinationChainMultichainContracts>,
}

impl ContractsConfig {
    /// Load all contract configurations for the given `chain_id` and deployment `env`.
    pub fn load(chain_id: u64, env: String) -> Result<Self, std::io::Error> {
        info!(
            "Loading contracts config for chain ID {} and environment {}",
            chain_id, env
        );

        Ok(Self {
            avs: NewtonAvsContractsConfig::load_src(chain_id, &env)?,
            eigenlayer: EigenlayerContractsConfig::load(chain_id, &env)?,
            policy: NewtonPolicyContractsConfig::load(chain_id, &env)?,
            destination_multichain: None,
        })
    }

    /// With the loaded source chain contracts, load additionally contracts on `chain_id` as a
    /// destination chain.
    ///
    /// Source-only fields (EigenLayer registries, AVS service manager) are absent from
    /// destination deployment JSON, so `merge_source_for_dest` copies them from the
    /// loaded source config. `validate_for_role` then asserts every required field is
    /// populated before the config leaves this method.
    pub fn with_destination_chain_id(self, chain_id: u64, env: &str) -> Result<Self, std::io::Error> {
        let eigenlayer = self.eigenlayer;
        let policy = self.policy;

        let mut avs = NewtonAvsContractsConfig::load_dest(chain_id, env)?;
        avs.merge_source_for_dest(&self.avs);
        avs.validate_for_role(ChainRole::Destination)?;

        let dest_multichain = DestinationChainMultichainContracts::load(chain_id, env)?;

        Ok(Self {
            avs,
            eigenlayer,
            policy,
            destination_multichain: Some(dest_multichain),
        })
    }

    /// Get the destination chain multichain contracts.
    ///
    /// Returns an error if not in multichain mode (destination_multichain is None).
    /// Use this to access operator_table_updater, bn254_certificate_verifier, and
    /// ecdsa_certificate_verifier on destination chains.
    pub fn destination_multichain(&self) -> Result<&DestinationChainMultichainContracts, std::io::Error> {
        self.destination_multichain.as_ref().ok_or_else(|| {
            std::io::Error::other(
                "destination_multichain_contracts not configured - use with_destination_chain_id() for multichain mode",
            )
        })
    }
}

/// The role this AVS plays for a given chain.
///
/// Source chains are home to EigenLayer middleware (BLS APK / index / stake registries,
/// strategy, operator state retriever) and the AVS service manager. Destination chains
/// receive cross-chain certificates and reuse the source's operator state — those fields
/// are populated via `merge_source_for_dest` because they are absent from destination
/// deployment JSON.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ChainRole {
    /// Source chain — owns EigenLayer registries, BLS signing, AVS service manager.
    Source,
    /// Destination chain — receives cross-chain task routing and certificate verification.
    Destination,
}

/// Newton AVS contract addresses for a single chain.
///
/// Fields typed `Address` are deployed on **both** source and destination chains
/// (overlap of the two deployment JSONs). Fields typed `Option<Address>` are deployed
/// on the source chain only and populated on destination configs by
/// `merge_source_for_dest` before `validate_for_role` asserts the invariant.
#[derive(Debug, Clone, Default)]
pub struct NewtonAvsContractsConfig {
    /// Address of the Newton prover service manager contract. Source chain only.
    /// `None` on destination configs until populated by `merge_source_for_dest`.
    pub newton_prover_service_manager: Option<Address>,
    /// Address of the Newton prover task manager contract.
    pub newton_prover_task_manager: Address,
    /// Address of the challenge verifier contract.
    pub challenge_verifier: Address,
    /// Address of the Rego verifier contract.
    pub rego_verifier: Address,
    /// Address of the attestation validator contract.
    pub attestation_validator: Address,
    /// Address of the operator registry contract.
    pub operator_registry: Address,
    /// Address of the operator state retriever contract. Source chain only.
    /// `None` on destination configs until populated by `merge_source_for_dest`.
    pub operator_state_retriever: Option<Address>,
    /// Address of the BLS aggregate public key registry contract. Source chain only.
    /// `None` on destination configs until populated by `merge_source_for_dest`.
    pub bls_apk_registry: Option<Address>,
    /// Address of the index registry contract. Source chain only.
    /// `None` on destination configs until populated by `merge_source_for_dest`.
    pub index_registry: Option<Address>,
    /// Address of the stake registry contract. Source chain only.
    /// `None` on destination configs until populated by `merge_source_for_dest`.
    pub stake_registry: Option<Address>,
    /// Address of the socket registry contract.
    pub socket_registry: Address,
    /// Address of the strategy contract. Source chain only.
    /// `None` on destination configs until populated by `merge_source_for_dest`.
    pub strategy: Option<Address>,
    /// Address of the identity registry contract.
    pub identity_registry: Address,
    /// Address of the ConfidentialDataRegistry contract for provider-managed datasets.
    pub confidential_data_registry: Address,
    /// Address of the BatchTaskManager contract for high-throughput batched submission.
    pub batch_task_manager: Address,
    /// Address of the EpochRegistry contract for privacy epoch lifecycle. Source chain only.
    /// `None` on destination configs (privacy services are source-side only).
    pub epoch_registry: Option<Address>,
    /// Address of the EnclaveVersionRegistry for PCR0 whitelist governance.
    /// None when not deployed (pre-TEE chains).
    pub enclave_version_registry: Option<Address>,
    /// Address of the `StateCommitRegistry` contract for per-chain JMT state-root
    /// commits (Phase 1 PDS). None when not yet deployed on the chain;
    /// `commit_state_root` runtime-guards against zero so callers degrade gracefully.
    pub state_commit_registry: Option<Address>,
    /// Address of the PolicyClientRegistry contract — on-chain directory of approved
    /// policy clients. Deployed on both source and destination chains.
    pub policy_client_registry: Address,
}

// TODO: Migrate address loading from deployment JSON to the on-chain
// `NewtonAddressesProvider` (Aave V3 `PoolAddressesProvider` pattern).
// On service startup, batch-read every well-known address from the
// provider via `getAddress(bytes32)` and populate this struct once —
// preserves today's no-runtime-RPC property (single startup read,
// in-memory thereafter) while making the on-chain registry the single
// source of truth across Solidity and Rust. The Solidity layer already
// deploys and consumes the provider via `INewtonAddressesProvider`; the
// alloy binding lives at `crates/core/src/generated/newton_addresses_provider.rs`.
// Post-migration, `load_src` / `load_dest` can be deleted and the
// `Option<Address>` fields collapsed back to plain `Address`.
impl NewtonAvsContractsConfig {
    /// Materializes `state_commit_registry` as an `Address`. Returns `Address::ZERO`
    /// when the registry is absent from the deployment JSON for this chain
    /// (`stateCommitRegistry` missing or set to `0x0`). `commit_state_root`
    /// runtime-guards against the zero address so callers degrade gracefully on
    /// chains without PDS wiring.
    pub fn state_commit_registry_or_zero(&self) -> Address {
        self.state_commit_registry.unwrap_or(Address::ZERO)
    }

    /// Validate that every field required for the given `role` is populated.
    ///
    /// Source-only fields are typed `Option<Address>` so destination configs can be
    /// built without their source-side EigenLayer counterparts present in JSON.
    /// `merge_source_for_dest` is what populates them on destination configs;
    /// this method asserts the invariant after construction.
    ///
    /// `epoch_registry` is intentionally excluded — privacy services are feature-gated
    /// and may legitimately be absent on a source chain that does not run privacy.
    /// `enclave_version_registry` and `state_commit_registry` are similarly
    /// pre-deployment-tolerant infrastructure fields.
    pub fn validate_for_role(&self, role: ChainRole) -> Result<(), std::io::Error> {
        let missing = |field: &str| {
            std::io::Error::other(format!(
                "{:?} role requires {} but it is None — check deployment JSON or merge_source_for_dest",
                role, field
            ))
        };
        if self.newton_prover_service_manager.is_none() {
            return Err(missing("newton_prover_service_manager"));
        }
        if self.bls_apk_registry.is_none() {
            return Err(missing("bls_apk_registry"));
        }
        if self.index_registry.is_none() {
            return Err(missing("index_registry"));
        }
        if self.stake_registry.is_none() {
            return Err(missing("stake_registry"));
        }
        if self.operator_state_retriever.is_none() {
            return Err(missing("operator_state_retriever"));
        }
        if self.strategy.is_none() {
            return Err(missing("strategy"));
        }
        Ok(())
    }

    /// Copy source-only fields from `src` into `self`. Used to build a destination
    /// config that retains source-side EigenLayer state (BLS APK, index, stake,
    /// strategy, operator state retriever) plus the AVS service manager — these
    /// are absent from destination deployment JSON by design.
    ///
    /// `epoch_registry` is NOT merged — privacy services run on the source chain
    /// only, and destination configs should keep `None` for that field.
    fn merge_source_for_dest(&mut self, src: &Self) {
        self.newton_prover_service_manager = src.newton_prover_service_manager;
        self.bls_apk_registry = src.bls_apk_registry;
        self.index_registry = src.index_registry;
        self.stake_registry = src.stake_registry;
        self.operator_state_retriever = src.operator_state_retriever;
        self.strategy = src.strategy;
        // TODO: socket_registry is deployed on both source and destination chains
        // with distinct addresses, but historical behavior overrides the dest
        // address with the source address. Investigate whether destination chains
        // should use their own socket_registry from load_dest() instead — kept
        // here to preserve pre-refactor behavior.
        self.socket_registry = src.socket_registry;
    }

    /// Load Newton AVS contract addresses for source chain from the deployment JSON
    pub fn load_src(chain_id: u64, env: &str) -> Result<Self, std::io::Error> {
        let cfg = Self::load(chain_id, false, env)?;
        cfg.validate_for_role(ChainRole::Source)?;
        Ok(cfg)
    }

    /// Load Newton AVS contract addresses for destination chain from the deployment JSON
    pub fn load_dest(chain_id: u64, env: &str) -> Result<Self, std::io::Error> {
        Self::load(chain_id, true, env)
    }

    /// Load Newton AVS contract addresses from the deployment JSON for the given `chain_id` and `env`.
    ///
    /// Overlap fields (deployed on both source and destination chains) parse via the
    /// `required` closure and surface a typed error if a JSON key is missing or
    /// malformed. Source-only fields parse via `optional`, returning `None` when the
    /// JSON key is absent (destination JSONs) or set to the zero address.
    fn load(chain_id: u64, is_destination_chain: bool, env: &str) -> Result<Self, std::io::Error> {
        let category = if !is_destination_chain {
            "newton-prover"
        } else {
            "newton-cross-chain"
        };

        info!(
            "[NewtonAvsContractsConfig] Loading deployment for {}/{}-{}",
            category, chain_id, env
        );
        let raw = load_deployment_json(category, chain_id, env)?;
        let json: serde_json::Value = serde_json::from_str(&raw)?;
        let addrs = &json["addresses"];

        let required = |key: &str| -> Result<Address, std::io::Error> {
            let s = addrs[key]
                .as_str()
                .ok_or_else(|| std::io::Error::other(format!("{} address not found", key)))?;
            Address::from_str(s).map_err(|e| std::io::Error::other(format!("Failed to parse {} address: {}", key, e)))
        };

        let optional = |key: &str| -> Option<Address> {
            addrs[key]
                .as_str()
                .and_then(|s| Address::from_str(s).ok())
                .filter(|addr| *addr != Address::ZERO)
        };

        Ok(Self {
            // Overlap fields — required on every chain.
            newton_prover_task_manager: required("newtonProverTaskManager")?,
            challenge_verifier: required("challengeVerifier")?,
            rego_verifier: required("regoVerifier")?,
            attestation_validator: required("attestationValidator")?,
            operator_registry: required("operatorRegistry")?,
            socket_registry: required("socketRegistry")?,
            identity_registry: required("identityRegistry")?,
            confidential_data_registry: required("confidentialDataRegistry")?,
            batch_task_manager: required("batchTaskManager")?,
            policy_client_registry: required("policyClientRegistry")?,

            // Source-only fields — `validate_for_role(Source)` enforces presence on
            // source chains; `merge_source_for_dest` populates these on destination
            // configs from the loaded source config before validation runs.
            newton_prover_service_manager: optional("newtonProverServiceManager"),
            operator_state_retriever: optional("operatorStateRetriever"),
            bls_apk_registry: optional("blsApkRegistry"),
            index_registry: optional("indexRegistry"),
            stake_registry: optional("stakeRegistry"),
            strategy: optional("strategy"),

            // Optional infrastructure — None when not yet deployed on the chain.
            epoch_registry: optional("epochRegistry"),
            enclave_version_registry: optional("enclaveVersionRegistry"),
            state_commit_registry: optional("stateCommitRegistry"),
        })
    }
}

/// Contract addresses for Newton policy-related components.
#[derive(Debug, Clone, Default)]
pub struct NewtonPolicyContractsConfig {
    /// Address of the policy factory contract.
    pub policy_factory: Address,
    /// Address of the policy data factory contract.
    pub policy_data_factory: Address,
}

impl NewtonPolicyContractsConfig {
    /// Load Newton policy contract addresses from the deployment JSON for the given `chain_id` and `env`.
    pub fn load(chain_id: u64, env: &str) -> Result<Self, std::io::Error> {
        info!(
            "[NewtonPolicyContractsConfig] Loading deployment for policy/{}-{}",
            chain_id, env
        );
        let raw = load_deployment_json("policy", chain_id, env)?;
        let newton_policy_deployment_data: serde_json::Value = serde_json::from_str(&raw)?;
        Ok(Self {
            policy_factory: Address::from_str(
                newton_policy_deployment_data["addresses"]["policyFactory"]
                    .as_str()
                    .ok_or_else(|| std::io::Error::other("policyFactory address not found"))?,
            )
            .map_err(|e| std::io::Error::other(format!("Failed to parse policy factory address: {}", e)))?,
            policy_data_factory: Address::from_str(
                newton_policy_deployment_data["addresses"]["policyDataFactory"]
                    .as_str()
                    .ok_or_else(|| std::io::Error::other("policyDataFactory address not found"))?,
            )
            .map_err(|e| std::io::Error::other(format!("Failed to parse policy data factory address: {}", e)))?,
        })
    }
}

/// Contract addresses for EigenLayer protocol integration.
///
/// Note: Destination chain multichain contracts (operator_table_updater, certificate verifiers)
/// are now accessed via `ContractsConfig::destination_multichain()`.
#[derive(Debug, Clone, Default)]
pub struct EigenlayerContractsConfig {
    /// Address of the delegation manager contract.
    pub delegation_manager: Address,
    /// Address of the AVS directory contract.
    pub avs_directory: Address,
    /// Address of the strategy manager contract.
    pub strategy_manager: Address,
    /// Address of the allocation manager contract.
    pub allocation_manager: Address,
    /// Address of the rewards coordinator contract.
    pub rewards_coordinator: Address,
    /// Address of the strategy factory contract.
    pub strategy_factory: Address,
    /// Address of the permission controller contract.
    pub permission_controller: Address,
    /// Address of the BN254 certificate verifier contract (source chain).
    pub bn254_certificate_verifier: Address,
    /// Address of the key registrar contract for cross-chain operator key management.
    pub key_registrar: Address,
}

impl EigenlayerContractsConfig {
    /// Load EigenLayer contract addresses from the deployment JSON for the given `chain_id` and `env`.
    pub fn load(chain_id: u64, env: &str) -> Result<Self, std::io::Error> {
        info!(
            "[EigenlayerContractsConfig] Loading deployment for core/{}-{}",
            chain_id, env
        );
        let raw = load_deployment_json("core", chain_id, env)?;
        let eigenlayer_deployment_data: serde_json::Value = serde_json::from_str(&raw)?;
        Ok(Self {
            delegation_manager: Address::from_str(
                eigenlayer_deployment_data["addresses"]["delegation"]
                    .as_str()
                    .unwrap_or_default(),
            )
            .unwrap_or_default(),
            avs_directory: Address::from_str(
                eigenlayer_deployment_data["addresses"]["avsDirectory"]
                    .as_str()
                    .unwrap_or_default(),
            )
            .unwrap_or_default(),
            strategy_manager: Address::from_str(
                eigenlayer_deployment_data["addresses"]["strategyManager"]
                    .as_str()
                    .unwrap_or_default(),
            )
            .unwrap_or_default(),
            allocation_manager: Address::from_str(
                eigenlayer_deployment_data["addresses"]["allocationManager"]
                    .as_str()
                    .unwrap_or_default(),
            )
            .unwrap_or_default(),
            rewards_coordinator: Address::from_str(
                eigenlayer_deployment_data["addresses"]["rewardsCoordinator"]
                    .as_str()
                    .unwrap_or_default(),
            )
            .unwrap_or_default(),
            strategy_factory: Address::from_str(
                eigenlayer_deployment_data["addresses"]["strategyFactory"]
                    .as_str()
                    .unwrap_or_default(),
            )
            .unwrap_or_default(),
            permission_controller: Address::from_str(
                eigenlayer_deployment_data["addresses"]["permissionController"]
                    .as_str()
                    .unwrap_or_default(),
            )
            .unwrap_or_default(),
            bn254_certificate_verifier: Address::from_str(
                eigenlayer_deployment_data["addresses"]["bn254CertificateVerifier"]
                    .as_str()
                    .unwrap_or_default(),
            )
            .unwrap_or_default(),
            key_registrar: Address::from_str(
                eigenlayer_deployment_data["addresses"]["keyRegistrar"]
                    .as_str()
                    .unwrap_or_default(),
            )
            .unwrap_or_default(),
        })
    }
}

/// Return the current deployment environment from `DEPLOYMENT_ENV`, defaulting to `"stagef"`.
pub fn get_deployment_env() -> String {
    std::env::var("DEPLOYMENT_ENV").unwrap_or_else(|_| "stagef".to_string())
}

/// Contract addresses for destination chain multichain contracts.
///
/// These contracts are deployed on L2 destination chains to enable cross-chain
/// operator table verification. They are loaded from the newton-cross-chain
/// deployment JSON for the destination chain.
#[derive(Debug, Clone)]
pub struct DestinationChainMultichainContracts {
    /// Address of the ECDSA operator table updater contract.
    pub operator_table_updater: Address,
    /// Address of the BN254 certificate verifier contract.
    pub bn254_certificate_verifier: Address,
    /// Address of the ECDSA certificate verifier contract.
    pub ecdsa_certificate_verifier: Address,
}

impl DestinationChainMultichainContracts {
    /// Load destination chain multichain contracts from the newton-cross-chain deployment JSON.
    pub fn load(chain_id: u64, env: &str) -> Result<Self, std::io::Error> {
        info!(
            "[DestinationChainMultichainContracts] Loading deployment for newton-cross-chain/{}-{}",
            chain_id, env
        );
        let raw = load_deployment_json("newton-cross-chain", chain_id, env)?;
        let deployment_data: serde_json::Value = serde_json::from_str(&raw)?;

        Ok(Self {
            operator_table_updater: Address::from_str(
                deployment_data["addresses"]["operatorTableUpdater"]
                    .as_str()
                    .ok_or_else(|| std::io::Error::other("operatorTableUpdater address not found"))?,
            )
            .map_err(|e| std::io::Error::other(format!("Failed to parse operator table updater address: {}", e)))?,
            bn254_certificate_verifier: Address::from_str(
                deployment_data["addresses"]["bn254CertificateVerifier"]
                    .as_str()
                    .ok_or_else(|| std::io::Error::other("bn254CertificateVerifier address not found"))?,
            )
            .map_err(|e| std::io::Error::other(format!("Failed to parse bn254 certificate verifier address: {}", e)))?,
            ecdsa_certificate_verifier: Address::from_str(
                deployment_data["addresses"]["ecdsaCertificateVerifier"]
                    .as_str()
                    .ok_or_else(|| std::io::Error::other("ecdsaCertificateVerifier address not found"))?,
            )
            .map_err(|e| std::io::Error::other(format!("Failed to parse ecdsa certificate verifier address: {}", e)))?,
        })
    }
}

/// Contract addresses for Newton cross-chain infrastructure on source chains.
///
/// These contracts are deployed on Ethereum mainnet/Sepolia to enable cross-chain transport
/// of operator tables to L2 destinations (e.g., Arbitrum, Optimism, Polygon, Base).
#[derive(Debug, Clone)]
pub struct SourceChainMultichainContracts {
    /// Address of the NewtonCrossChainRegistry contract.
    /// Manages operator set registrations and configurations for cross-chain transport.
    pub cross_chain_registry: Address,
    /// Address of the BN254TableCalculator contract.
    /// Computes BN254 operator table info with owner-based access control.
    pub table_calculator: Address,
}

impl SourceChainMultichainContracts {
    /// Load Newton cross-chain contract addresses from the deployment JSON for the given `chain_id` and `env`.
    pub fn load(chain_id: u64, env: &str) -> Result<Self, std::io::Error> {
        info!(
            "[SourceChainMultichainContracts] Loading deployment for newton-cross-chain/{}-{}",
            chain_id, env
        );
        let raw = load_deployment_json("newton-cross-chain", chain_id, env)?;
        let deployment_data: serde_json::Value = serde_json::from_str(&raw)?;

        Ok(Self {
            cross_chain_registry: Address::from_str(
                deployment_data["addresses"]["crossChainRegistry"]
                    .as_str()
                    .ok_or_else(|| std::io::Error::other("crossChainRegistry address not found"))?,
            )
            .map_err(|e| std::io::Error::other(format!("Failed to parse cross chain registry address: {}", e)))?,
            table_calculator: Address::from_str(
                deployment_data["addresses"]["operatorTableCalculator"]
                    .as_str()
                    .ok_or_else(|| std::io::Error::other("operatorTableCalculator address not found"))?,
            )
            .map_err(|e| std::io::Error::other(format!("Failed to parse table calculator address: {}", e)))?,
        })
    }

    /// Try to load Newton cross-chain contract addresses, returning None if the deployment file doesn't exist.
    /// This is useful for chains that may or may not have cross-chain contracts deployed.
    pub fn try_load(chain_id: u64, env: &str) -> Option<Self> {
        Self::load(chain_id, env).ok()
    }
}