dynamic_grounding_for_github_copilot 0.1.0

use crate::error::{Error, Result};
use chrono::{DateTime, Duration, Utc};
use parking_lot::RwLock;
use std::collections::VecDeque;

/// Rolling window for tracking events within a time period
#[derive(Debug)]
pub struct RollingWindow {
    events: VecDeque<DateTime<Utc>>,
    window_duration: Duration,
    limit: usize,
}

impl RollingWindow {
    pub fn new(window_duration: Duration, limit: usize) -> Self {
        Self {
            events: VecDeque::new(),
            window_duration,
            limit,
        }
    }

    /// Remove events outside the rolling window
    fn cleanup(&mut self) {
        let cutoff = Utc::now() - self.window_duration;
        while let Some(&earliest) = self.events.front() {
            if earliest < cutoff {
                self.events.pop_front();
            } else {
                break;
            }
        }
    }

    /// Add a new event to the window
    pub fn add_event(&mut self) {
        self.cleanup();
        self.events.push_back(Utc::now());
    }

    /// Get current count of events in the window
    pub fn count(&mut self) -> usize {
        self.cleanup();
        self.events.len()
    }

    /// Get remaining capacity before hitting limit
    pub fn remaining(&mut self) -> usize {
        self.limit.saturating_sub(self.count())
    }

    /// Get usage percentage (0.0 to 1.0)
    pub fn usage_percentage(&mut self) -> f64 {
        let count = self.count();
        if self.limit == 0 {
            0.0
        } else {
            count as f64 / self.limit as f64
        }
    }

    /// Check if adding an event would exceed the limit
    pub fn would_exceed(&mut self) -> bool {
        self.count() >= self.limit
    }
}

/// Tracks API usage across multiple quota dimensions
#[derive(Debug)]
pub struct QuotaTracker {
    requests_per_minute: RwLock<RollingWindow>,
    requests_per_day: RwLock<RollingWindow>,
    tokens_per_minute: RwLock<RollingWindow>,
}

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

impl QuotaTracker {
    /// Create a new quota tracker with Gemini free tier limits
    pub fn new() -> Self {
        Self {
            // 15 requests per minute
            requests_per_minute: RwLock::new(RollingWindow::new(Duration::minutes(1), 15)),
            // 1500 requests per day
            requests_per_day: RwLock::new(RollingWindow::new(Duration::days(1), 1500)),
            // 1,000,000 tokens per minute (track requests as proxy)
            tokens_per_minute: RwLock::new(RollingWindow::new(Duration::minutes(1), 1_000_000)),
        }
    }

    /// Check if a request can be made without exceeding quotas
    pub fn check_quota(&self) -> Result<QuotaStatus> {
        let rpm = self.requests_per_minute.write().usage_percentage();
        let rpd = self.requests_per_day.write().usage_percentage();
        let tpm = self.tokens_per_minute.write().usage_percentage();

        let max_usage = rpm.max(rpd).max(tpm);

        // Block at 95% usage
        if max_usage >= 0.95 {
            return Err(Error::QuotaExceeded(format!(
                "Quota usage at {:.1}% (RPM: {:.1}%, RPD: {:.1}%, TPM: {:.1}%)",
                max_usage * 100.0,
                rpm * 100.0,
                rpd * 100.0,
                tpm * 100.0
            )));
        }

        let status = QuotaStatus {
            rpm_usage: rpm,
            rpd_usage: rpd,
            tpm_usage: tpm,
            warning: max_usage >= 0.80,
            remaining_rpm: self.requests_per_minute.write().remaining(),
            remaining_rpd: self.requests_per_day.write().remaining(),
        };

        Ok(status)
    }

    /// Record a request and token usage
    pub fn record_request(&self, token_count: usize) {
        self.requests_per_minute.write().add_event();
        self.requests_per_day.write().add_event();

        // Add token count as individual events (simplified tracking)
        let mut tpm = self.tokens_per_minute.write();
        for _ in 0..token_count.min(1000) {
            // Cap at 1000 to prevent extreme memory usage
            tpm.add_event();
        }
    }

    /// Get current quota status without checking limits
    pub fn get_status(&self) -> QuotaStatus {
        let rpm = self.requests_per_minute.write().usage_percentage();
        let rpd = self.requests_per_day.write().usage_percentage();
        let tpm = self.tokens_per_minute.write().usage_percentage();

        QuotaStatus {
            rpm_usage: rpm,
            rpd_usage: rpd,
            tpm_usage: tpm,
            warning: rpm.max(rpd).max(tpm) >= 0.80,
            remaining_rpm: self.requests_per_minute.write().remaining(),
            remaining_rpd: self.requests_per_day.write().remaining(),
        }
    }
}

/// Current quota usage status
#[derive(Debug, Clone)]
pub struct QuotaStatus {
    pub rpm_usage: f64,
    pub rpd_usage: f64,
    pub tpm_usage: f64,
    pub warning: bool,
    pub remaining_rpm: usize,
    pub remaining_rpd: usize,
}

impl QuotaStatus {
    /// Format quota status for display in MCP responses
    pub fn format_message(&self) -> String {
        let warning_prefix = if self.warning {
            "⚠️ WARNING: High quota usage! "
        } else {
            ""
        };

        format!(
            "{}Gemini API Quota: {:.1}% used (RPM: {}/{}, RPD: {}/1500)",
            warning_prefix,
            self.rpm_usage.max(self.rpd_usage).max(self.tpm_usage) * 100.0,
            15 - self.remaining_rpm,
            15,
            1500 - self.remaining_rpd
        )
    }
}

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

    #[test]
    fn test_rolling_window_basic() {
        let mut window = RollingWindow::new(Duration::seconds(1), 5);

        assert_eq!(window.count(), 0);
        assert_eq!(window.remaining(), 5);
        assert!(!window.would_exceed());

        window.add_event();
        assert_eq!(window.count(), 1);
        assert_eq!(window.remaining(), 4);
    }

    #[test]
    fn test_rolling_window_limit() {
        let mut window = RollingWindow::new(Duration::seconds(1), 3);

        window.add_event();
        window.add_event();
        window.add_event();

        assert_eq!(window.count(), 3);
        assert!(window.would_exceed());
    }

    #[test]
    fn test_quota_tracker_initialization() {
        let tracker = QuotaTracker::new();
        let status = tracker.get_status();

        assert_eq!(status.rpm_usage, 0.0);
        assert_eq!(status.rpd_usage, 0.0);
        assert!(!status.warning);
    }

    #[test]
    fn test_quota_tracker_record() {
        let tracker = QuotaTracker::new();

        tracker.record_request(100);
        let status = tracker.get_status();

        assert!(status.rpm_usage > 0.0);
        assert!(status.rpd_usage > 0.0);
    }

    #[test]
    fn test_quota_status_formatting() {
        let status = QuotaStatus {
            rpm_usage: 0.45,
            rpd_usage: 0.30,
            tpm_usage: 0.20,
            warning: false,
            remaining_rpm: 8,
            remaining_rpd: 1050,
        };

        let message = status.format_message();
        assert!(message.contains("45.0%"));
        assert!(message.contains("7/15"));
        assert!(message.contains("450/1500"));
    }

    #[test]
    fn test_quota_warning() {
        let status = QuotaStatus {
            rpm_usage: 0.85,
            rpd_usage: 0.70,
            tpm_usage: 0.60,
            warning: true,
            remaining_rpm: 2,
            remaining_rpd: 450,
        };

        let message = status.format_message();
        assert!(message.contains("⚠️ WARNING"));
        assert!(message.contains("85.0%"));
    }
}