zc2 0.0.4

//! Credit management system for compute billing.

use chrono::{DateTime, Utc};
use dashmap::DashMap;
use serde::{Deserialize, Serialize};

/// Balance data source — indicates how current the cached balance is
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum BalanceStatus {
    /// This broker is the authoritative owner of this user's balance
    Authoritative,
    /// Balance was fetched from a peer broker; may be slightly stale
    Prefetched,
    /// Currently re-fetching balance from peer for reconciliation
    Reconciling,
}

impl BalanceStatus {
    pub fn as_str(&self) -> &'static str {
        match self {
            BalanceStatus::Authoritative => "authoritative",
            BalanceStatus::Prefetched => "prefetched",
            BalanceStatus::Reconciling => "reconciling",
        }
    }
}

/// User credit balance and usage tracking
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserCredits {
    /// User ID
    pub user_id: String,
    /// Current credit balance
    pub balance: f64,
    /// Total credits ever added
    pub total_added: f64,
    /// Total credits spent
    pub total_spent: f64,
    /// Last transaction timestamp
    pub last_updated: DateTime<Utc>,
    /// Rate limit: max requests per minute (legacy, kept for compat)
    pub rate_limit: u32,
    /// Current requests this minute
    pub requests_this_minute: u32,
    /// Minute window start
    pub minute_window_start: DateTime<Utc>,
    /// Rate limit: max requests per second (from dashboard)
    pub rate_limit_per_second: Option<u32>,
    /// Rate limit: max requests per day (from dashboard)
    pub rate_limit_per_day: Option<u32>,
    /// Rate limit: max requests per month (from dashboard)
    pub rate_limit_per_month: Option<u32>,
    /// Balance data source (authoritative / prefetched / reconciling)
    pub balance_status: BalanceStatus,
    /// When balance was last fetched from a peer (None if authoritative)
    pub last_prefetched: Option<DateTime<Utc>>,
    /// Requests in current second window
    pub requests_this_second: u32,
    /// Second window start
    pub second_window_start: DateTime<Utc>,
    /// Requests in current day window
    pub requests_this_day: u32,
    /// Day window start
    pub day_window_start: DateTime<Utc>,
    /// Requests in current month window
    pub requests_this_month: u32,
    /// Month window start
    pub month_window_start: DateTime<Utc>,
}

impl UserCredits {
    /// Create a new user with initial credits
    pub fn new(user_id: String, initial_credits: f64) -> Self {
        let now = Utc::now();
        Self {
            user_id,
            balance: initial_credits,
            total_added: initial_credits,
            total_spent: 0.0,
            last_updated: now,
            rate_limit: 0, // Default: 0 = unlimited (per-minute; only applied when > 0)
            requests_this_minute: 0,
            minute_window_start: now,
            rate_limit_per_second: None,
            rate_limit_per_day: None,
            rate_limit_per_month: None,
            balance_status: BalanceStatus::Authoritative,
            last_prefetched: None,
            requests_this_second: 0,
            second_window_start: now,
            requests_this_day: 0,
            day_window_start: now,
            requests_this_month: 0,
            month_window_start: now,
        }
    }

    /// Check if user has enough credits
    pub fn has_credits(&self, amount: f64) -> bool {
        self.balance >= amount
    }

    /// Deduct credits (returns false if insufficient)
    pub fn deduct(&mut self, amount: f64) -> bool {
        if self.balance >= amount {
            self.balance -= amount;
            self.total_spent += amount;
            self.last_updated = Utc::now();
            true
        } else {
            false
        }
    }

    /// Add credits
    pub fn add(&mut self, amount: f64) {
        self.balance += amount;
        self.total_added += amount;
        self.last_updated = Utc::now();
    }

    /// Check and update rate limit (per-minute legacy + per-sec/day/month)
    pub fn check_rate_limit(&mut self) -> bool {
        let now = Utc::now();

        // Per-minute (legacy) — only enforced when rate_limit > 0
        if self.rate_limit > 0 {
            let elapsed_min = now.signed_duration_since(self.minute_window_start);
            if elapsed_min.num_seconds() >= 60 {
                self.requests_this_minute = 0;
                self.minute_window_start = now;
            }
            if self.requests_this_minute >= self.rate_limit {
                return false;
            }
        }

        // Per-second
        if let Some(limit) = self.rate_limit_per_second {
            let elapsed_sec = now.signed_duration_since(self.second_window_start);
            if elapsed_sec.num_milliseconds() >= 1000 {
                self.requests_this_second = 0;
                self.second_window_start = now;
            }
            if self.requests_this_second >= limit {
                return false;
            }
        }

        // Per-day
        if let Some(limit) = self.rate_limit_per_day {
            let elapsed_day = now.signed_duration_since(self.day_window_start);
            if elapsed_day.num_seconds() >= 86400 {
                self.requests_this_day = 0;
                self.day_window_start = now;
            }
            if self.requests_this_day >= limit {
                return false;
            }
        }

        // Per-month
        if let Some(limit) = self.rate_limit_per_month {
            let elapsed_month = now.signed_duration_since(self.month_window_start);
            if elapsed_month.num_seconds() >= 2_592_000 { // 30 days
                self.requests_this_month = 0;
                self.month_window_start = now;
            }
            if self.requests_this_month >= limit {
                return false;
            }
        }

        // All checks passed — increment all counters
        self.requests_this_minute += 1;
        self.requests_this_second += 1;
        self.requests_this_day += 1;
        self.requests_this_month += 1;
        true
    }
}

/// Credit manager: per-user balance cache and rate-limit tracking.
///
/// Reservations and transaction history are managed by the `Ledger`; this
/// struct is intentionally lightweight — it only caches what the router needs
/// on the hot path (balance, rate limits, reconciliation status).
#[derive(Debug)]
pub struct CreditManager {
    /// User balances
    balances: DashMap<String, UserCredits>,
}

impl CreditManager {
    /// Create a new credit manager
    pub fn new() -> Self {
        Self {
            balances: DashMap::new(),
        }
    }

    /// Get or create user credits
    pub fn get_or_create(&self, user_id: &str, initial_credits: f64) -> UserCredits {
        self.balances
            .entry(user_id.to_string())
            .or_insert_with(|| UserCredits::new(user_id.to_string(), initial_credits))
            .clone()
    }

    /// Get user credits
    pub fn get(&self, user_id: &str) -> Option<UserCredits> {
        self.balances.get(user_id).map(|c| c.clone())
    }

    /// Set user balance (sync from ledger — marks as authoritative)
    pub fn set_balance(&self, user_id: &str, balance: f64) {
        if let Some(mut entry) = self.balances.get_mut(user_id) {
            entry.balance = balance;
        }
    }

    /// Set user balance fetched from a peer broker (marks as prefetched)
    pub fn set_prefetched_balance(&self, user_id: &str, balance: f64) {
        if let Some(mut entry) = self.balances.get_mut(user_id) {
            entry.balance = balance;
            entry.balance_status = BalanceStatus::Prefetched;
            entry.last_prefetched = Some(Utc::now());
        }
    }

    /// Mark balance as reconciling (re-fetch in progress)
    pub fn set_reconciling(&self, user_id: &str) {
        if let Some(mut entry) = self.balances.get_mut(user_id) {
            entry.balance_status = BalanceStatus::Reconciling;
        }
    }

    /// Mark balance as authoritative (this broker owns it)
    pub fn set_authoritative(&self, user_id: &str) {
        if let Some(mut entry) = self.balances.get_mut(user_id) {
            entry.balance_status = BalanceStatus::Authoritative;
            entry.last_prefetched = None;
        }
    }

    /// Check if a prefetched balance is stale and needs reconciliation
    pub fn needs_reconciliation(&self, user_id: &str, max_age_secs: i64) -> bool {
        self.balances.get(user_id).map(|c| {
            c.balance_status == BalanceStatus::Prefetched
                && c.last_prefetched
                    .map(|t| Utc::now().signed_duration_since(t).num_seconds() >= max_age_secs)
                    .unwrap_or(true)
        }).unwrap_or(false)
    }

    /// Return all known user IDs (for reconciliation sweep)
    pub fn get_all_user_ids(&self) -> Vec<String> {
        self.balances.iter().map(|e| e.key().clone()).collect()
    }

    /// Check if user has sufficient credits
    pub fn has_credits(&self, user_id: &str, amount: f64) -> bool {
        self.balances
            .get(user_id)
            .map(|c| c.has_credits(amount))
            .unwrap_or(false)
    }

    /// Check rate limit for user
    pub fn check_rate_limit(&self, user_id: &str) -> bool {
        self.balances
            .get_mut(user_id)
            .map(|mut c| c.check_rate_limit())
            .unwrap_or(false)
    }

    /// Set rate limit for user (legacy per-minute)
    pub fn set_rate_limit(&self, user_id: &str, limit: u32) {
        if let Some(mut credits) = self.balances.get_mut(user_id) {
            credits.rate_limit = limit;
        }
    }

    /// Sync per-second/day/month rate limits from dashboard DB
    pub fn set_rate_limits(
        &self,
        user_id: &str,
        per_second: Option<u32>,
        per_day: Option<u32>,
        per_month: Option<u32>,
    ) {
        if let Some(mut credits) = self.balances.get_mut(user_id) {
            credits.rate_limit_per_second = per_second;
            credits.rate_limit_per_day = per_day;
            credits.rate_limit_per_month = per_month;
        }
    }
}

impl Default for CreditManager {
    fn default() -> Self {
        Self::new()
    }
}

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

    // --- UserCredits ---

    #[test]
    fn test_new_user_initial_balance() {
        let u = UserCredits::new("alice".to_string(), 100.0);
        assert_eq!(u.balance, 100.0);
        assert_eq!(u.total_added, 100.0);
        assert_eq!(u.total_spent, 0.0);
    }

    #[test]
    fn test_deduct_success_updates_balance_and_spent() {
        let mut u = UserCredits::new("alice".to_string(), 50.0);
        assert!(u.deduct(20.0));
        assert_eq!(u.balance, 30.0);
        assert_eq!(u.total_spent, 20.0);
    }

    #[test]
    fn test_deduct_fails_on_insufficient_balance() {
        let mut u = UserCredits::new("alice".to_string(), 10.0);
        assert!(!u.deduct(20.0));
        assert_eq!(u.balance, 10.0); // unchanged
        assert_eq!(u.total_spent, 0.0);
    }

    #[test]
    fn test_add_increases_balance_and_total_added() {
        let mut u = UserCredits::new("alice".to_string(), 50.0);
        u.add(25.0);
        assert_eq!(u.balance, 75.0);
        assert_eq!(u.total_added, 75.0);
    }

    #[test]
    fn test_has_credits() {
        let u = UserCredits::new("alice".to_string(), 10.0);
        assert!(u.has_credits(10.0));
        assert!(u.has_credits(5.0));
        assert!(!u.has_credits(10.001));
    }

    #[test]
    fn test_rate_limit_per_minute_blocks_after_limit() {
        let mut u = UserCredits::new("alice".to_string(), 100.0);
        u.rate_limit = 3;
        assert!(u.check_rate_limit()); // 1
        assert!(u.check_rate_limit()); // 2
        assert!(u.check_rate_limit()); // 3
        assert!(!u.check_rate_limit()); // 4 — blocked
    }

    #[test]
    fn test_rate_limit_per_second_blocks_after_limit() {
        let mut u = UserCredits::new("alice".to_string(), 100.0);
        u.rate_limit = 1000; // high per-minute limit
        u.rate_limit_per_second = Some(2);
        assert!(u.check_rate_limit()); // 1
        assert!(u.check_rate_limit()); // 2
        assert!(!u.check_rate_limit()); // 3 — blocked by per-second
    }

    #[test]
    fn test_rate_limit_per_day_blocks_after_limit() {
        let mut u = UserCredits::new("alice".to_string(), 100.0);
        u.rate_limit = 1000;
        u.rate_limit_per_day = Some(2);
        assert!(u.check_rate_limit()); // 1
        assert!(u.check_rate_limit()); // 2
        assert!(!u.check_rate_limit()); // 3 — blocked by per-day
    }

    // --- CreditManager ---

    #[test]
    fn test_get_or_create_initialises_balance() {
        let mgr = CreditManager::new();
        let u = mgr.get_or_create("bob", 200.0);
        assert_eq!(u.balance, 200.0);
        // Second call returns existing user (same balance)
        let u2 = mgr.get_or_create("bob", 999.0);
        assert_eq!(u2.balance, 200.0);
    }

    #[test]
    fn test_get_returns_none_for_unknown_user() {
        let mgr = CreditManager::new();
        assert!(mgr.get("ghost").is_none());
    }

    #[test]
    fn test_set_balance_syncs_from_ledger() {
        let mgr = CreditManager::new();
        mgr.get_or_create("alice", 10.0);
        mgr.set_balance("alice", 500.0);
        assert_eq!(mgr.get("alice").unwrap().balance, 500.0);
    }

    #[test]
    fn test_has_credits_check() {
        let mgr = CreditManager::new();
        mgr.get_or_create("alice", 50.0);
        assert!(mgr.has_credits("alice", 50.0));
        assert!(!mgr.has_credits("alice", 50.001));
        assert!(!mgr.has_credits("nobody", 1.0));
    }

    #[test]
    fn test_rate_limit_manager_check() {
        let mgr = CreditManager::new();
        mgr.get_or_create("alice", 100.0);
        mgr.set_rate_limit("alice", 2);

        assert!(mgr.check_rate_limit("alice")); // 1
        assert!(mgr.check_rate_limit("alice")); // 2
        assert!(!mgr.check_rate_limit("alice")); // blocked
    }

    #[test]
    fn test_check_rate_limit_unknown_user_returns_false() {
        let mgr = CreditManager::new();
        assert!(!mgr.check_rate_limit("ghost"));
    }
}