use std::collections::{BTreeMap, LinkedList};

pub struct TimeoutSet<T> {
    time_list: BTreeMap<u64, LinkedList<T>>,
}

impl<T> Default for TimeoutSet<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> TimeoutSet<T> {
    pub fn new() -> Self {
        Self {
            time_list: BTreeMap::new(),
        }
    }

    pub fn add(&mut self, time: u64, val: T) {
        match self.time_list.get_mut(&time) {
            Some(list) => {
                list.push_back(val);
            }
            None => {
                let mut list = LinkedList::new();
                list.push_back(val);
                self.time_list.insert(time, list);
            }
        }
    }

    fn get_timeout_keys(&self, time: u64) -> LinkedList<u64> {
        let mut times = LinkedList::new();
        for (t, _) in self.time_list.iter() {
            if *t <= time {
                times.push_back(*t)
            } else {
                break;
            }
        }
        times
    }

    /// only get the timeout value list
    pub fn get_timeout_values(&self, time: u64) -> LinkedList<&T> {
        let mut list = LinkedList::new();
        for (t, values) in self.time_list.iter() {
            if *t <= time {
                for item in values {
                    list.push_back(item);
                }
            } else {
                break;
            }
        }
        list
    }

    /// remove timeout value and return the value list
    pub fn timeout(&mut self, time: u64) -> LinkedList<T> {
        let mut list = LinkedList::new();
        let times = self.get_timeout_keys(time);
        let mut time_list = &mut self.time_list;
        for t in times {
            if let Some(mut values) = time_list.remove(&t) {
                list.append(&mut values)
            }
        }
        list
    }

    pub fn clear(&mut self) {
        self.time_list.clear();
    }
}