use rudo_gc::heap::{page_size, with_heap};
use rudo_gc::Gc;
use std::thread;

#[test]
fn test_tlab_thread_isolation() {
    use std::sync::{Arc, Barrier};
    let barrier = Arc::new(Barrier::new(2));

    let b1 = Arc::clone(&barrier);
    let t1 = thread::spawn(move || {
        // Allocate a small object to trigger TLAB initialization/page allocation
        let _g1 = Gc::new(42i32);
        let addr = with_heap(|h| {
            assert!(!h.pages.is_empty());
            // Return the first page address
            h.pages[0].as_ptr() as usize
        });
        b1.wait();
        addr
    });

    let b2 = Arc::clone(&barrier);
    let t2 = thread::spawn(move || {
        // Allocate a small object
        let _g2 = Gc::new(43i32);
        let addr = with_heap(|h| {
            assert!(!h.pages.is_empty());
            // Return the first page address
            h.pages[0].as_ptr() as usize
        });
        b2.wait();
        addr
    });

    let addr1 = t1.join().unwrap();
    let addr2 = t2.join().unwrap();

    assert_ne!(
        addr1, addr2,
        "Each thread should have its own TLAB and unique pages"
    );
}

#[test]
fn test_tlab_exhaustion_allocates_new_page() {
    // We want to exhaust a TLAB for a specific size class.
    // GcBox<i32> fits in the 64-byte size class.
    // We calculate the number of objects dynamically based on the page size.

    let size_class = 64;
    let header_size = rudo_gc::heap::PageHeader::header_size(size_class);
    let objects_per_page = (page_size() - header_size) / size_class;
    let count = objects_per_page + 10; // Ensure we need more than one page

    thread::spawn(move || {
        let initial_pages = with_heap(|h| h.pages.len());
        assert_eq!(initial_pages, 0, "New thread should have 0 pages initially");

        let mut pointers = Vec::new();
        for i in 0..count {
            #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
            pointers.push(Gc::new(i as i32));
        }

        let final_pages = with_heap(|h| h.pages.len());
        assert!(
            final_pages >= 2,
            "TLAB exhaustion should have forced at least one new page allocation (total >= 2), got {final_pages}"
        );

        // Verify they are not all in the same page
        let p1 = Gc::as_ptr(&pointers[0]) as usize;
        let p_last = Gc::as_ptr(&pointers[count - 1]) as usize;
        assert_ne!(
            p1 & !(page_size() - 1),
            p_last & !(page_size() - 1),
            "First and last objects should be in different pages"
        );
    }).join().unwrap();
}

#[test]
fn test_tlab_bump_pointer_contiguity() {
    // Sequential allocations in the same TLAB should be contiguous (within the same size class)
    thread::spawn(|| {
        let g1 = Gc::new(1i32);
        let g2 = Gc::new(2i32);
        let g3 = Gc::new(3i32);

        let p1 = Gc::as_ptr(&g1) as usize;
        let p2 = Gc::as_ptr(&g2) as usize;
        let p3 = Gc::as_ptr(&g3) as usize;

        // All should be in the same page
        let page = p1 & !(page_size() - 1);
        assert_eq!(p2 & !(page_size() - 1), page);
        assert_eq!(p3 & !(page_size() - 1), page);

        // Distance between them should be constant (the size class)
        let diff1 = p2 - p1;
        let diff2 = p3 - p2;

        assert_eq!(
            diff1, diff2,
            "Bump pointer should allocate with constant stride"
        );
        // For i32, it's 64 bytes as calculated before.
        assert_eq!(diff1, 64, "Expected 64-byte size class for GcBox<i32>");
    })
    .join()
    .unwrap();
}

#[test]
fn test_mixed_size_class_tlabs() {
    // Different size classes should use different TLABs (and different pages)
    thread::spawn(|| {
        use rudo_gc::Trace;

        #[derive(Trace)]
        struct Large {
            data: [u8; 100],
        }

        let g_small = Gc::new(()); // 32 bytes -> 32 size class
        let g_large = Gc::new(Large { data: [0; 100] }); // 32 + 100 = 132 bytes -> 256 size class

        let p_small = Gc::as_ptr(&g_small) as usize;
        let p_large = Gc::as_ptr(&g_large) as usize;

        assert_ne!(
            p_small & !(page_size() - 1),
            p_large & !(page_size() - 1),
            "Different size classes must be in different pages"
        );

        with_heap(|h| {
            assert!(
                h.pages.len() >= 2,
                "Should have allocated at least 2 pages for 2 different size classes"
            );
        });
    })
    .join()
    .unwrap();
}

#[test]
fn test_free_slot_reuse() {
    #[inline(never)]
    fn clear_stack() {
        let mut x = [0u64; 1024];
        std::hint::black_box(&mut x);
    }

    // We want to verify that slots reclaimed by GC are reused for new allocations
    // instead of always allocating new pages.

    let size_class = 64;
    let header_size = rudo_gc::heap::PageHeader::header_size(size_class);
    let objects_per_page = (page_size() - header_size) / size_class;
    let count = objects_per_page + 10;

    thread::spawn(move || {
        let initial_pages = {
            let mut pointers = Vec::new();
            for i in 0..count {
                pointers.push(Gc::new(i));
            }
            let len = with_heap(|h| h.pages.len());
            drop(pointers);
            len
        };

        assert!(
            initial_pages >= 2,
            "Should have at least 2 pages for {count} objects of {size_class} bytes"
        );

        // Clear stack to remove stale Gc pointers
        clear_stack();

        // Force collection.
        // The first set of 'pointers' is already dropped as it was local to the block above.
        rudo_gc::collect_full();

        // Allocate more objects
        let mut pointers = Vec::new();
        for i in 0..count {
            pointers.push(Gc::new(i));
        }

        let final_pages = with_heap(|h| h.pages.len());

        assert!(
            final_pages <= initial_pages + 1,
            "Page count increased from {initial_pages} to {final_pages} indicating no reuse of free slots"
        );
        drop(pointers);
    })
    .join()
    .unwrap();
}
