use iris::prelude::*;
use std::{
    alloc::{GlobalAlloc, Layout, System},
    cell::Cell,
    time::{Duration, Instant},
};

struct CountingAllocator;

thread_local! {
    static TRACKING: Cell<bool> = const { Cell::new(false) };
    static ALLOCATIONS: Cell<usize> = const { Cell::new(0) };
}

fn note_allocation() {
    TRACKING.with(|tracking| {
        if tracking.get() {
            ALLOCATIONS.with(|allocations| allocations.set(allocations.get() + 1));
        }
    });
}

unsafe impl GlobalAlloc for CountingAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let ptr = unsafe { System.alloc(layout) };
        note_allocation();
        ptr
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        let ptr = unsafe { System.alloc_zeroed(layout) };
        note_allocation();
        ptr
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { System.dealloc(ptr, layout) }
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        let ptr = unsafe { System.realloc(ptr, layout, new_size) };
        note_allocation();
        ptr
    }
}

#[global_allocator]
static ALLOCATOR: CountingAllocator = CountingAllocator;

fn allocations_during(run: impl FnOnce()) -> usize {
    ALLOCATIONS.with(|allocations| allocations.set(0));
    TRACKING.with(|tracking| tracking.set(true));
    run();
    TRACKING.with(|tracking| tracking.set(false));
    ALLOCATIONS.with(Cell::get)
}

struct TestRsc {
    ui: Ui,
}

impl UiRsc for TestRsc {
    fn ui(&self) -> &Ui {
        &self.ui
    }

    fn ui_mut(&mut self) -> &mut Ui {
        &mut self.ui
    }
}

struct FrameRequester;

impl Widget for FrameRequester {
    fn draw(&mut self, painter: &mut Painter) {
        let paint = painter.paint(&PaintId::WHITE);
        painter.primitive(RectPrimitive::color(paint));
        painter.request_next_frame();
        painter.set_size(Size::REST);
    }
}

#[test]
fn a_warmed_self_redraw_allocates_nothing() {
    let mut rsc = TestRsc { ui: Ui::default() };
    let root = rsc.ui.widgets.add_strong(FrameRequester).any();
    let mut render = UiRenderState::new();
    render.resize((800.0, 600.0));
    let start = Instant::now();

    for frame in 0..3 {
        assert!(render.update_at(&root, &mut rsc, start + Duration::from_millis(frame * 16),));
    }

    let allocations = allocations_during(|| {
        assert!(render.update_at(&root, &mut rsc, start + Duration::from_millis(48),));
    });
    assert_eq!(allocations, 0);
}

#[test]
fn refreshing_retained_accessibility_allocates_nothing() {
    let mut rsc = TestRsc { ui: Ui::default() };
    let leaf = rsc.ui.widgets.add_strong(Rect::new(PaintId::WHITE));
    rsc.ui.widgets.set_label(&leaf, "named".to_owned());
    let offset = rsc.ui.widgets.add_strong(Offset {
        inner: leaf.any(),
        amt: UiVec2::ZERO,
    });
    let offset_id = offset.weak();
    let root = offset.any();
    let mut render = UiRenderState::new();
    render.resize((800.0, 600.0));
    render.update(&root, &mut rsc);
    let mut access = AccessTree::new();
    assert!(access.refresh(rsc.widgets(), &render, &rsc));

    rsc.ui.widgets.get_mut(&offset_id).unwrap().amt = UiVec2::abs(Vec2::new(20.0, 0.0));
    render.update(&root, &mut rsc);

    let moved = allocations_during(|| assert!(access.refresh(rsc.widgets(), &render, &rsc)));
    assert_eq!(moved, 0);

    let unchanged = allocations_during(|| assert!(!access.refresh(rsc.widgets(), &render, &rsc)));
    assert_eq!(unchanged, 0);
}