newton-core 0.4.16

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

use crate::{
    common::chain::{
        is_supported_chain, ARBITRUM_ONE, ARBITRUM_SEPOLIA, BASE_MAINNET, BASE_SEPOLIA, ETHEREUM_MAINNET, LOCAL_ANVIL,
        LOCAL_ANVIL_DESTINATION, OPTIMISM_MAINNET, OPTIMISM_SEPOLIA, POLYGON_AMOY, POLYGON_MAINNET, SEPOLIA,
    },
    config::error::ConfigError,
};

// =============================================================================
// Chain ID Constants (kept for backward compatibility)
// =============================================================================

/// Local Anvil chain id (source chain for multichain testing).
pub const ANVIL_CHAIN_ID: u64 = LOCAL_ANVIL;

/// Local Anvil destination chain id (for dual-anvil multichain testing).
pub const ANVIL_DESTINATION_CHAIN_ID: u64 = LOCAL_ANVIL_DESTINATION;

/// Sepolia chain id.
pub const SEPOLIA_CHAIN_ID: u64 = SEPOLIA;

/// Ethereum mainnet chain id.
pub const MAINNET_CHAIN_ID: u64 = ETHEREUM_MAINNET;

// =============================================================================
// Default RPC URLs - Local Anvil
// =============================================================================

/// Local Anvil RPC HTTP URL (source chain, port 8545).
pub const ANVIL_HTTP_URL: &str = "http://127.0.0.1:8545";

/// Local Anvil RPC WebSocket URL (source chain, port 8545).
pub const ANVIL_WS_URL: &str = "ws://127.0.0.1:8545";

/// Local Anvil destination chain RPC HTTP URL (port 8546).
pub const ANVIL_DESTINATION_HTTP_URL: &str = "http://127.0.0.1:8546";

/// Local Anvil destination chain RPC WebSocket URL (port 8546).
pub const ANVIL_DESTINATION_WS_URL: &str = "ws://127.0.0.1:8546";

// =============================================================================
// Default RPC URLs - Ethereum
// =============================================================================

/// Ethereum Mainnet default RPC HTTP URL.
pub const MAINNET_HTTP_URL: &str = "https://mainnet.gateway.tenderly.co";

/// Ethereum Mainnet default RPC WebSocket URL.
pub const MAINNET_WS_URL: &str = "wss://mainnet.gateway.tenderly.co";

/// Sepolia default RPC HTTP URL.
pub const SEPOLIA_HTTP_URL: &str = "https://sepolia.gateway.tenderly.co";

/// Sepolia default RPC WebSocket URL.
pub const SEPOLIA_WS_URL: &str = "wss://sepolia.gateway.tenderly.co";

// =============================================================================
// Default RPC URLs - Base
// =============================================================================

/// Base Mainnet default RPC HTTP URL.
pub const BASE_MAINNET_HTTP_URL: &str = "https://base.gateway.tenderly.co";

/// Base Mainnet default RPC WebSocket URL.
pub const BASE_MAINNET_WS_URL: &str = "wss://base.gateway.tenderly.co";

/// Base Sepolia default RPC HTTP URL.
pub const BASE_SEPOLIA_HTTP_URL: &str = "https://base-sepolia.gateway.tenderly.co";

/// Base Sepolia default RPC WebSocket URL.
pub const BASE_SEPOLIA_WS_URL: &str = "wss://base-sepolia.gateway.tenderly.co";

// =============================================================================
// Default RPC URLs - Optimism
// =============================================================================

/// Optimism Mainnet default RPC HTTP URL.
pub const OPTIMISM_MAINNET_HTTP_URL: &str = "https://optimism.gateway.tenderly.co";

/// Optimism Mainnet default RPC WebSocket URL.
pub const OPTIMISM_MAINNET_WS_URL: &str = "wss://optimism.gateway.tenderly.co";

/// Optimism Sepolia default RPC HTTP URL.
pub const OPTIMISM_SEPOLIA_HTTP_URL: &str = "https://optimism-sepolia.gateway.tenderly.co";

/// Optimism Sepolia default RPC WebSocket URL.
pub const OPTIMISM_SEPOLIA_WS_URL: &str = "wss://optimism-sepolia.gateway.tenderly.co";

// =============================================================================
// Default RPC URLs - Arbitrum
// =============================================================================

/// Arbitrum One default RPC HTTP URL.
pub const ARBITRUM_ONE_HTTP_URL: &str = "https://arbitrum.gateway.tenderly.co";

/// Arbitrum One default RPC WebSocket URL.
pub const ARBITRUM_ONE_WS_URL: &str = "wss://arbitrum.gateway.tenderly.co";

/// Arbitrum Sepolia default RPC HTTP URL.
pub const ARBITRUM_SEPOLIA_HTTP_URL: &str = "https://arbitrum-sepolia.gateway.tenderly.co";

/// Arbitrum Sepolia default RPC WebSocket URL.
pub const ARBITRUM_SEPOLIA_WS_URL: &str = "wss://arbitrum-sepolia.gateway.tenderly.co";

// =============================================================================
// Default RPC URLs - Polygon
// =============================================================================

/// Polygon Mainnet default RPC HTTP URL.
pub const POLYGON_MAINNET_HTTP_URL: &str = "https://polygon.gateway.tenderly.co";

/// Polygon Mainnet default RPC WebSocket URL.
pub const POLYGON_MAINNET_WS_URL: &str = "wss://polygon.gateway.tenderly.co";

/// Polygon Amoy (testnet) default RPC HTTP URL.
pub const POLYGON_AMOY_HTTP_URL: &str = "https://polygon-amoy.gateway.tenderly.co";

/// Polygon Amoy (testnet) default RPC WebSocket URL.
pub const POLYGON_AMOY_WS_URL: &str = "wss://polygon-amoy.gateway.tenderly.co";

// =============================================================================
// Websocket Constants
// =============================================================================

/// If websocket is dropped this is how many times to retry connection before
/// erroring out
pub const WEBSOCKET_DROPPED_MAX_RETRY: u64 = 6;

/// If websocket is dropped this is how many seconds to wait before retrying
/// again up the WEBSOCKET_DROPPED_MAX_RETRY
pub const WEBSOCKET_DROPPED_RETRY_DELAY: u64 = 10;

// =============================================================================
// WebSocket Reconnection Loop
// =============================================================================

/// Callback hooks for WebSocket reconnection lifecycle events.
///
/// Implement this trait to integrate metrics or custom logging into the
/// generic reconnection loop. All methods have default no-op implementations.
#[cfg(any(feature = "rpc", feature = "database", feature = "websocket"))]
pub trait WsReconnectCallbacks: Send + 'static {
    /// Called when the WebSocket stream ends or encounters an error.
    fn on_disconnected(&self) {}
    /// Called on each reconnection attempt (both graceful and error).
    fn on_reconnect_attempt(&self) {}
}

/// No-op implementation for listeners that don't need metric hooks.
#[derive(Debug)]
#[cfg(any(feature = "rpc", feature = "database", feature = "websocket"))]
pub struct NoOpWsCallbacks;
#[cfg(any(feature = "rpc", feature = "database", feature = "websocket"))]
impl WsReconnectCallbacks for NoOpWsCallbacks {}

/// Run a WebSocket listener in a reconnection loop with retry.
///
/// `listen_fn` is called repeatedly. Each invocation should connect to the
/// WebSocket, subscribe to events, and process them until the stream ends
/// (returns `Ok(())`) or an error occurs (returns `Err`).
///
/// On graceful stream end (`Ok`), the retry counter resets and the loop
/// reconnects immediately. On error, the loop retries up to `max_retries`
/// times with a `retry_delay` sleep between attempts. When retries are
/// exhausted, the loop exits.
///
/// Returns `true` if retries were exhausted (fatal), `false` if the loop
/// was terminated by other means (shouldn't happen with current callers).
#[cfg(any(feature = "rpc", feature = "database", feature = "websocket"))]
pub async fn ws_reconnect_loop<F, Fut, C>(
    label: &str,
    max_retries: u64,
    retry_delay: std::time::Duration,
    mut listen_fn: F,
    callbacks: C,
) -> bool
where
    F: FnMut() -> Fut,
    Fut: std::future::Future<Output = eyre::Result<()>>,
    C: WsReconnectCallbacks,
{
    let mut retry_count: u64 = 0;
    loop {
        match listen_fn().await {
            Ok(()) => {
                // Stream ended gracefully — reset retry counter and reconnect
                callbacks.on_disconnected();
                callbacks.on_reconnect_attempt();
                retry_count = 0;
                tracing::info!("{label} ws stream ended, reconnecting");
            }
            Err(e) => {
                callbacks.on_disconnected();
                if retry_count >= max_retries {
                    tracing::error!("{label} websocket failed to reconnect after {retry_count} retries: {e}");
                    return true;
                }
                callbacks.on_reconnect_attempt();
                tracing::error!("{label} websocket error: {e}, retrying");
                tokio::time::sleep(retry_delay).await;
                retry_count += 1;
            }
        }
    }
}

// =============================================================================
// Utility Functions
// =============================================================================

/// Check if the chain is anvil (either source 31337 or destination 31338)
pub fn is_anvil() -> bool {
    match std::env::var("CHAIN_ID") {
        Ok(chain_id_str) => {
            let chain_id = chain_id_str.parse::<u64>().unwrap_or(0);
            chain_id == ANVIL_CHAIN_ID || chain_id == ANVIL_DESTINATION_CHAIN_ID
        }
        Err(_) => true,
    }
}

// =============================================================================
// RpcProviderConfig - Single Chain Configuration
// =============================================================================

/// RPC provider configuration containing HTTP and WebSocket endpoints for a single chain.
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
pub struct RpcProviderConfig {
    /// HTTP RPC endpoint URL.
    pub http: String,
    /// WebSocket RPC endpoint URL.
    pub ws: String,
}

impl RpcProviderConfig {
    /// Create a new RpcProviderConfig with the given HTTP and WebSocket URLs.
    pub fn new(http: String, ws: String) -> Self {
        Self { http, ws }
    }

    /// Load an `RpcProviderConfig` for the given `chain_id`, optionally overridden by environment variables.
    /// This is the legacy single-chain loading method. Prefer `ChainRpcProviderConfig::load()` for multichain.
    pub fn load(chain_id: u64) -> Result<Self, std::io::Error> {
        let default_config = get_default_rpc_config(chain_id).ok_or_else(|| {
            std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                format!("Unsupported chain id: {}", chain_id),
            )
        })?;

        Ok(Self {
            http: get_env_or_default(
                &format!("RPC_HTTP_URL_{}", chain_id),
                "RPC_HTTP_URL",
                &default_config.http,
            ),
            ws: get_env_or_default(&format!("RPC_WS_URL_{}", chain_id), "RPC_WS_URL", &default_config.ws),
        })
    }
}

// =============================================================================
// ChainRpcProviderConfig - Multi-Chain Configuration
// =============================================================================

/// RPC provider configuration for multiple chains.
///
/// This struct maintains a mapping from chain ID to RPC configuration,
/// pre-populated with defaults for all supported chains. Environment variables
/// can override specific chains using the pattern `RPC_HTTP_URL_{CHAIN_ID}`.
#[derive(Debug, Clone)]
pub struct ChainRpcProviderConfig {
    configs: HashMap<u64, RpcProviderConfig>,
}

impl ChainRpcProviderConfig {
    /// Load RPC configurations for all supported chains with defaults and env var overrides.
    ///
    /// Environment variable pattern:
    /// - Chain-specific: `RPC_HTTP_URL_{CHAIN_ID}` and `RPC_WS_URL_{CHAIN_ID}` (e.g., `RPC_HTTP_URL_84532`)
    /// - Legacy global: `RPC_HTTP_URL` and `RPC_WS_URL` (applies only if chain-specific not set)
    pub fn load() -> Self {
        let mut configs = HashMap::new();

        let supported_chains = [
            LOCAL_ANVIL,
            LOCAL_ANVIL_DESTINATION,
            ETHEREUM_MAINNET,
            SEPOLIA,
            BASE_MAINNET,
            BASE_SEPOLIA,
            OPTIMISM_MAINNET,
            OPTIMISM_SEPOLIA,
            ARBITRUM_ONE,
            ARBITRUM_SEPOLIA,
            POLYGON_MAINNET,
            POLYGON_AMOY,
        ];

        for chain_id in supported_chains {
            if let Some(default_config) = get_default_rpc_config(chain_id) {
                let http = get_env_or_default(
                    &format!("RPC_HTTP_URL_{}", chain_id),
                    "RPC_HTTP_URL",
                    &default_config.http,
                );
                let ws = get_env_or_default(&format!("RPC_WS_URL_{}", chain_id), "RPC_WS_URL", &default_config.ws);
                configs.insert(chain_id, RpcProviderConfig { http, ws });
            }
        }

        Self { configs }
    }

    /// Get the RPC configuration for a specific chain ID.
    pub fn get(&self, chain_id: u64) -> Option<&RpcProviderConfig> {
        self.configs.get(&chain_id)
    }

    /// Get the RPC configuration for a specific chain ID, returning an error if not found.
    pub fn get_or_err(&self, chain_id: u64) -> Result<&RpcProviderConfig, ConfigError> {
        self.configs
            .get(&chain_id)
            .ok_or(ConfigError::MissingRpcForChain { chain_id })
    }

    /// Insert or update an RPC configuration for a specific chain ID.
    ///
    /// Returns an error if the chain ID is not supported.
    pub fn insert(&mut self, chain_id: u64, config: RpcProviderConfig) -> Result<(), std::io::Error> {
        if !is_supported_chain(chain_id) {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                format!("Chain ID {} is not supported", chain_id),
            ));
        }
        self.configs.insert(chain_id, config);
        Ok(())
    }

    /// Check if a configuration exists for the given chain ID.
    pub fn contains(&self, chain_id: u64) -> bool {
        self.configs.contains_key(&chain_id)
    }

    /// Get all configured chain IDs.
    pub fn chain_ids(&self) -> impl Iterator<Item = &u64> {
        self.configs.keys()
    }
}

// =============================================================================
// Helper Functions
// =============================================================================

/// Get the default RPC configuration for a chain ID.
fn get_default_rpc_config(chain_id: u64) -> Option<RpcProviderConfig> {
    match chain_id {
        LOCAL_ANVIL => Some(RpcProviderConfig {
            http: ANVIL_HTTP_URL.to_string(),
            ws: ANVIL_WS_URL.to_string(),
        }),
        LOCAL_ANVIL_DESTINATION => Some(RpcProviderConfig {
            http: ANVIL_DESTINATION_HTTP_URL.to_string(),
            ws: ANVIL_DESTINATION_WS_URL.to_string(),
        }),
        ETHEREUM_MAINNET => Some(RpcProviderConfig {
            http: MAINNET_HTTP_URL.to_string(),
            ws: MAINNET_WS_URL.to_string(),
        }),
        SEPOLIA => Some(RpcProviderConfig {
            http: SEPOLIA_HTTP_URL.to_string(),
            ws: SEPOLIA_WS_URL.to_string(),
        }),
        BASE_MAINNET => Some(RpcProviderConfig {
            http: BASE_MAINNET_HTTP_URL.to_string(),
            ws: BASE_MAINNET_WS_URL.to_string(),
        }),
        BASE_SEPOLIA => Some(RpcProviderConfig {
            http: BASE_SEPOLIA_HTTP_URL.to_string(),
            ws: BASE_SEPOLIA_WS_URL.to_string(),
        }),
        OPTIMISM_MAINNET => Some(RpcProviderConfig {
            http: OPTIMISM_MAINNET_HTTP_URL.to_string(),
            ws: OPTIMISM_MAINNET_WS_URL.to_string(),
        }),
        OPTIMISM_SEPOLIA => Some(RpcProviderConfig {
            http: OPTIMISM_SEPOLIA_HTTP_URL.to_string(),
            ws: OPTIMISM_SEPOLIA_WS_URL.to_string(),
        }),
        ARBITRUM_ONE => Some(RpcProviderConfig {
            http: ARBITRUM_ONE_HTTP_URL.to_string(),
            ws: ARBITRUM_ONE_WS_URL.to_string(),
        }),
        ARBITRUM_SEPOLIA => Some(RpcProviderConfig {
            http: ARBITRUM_SEPOLIA_HTTP_URL.to_string(),
            ws: ARBITRUM_SEPOLIA_WS_URL.to_string(),
        }),
        POLYGON_MAINNET => Some(RpcProviderConfig {
            http: POLYGON_MAINNET_HTTP_URL.to_string(),
            ws: POLYGON_MAINNET_WS_URL.to_string(),
        }),
        POLYGON_AMOY => Some(RpcProviderConfig {
            http: POLYGON_AMOY_HTTP_URL.to_string(),
            ws: POLYGON_AMOY_WS_URL.to_string(),
        }),
        _ => None,
    }
}

/// Get environment variable value with chain-specific and global fallback.
///
/// Priority: chain-specific env var -> global env var -> default value
fn get_env_or_default(chain_specific_key: &str, global_key: &str, default: &str) -> String {
    if let Ok(url) = std::env::var(chain_specific_key) {
        if !url.is_empty() {
            return url;
        }
    }
    if let Ok(url) = std::env::var(global_key) {
        if !url.is_empty() {
            return url;
        }
    }
    default.to_string()
}