use core::{any::TypeId, mem::MaybeUninit};
use std::ptr::{addr_of_mut, addr_of};

use crate::{never::Never, ValidAlignment, Alignment, align::AlignedBuffer};

fn never_type_id() -> TypeId {
    core::any::TypeId::of::<Never>()
}

/// Core of a dungeon primitive.
///
/// Any static type with size less than or equal to `SIZE` bytes and alignment
/// requirements less than or equal to `2^ALIGN` bytes can be stored in the core.
///
/// By default, this type uses a `8` byte alignment as that
/// is the most common alignment for x64 machines.
pub struct DungeonCore<const SIZE: usize, const ALIGN: u8 = 3>
where
    Alignment<ALIGN>: ValidAlignment,
{
    /// The type of the value `data` contains.
    ///
    /// If `current` is the TypeId of a uninhabited type
    /// then `data` must not be used to make an instance of it.
    current: TypeId,

    /// A function to cause the value in `data` to be dropped.
    drop_impl: fn(&mut DungeonCore<SIZE, ALIGN>),

    /// Value of whatever type `current` says.
    ///
    /// If `current` is a uninhabited type then `data` cannot be used
    /// to make an instance of it.
    buffer: AlignedBuffer<SIZE, ALIGN>,
}

impl<const SIZE: usize, const ALIGN: u8> core::fmt::Debug for DungeonCore<SIZE, ALIGN> 
where
    Alignment<ALIGN>: ValidAlignment,
{
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("DungeonCore")
         .field("current_type", &if self.is_empty() { None } else { Some(self.current) })
         .finish()
    }
}


/// Drop the value inside a `DungeonCore` by using the `take` method.
/// Then immediately dropping the return. If the dungeon core was empty
/// then this does nothing.
fn drop_t<T, const SIZE: usize, const ALIGN: u8>(core: &mut DungeonCore<SIZE, ALIGN>)
where
    T: 'static,
    Alignment<ALIGN>: ValidAlignment,
{
    let _ = core.take::<T>();
}

impl<const SIZE: usize, const ALIGN: u8> DungeonCore<SIZE, ALIGN> 
where
    Alignment<ALIGN>: ValidAlignment,
{
    /// Create a empty dungeon core.
    pub fn new() -> Self {
        Self {
            current: never_type_id(),
            drop_impl: drop_t::<Never, SIZE, ALIGN>,
            buffer: AlignedBuffer::new(),
        }
    }

    /// Store a value of type `T`.
    ///
    /// If `T` can not be stored by this dungeon core do to
    /// size or alignment requirements a compiler error or
    /// runtime panic will be generated. This behavior is
    /// controlled by the `assert_runtime` feature.
    ///
    /// If the dungeon core contains a value already it will
    /// be dropped.
    #[track_caller]
    pub fn store<T: 'static>(&mut self, value: T) {
        // Assert that this dungeon core can store a value of `T`.
        crate::can_store::assert_can_store::<SIZE, ALIGN, T>();

        // Drop the current value.
        // Does nothing if already empty.
        let drop_old = self.drop_impl;
        drop_old(self);

        // SAFETY:
        // The call do `assert_can_store` makes the `store_unchecked` requirements
        // satisfied if this code runs.
        unsafe { self.store_unchecked::<T>(value) };
    }

    /// Gain ownership of the value inside the dungeon core.
    ///
    /// This will leave the dungeon core empty if the value is of 
    /// type `T`. If the dungeon core is empty or has another type's
    /// value stored in it then a `None` will be returned.
    pub fn take<T: 'static>(&mut self) -> Option<T> {
        if self.is_type::<T>() {
            // SAFETY:
            // We just checked that this dungeon core is storing a value
            // of `T`.
            Some(unsafe { self.take_unchecked::<T>() })
        } else {
            None
        }
    }

    /// Replace the current value with a `T`.
    ///
    /// The current value is returned if the dungeon core is not empty,
    /// and the current stored value is of type `U`.
    ///
    /// If `T` can not be stored by this dungeon core do to
    /// size or alignment requirements a compiler error or
    /// runtime panic will be generated. This behavior is
    /// controlled by the `assert_runtime` feature.
    ///
    /// This is slightly more efficient than a `take()` followed be
    /// a `store()
    #[track_caller]
    pub fn replace<T: 'static, U: 'static>(&mut self, value: T) -> Option<U> {
        // Assert `T` can be stored.
        crate::can_store::assert_can_store::<SIZE, ALIGN, T>();

        // Take the old value.
        let old_value = self.take();

        // SAFETY:
        // The call to `assert_can_store()` satisfies the safety requirements.
        // 
        // We also won't forget a value because we just took it.
        unsafe { self.store_unchecked::<T>(value) };

        old_value
    }

    /// Try to borrow the stored value.
    ///
    /// If the dungeon core is empty or the stored value is
    /// not of type `T` then `None` will be returned.
    pub fn try_borrow<T: 'static>(&self) -> Option<&T> {
        if !self.is_type::<T>() {
            return None;
        }

        // SAFETY:
        // We just checked that the current stored type is a `T`.
        let value = unsafe { self.try_borrow_unchecked::<T>() };

        Some(value)
    }

    /// Try to mutably borrow the stored value.
    ///
    /// If the dungeon core is empty or the stored value is
    /// not of type `T` then `None` will be returned.
    pub fn try_borrow_mut<T: 'static>(&mut self) -> Option<&mut T> {
        if !self.is_type::<T>() {
            return None;
        }

        // SAFETY:
        // We just checked that the current stored type is a `T`.
        let value = unsafe { self.try_borrow_mut_unchecked::<T>() };

        Some(value)
    }

    /// Check if a value is being stored.
    pub fn is_empty(&self) -> bool {
        self.current == never_type_id()
    }

    /// Check if a value being stored is of type `T`.
    ///
    /// Will always return `false` if empty.
    pub fn is_type<T: 'static>(&self) -> bool {
        if self.is_empty() {
            // A dungeon core never stores a value of `!`.
            // It is used as a sentinel value.
            false
        } else {
            self.current == core::any::TypeId::of::<T>()
        }
    }
}


/// Unsafe methods.
impl<const SIZE: usize, const ALIGN: u8> DungeonCore<SIZE, ALIGN> 
where
    Alignment<ALIGN>: ValidAlignment,
{
    /// Unchecked form of `store`.
    ///
    /// If the dungeon core is not empty, then the stored value will 
    /// be forgotten.
    ///
    /// # Safety
    /// `T` must be storable in `SIZE` bytes with alignment of `2^ALIGN` bytes.
    pub unsafe fn store_unchecked<T: 'static>(&mut self, value: T) {
        // Wrap the value so we can manually move it.
        // `value` will be considered uninit after the read below.
        let value = MaybeUninit::new(value);

        // Get a pointer to the aligned buffer.
        // This is done in long form to make it easier to read.
        let data: &mut [MaybeUninit<u8>; SIZE] = &mut self.buffer.data;
        let data: *mut [MaybeUninit<u8>; SIZE] = data as _;
        let data: *mut MaybeUninit<T> = data.cast();

        // SAFETY:
        // `value` is valid for a read because we just made it above 
        // from a valid value.
        //
        // `data` is valid for a write because of the safety requirements
        // for the outer function.
        //
        // They do not overlap.
        unsafe {
            core::ptr::copy_nonoverlapping(addr_of!(value), data, 1);
        }

        // Set the drop function to drop a value of `T`.
        self.drop_impl = drop_t::<T, SIZE, ALIGN>;

        // Set the current type to `T`.
        self.current = core::any::TypeId::of::<T>();
    }

    /// Unchecked form of [take].
    ///
    /// # Safety
    /// The dungeon core must be storing a valid value of type `T`.
    pub unsafe fn take_unchecked<T: 'static>(&mut self) -> T {
        // Buffer to move the stored value into.
        let mut value = MaybeUninit::uninit();

        // Get a pointer to the aligned buffer.
        // This is done in long form to make it easier to read.
        let data: &[MaybeUninit<u8>; SIZE] = &self.buffer.data;
        let data: *const [MaybeUninit<u8>; SIZE] = data as _;
        let data: *const MaybeUninit<T> = data.cast();

        // SAFETY:
        // `data` is valid for a read because we borrowed it.
        // And we must be storing a `T` because of the outer safety requirements.
        //
        // `value` is valid for a write because we just made it above.
        //
        // They do not overlap.
        unsafe {
            core::ptr::copy_nonoverlapping(data, addr_of_mut!(value), 1);
        }

        // Mark the dungeon core as empty.
        self.current = never_type_id();

        // SAFETY:
        // The requirements of the outer function require that the
        // value we just copied to be a valid `T`; And, by marking 
        // the dungeon core as empty the value will not be aliased.
        unsafe { value.assume_init() }
    }

    /// Unchecked form of `try_borrow`.
    ///
    /// # Safety
    /// The dungeon core must be storing a valid value of type `T`.
    pub unsafe fn try_borrow_unchecked<T: 'static>(&self) -> &T {
        // Get a pointer to the aligned buffer.
        // This is done in long form to make it easier to read.
        let data: &[MaybeUninit<u8>; SIZE] = &self.buffer.data;
        let data: *const [MaybeUninit<u8>; SIZE] = data as _;
        let data: *const MaybeUninit<T> = data.cast();

        // SAFETY:
        // We just borrowed `data` so it's valid to dereference.
        let data: &MaybeUninit<T> = unsafe { &*data };

        // SAFETY:
        // The outer safety requirements makes this safe.
        // `data` is a valid value of `T`.
        unsafe { data.assume_init_ref() }
    }

    /// Unchecked form of `try_borrow_mut`.
    ///
    /// # Safety
    /// The dungeon core must be storing a valid value of type `T`.
    pub unsafe fn try_borrow_mut_unchecked<T: 'static>(&mut self) -> &mut T {
        // Get a pointer to the aligned buffer.
        // This is done in long form to make it easier to read.
        let data: &mut [MaybeUninit<u8>; SIZE] = &mut self.buffer.data;
        let data: *mut [MaybeUninit<u8>; SIZE] = data as _;
        let data: *mut MaybeUninit<T> = data.cast();
        
        // SAFETY:
        // We just borrowed `data` so it's valid to dereference.
        let data: &mut MaybeUninit<T> = unsafe { &mut *data };

        // SAFETY:
        // The outer safety requirements makes this safe.
        // `data` is a valid value of `T`.
        unsafe { data.assume_init_mut() }
    }
}