brokecrew/src/budget.rs

use chrono::{DateTime, Duration, Utc};
use serde::{Deserialize, Serialize};
use std::collections::VecDeque;
use std::fs;
use std::io::Write;
use std::path::{Path, PathBuf};

const MINUTE_LIMIT: usize = 20;

/// Estimates token count from text using standard heuristic (~4 characters per token).
pub fn estimate_tokens(text: &str) -> usize {
    if text.is_empty() {
        return 0;
    }
    text.len().div_ceil(4).max(1)
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum CreditTier {
    LowBalance, // < $10 lifetime purchased: 50 requests/day
    Funded,     // >= $10 lifetime purchased: 1,000 requests/day
}

impl CreditTier {
    pub fn daily_limit(&self) -> usize {
        match self {
            CreditTier::LowBalance => 50,
            CreditTier::Funded => 1_000,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BudgetStatus {
    pub minute_requests: usize,
    pub minute_limit: usize,
    pub daily_requests: usize,
    pub daily_limit: usize,
    pub approaching_daily_limit: bool,
}

#[derive(thiserror::Error, Debug, PartialEq, Eq)]
pub enum BudgetExceeded {
    #[error("Rate limited: exceeded {0} requests in 1 minute. Please wait {1}s before retrying.")]
    MinuteLimit(usize, u64),
    #[error("Daily budget exhausted: reached {0} of {1} requests today.")]
    DailyQuota(usize, usize),
}

#[derive(Debug, Serialize, Deserialize, Default)]
struct StoredBudgetState {
    daily_count: usize,
    last_day_utc: String, // YYYY-MM-DD
    recent_calls: VecDeque<DateTime<Utc>>,
}

pub struct BudgetTracker {
    storage_path: Option<PathBuf>,
    credit_tier: CreditTier,
    daily_count: usize,
    last_day_utc: String,
    recent_calls: VecDeque<DateTime<Utc>>,
}

impl BudgetTracker {
    pub fn new(storage_path: Option<PathBuf>, credit_tier: CreditTier) -> Self {
        let mut tracker = Self {
            storage_path,
            credit_tier,
            daily_count: 0,
            last_day_utc: Utc::now().format("%Y-%m-%d").to_string(),
            recent_calls: VecDeque::new(),
        };

        if let Some(ref path) = tracker.storage_path.clone() {
            let _ = tracker.load_from_disk(path);
        }

        tracker.refresh_window();
        tracker
    }

    fn today_utc() -> String {
        Utc::now().format("%Y-%m-%d").to_string()
    }

    fn refresh_window(&mut self) {
        let now = Utc::now();
        let current_day = Self::today_utc();

        // Day rollover check
        if self.last_day_utc != current_day {
            self.daily_count = 0;
            self.last_day_utc = current_day;
        }

        // Sliding 60-second window eviction
        let cutoff = now - Duration::seconds(60);
        while let Some(&first) = self.recent_calls.front() {
            if first < cutoff {
                self.recent_calls.pop_front();
            } else {
                break;
            }
        }
    }

    /// Checks the current budget status against OpenRouter rate limits.
    /// Returns `BudgetStatus` if permitted, or `BudgetExceeded` if limit reached.
    pub fn check_budget(&mut self) -> Result<BudgetStatus, BudgetExceeded> {
        self.refresh_window();

        let daily_limit = self.credit_tier.daily_limit();

        // Check daily ceiling
        if self.daily_count >= daily_limit {
            return Err(BudgetExceeded::DailyQuota(self.daily_count, daily_limit));
        }

        // Check 1-minute rolling ceiling
        if self.recent_calls.len() >= MINUTE_LIMIT {
            let oldest = self.recent_calls.front().copied().unwrap_or_else(Utc::now);
            let elapsed = (Utc::now() - oldest).num_seconds().max(0);
            let wait_secs = (60 - elapsed).max(1) as u64;
            return Err(BudgetExceeded::MinuteLimit(MINUTE_LIMIT, wait_secs));
        }

        let approaching_daily_limit = self.daily_count >= (daily_limit * 80) / 100;

        Ok(BudgetStatus {
            minute_requests: self.recent_calls.len(),
            minute_limit: MINUTE_LIMIT,
            daily_requests: self.daily_count,
            daily_limit,
            approaching_daily_limit,
        })
    }

    /// Records an outbound call and persists updated state to disk atomically.
    pub fn record_call(&mut self, _is_paid: bool) -> Result<(), std::io::Error> {
        self.refresh_window();
        let now = Utc::now();
        self.recent_calls.push_back(now);
        self.daily_count += 1;

        if let Some(ref path) = self.storage_path.clone() {
            self.save_to_disk(path)?;
        }
        Ok(())
    }

    fn load_from_disk(&mut self, path: &Path) -> Result<(), std::io::Error> {
        if !path.exists() {
            return Ok(());
        }
        let data = fs::read_to_string(path)?;
        if let Ok(state) = serde_json::from_str::<StoredBudgetState>(&data) {
            self.daily_count = state.daily_count;
            self.last_day_utc = state.last_day_utc;
            self.recent_calls = state.recent_calls;
        }
        Ok(())
    }

    fn save_to_disk(&self, path: &Path) -> Result<(), std::io::Error> {
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)?;
        }
        let state = StoredBudgetState {
            daily_count: self.daily_count,
            last_day_utc: self.last_day_utc.clone(),
            recent_calls: self.recent_calls.clone(),
        };
        let serialized = serde_json::to_string_pretty(&state).map_err(std::io::Error::other)?;

        // Atomic write via tempfile in the same parent directory
        let parent_dir = path.parent().unwrap_or_else(|| Path::new("."));
        let mut temp_file = tempfile::NamedTempFile::new_in(parent_dir)?;
        temp_file.write_all(serialized.as_bytes())?;
        temp_file.persist(path).map_err(|e| e.error)?;

        Ok(())
    }
}

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

    #[test]
    fn test_budget_minute_window_and_daily_limit() {
        let mut tracker = BudgetTracker::new(None, CreditTier::LowBalance);

        // Check budget initially
        let status = tracker.check_budget().expect("initial check should pass");
        assert_eq!(status.minute_requests, 0);
        assert_eq!(status.minute_limit, 20);
        assert_eq!(status.daily_requests, 0);
        assert_eq!(status.daily_limit, 50);
        assert!(!status.approaching_daily_limit);

        // Record 19 calls
        for _ in 0..19 {
            tracker.record_call(false).unwrap();
        }
        let status = tracker.check_budget().expect("19 calls should pass");
        assert_eq!(status.minute_requests, 19);

        // 20th call
        tracker.record_call(false).unwrap();
        // 21st check should fail minute limit
        let err = tracker.check_budget().unwrap_err();
        match err {
            BudgetExceeded::MinuteLimit(limit, _) => assert_eq!(limit, 20),
            _ => panic!("Expected MinuteLimit error, got {:?}", err),
        }
    }

    #[test]
    fn test_budget_daily_ceiling_and_approaching_warning() {
        let mut tracker = BudgetTracker::new(None, CreditTier::LowBalance);
        tracker.daily_count = 40; // 80% of 50
        let status = tracker.check_budget().unwrap();
        assert!(status.approaching_daily_limit);

        tracker.daily_count = 50;
        let err = tracker.check_budget().unwrap_err();
        assert_eq!(err, BudgetExceeded::DailyQuota(50, 50));
    }

    #[test]
    fn test_budget_persistence() {
        let temp_dir = tempfile::tempdir().unwrap();
        let file_path = temp_dir.path().join("budget.json");

        {
            let mut tracker = BudgetTracker::new(Some(file_path.clone()), CreditTier::LowBalance);
            tracker.record_call(false).unwrap();
            tracker.record_call(false).unwrap();
        }

        // Reload in new tracker instance
        let mut reloaded = BudgetTracker::new(Some(file_path), CreditTier::LowBalance);
        let status = reloaded.check_budget().unwrap();
        assert_eq!(status.daily_requests, 2);
        assert_eq!(status.minute_requests, 2);
    }
}