use fxhash::FxBuildHasher;
use leaf_assembly::types::{Type, derivations::PtrT, intrinsics::IntT};
use scc::HashMap;
use std::{
	alloc::Layout,
	any::Any,
	marker::PhantomData,
	ops::Range,
	sync::{Arc, OnceLock},
};

#[derive(Debug)]
pub struct LayoutElement {
	pub offset: usize,
	pub layout: Arc<TypeLayout>,
}

impl LayoutElement {
	pub fn range(&self) -> Range<usize> {
		self.offset..self.offset + self.layout.layout.size()
	}
}

#[derive(Debug)]
pub struct TypeLayout {
	pub layout: Layout,
	pub elements: Vec<(usize, Arc<TypeLayout>)>,
}

#[derive(Default)]
pub struct LayoutCache<'l> {
	phantom: PhantomData<&'l ()>,
	layouts: HashMap<*const u8, Arc<dyn Any>, FxBuildHasher>,
}

pub trait GetLayout<T> {
	type LayoutT;
	fn get_layout(&self, value: T) -> Arc<Self::LayoutT>;
}

impl<'l> GetLayout<Type<'_>> for LayoutCache<'l> {
	type LayoutT = TypeLayout;

	fn get_layout(&self, ty: Type<'_>) -> Arc<Self::LayoutT> {
		match ty {
			Type::Void(_) => {
				static LAYOUT: OnceLock<Arc<TypeLayout>> = OnceLock::new();
				LAYOUT
					.get_or_init(|| {
						Arc::new(TypeLayout {
							layout: Layout::new::<()>(),
							elements: Vec::new(),
						})
					})
					.clone()
			}

			Type::Bool(_) => {
				static LAYOUT: OnceLock<Arc<TypeLayout>> = OnceLock::new();
				LAYOUT
					.get_or_init(|| {
						Arc::new(TypeLayout {
							layout: Layout::new::<bool>(),
							elements: Vec::new(),
						})
					})
					.clone()
			}

			Type::Int(IntT {
				signed: false,
				precision: 32,
				..
			}) => {
				static LAYOUT: OnceLock<Arc<TypeLayout>> = OnceLock::new();
				LAYOUT
					.get_or_init(|| {
						Arc::new(TypeLayout {
							layout: Layout::new::<u32>(),
							elements: Vec::new(),
						})
					})
					.clone()
			}

			Type::Ptr(PtrT { .. }) => {
				static LAYOUT: OnceLock<Arc<TypeLayout>> = OnceLock::new();
				LAYOUT
					.get_or_init(|| {
						Arc::new(TypeLayout {
							layout: Layout::new::<*const u8>(),
							elements: Vec::new(),
						})
					})
					.clone()
			}

			_ => todo!("Unsupported type {ty}"),
		}
	}
}