#![feature(const_index)]
#![feature(const_trait_impl)]
#![feature(type_info)]
#![feature(const_heap)]
#![feature(str_from_raw_parts)]
#![feature(ptr_metadata)]
// needed as of now
#![allow(internal_features)]
#![feature(core_intrinsics)]

use std::{
    any::TypeId,
    collections::BTreeSet,
    fmt::Debug,
    intrinsics::{const_allocate, const_make_global, type_id},
    mem::{
        transmute,
        type_info::{DynTrait, DynTraitPredicate, Float, Int, Str, Type, TypeKind},
    },
    ptr::{self, Pointee},
    slice, str,
};

const fn get_rec<T: 'static>() -> TyKindRecursive {
    const {
        let info = Type::of::<T>();
        get_ty_recursive_impl(info)
    }
}

pub trait ReflectionRecursive: 'static {
    fn get_ty_recursive<'a>(&'a self) -> TyKindMapped<'a>
    where
        Self: Sized,
    {
        let rec = get_rec::<Self>();

        rec.map(unsafe { transmute::<_, *const u8>(self) })
    }
}

/// usize is offset
const fn get_ty_recursive_impl(info: Type) -> TyKindRecursive {
    match info.kind {
        TypeKind::Tuple(fields) => {
            let arr: &'static mut [(TyKindRecursive, usize)] = unsafe {
                slice::from_raw_parts_mut(
                    const_allocate(
                        size_of::<(TyKindRecursive, usize)>() * fields.fields.len(),
                        align_of::<(TyKindRecursive, usize)>(),
                    ) as *mut _ as _,
                    fields.fields.len(),
                )
            };

            let mut i = 0;
            while i < fields.fields.len() {
                let field = &fields.fields[i];
                arr[i] = (get_ty_recursive_impl(field.ty.info()), field.offset);
                i += 1;
            }
            unsafe {
                const_make_global(arr as *mut _ as _);
            }
            TyKindRecursive::Tuple(&arr[..fields.fields.len()])
        }
        TypeKind::Slice(slice) => {
            let res: &'static mut TyKindRecursive = unsafe {
                transmute(const_allocate(
                    size_of::<TyKindRecursive>(),
                    align_of::<TyKindRecursive>(),
                ))
            };
            *res = get_ty_recursive_impl(slice.element_ty.info());
            unsafe {
                const_make_global(res as *mut _ as _);
            }
            TyKindRecursive::Slice((res, slice.element_ty.info().size.unwrap()))
        }
        TypeKind::Array(array) => {
            let arr: &'static mut [(TyKindRecursive, usize)] = unsafe {
                slice::from_raw_parts_mut(
                    const_allocate(
                        size_of::<(TyKindRecursive, usize)>() * array.len,
                        align_of::<(TyKindRecursive, usize)>(),
                    ) as *mut _ as _,
                    array.len,
                )
            };

            let mut i = 0;
            let mut offset = 0;
            while i < array.len {
                let field = &array.element_ty;
                let info = field.info();
                let size = info.size.unwrap();
                arr[i] = (get_ty_recursive_impl(info), offset);
                offset += size;
                i += 1;
            }
            unsafe {
                const_make_global(arr as *mut _ as _);
            }
            TyKindRecursive::Array(&arr[..array.len])
        }
        TypeKind::Str(str) => TyKindRecursive::Str(str),
        TypeKind::Float(float) => TyKindRecursive::Float(float),
        TypeKind::Int(int) => TyKindRecursive::Int(int),
        TypeKind::DynTrait(dt) => TyKindRecursive::DynTrait(dt),
        TypeKind::Reference(re) => {
            let res: &'static mut TyKindRecursive = unsafe {
                transmute(const_allocate(
                    size_of::<TyKindRecursive>(),
                    align_of::<TyKindRecursive>(),
                ))
            };
            *res = get_ty_recursive_impl(re.pointee.info());
            unsafe {
                const_make_global(res as *mut _ as _);
            }
            TyKindRecursive::Reference(re.mutable, &*res)
        }
        other => TyKindRecursive::Other(other),
    }
}

#[derive(Debug)]
pub enum TyKindRecursive {
    // usize is offset
    Tuple(&'static [(TyKindRecursive, usize)]),
    // usize is offset
    Array(&'static [(TyKindRecursive, usize)]),
    Int(Int),
    Str(Str),
    DynTrait(DynTrait),
    // usize is size per elem
    Slice((&'static TyKindRecursive, usize)),
    // bool is mutable
    Reference(bool, &'static TyKindRecursive),
    Float(Float),
    Other(TypeKind),
}

impl TyKindRecursive {
    // second param is align
    fn map<'a>(&self, val: *const u8) -> TyKindMapped<'a> {
        match self {
            TyKindRecursive::Tuple(ty_kind_recursives) => {
                let tys: Vec<TyKindMapped> = ty_kind_recursives
                    .into_iter()
                    .map(|field| {
                        let res = field
                            .0
                            .map(unsafe { transmute(val.byte_offset(field.1 as _)) });
                        res
                    })
                    .collect();

                TyKindMapped::Tuple(tys)
            }
            TyKindRecursive::Array(ty_kind_recursives) => {
                let tys: Vec<TyKindMapped> = ty_kind_recursives
                    .into_iter()
                    .map(|field| {
                        let res = field
                            .0
                            .map(unsafe { transmute(val.byte_offset(field.1 as _)) });
                        res
                    })
                    .collect();

                TyKindMapped::Array(tys)
            }
            TyKindRecursive::Int(int) => match (int.signed, int.bits) {
                (false, 8) => {
                    TyKindMapped::IntMapped(IntMapped::U8(*unsafe { transmute::<_, &_>(val) }))
                }

                (false, 16) => {
                    TyKindMapped::IntMapped(IntMapped::U16(*unsafe { transmute::<_, &_>(val) }))
                }
                (false, 32) => {
                    TyKindMapped::IntMapped(IntMapped::U32(*unsafe { transmute::<_, &_>(val) }))
                }

                (false, 64) => {
                    TyKindMapped::IntMapped(IntMapped::U64(*unsafe { transmute::<_, &_>(val) }))
                }
                (false, 128) => {
                    TyKindMapped::IntMapped(IntMapped::U128(*unsafe { transmute::<_, &_>(val) }))
                }
                (true, 8) => {
                    TyKindMapped::IntMapped(IntMapped::I8(*unsafe { transmute::<_, &_>(val) }))
                }
                (true, 16) => {
                    TyKindMapped::IntMapped(IntMapped::I16(*unsafe { transmute::<_, &_>(val) }))
                }
                (true, 32) => {
                    TyKindMapped::IntMapped(IntMapped::I32(*unsafe { transmute::<_, &_>(val) }))
                }
                (true, 64) => {
                    TyKindMapped::IntMapped(IntMapped::I64(*unsafe { transmute::<_, &_>(val) }))
                }
                (true, 128) => {
                    TyKindMapped::IntMapped(IntMapped::I128(*unsafe { transmute::<_, &_>(val) }))
                }
                _ => unreachable!(),
            },
            TyKindRecursive::Float(float) => match float.bits {
                32 => {
                    TyKindMapped::FloatMapped(FloatMapped::F32(*unsafe { transmute::<_, &_>(val) }))
                }
                64 => {
                    TyKindMapped::FloatMapped(FloatMapped::F64(*unsafe { transmute::<_, &_>(val) }))
                }
                _ => unimplemented!(),
            },
            TyKindRecursive::DynTrait(_) => {
                unreachable!()
            }
            TyKindRecursive::Reference(mutable, TyKindRecursive::DynTrait(dt)) => {
                TyKindMapped::DynTraitReference(*mutable, DynTraitMapped(dt.predicates, val))
            }

            TyKindRecursive::Reference(mutable, TyKindRecursive::Slice((pointee, size))) => {
                let mut slice_start = unsafe { *(val as *const *const u8) };
                let length = unsafe { *(val.byte_offset(8) as *const usize) };

                let mut res = vec![];

                for _ in 0..length {
                    res.push(pointee.map(slice_start));
                    slice_start = unsafe { slice_start.byte_offset(*size as _) };
                }

                TyKindMapped::SliceReference(*mutable, res)
            }
            TyKindRecursive::Reference(mutable, TyKindRecursive::Str(..)) => {
                let slice_start = unsafe { *(val as *const *const u8) };
                let length = unsafe { *(val.byte_offset(8) as *const usize) };
                let str = unsafe { str::from_raw_parts(slice_start, length) };
                TyKindMapped::StrReference(*mutable, str)
            }
            TyKindRecursive::Reference(mutable, pointee) => {
                let val = val as *const *const u8;
                TyKindMapped::Reference(*mutable, Box::new((**pointee).map(unsafe { *val })))
            }
            TyKindRecursive::Str(_str) => unreachable!(),
            TyKindRecursive::Slice(_) => unreachable!(),
            TyKindRecursive::Other(kind) => todo!("{kind:?}"),
        }
    }
}

#[derive(Debug, PartialEq)]
pub enum TyKindMapped<'a> {
    IntMapped(IntMapped),
    FloatMapped(FloatMapped),
    Tuple(Vec<TyKindMapped<'a>>),
    Array(Vec<TyKindMapped<'a>>),
    // bool is Mutability
    Reference(bool, Box<TyKindMapped<'a>>),
    // bool is Mutability
    DynTraitReference(bool, DynTraitMapped),
    // bool is Mutability
    StrReference(bool, &'a str),
    SliceReference(bool, Vec<TyKindMapped<'a>>),
}

#[derive(Debug)]
pub struct DynTraitMapped(pub &'static [DynTraitPredicate], *const u8);

fn to_pred_set(dt: &'static [DynTraitPredicate]) -> BTreeSet<TypeId> {
    dt.iter().map(|pred| pred.trait_ty.ty).collect()
}

impl PartialEq for DynTraitMapped {
    fn eq(&self, other: &Self) -> bool {
        to_pred_set(self.0) == to_pred_set(other.0)
    }
}

impl DynTraitMapped {
    pub fn try_as<U: Pointee<Metadata = ptr::DynMetadata<U>> + ?Sized + 'static>(
        &self,
    ) -> Option<&U> {
        if to_pred_set(self.0).contains(&const { type_id::<U>() }) {
            Some(unsafe { *(self.1 as *const &U) })
        } else {
            None
        }
    }
}

#[derive(PartialEq)]
pub enum IntMapped {
    U8(u8),
    U16(u16),
    U32(u32),
    U64(u64),
    U128(u128),
    I8(i8),
    I16(i16),
    I32(i32),
    I64(i64),
    I128(i128),
}

#[derive(PartialEq)]
pub enum FloatMapped {
    F32(f32),
    F64(f64),
}

impl Debug for FloatMapped {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::F32(arg0) => write!(f, "{}", *arg0),
            Self::F64(arg0) => write!(f, "{}", *arg0),
        }
    }
}

impl Debug for IntMapped {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::U8(arg0) => write!(f, "{}", *arg0),
            Self::U16(arg0) => write!(f, "{}", *arg0),
            Self::U32(arg0) => write!(f, "{}", *arg0),
            Self::U64(arg0) => write!(f, "{}", *arg0),
            Self::U128(arg0) => write!(f, "{}", *arg0),
            Self::I8(arg0) => write!(f, "{}", *arg0),
            Self::I16(arg0) => write!(f, "{}", *arg0),
            Self::I32(arg0) => write!(f, "{}", *arg0),
            Self::I64(arg0) => write!(f, "{}", *arg0),
            Self::I128(arg0) => write!(f, "{}", *arg0),
        }
    }
}

impl<T: ?Sized + 'static> ReflectionRecursive for T {}
