use std::sync::{Condvar, Mutex};

/// Format a byte count with a binary unit (B/KiB/MiB/GiB), one decimal place.
pub fn human_bytes(bytes: u64) -> String {
    const KIB: u64 = 1024;
    const MIB: u64 = KIB * 1024;
    const GIB: u64 = MIB * 1024;
    if bytes >= GIB {
        format!("{:.1} GiB", bytes as f64 / GIB as f64)
    } else if bytes >= MIB {
        format!("{:.1} MiB", bytes as f64 / MIB as f64)
    } else if bytes >= KIB {
        format!("{:.1} KiB", bytes as f64 / KIB as f64)
    } else {
        format!("{bytes} B")
    }
}

/// A counting semaphore measured in bytes. `acquire(n)` blocks until `n` bytes
/// are free, then reserves them; `release(n)` returns them and wakes waiters.
/// A single request larger than the total is clamped to the total so it can
/// still proceed (callers gate truly-oversized work with a skip rule upstream).
pub struct MemoryBudget {
    total: u64,
    used: Mutex<u64>,
    cv: Condvar,
}

impl MemoryBudget {
    pub fn new(total: u64) -> Self {
        MemoryBudget { total: total.max(1), used: Mutex::new(0), cv: Condvar::new() }
    }

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

    pub fn used(&self) -> u64 {
        *self.used.lock().unwrap()
    }

    /// Block until `n` bytes are available, then reserve them.
    pub fn acquire(&self, n: u64) {
        let n = n.min(self.total);
        let mut used = self.used.lock().unwrap();
        while *used + n > self.total {
            used = self.cv.wait(used).unwrap();
        }
        *used += n;
    }

    /// Return `n` previously-acquired bytes and wake any waiters.
    pub fn release(&self, n: u64) {
        let n = n.min(self.total);
        let mut used = self.used.lock().unwrap();
        *used = used.saturating_sub(n);
        drop(used);
        self.cv.notify_all();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::mpsc::channel;
    use std::sync::Arc;
    use std::thread;

    #[test]
    fn human_bytes_picks_a_sensible_unit() {
        assert_eq!(human_bytes(512), "512 B");
        assert_eq!(human_bytes(2048), "2.0 KiB");
        assert_eq!(human_bytes(3 * 1024 * 1024), "3.0 MiB");
        assert_eq!(human_bytes(2 * 1024 * 1024 * 1024), "2.0 GiB");
    }

    #[test]
    fn tracks_used_bytes() {
        let mb = MemoryBudget::new(100);
        mb.acquire(30);
        assert_eq!(mb.used(), 30);
        mb.acquire(20);
        assert_eq!(mb.used(), 50);
        mb.release(30);
        assert_eq!(mb.used(), 20);
    }

    #[test]
    fn oversized_request_is_clamped_to_total() {
        let mb = MemoryBudget::new(100);
        mb.acquire(1000); // clamped to 100, must not hang
        assert_eq!(mb.used(), 100);
    }

    #[test]
    fn release_wakes_a_blocked_acquire() {
        let mb = Arc::new(MemoryBudget::new(100));
        mb.acquire(100); // budget now full

        let mb2 = mb.clone();
        let (tx, rx) = channel();
        let h = thread::spawn(move || {
            mb2.acquire(50); // blocks until main releases
            tx.send(()).unwrap();
            mb2.release(50);
        });

        // Free the budget; the blocked thread must now proceed.
        mb.release(100);
        rx.recv().unwrap(); // proves the waiter was woken (no deadlock)
        h.join().unwrap();
    }
}
