Vapora/crates/vapora-llm-router/src/budget.rs

use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;

use chrono::{Datelike, Utc};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use tokio::sync::RwLock;

/// Budget configuration errors
#[derive(Debug, Error)]
pub enum BudgetConfigError {
    #[error("Failed to read budget config file: {0}")]
    ReadError(#[from] std::io::Error),

    #[error("Failed to parse TOML: {0}")]
    ParseError(#[from] toml::de::Error),

    #[error("Invalid budget configuration: {0}")]
    ValidationError(String),
}

/// Role-specific budget allocation and spending limits.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoleBudget {
    /// Agent role (e.g., "developer", "architect", "reviewer")
    pub role: String,
    /// Monthly spending limit in cents ($500 = 50000 cents)
    pub monthly_limit_cents: u32,
    /// Weekly spending limit in cents ($125 = 12500 cents)
    pub weekly_limit_cents: u32,
    /// Fallback provider when budget is exceeded (cheaper alternative)
    pub fallback_provider: String,
    /// Alert threshold as fraction (0.8 = alert at 80% utilization)
    pub alert_threshold: f32,
}

/// Current budget status for a role.
#[derive(Debug, Clone)]
pub struct BudgetStatus {
    pub role: String,
    pub monthly_remaining_cents: u32,
    pub weekly_remaining_cents: u32,
    pub monthly_utilization: f32,
    pub weekly_utilization: f32,
    pub exceeded: bool,
    pub near_threshold: bool,
    pub fallback_provider: String,
}

/// Budget manager for enforcing role-based spending limits.
/// Tracks monthly and weekly spend, enforces caps, and triggers alerts.
pub struct BudgetManager {
    budgets: Arc<RwLock<HashMap<String, RoleBudget>>>,
    spending: Arc<RwLock<HashMap<String, RoleSpending>>>,
}

#[derive(Debug, Clone)]
struct RoleSpending {
    #[allow(dead_code)]
    role: String,
    current_month: MonthBudget,
    current_week: WeekBudget,
}

#[derive(Debug, Clone)]
struct MonthBudget {
    year: i32,
    month: u32,
    spent_cents: u32,
}

#[derive(Debug, Clone)]
struct WeekBudget {
    year: i32,
    week: u32,
    spent_cents: u32,
}

/// Budget configuration file structure
#[derive(Debug, Serialize, Deserialize)]
pub struct BudgetConfig {
    pub budgets: HashMap<String, RoleBudget>,
}

impl BudgetConfig {
    /// Load budget configuration from TOML file
    pub fn load<P: AsRef<Path>>(path: P) -> Result<Self, BudgetConfigError> {
        let content = std::fs::read_to_string(path)?;
        let config: Self = toml::from_str(&content)?;
        config.validate()?;
        Ok(config)
    }

    /// Load from TOML with default fallback if file doesn't exist
    pub fn load_or_default<P: AsRef<Path>>(path: P) -> Result<Self, BudgetConfigError> {
        match Self::load(&path) {
            Ok(config) => Ok(config),
            Err(BudgetConfigError::ReadError(_)) => {
                // File doesn't exist, use defaults
                Ok(BudgetConfig {
                    budgets: HashMap::new(),
                })
            }
            Err(e) => Err(e),
        }
    }

    /// Validate budget configuration
    fn validate(&self) -> Result<(), BudgetConfigError> {
        if self.budgets.is_empty() {
            return Err(BudgetConfigError::ValidationError(
                "At least one role budget must be configured".to_string(),
            ));
        }

        for (role, budget) in &self.budgets {
            if budget.monthly_limit_cents == 0 {
                return Err(BudgetConfigError::ValidationError(format!(
                    "Role {} has zero monthly limit",
                    role
                )));
            }
            if budget.weekly_limit_cents == 0 {
                return Err(BudgetConfigError::ValidationError(format!(
                    "Role {} has zero weekly limit",
                    role
                )));
            }
            if budget.alert_threshold < 0.0 || budget.alert_threshold > 1.0 {
                return Err(BudgetConfigError::ValidationError(format!(
                    "Role {} has invalid alert_threshold: {}",
                    role, budget.alert_threshold
                )));
            }
        }

        Ok(())
    }
}

impl BudgetManager {
    /// Create new budget manager with role budgets.
    pub fn new(budgets: HashMap<String, RoleBudget>) -> Self {
        let spending = budgets
            .keys()
            .map(|role| {
                (
                    role.clone(),
                    RoleSpending {
                        role: role.clone(),
                        current_month: MonthBudget {
                            year: Utc::now().year(),
                            month: Utc::now().month(),
                            spent_cents: 0,
                        },
                        current_week: WeekBudget {
                            year: Utc::now().year(),
                            week: Utc::now().iso_week().week(),
                            spent_cents: 0,
                        },
                    },
                )
            })
            .collect();

        Self {
            budgets: Arc::new(RwLock::new(budgets)),
            spending: Arc::new(RwLock::new(spending)),
        }
    }

    /// Check budget status for role.
    /// Returns BudgetStatus with remaining balance, utilization %, and alert
    /// flags.
    pub async fn check_budget(&self, role: &str) -> Result<BudgetStatus, String> {
        let budgets = self.budgets.read().await;
        let mut spending = self.spending.write().await;

        let budget = budgets
            .get(role)
            .ok_or_else(|| format!("Unknown role: {}", role))?;
        let spending_entry = spending
            .entry(role.to_string())
            .or_insert_with(|| RoleSpending {
                role: role.to_string(),
                current_month: MonthBudget {
                    year: Utc::now().year(),
                    month: Utc::now().month(),
                    spent_cents: 0,
                },
                current_week: WeekBudget {
                    year: Utc::now().year(),
                    week: Utc::now().iso_week().week(),
                    spent_cents: 0,
                },
            });

        // Reset month if new month
        let now = Utc::now();
        if now.year() != spending_entry.current_month.year
            || now.month() != spending_entry.current_month.month
        {
            spending_entry.current_month = MonthBudget {
                year: now.year(),
                month: now.month(),
                spent_cents: 0,
            };
        }

        // Reset week if new week
        let current_week = now.iso_week().week();
        if now.year() != spending_entry.current_week.year
            || current_week != spending_entry.current_week.week
        {
            spending_entry.current_week = WeekBudget {
                year: now.year(),
                week: current_week,
                spent_cents: 0,
            };
        }

        let monthly_remaining = budget
            .monthly_limit_cents
            .saturating_sub(spending_entry.current_month.spent_cents);
        let weekly_remaining = budget
            .weekly_limit_cents
            .saturating_sub(spending_entry.current_week.spent_cents);

        let monthly_utilization = if budget.monthly_limit_cents > 0 {
            spending_entry.current_month.spent_cents as f32 / budget.monthly_limit_cents as f32
        } else {
            1.0
        };

        let weekly_utilization = if budget.weekly_limit_cents > 0 {
            spending_entry.current_week.spent_cents as f32 / budget.weekly_limit_cents as f32
        } else {
            1.0
        };

        let exceeded = monthly_remaining == 0 || weekly_remaining == 0;
        let near_threshold = monthly_utilization >= budget.alert_threshold
            || weekly_utilization >= budget.alert_threshold;

        Ok(BudgetStatus {
            role: role.to_string(),
            monthly_remaining_cents: monthly_remaining,
            weekly_remaining_cents: weekly_remaining,
            monthly_utilization,
            weekly_utilization,
            exceeded,
            near_threshold,
            fallback_provider: budget.fallback_provider.clone(),
        })
    }

    /// Record spending against role budget.
    /// Automatically resets month/week tracking if boundaries crossed.
    pub async fn record_spend(&self, role: &str, cost_cents: u32) -> Result<(), String> {
        let mut spending = self.spending.write().await;
        let entry = spending
            .get_mut(role)
            .ok_or_else(|| format!("Unknown role: {}", role))?;

        let now = Utc::now();

        // Reset month if needed
        if now.year() != entry.current_month.year || now.month() != entry.current_month.month {
            entry.current_month = MonthBudget {
                year: now.year(),
                month: now.month(),
                spent_cents: 0,
            };
        }

        // Reset week if needed
        let current_week = now.iso_week().week();
        if now.year() != entry.current_week.year || current_week != entry.current_week.week {
            entry.current_week = WeekBudget {
                year: now.year(),
                week: current_week,
                spent_cents: 0,
            };
        }

        entry.current_month.spent_cents =
            entry.current_month.spent_cents.saturating_add(cost_cents);
        entry.current_week.spent_cents = entry.current_week.spent_cents.saturating_add(cost_cents);

        Ok(())
    }

    /// Get fallback provider for role when budget exceeded.
    pub async fn get_fallback_provider(&self, role: &str) -> Result<String, String> {
        let budgets = self.budgets.read().await;
        budgets
            .get(role)
            .map(|b| b.fallback_provider.clone())
            .ok_or_else(|| format!("Unknown role: {}", role))
    }

    /// Get all budget statuses (for monitoring dashboards).
    pub async fn get_all_budgets(&self) -> Vec<BudgetStatus> {
        let budgets = self.budgets.read().await;
        let spending = self.spending.read().await;

        budgets
            .iter()
            .filter_map(|(role, budget)| {
                spending.get(role).map(|sp| {
                    let monthly_remaining = budget
                        .monthly_limit_cents
                        .saturating_sub(sp.current_month.spent_cents);
                    let weekly_remaining = budget
                        .weekly_limit_cents
                        .saturating_sub(sp.current_week.spent_cents);

                    let monthly_utilization = if budget.monthly_limit_cents > 0 {
                        sp.current_month.spent_cents as f32 / budget.monthly_limit_cents as f32
                    } else {
                        1.0
                    };

                    let weekly_utilization = if budget.weekly_limit_cents > 0 {
                        sp.current_week.spent_cents as f32 / budget.weekly_limit_cents as f32
                    } else {
                        1.0
                    };

                    let exceeded = monthly_remaining == 0 || weekly_remaining == 0;
                    let near_threshold = monthly_utilization >= budget.alert_threshold
                        || weekly_utilization >= budget.alert_threshold;

                    BudgetStatus {
                        role: role.clone(),
                        monthly_remaining_cents: monthly_remaining,
                        weekly_remaining_cents: weekly_remaining,
                        monthly_utilization,
                        weekly_utilization,
                        exceeded,
                        near_threshold,
                        fallback_provider: budget.fallback_provider.clone(),
                    }
                })
            })
            .collect()
    }

    /// List all configured role budgets
    pub async fn list_budgets(&self) -> Vec<RoleBudget> {
        let budgets = self.budgets.read().await;
        budgets.values().cloned().collect()
    }
}

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

    fn create_test_budgets() -> HashMap<String, RoleBudget> {
        let mut budgets = HashMap::new();
        budgets.insert(
            "developer".to_string(),
            RoleBudget {
                role: "developer".to_string(),
                monthly_limit_cents: 30000, // $300
                weekly_limit_cents: 7500,   // $75
                fallback_provider: "ollama".to_string(),
                alert_threshold: 0.8,
            },
        );
        budgets
    }

    #[tokio::test]
    async fn test_budget_manager_creation() {
        let budgets = create_test_budgets();
        let manager = BudgetManager::new(budgets);

        let status = manager.check_budget("developer").await.unwrap();
        assert_eq!(status.role, "developer");
        assert_eq!(status.monthly_remaining_cents, 30000);
        assert!(!status.exceeded);
        assert!(!status.near_threshold);
    }

    #[tokio::test]
    async fn test_record_spend() {
        let budgets = create_test_budgets();
        let manager = BudgetManager::new(budgets);

        manager.record_spend("developer", 5000).await.unwrap();
        let status = manager.check_budget("developer").await.unwrap();
        assert_eq!(status.monthly_remaining_cents, 25000);
        assert!((status.monthly_utilization - 0.1667).abs() < 0.01);
    }

    #[tokio::test]
    async fn test_threshold_alert() {
        let budgets = create_test_budgets();
        let manager = BudgetManager::new(budgets);

        // Spend 81% of weekly budget (7500 * 0.81 = 6075)
        // This triggers near_threshold (> 80%) but not exceeded
        manager.record_spend("developer", 6075).await.unwrap();
        let status = manager.check_budget("developer").await.unwrap();
        assert!(!status.exceeded);
        assert!(status.near_threshold); // 81% > 80% threshold
    }

    #[tokio::test]
    async fn test_budget_exceeded() {
        let budgets = create_test_budgets();
        let manager = BudgetManager::new(budgets);

        // Exceed monthly budget
        manager.record_spend("developer", 30000).await.unwrap();
        let status = manager.check_budget("developer").await.unwrap();
        assert!(status.exceeded);
        assert_eq!(status.monthly_remaining_cents, 0);
    }

    #[tokio::test]
    async fn test_fallback_provider() {
        let budgets = create_test_budgets();
        let manager = BudgetManager::new(budgets);

        let fallback = manager.get_fallback_provider("developer").await.unwrap();
        assert_eq!(fallback, "ollama");
    }

    #[tokio::test]
    async fn test_unknown_role() {
        let budgets = create_test_budgets();
        let manager = BudgetManager::new(budgets);

        let result = manager.check_budget("unknown_role").await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_get_all_budgets() {
        let budgets = create_test_budgets();
        let manager = BudgetManager::new(budgets);

        manager.record_spend("developer", 3000).await.unwrap();
        let all_statuses = manager.get_all_budgets().await;
        assert_eq!(all_statuses.len(), 1);
        assert_eq!(all_statuses[0].monthly_remaining_cents, 27000);
    }
}