use std::ops::Range;

#[derive(Default)]
pub struct Dirty {
    words: Vec<u64>,
    /// Everything is dirty regardless of the bits -- the state after a
    /// buffer reallocation, whose contents are undefined, and the state a
    /// freshly built arena starts in. Kept as a flag rather than by
    /// setting every bit so that it costs nothing to say and cannot be
    /// half-applied as the array grows.
    all: bool,
}

impl Dirty {
    pub fn new_all() -> Self {
        Self {
            words: Vec::new(),
            all: true,
        }
    }

    pub fn mark(&mut self, i: usize) {
        if self.all {
            return;
        }
        let word = i / 64;
        if word >= self.words.len() {
            self.words.resize(word + 1, 0);
        }
        self.words[word] |= 1 << (i % 64);
    }

    pub fn contains(&self, i: usize) -> bool {
        self.all
            || self
                .words
                .get(i / 64)
                .is_some_and(|word| word & (1 << (i % 64)) != 0)
    }

    /// Clear one entry that was restored to the value already on the GPU.
    /// `all` has no per-entry representation and is used only when every
    /// byte must be uploaded regardless of later writes, so it stays set.
    pub fn unmark(&mut self, i: usize) {
        if self.all {
            return;
        }
        if let Some(word) = self.words.get_mut(i / 64) {
            *word &= !(1 << (i % 64));
        }
    }

    /// Everything must be written: the buffer was reallocated (its
    /// contents are undefined), or the array was cleared.
    pub fn mark_all(&mut self) {
        self.all = true;
        self.words.clear();
    }

    pub fn is_clean(&self) -> bool {
        !self.all && self.words.iter().all(|w| *w == 0)
    }

    pub fn ranges(&self, len: usize, gap: usize) -> Vec<Range<usize>> {
        if self.all {
            return Vec::from_iter((len > 0).then_some(0..len));
        }
        let mut ranges: Vec<Range<usize>> = Vec::new();
        for (w, word) in self.words.iter().enumerate() {
            let mut bits = *word;
            while bits != 0 {
                let start = w * 64 + bits.trailing_zeros() as usize;
                let run = (bits >> (start - w * 64)).trailing_ones() as usize;
                let end = (start + run).min(len);
                if start >= len {
                    break;
                }
                match ranges.last_mut() {
                    Some(last) if start - last.end <= gap => last.end = end,
                    _ => ranges.push(start..end),
                }
                bits &= !(((1u128 << run) - 1) as u64) << (start - w * 64);
            }
        }
        ranges
    }

    pub fn clear(&mut self) {
        self.all = false;
        self.words.clear();
    }
}

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

    fn marked(indices: &[usize], len: usize, gap: usize) -> Vec<Range<usize>> {
        let mut d = Dirty::default();
        for &i in indices {
            d.mark(i);
        }
        d.ranges(len, gap)
    }

    #[test]
    fn adjacent_entries_are_one_range() {
        assert_eq!(marked(&[3, 4, 5], 64, 0), vec![3..6]);
    }

    #[test]
    fn a_restored_entry_can_be_unmarked() {
        let mut dirty = Dirty::default();
        dirty.mark(3);
        dirty.mark(5);
        assert!(dirty.contains(3));
        dirty.unmark(3);
        assert!(!dirty.contains(3));
        assert_eq!(dirty.ranges(8, 0), vec![5..6]);
    }

    #[test]
    fn a_run_that_crosses_a_word_boundary_is_one_range() {
        assert_eq!(marked(&[62, 63, 64, 65], 128, 0), vec![62..66]);
    }

    #[test]
    fn a_gap_wider_than_the_threshold_stays_two_ranges() {
        assert_eq!(marked(&[0, 10], 64, 4), vec![0..1, 10..11]);
        assert_eq!(marked(&[0, 10], 64, 16), vec![0..11]);
    }

    #[test]
    fn ranges_stop_at_the_length() {
        assert_eq!(marked(&[1, 2, 40], 3, 0), vec![1..3]);
    }

    #[test]
    fn mark_all_covers_everything_and_survives_later_marks() {
        let mut d = Dirty::new_all();
        d.mark(2);
        assert_eq!(d.ranges(9, 0), vec![0..9]);
        assert!(!d.is_clean());
        d.clear();
        assert!(d.is_clean());
        assert!(d.ranges(9, 0).is_empty());
    }

    #[test]
    fn an_empty_array_has_nothing_to_upload_even_when_all_is_set() {
        assert!(Dirty::new_all().ranges(0, 0).is_empty());
    }
}