use std::{alloc::Layout, cell::UnsafeCell, ops::Deref, ptr::NonNull, sync::atomic::AtomicUsize};

use crate::vec3::Color;

pub struct RenderContextInner {
    width: usize,
    render_id: UnsafeCell<usize>,
    max: usize,
    layout: Layout,
    buffer: UnsafeCell<*mut Color>,
    refs: AtomicUsize,
}

impl RenderContextInner {
    pub fn new(width: usize, height: usize, id: usize) -> Self {
		let layout = Layout::array::<Color>(width * height).unwrap();
		let buffer = UnsafeCell::new(unsafe { std::alloc::alloc_zeroed(layout) } as *mut Color);
		Self {
			width,
			max: width * height,
			layout,
			buffer,
			refs: AtomicUsize::new(1),
			render_id: UnsafeCell::new(id),
		}
	}

	pub fn render_id(&self) -> usize {
		let ptr = self.render_id.get();
		unsafe { std::ptr::read_volatile(ptr) }
	}

	pub fn next(&self) {
		let ptr = self.render_id.get();
		unsafe {
			std::ptr::write_volatile(ptr, std::ptr::read_volatile(ptr) + 1);
		}
	}

	pub fn set_color(&self, x: usize, y: usize, color: Color) {
		let offset = (y * self.width) + x;

		if offset >= self.max {
			panic!("Bad index: {offset} >= {}", self.max);
		}

		unsafe {
			let ptr = *self.buffer.get();
			std::ptr::write(ptr.add(offset), color);
		}
	}

	pub fn get_color_raw_idx(&self, offset: usize) -> Color {
		if offset >= self.max {
			panic!("Bad index: {offset} >= {}", self.max);
		}

		unsafe {
			let ptr = *self.buffer.get();
			std::ptr::read(ptr.add(offset))
		}
	}

	pub fn get_color(&self, x: usize, y: usize) -> Color {
		let offset = (y * self.width) + x;
		self.get_color_raw_idx(offset)
	}
}

#[derive(Debug)]
pub struct RenderContext {
    inner: NonNull<RenderContextInner>,
}

impl Deref for RenderContext {
    type Target = RenderContextInner;
    
    fn deref(&self) -> &Self::Target {
        unsafe { self.inner.as_ref() }
    }
}

impl RenderContext {
	pub fn new(width: usize, height: usize, id: usize) -> Self {
        const LAYOUT: Layout = Layout::new::<RenderContextInner>();
        let inner = NonNull::new(unsafe { std::alloc::alloc_zeroed(LAYOUT) as *mut RenderContextInner }).unwrap();
        unsafe { inner.write_volatile(RenderContextInner::new(width, height, id)); }
		Self { inner }
	}
}

impl Clone for RenderContext {
	fn clone(&self) -> Self {
		if unsafe { self.inner.as_ref() }.refs.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == usize::MAX {
			panic!("RenderContextInner refcount overflowed??");
		}

		unsafe { std::ptr::read_volatile(self as *const Self) }
	}
}

impl Drop for RenderContext {
	fn drop(&mut self) {
        let r = unsafe { self.inner.as_ref() }.refs.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
		if r <= 1 {
            let n = unsafe { self.inner.as_ref() }.refs.load(std::sync::atomic::Ordering::SeqCst);
            println!("Freeing... old={r} new={n}");
            unsafe {
				let ptr_ptr: *mut *mut Color = self.inner.as_ref().buffer.get();
				let ptr = *ptr_ptr;
				std::alloc::dealloc(ptr as *mut u8, self.inner.as_ref().layout);
			}
		}
	}
}

unsafe impl Send for RenderContext {}
unsafe impl Sync for RenderContext {}
