use std::{
    collections::{BTreeMap, BTreeSet},
    ops::Bound,
};

#[derive(Debug, Default)]
pub struct AddressSet {
    ranges: BTreeSet<(u32, usize)>,
    bounds: BTreeMap<u32, usize>,
}

impl AddressSet {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn add(&mut self, address: u32, length: usize) {
        if length == 0 || !self.ranges.insert((address, length)) {
            return;
        }
        let end = (address as usize)
            .checked_add(length)
            .and_then(|e| u32::try_from(e).ok());

        if let Some(end) = end {
            let val_before = self.bounds.range(..=end).next_back().map_or(0, |(_, &v)| v);
            self.bounds.insert(end, val_before);
        }

        let val_before = self
            .bounds
            .range(..address)
            .next_back()
            .map_or(0, |(_, &v)| v);
        if let Some(&val_at) = self.bounds.get(&address) {
            if val_before == val_at + 1 {
                self.bounds.remove(&address);
            }
        } else {
            self.bounds.insert(address, val_before);
        }

        let start_bound = Bound::Included(address);
        let end_bound = match end {
            Some(e) => Bound::Excluded(e),
            None => Bound::Unbounded,
        };

        for (_, val) in self.bounds.range_mut((start_bound, end_bound)) {
            *val += 1;
        }
    }

    pub fn remove(&mut self, address: u32, length: usize) {
        if !self.ranges.remove(&(address, length)) {
            return;
        }
        let end = (address as usize)
            .checked_add(length)
            .and_then(|e| u32::try_from(e).ok());

        if let Some(end) = end {
            let val_before = self.bounds.range(..end).next_back().map_or(0, |(_, &v)| v);
            if let Some(&val_at) = self.bounds.get(&end) {
                if val_at + 1 == val_before {
                    self.bounds.remove(&end);
                }
            } else {
                self.bounds.insert(end, val_before);
            }
        }

        let val_before = self
            .bounds
            .range(..address)
            .next_back()
            .map_or(0, |(_, &v)| v);
        if let Some(&val_at) = self.bounds.get(&address) {
            if val_before + 1 == val_at {
                self.bounds.remove(&address);
            }
        } else {
            self.bounds.insert(address, val_before);
        }

        let start_bound = Bound::Included(address);
        let end_bound = match end {
            Some(e) => Bound::Excluded(e),
            None => Bound::Unbounded,
        };

        for (_, val) in self.bounds.range_mut((start_bound, end_bound)) {
            *val -= 1;
        }
    }

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

    pub fn is_empty(&self) -> bool {
        self.bounds.is_empty()
    }

    pub fn contains(&self, address: u32) -> bool {
        self.bounds
            .range(..=address)
            .next_back()
            .is_some_and(|(_, &val)| val > 0)
    }

    pub fn start_of_range_containing(&self, address: u32) -> Option<u32> {
        if !self.contains(address) {
            return None;
        }
        self.ranges
            .range(..=(address, usize::MAX))
            .rev()
            .find_map(|&(start, length)| {
                let contains = start <= address
                    && (start as usize)
                        .checked_add(length)
                        .is_none_or(|end| end > address as usize);
                contains.then_some(start)
            })
    }
}

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

    #[test]
    fn should_not_include_addresses_when_empty() {
        let set = AddressSet::new();
        assert!(set.is_empty());
        assert!(!set.contains(0x13374200));
        assert_eq!(set.start_of_range_containing(0x13374200), None);
    }

    #[test]
    fn should_include_addresses_when_full() {
        let mut set = AddressSet::new();
        set.add(0x00000000, 0x100000000);
        assert!(set.contains(0x13374200));
        assert_eq!(set.start_of_range_containing(0x13374200), Some(0x00000000));
    }

    #[test]
    fn should_ignore_empty_address_ranges() {
        let mut set = AddressSet::new();
        set.add(0x13374200, 0);
        assert!(set.is_empty());
        assert!(!set.contains(0x13374200));
    }

    #[test]
    fn should_add_addresses_idempotently() {
        let mut set = AddressSet::new();
        set.add(0x13374200, 1);
        set.add(0x13374200, 1);
        set.remove(0x13374200, 1);
        assert!(set.is_empty());
        assert!(!set.contains(0x13374200));
    }

    #[test]
    fn should_remove_addresses_idempotently() {
        let mut set = AddressSet::new();
        set.add(0x13374200, 1);
        set.remove(0x13374200, 1);
        set.remove(0x13374200, 1);
        assert!(set.is_empty());
        assert!(!set.contains(0x13374200));
    }

    #[test]
    fn should_report_address_in_range() {
        let mut set = AddressSet::new();
        set.add(0x13374200, 4);
        assert!(!set.contains(0x133741ff));
        for address in 0x13374200..0x13374204 {
            assert!(set.contains(address));
        }
        assert!(!set.contains(0x13374204));
    }

    #[test]
    fn should_allow_overlapping_addresses() {
        let mut set = AddressSet::new();
        set.add(0x13374200, 4);
        set.add(0x13374201, 1);
        set.add(0x13374202, 2);

        assert!(!set.contains(0x133741ff));
        for address in 0x13374200..0x13374204 {
            assert!(set.contains(address));
        }
        assert!(!set.contains(0x13374204));

        set.remove(0x13374200, 4);
        assert!(!set.contains(0x13374200));
        for address in 0x13374201..0x13374204 {
            assert!(set.contains(address));
        }
        assert!(!set.contains(0x13374204));
    }

    #[test]
    fn should_allow_removing_overlapped_address_ranges() {
        let mut set = AddressSet::new();
        set.add(0x13374200, 8);
        set.add(0x13374204, 8);
        set.remove(0x13374204, 8);

        for address in 0x13374200..0x13374208 {
            assert!(set.contains(address));
        }
        for address in 0x13374208..0x1337420c {
            assert!(!set.contains(address));
        }
    }

    #[test]
    fn should_merge_adjacent_ranges() {
        let mut set = AddressSet::new();
        set.add(0x13374200, 4);
        set.add(0x13374204, 4);
        set.add(0x13374208, 4);
        set.add(0x1337420c, 4);

        assert!(!set.contains(0x133741ff));
        for address in 0x13374200..0x13374210 {
            assert!(set.contains(address));
        }
        assert!(!set.contains(0x13374210));

        set.remove(0x13374200, 4);
        set.remove(0x13374204, 4);
        set.remove(0x13374208, 4);
        set.remove(0x1337420c, 4);

        assert!(set.is_empty());
        for address in 0x133741ff..=0x13374210 {
            assert!(!set.contains(address));
        }
    }

    #[test]
    fn should_find_start_of_range() {
        let mut set = AddressSet::new();
        set.add(0x13374200, 4);
        assert_eq!(set.start_of_range_containing(0x133741ff), None);
        for address in 0x13374200..0x13374204 {
            assert_eq!(set.start_of_range_containing(address), Some(0x13374200));
        }
        assert_eq!(set.start_of_range_containing(0x13374204), None);
    }

    #[test]
    fn should_ignore_ranges_not_containing_address() {
        let mut set = AddressSet::new();
        set.add(0x10000000, 1024);
        set.add(0x30000000, 1024);

        assert!(!set.contains(0x20000000));
        assert_eq!(set.start_of_range_containing(0x20000000), None);
    }
}