//! Timer wheel for scheduling periodic and one-shot
//! callbacks without threads or async.
//!
//! The event loop calls `tick()` each iteration to
//! fire expired timers.

use std::collections::{BTreeMap, HashMap};
use std::time::{Duration, Instant};

/// Opaque timer identifier.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct TimerId(u64);

/// A simple timer wheel driven by the event loop.
///
/// Timers are stored in a BTreeMap keyed by deadline,
/// which gives O(log n) insert/cancel and efficient
/// scanning of expired timers.
pub struct TimerWheel {
    deadlines: BTreeMap<Instant, Vec<TimerId>>,
    intervals: HashMap<TimerId, Duration>,
    next_id: u64,
}

impl TimerWheel {
    pub fn new() -> Self {
        Self {
            deadlines: BTreeMap::new(),
            intervals: HashMap::new(),
            next_id: 0,
        }
    }

    /// Schedule a one-shot timer that fires after `delay`.
    pub fn schedule(&mut self, delay: Duration) -> TimerId {
        let id = self.alloc_id();
        let deadline = Instant::now() + delay;
        self.deadlines.entry(deadline).or_default().push(id);
        id
    }

    /// Schedule a repeating timer that fires every
    /// `interval`, starting after one interval.
    pub fn schedule_repeating(&mut self, interval: Duration) -> TimerId {
        let id = self.alloc_id();
        let deadline = Instant::now() + interval;
        self.deadlines.entry(deadline).or_default().push(id);
        self.intervals.insert(id, interval);
        id
    }

    /// Cancel a pending timer.
    ///
    /// Returns `true` if the timer was found and removed.
    pub fn cancel(&mut self, id: TimerId) -> bool {
        self.intervals.remove(&id);

        let mut found = false;

        // Remove from deadline map
        let mut empty_keys = Vec::new();
        for (key, ids) in self.deadlines.iter_mut() {
            if let Some(pos) = ids.iter().position(|i| *i == id) {
                ids.swap_remove(pos);
                found = true;
                if ids.is_empty() {
                    empty_keys.push(*key);
                }
                break;
            }
        }
        for k in empty_keys {
            self.deadlines.remove(&k);
        }
        found
    }

    /// Fire all expired timers, returning their IDs.
    ///
    /// Repeating timers are automatically rescheduled.
    /// Call this once per event loop iteration.
    pub fn tick(&mut self) -> Vec<TimerId> {
        let now = Instant::now();
        let mut fired = Vec::new();

        // Collect all deadlines <= now
        let expired: Vec<Instant> =
            self.deadlines.range(..=now).map(|(k, _)| *k).collect();

        for deadline in expired {
            if let Some(ids) = self.deadlines.remove(&deadline) {
                for id in ids {
                    fired.push(id);

                    // Reschedule if repeating
                    if let Some(&interval) = self.intervals.get(&id) {
                        let next = now + interval;
                        self.deadlines.entry(next).or_default().push(id);
                    }
                }
            }
        }

        fired
    }

    /// Duration until the next timer fires, or `None`
    /// if no timers are scheduled.
    ///
    /// Useful for determining the poll timeout.
    pub fn next_deadline(&self) -> Option<Duration> {
        self.deadlines.keys().next().map(|deadline| {
            let now = Instant::now();
            if *deadline <= now {
                Duration::ZERO
            } else {
                *deadline - now
            }
        })
    }

    /// Number of pending timers.
    pub fn pending_count(&self) -> usize {
        self.deadlines.values().map(|v| v.len()).sum()
    }

    /// Check if there are no pending timers.
    pub fn is_empty(&self) -> bool {
        self.deadlines.is_empty()
    }

    fn alloc_id(&mut self) -> TimerId {
        let id = TimerId(self.next_id);
        self.next_id += 1;
        id
    }
}

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

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

    #[test]
    fn schedule_and_tick() {
        let mut tw = TimerWheel::new();
        let id = tw.schedule(Duration::from_millis(10));

        // Not yet expired
        assert!(tw.tick().is_empty());
        assert_eq!(tw.pending_count(), 1);

        // Wait and tick
        thread::sleep(Duration::from_millis(15));
        let fired = tw.tick();
        assert_eq!(fired, vec![id]);
        assert!(tw.is_empty());
    }

    #[test]
    fn cancel_timer() {
        let mut tw = TimerWheel::new();
        let id = tw.schedule(Duration::from_millis(100));
        assert!(tw.cancel(id));
        assert!(tw.is_empty());

        // Cancel non-existent returns false
        assert!(!tw.cancel(TimerId(999)));
    }

    #[test]
    fn repeating_timer() {
        let mut tw = TimerWheel::new();
        let id = tw.schedule_repeating(Duration::from_millis(10));

        thread::sleep(Duration::from_millis(15));
        let fired = tw.tick();
        assert_eq!(fired, vec![id]);

        // Should be rescheduled, not empty
        assert_eq!(tw.pending_count(), 1);

        // Cancel the repeating timer
        tw.cancel(id);
        assert!(tw.is_empty());
    }

    #[test]
    fn next_deadline_empty() {
        let tw = TimerWheel::new();
        assert!(tw.next_deadline().is_none());
    }

    #[test]
    fn next_deadline_returns_duration() {
        let mut tw = TimerWheel::new();
        tw.schedule(Duration::from_secs(10));
        let d = tw.next_deadline().unwrap();
        assert!(d <= Duration::from_secs(10));
        assert!(d > Duration::from_secs(9));
    }

    #[test]
    fn multiple_timers_fire_in_order() {
        let mut tw = TimerWheel::new();
        let a = tw.schedule(Duration::from_millis(5));
        let b = tw.schedule(Duration::from_millis(10));

        thread::sleep(Duration::from_millis(15));
        let fired = tw.tick();
        assert!(fired.contains(&a));
        assert!(fired.contains(&b));
        assert!(tw.is_empty());
    }
}