use lancer_core::bitmap::BitmapAllocator;
use proptest::prelude::*;

#[test]
fn allocate_all_then_free_all() {
    let mut bm: BitmapAllocator<[u64; 2]> = BitmapAllocator::new();
    bm.init(100);
    bm.mark_range_free(0..100);
    assert_eq!(bm.free_items(), 100);

    let indices: Vec<usize> = (0..100)
        .map(|_| bm.allocate().expect("should allocate"))
        .collect();
    assert_eq!(bm.free_items(), 0);
    assert!(bm.allocate().is_none());

    indices.iter().for_each(|&idx| {
        bm.deallocate(idx).expect("should deallocate");
    });
    assert_eq!(bm.free_items(), 100);
}

#[test]
fn double_free_returns_error() {
    let mut bm: BitmapAllocator<[u64; 1]> = BitmapAllocator::new();
    bm.init(64);
    bm.mark_range_free(0..64);

    let idx = bm.allocate().unwrap();
    bm.deallocate(idx).unwrap();
    assert!(bm.deallocate(idx).is_err());
}

#[test]
fn invariant_holds_after_mixed_ops() {
    let mut bm: BitmapAllocator<[u64; 2]> = BitmapAllocator::new();
    bm.init(50);
    bm.mark_range_free(0..50);

    (0..25).for_each(|_| {
        bm.allocate().unwrap();
    });
    assert_eq!(bm.used_items() + bm.free_items(), bm.total_items());
}

#[test]
fn allocated_index_in_range() {
    let mut bm: BitmapAllocator<[u64; 2]> = BitmapAllocator::new();
    bm.init(80);
    bm.mark_range_free(0..80);

    (0..80).for_each(|_| {
        let idx = bm.allocate().unwrap();
        assert!(idx < 80, "index {} out of range", idx);
    });
}

proptest! {
    #[test]
    fn alloc_n_free_permuted(seed in 0u64..10000) {
        let total = 64usize;
        let mut bm: BitmapAllocator<[u64; 1]> = BitmapAllocator::new();
        bm.init(total);
        bm.mark_range_free(0..total);

        let mut indices: Vec<usize> = (0..total)
            .map(|_| bm.allocate().expect("should allocate"))
            .collect();

        let mut rng_state = seed;
        let len = indices.len();
        (0..len).rev().for_each(|i| {
            rng_state = rng_state.wrapping_mul(6364136223846793005).wrapping_add(1);
            let j = (rng_state as usize) % (i + 1);
            indices.swap(i, j);
        });

        indices.iter().for_each(|&idx| {
            bm.deallocate(idx).expect("should deallocate");
        });
        prop_assert_eq!(bm.used_items(), 0);
        prop_assert_eq!(bm.free_items(), total);
    }

    #[test]
    fn alloc_until_none_matches_free_count(free_count in 1usize..65) {
        let mut bm: BitmapAllocator<[u64; 1]> = BitmapAllocator::new();
        bm.init(64);
        bm.mark_range_free(0..free_count);

        let initial_free = bm.free_items();
        let mut allocated = 0usize;
        while bm.allocate().is_some() {
            allocated += 1;
        }
        prop_assert_eq!(allocated, initial_free);
    }

    #[test]
    fn invariant_always_holds(ops in proptest::collection::vec(0u8..3, 1..200)) {
        let mut bm: BitmapAllocator<[u64; 1]> = BitmapAllocator::new();
        bm.init(64);
        bm.mark_range_free(0..64);
        let mut live: Vec<usize> = Vec::new();

        ops.iter().for_each(|&op| {
            match op {
                0 => {
                    bm.allocate().into_iter().for_each(|idx| live.push(idx));
                }
                1 => {
                    live.pop().into_iter().for_each(|idx| {
                        bm.deallocate(idx).unwrap();
                    });
                }
                _ => {}
            }
            assert_eq!(bm.used_items() + bm.free_items(), bm.total_items());
        });
    }
}