// VRAM Budget Tracker — enforces GPU memory allocation limits.
//
// Every buffer and texture allocation goes through the tracker.
// Allocations that would exceed the budget are refused gracefully.
// Per-allocation labels enable debugging VRAM usage breakdowns.

/// VRAM allocation tracker with budget enforcement.
pub struct VramTracker {
    allocated: u64,
    budget: u64,
    peak: u64,
    allocation_count: u32,
}

impl VramTracker {
    /// Create a tracker with the given budget in bytes.
    pub fn new(budget_bytes: u64) -> Self {
        Self {
            allocated: 0,
            budget: budget_bytes,
            peak: 0,
            allocation_count: 0,
        }
    }

    /// Try to allocate `bytes`. Returns true if within budget.
    pub fn try_allocate(&mut self, label: &str, bytes: u64) -> bool {
        if self.allocated + bytes > self.budget {
            log::warn!(
                "VRAM budget refused: {} ({} MB) — would exceed {:.1} GB budget (used: {:.1} MB)",
                label,
                bytes / (1024 * 1024),
                self.budget as f64 / (1024.0 * 1024.0 * 1024.0),
                self.allocated as f64 / (1024.0 * 1024.0)
            );
            return false;
        }
        self.allocated += bytes;
        self.peak = self.peak.max(self.allocated);
        self.allocation_count += 1;
        true
    }

    /// Record an allocation without budget check (for tracking already-made allocations).
    pub fn record(&mut self, bytes: u64) {
        self.allocated += bytes;
        self.peak = self.peak.max(self.allocated);
        self.allocation_count += 1;
    }

    /// Free bytes from tracked allocation.
    pub fn free(&mut self, bytes: u64) {
        self.allocated = self.allocated.saturating_sub(bytes);
    }

    /// Current allocated bytes.
    pub fn allocated(&self) -> u64 {
        self.allocated
    }

    /// Budget in bytes.
    pub fn budget(&self) -> u64 {
        self.budget
    }

    /// Peak allocation in bytes.
    pub fn peak(&self) -> u64 {
        self.peak
    }

    /// Remaining budget in bytes.
    pub fn remaining(&self) -> u64 {
        self.budget.saturating_sub(self.allocated)
    }

    /// Number of tracked allocations.
    pub fn allocation_count(&self) -> u32 {
        self.allocation_count
    }

    /// Format current usage as human-readable string.
    pub fn report(&self) -> String {
        format!(
            "{:.1} MB / {:.1} GB (peak: {:.1} MB, {} allocs)",
            self.allocated as f64 / (1024.0 * 1024.0),
            self.budget as f64 / (1024.0 * 1024.0 * 1024.0),
            self.peak as f64 / (1024.0 * 1024.0),
            self.allocation_count
        )
    }

    /// Format as fraction string for profiler display.
    pub fn usage_display(&self) -> String {
        let used_gb = self.allocated as f64 / (1024.0 * 1024.0 * 1024.0);
        let budget_gb = self.budget as f64 / (1024.0 * 1024.0 * 1024.0);
        format!("{:.2} / {:.1} GB", used_gb, budget_gb)
    }
}

impl Default for VramTracker {
    fn default() -> Self {
        Self::new(8 * 1024 * 1024 * 1024) // 8 GB default
    }
}

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

    #[test]
    fn budget_enforcement() {
        let mut tracker = VramTracker::new(1024);
        assert!(tracker.try_allocate("small", 500));
        assert!(tracker.try_allocate("medium", 500));
        assert!(!tracker.try_allocate("overflow", 100)); // over budget
        assert_eq!(tracker.allocated(), 1000);
    }

    #[test]
    fn free_releases() {
        let mut tracker = VramTracker::new(1024);
        tracker.try_allocate("a", 500);
        tracker.free(500);
        assert!(tracker.try_allocate("b", 800)); // fits after free
    }

    #[test]
    fn peak_tracking() {
        let mut tracker = VramTracker::new(2048);
        tracker.try_allocate("a", 1000);
        tracker.try_allocate("b", 500);
        tracker.free(500);
        assert_eq!(tracker.peak(), 1500);
        assert_eq!(tracker.allocated(), 1000);
    }

    #[test]
    fn remaining_correct() {
        let mut tracker = VramTracker::new(1024);
        tracker.try_allocate("a", 300);
        assert_eq!(tracker.remaining(), 724);
    }

    #[test]
    fn report_format() {
        let tracker = VramTracker::new(8 * 1024 * 1024 * 1024);
        let report = tracker.report();
        assert!(report.contains("GB"));
        assert!(report.contains("allocs"));
    }
}
