use super::tier::{DEFAULT_CRITICAL_THRESHOLD, DEFAULT_WARNING_THRESHOLD};
use crate::models::{Capabilities, ModelSpec};

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum WindowStatus {
    Ok { utilization: f64, remaining: u64 },
    Warning { utilization: f64, remaining: u64 },
    Critical { utilization: f64, remaining: u64 },
    Exceeded { overage: u64 },
}

impl WindowStatus {
    pub fn should_proceed(&self) -> bool {
        !matches!(self, Self::Exceeded { .. })
    }

    pub fn utilization(&self) -> Option<f64> {
        match self {
            Self::Ok { utilization, .. }
            | Self::Warning { utilization, .. }
            | Self::Critical { utilization, .. } => Some(*utilization),
            Self::Exceeded { .. } => None,
        }
    }
}

#[derive(Debug, Clone)]
pub struct ContextWindow {
    capabilities: Capabilities,
    extended_enabled: bool,
    current_usage: u64,
    peak_usage: u64,
    warning_threshold: f64,
    critical_threshold: f64,
}

impl ContextWindow {
    pub fn new(spec: &ModelSpec, extended_enabled: bool) -> Self {
        Self {
            capabilities: spec.capabilities,
            extended_enabled,
            current_usage: 0,
            peak_usage: 0,
            warning_threshold: DEFAULT_WARNING_THRESHOLD,
            critical_threshold: DEFAULT_CRITICAL_THRESHOLD,
        }
    }

    pub fn limit(&self) -> u64 {
        self.capabilities.effective_context(self.extended_enabled)
    }

    pub fn usage(&self) -> u64 {
        self.current_usage
    }

    pub fn remaining(&self) -> u64 {
        self.limit().saturating_sub(self.current_usage)
    }

    pub fn utilization(&self) -> f64 {
        let limit = self.limit();
        if limit == 0 {
            return 0.0;
        }
        self.current_usage as f64 / limit as f64
    }

    pub fn status(&self) -> WindowStatus {
        let limit = self.limit();
        let utilization = self.utilization();

        if self.current_usage > limit {
            WindowStatus::Exceeded {
                overage: self.current_usage - limit,
            }
        } else if utilization >= self.critical_threshold {
            WindowStatus::Critical {
                utilization,
                remaining: self.remaining(),
            }
        } else if utilization >= self.warning_threshold {
            WindowStatus::Warning {
                utilization,
                remaining: self.remaining(),
            }
        } else {
            WindowStatus::Ok {
                utilization,
                remaining: self.remaining(),
            }
        }
    }

    pub fn can_fit(&self, additional: u64) -> bool {
        self.current_usage + additional <= self.limit()
    }

    pub fn update(&mut self, new_usage: u64) {
        self.current_usage = new_usage;
        if new_usage > self.peak_usage {
            self.peak_usage = new_usage;
        }
    }

    pub fn add(&mut self, tokens: u64) {
        self.update(self.current_usage.saturating_add(tokens));
    }

    pub fn reset(&mut self, new_usage: u64) {
        self.current_usage = new_usage;
    }

    pub fn peak(&self) -> u64 {
        self.peak_usage
    }

    pub fn warning_threshold(&self) -> f64 {
        self.warning_threshold
    }

    pub fn thresholds(mut self, warning: f64, critical: f64) -> Self {
        self.warning_threshold = warning;
        self.critical_threshold = critical;
        self
    }
}

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

    #[test]
    fn test_context_window_status() {
        let reg = registry();
        let spec = reg.resolve("sonnet").unwrap();
        let mut window = ContextWindow::new(spec, false);

        window.update(100_000);
        assert!(matches!(window.status(), WindowStatus::Ok { .. }));

        window.update(180_000);
        assert!(matches!(window.status(), WindowStatus::Warning { .. }));

        window.update(195_000);
        assert!(matches!(window.status(), WindowStatus::Critical { .. }));

        window.update(250_000);
        assert!(matches!(window.status(), WindowStatus::Exceeded { .. }));
    }

    #[test]
    fn test_extended_context() {
        let reg = registry();
        let spec = reg.resolve("sonnet").unwrap();

        let standard = ContextWindow::new(spec, false);
        assert_eq!(standard.limit(), 200_000);

        let extended = ContextWindow::new(spec, true);
        assert_eq!(extended.limit(), 1_000_000);
    }
}
