hypercall/read_cache/

mmp.rs

use crate::shared::order_types::ParsedSymbol;
use anyhow::{Context, Result};
use hypercall_db::MmpConfigRecord as MmpConfig;
use hypercall_db::MmpConfigWriter;
use hypercall_types::Fill;
use hypercall_types::{to_contract_units_decimal, WalletAddress};
use rust_decimal::prelude::ToPrimitive;
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, info, warn};

/// Represents a single fill record with calculated metrics
#[derive(Debug, Clone)]
struct FillRecord {
    timestamp: u64,
    quantity: u64,
    delta: f64,
    vega: f64,
}

/// Per-currency fill tracking window for one wallet
#[derive(Debug, Clone)]
struct CurrencyFillWindow {
    fills: VecDeque<FillRecord>,
    cumulative_qty: u64,
    cumulative_delta: f64,
    cumulative_vega: f64,
    frozen_until: Option<u64>, // Timestamp when this currency becomes unfrozen
}

impl CurrencyFillWindow {
    fn new() -> Self {
        Self {
            fills: VecDeque::new(),
            cumulative_qty: 0,
            cumulative_delta: 0.0,
            cumulative_vega: 0.0,
            frozen_until: None,
        }
    }

    /// Add the fill to tracking and check if limits are now breached
    /// ALWAYS adds the fill first, then checks if cumulative metrics breach limits
    /// Returns true if limits are breached after adding, false otherwise
    fn eval_fill(
        &mut self,
        record: FillRecord,
        qty_limit: Option<i64>,
        delta_limit: Option<f64>,
        vega_limit: Option<f64>,
    ) -> bool {
        // ALWAYS add the fill first
        self.cumulative_qty += record.quantity;
        self.cumulative_delta += record.delta;
        self.cumulative_vega += record.vega;
        self.fills.push_back(record);

        // Then check if the new cumulative values breach any limits
        if let Some(limit) = qty_limit {
            if self.cumulative_qty > limit as u64 {
                return true; // Breach detected
            }
        }

        if let Some(limit) = delta_limit {
            if self.cumulative_delta.abs() > limit {
                return true; // Breach detected
            }
        }

        if let Some(limit) = vega_limit {
            if self.cumulative_vega.abs() > limit {
                return true; // Breach detected
            }
        }

        false // No breach
    }

    /// Remove fills older than the window interval
    fn evict_old_fills(&mut self, current_time: u64, interval_ms: u64) {
        let cutoff_time = current_time.saturating_sub(interval_ms);

        while let Some(front) = self.fills.front() {
            if front.timestamp < cutoff_time {
                // Remove from cumulative counts
                self.cumulative_qty = self.cumulative_qty.saturating_sub(front.quantity);
                self.cumulative_delta -= front.delta;
                self.cumulative_vega -= front.vega;
                self.fills.pop_front();
            } else {
                break;
            }
        }
    }

    /// Check if currently frozen
    fn is_frozen(&self, current_time: u64) -> bool {
        if let Some(frozen_until) = self.frozen_until {
            current_time < frozen_until
        } else {
            false
        }
    }

    /// Set frozen state
    fn set_frozen(&mut self, current_time: u64, frozen_duration_ms: u64) {
        self.frozen_until = Some(current_time + frozen_duration_ms);
    }

    /// Clear all fills and unfreeze
    fn reset(&mut self) {
        self.fills.clear();
        self.cumulative_qty = 0;
        self.cumulative_delta = 0.0;
        self.cumulative_vega = 0.0;
        self.frozen_until = None;
    }
}

/// Aggregate state for one wallet across all currencies
#[derive(Debug, Clone)]
struct WalletMmpState {
    currencies: HashMap<String, CurrencyFillWindow>,
}

impl WalletMmpState {
    fn new() -> Self {
        Self {
            currencies: HashMap::new(),
        }
    }

    fn get_or_create_currency(&mut self, currency: &str) -> &mut CurrencyFillWindow {
        self.currencies
            .entry(currency.to_string())
            .or_insert_with(CurrencyFillWindow::new)
    }
}

/// Main MMP cache for tracking fills and checking limits
pub struct MmpCache {
    /// MMP configurations indexed by (wallet_address, currency)
    configs: Arc<RwLock<HashMap<(WalletAddress, String), MmpConfig>>>,

    /// Per-wallet state tracking
    state: Arc<RwLock<HashMap<WalletAddress, WalletMmpState>>>,

    /// Database handler for persistence
    db: Arc<dyn MmpConfigWriter>,

    /// Greeks cache for calculating delta and vega
    greeks_cache: Arc<crate::read_cache::greeks::GreeksCache>,
}

impl MmpCache {
    /// Create a new MMP cache
    pub fn new(
        db: Arc<dyn MmpConfigWriter>,
        greeks_cache: Arc<crate::read_cache::greeks::GreeksCache>,
    ) -> Result<Self> {
        info!("Initializing MmpCache");

        Ok(Self {
            configs: Arc::new(RwLock::new(HashMap::new())),
            state: Arc::new(RwLock::new(HashMap::new())),
            db: db,
            greeks_cache,
        })
    }

    /// Load all MMP configs from database
    pub async fn load_configs_from_db(&self) -> Result<()> {
        let configs_list = self
            .db
            .get_all_mmp_configs_sync()
            .context("Failed to load MMP configs from database")?;

        let mut configs = self.configs.write().await;
        for config in configs_list {
            let key = (config.wallet_address, config.currency.clone());
            info!(
                "Loaded MMP config: wallet={}, currency={}, enabled={}",
                config.wallet_address, config.currency, config.enabled
            );
            configs.insert(key, config);
        }

        info!("Loaded {} MMP configs from database", configs.len());
        Ok(())
    }

    /// Set or update MMP config
    pub async fn set_config(&self, config: MmpConfig) -> Result<()> {
        // Save to database
        self.db
            .save_mmp_config_sync(&config)
            .context("Failed to save MMP config to database")?;

        // Load the saved config to get timestamps
        let saved_config = self
            .db
            .get_mmp_config_sync(&config.wallet_address, &config.currency)?
            .ok_or_else(|| anyhow::anyhow!("Config not found after save"))?;

        // Update in-memory cache
        let key = (config.wallet_address, config.currency.clone());
        let mut configs = self.configs.write().await;
        configs.insert(key, saved_config);

        info!(
            "MMP config updated: wallet={}, currency={}",
            config.wallet_address, config.currency
        );
        Ok(())
    }

    /// Get MMP config for wallet and currency
    pub async fn get_config(&self, wallet: &WalletAddress, currency: &str) -> Option<MmpConfig> {
        let configs = self.configs.read().await;
        configs.get(&(*wallet, currency.to_string())).cloned()
    }

    /// Fast path: Check if MMP is enabled for wallet and currency
    /// Returns true if MMP config exists and is enabled, false otherwise
    /// This is faster than calling process_fill() for wallets without MMP
    pub async fn is_mmp_enabled(&self, wallet: &WalletAddress, currency: &str) -> bool {
        let configs = self.configs.read().await;
        if let Some(config) = configs.get(&(*wallet, currency.to_string())) {
            config.enabled
        } else {
            false
        }
    }

    /// Get all MMP configs for a wallet
    pub async fn get_configs_for_wallet(&self, wallet: &WalletAddress) -> Vec<MmpConfig> {
        let configs = self.configs.read().await;
        configs
            .iter()
            .filter(|(key, _)| &key.0 == wallet)
            .map(|(_, config)| config.clone())
            .collect()
    }

    /// Delete (disable) MMP config
    pub async fn delete_config(&self, wallet: &WalletAddress, currency: &str) -> Result<()> {
        // Update database
        self.db
            .delete_mmp_config_sync(wallet, currency)
            .context("Failed to delete MMP config from database")?;

        // Remove from in-memory cache
        let mut configs = self.configs.write().await;
        let key = (*wallet, currency.to_string());
        configs.remove(&key);

        info!(
            "MMP config deleted: wallet={}, currency={}",
            wallet, currency
        );
        Ok(())
    }

    /// Reset MMP state for a specific currency (clear fills and unfreeze)
    pub async fn reset_mmp(&self, wallet: &WalletAddress, currency: &str) {
        let mut state = self.state.write().await;
        if let Some(wallet_state) = state.get_mut(wallet) {
            if let Some(currency_window) = wallet_state.currencies.get_mut(currency) {
                currency_window.reset();
                info!("MMP state reset: wallet={}, currency={}", wallet, currency);
            }
        }
    }

    /// Check if wallet+currency is currently frozen
    pub async fn is_frozen(
        &self,
        wallet: &WalletAddress,
        currency: &str,
        current_time: u64,
    ) -> bool {
        let state = self.state.read().await;
        if let Some(wallet_state) = state.get(wallet) {
            if let Some(currency_window) = wallet_state.currencies.get(currency) {
                return currency_window.is_frozen(current_time);
            }
        }
        false
    }

    /// Process a fill and check if MMP limits are triggered
    /// Returns Ok(()) if limits not exceeded, Err with reason if triggered
    pub async fn process_fill(
        &self,
        wallet: &WalletAddress,
        currency: &str,
        fill: &Fill,
        current_time: u64,
    ) -> Result<(), String> {
        // Skip perps - MMP only applies to options
        if ParsedSymbol::from_symbol(&fill.symbol).is_err() {
            return Ok(());
        }

        // Get config for this wallet+currency
        let config = match self.get_config(wallet, currency).await {
            Some(cfg) if cfg.enabled => cfg,
            _ => {
                // No config or disabled, allow fill
                return Ok(());
            }
        };

        // Calculate fill metrics (delta and vega)
        let (delta, vega) = self
            .calculate_fill_metrics(fill)
            .await
            .map_err(|e| format!("Failed to calculate fill metrics: {}", e))?;

        // Convert Decimal size to u64 contract units for MMP tracking
        let size_contract_units = to_contract_units_decimal(&fill.symbol, fill.size)
            .to_u64()
            .unwrap_or(0);
        let record = FillRecord {
            timestamp: fill.timestamp,
            quantity: size_contract_units,
            delta,
            vega,
        };

        // Get or create wallet state
        let mut state = self.state.write().await;
        let wallet_state = state.entry(*wallet).or_insert_with(WalletMmpState::new);

        let currency_window = wallet_state.get_or_create_currency(currency);

        // Check if frozen
        if currency_window.is_frozen(current_time) {
            debug!(
                "MMP: Fill rejected - wallet {} currency {} is frozen",
                wallet, currency
            );
            return Err("MMP triggered - currently frozen".to_string());
        }

        // Evict old fills from rolling window
        currency_window.evict_old_fills(current_time, config.interval_ms as u64);

        // Evaluate if adding this fill would breach limits
        // Convert Decimal limits to i64/f64 for internal MMP logic
        let breached = currency_window.eval_fill(
            record,
            config.qty_limit.and_then(|d| d.to_i64()),
            config.delta_limit.and_then(|d| d.to_f64()),
            config.vega_limit.and_then(|d| d.to_f64()),
        );

        if breached {
            // Freeze this currency
            currency_window.set_frozen(current_time, config.frozen_time_ms as u64);

            warn!(
                "🚨 MMP TRIGGERED: Limit exceeded for wallet={}, currency={} (qty={}, delta={:.2}, vega={:.2})",
                wallet, currency, currency_window.cumulative_qty, currency_window.cumulative_delta, currency_window.cumulative_vega
            );

            return Err("MMP limit exceeded".to_string());
        }

        debug!(
            "MMP: Fill accepted - wallet={}, currency={}, qty={}, delta={:.2}, vega={:.2}",
            wallet,
            currency,
            currency_window.cumulative_qty,
            currency_window.cumulative_delta,
            currency_window.cumulative_vega
        );

        Ok(())
    }

    /// Start the automatic eviction task (legacy, no shutdown support).
    /// This spawns a background task that periodically evicts old fills from all currency windows.
    /// For graceful shutdown support, use `start_with_shutdown` instead.
    pub fn start(&self) {
        // Create a dummy shutdown channel that never fires.
        // IMPORTANT: We must keep the sender alive, otherwise recv() returns Closed immediately.
        let (tx, rx) = tokio::sync::broadcast::channel::<()>(1);
        std::mem::forget(tx);
        self.start_with_shutdown(rx);
    }

    /// Start the automatic eviction task with shutdown support.
    /// This spawns a background task that periodically evicts old fills from all currency windows.
    /// The task will exit gracefully when the shutdown signal is received.
    pub fn start_with_shutdown(
        &self,
        mut shutdown_rx: tokio::sync::broadcast::Receiver<()>,
    ) -> tokio::task::JoinHandle<()> {
        let state = Arc::clone(&self.state);
        let configs = Arc::clone(&self.configs);

        tokio::spawn(async move {
            info!("MMP automatic eviction task started");
            let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(1000));

            loop {
                tokio::select! {
                    _ = shutdown_rx.recv() => {
                        info!("MMP eviction task received shutdown signal");
                        break;
                    }
                    _ = interval.tick() => {
                        let current_time = std::time::SystemTime::now()
                            .duration_since(std::time::UNIX_EPOCH)
                            .unwrap()
                            .as_millis() as u64;

                        // Get read lock on configs to find all active intervals
                        let configs_snapshot = configs.read().await;
                        let mut wallet_currency_intervals: HashMap<(WalletAddress, String), u64> =
                            HashMap::new();

                        for ((wallet, currency), config) in configs_snapshot.iter() {
                            if config.enabled {
                                wallet_currency_intervals
                                    .insert((*wallet, currency.clone()), config.interval_ms as u64);
                            }
                        }
                        drop(configs_snapshot);

                        // Get write lock on state to evict old fills
                        let mut state_guard = state.write().await;

                        for ((wallet, currency), interval_ms) in wallet_currency_intervals.iter() {
                            if let Some(wallet_state) = state_guard.get_mut(wallet) {
                                if let Some(currency_window) = wallet_state.currencies.get_mut(currency) {
                                    let before_count = currency_window.fills.len();
                                    currency_window.evict_old_fills(current_time, *interval_ms);
                                    let after_count = currency_window.fills.len();

                                    if before_count > after_count {
                                        debug!(
                                            "MMP eviction: wallet={}, currency={}, evicted {} fills",
                                            wallet,
                                            currency,
                                            before_count - after_count
                                        );
                                    }
                                }
                            }
                        }
                    }
                }
            }
            info!("MMP eviction task stopped");
        })
    }

    /// Calculate delta and vega for a fill using the Greeks cache
    /// Returns error if Greeks data is not available for the symbol
    async fn calculate_fill_metrics(&self, fill: &Fill) -> Result<(f64, f64)> {
        // Get Greeks from cache
        let greeks = self
            .greeks_cache
            .get_greeks(&fill.symbol)
            .await
            .context("Greeks not available for symbol")?;

        let quantity_in_contracts = fill.size.to_f64().unwrap_or(0.0);
        let side_multiplier: f64 = match fill.taker_side {
            hypercall_types::Side::Buy => 1.0,
            hypercall_types::Side::Sell => -1.0,
        };

        // Apply side multiplier to delta and vega
        let delta = greeks.delta * quantity_in_contracts * side_multiplier;
        let vega = greeks.vega * quantity_in_contracts * side_multiplier.abs();

        Ok((delta, vega))
    }
}

#[async_trait::async_trait]
impl crate::shared::service::Service for MmpCache {
    fn name(&self) -> &'static str {
        "MmpCache"
    }

    fn owner(&self) -> crate::shared::service::ServiceOwner {
        crate::shared::service::ServiceOwner::Engine
    }

    async fn run(
        self: std::sync::Arc<Self>,
        shutdown: crate::shared::ShutdownRx,
    ) -> anyhow::Result<()> {
        let handle = self.start_with_shutdown(shutdown);
        handle
            .await
            .map_err(|e| anyhow::anyhow!("MmpCache task task failed: {:?}", e))
    }
}

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

    #[test]
    fn test_currency_fill_window_basic() {
        let mut window = CurrencyFillWindow::new();

        let record = FillRecord {
            timestamp: 1000,
            quantity: 100,
            delta: 50.0,
            vega: 10.0,
        };

        // No limits, should accept
        let breached = window.eval_fill(record, None, None, None);
        assert!(!breached);

        assert_eq!(window.cumulative_qty, 100);
        assert_eq!(window.cumulative_delta, 50.0);
        assert_eq!(window.cumulative_vega, 10.0);
    }

    #[test]
    fn test_evict_old_fills() {
        let mut window = CurrencyFillWindow::new();

        window.eval_fill(
            FillRecord {
                timestamp: 1000,
                quantity: 100,
                delta: 50.0,
                vega: 10.0,
            },
            None,
            None,
            None,
        );

        window.eval_fill(
            FillRecord {
                timestamp: 2000,
                quantity: 200,
                delta: 100.0,
                vega: 20.0,
            },
            None,
            None,
            None,
        );

        // Evict fills older than 1500ms ago (cutoff at 3500 - 1500 = 2000)
        window.evict_old_fills(3500, 1500);

        // First fill (timestamp 1000) should be evicted
        assert_eq!(window.cumulative_qty, 200);
        assert_eq!(window.cumulative_delta, 100.0);
        assert_eq!(window.cumulative_vega, 20.0);
        assert_eq!(window.fills.len(), 1);
    }

    #[test]
    fn test_frozen_state() {
        let mut window = CurrencyFillWindow::new();

        // Not frozen initially
        assert!(!window.is_frozen(1000));

        // Set frozen for 5000ms from time 1000
        window.set_frozen(1000, 5000);

        // Should be frozen at time 3000
        assert!(window.is_frozen(3000));

        // Should not be frozen at time 7000
        assert!(!window.is_frozen(7000));
    }

    #[test]
    fn test_reset() {
        let mut window = CurrencyFillWindow::new();

        window.eval_fill(
            FillRecord {
                timestamp: 1000,
                quantity: 100,
                delta: 50.0,
                vega: 10.0,
            },
            None,
            None,
            None,
        );

        window.set_frozen(1000, 5000);

        window.reset();

        assert_eq!(window.cumulative_qty, 0);
        assert_eq!(window.cumulative_delta, 0.0);
        assert_eq!(window.cumulative_vega, 0.0);
        assert!(window.fills.is_empty());
        assert!(!window.is_frozen(3000));
    }

    #[test]
    fn test_eval_fill_with_qty_limit() {
        let mut window = CurrencyFillWindow::new();

        // Add a fill that's within limits
        let breached = window.eval_fill(
            FillRecord {
                timestamp: 1000,
                quantity: 100,
                delta: 50.0,
                vega: 10.0,
            },
            Some(200), // qty_limit
            None,
            None,
        );
        assert!(!breached);
        assert_eq!(window.cumulative_qty, 100);

        // Try to add a fill that would exceed the limit
        let breached = window.eval_fill(
            FillRecord {
                timestamp: 2000,
                quantity: 150,
                delta: 75.0,
                vega: 15.0,
            },
            Some(200), // qty_limit
            None,
            None,
        );
        assert!(breached);
        // Fill SHOULD have been added (new behavior: always add first, then check)
        assert_eq!(window.cumulative_qty, 250);
    }

    #[test]
    fn test_eval_fill_with_delta_limit() {
        let mut window = CurrencyFillWindow::new();

        // Add a fill within delta limit
        let breached = window.eval_fill(
            FillRecord {
                timestamp: 1000,
                quantity: 100,
                delta: 50.0,
                vega: 10.0,
            },
            None,
            Some(100.0), // delta_limit
            None,
        );
        assert!(!breached);

        // Try to add a fill that would exceed delta limit
        let breached = window.eval_fill(
            FillRecord {
                timestamp: 2000,
                quantity: 100,
                delta: 60.0,
                vega: 10.0,
            },
            None,
            Some(100.0), // delta_limit
            None,
        );
        assert!(breached);
        assert_eq!(window.cumulative_delta, 110.0); // Fill was added (new behavior)
    }

    #[test]
    fn test_eval_fill_with_vega_limit() {
        let mut window = CurrencyFillWindow::new();

        // Add a fill within vega limit
        let breached = window.eval_fill(
            FillRecord {
                timestamp: 1000,
                quantity: 100,
                delta: 50.0,
                vega: 10.0,
            },
            None,
            None,
            Some(20.0), // vega_limit
        );
        assert!(!breached);

        // Try to add a fill that would exceed vega limit
        let breached = window.eval_fill(
            FillRecord {
                timestamp: 2000,
                quantity: 100,
                delta: 50.0,
                vega: 15.0,
            },
            None,
            None,
            Some(20.0), // vega_limit
        );
        assert!(breached);
        assert_eq!(window.cumulative_vega, 25.0); // Fill was added (new behavior)
    }
}