use core::{
    cell::UnsafeCell,
    mem::MaybeUninit,
    panic,
    sync::atomic::{AtomicUsize, Ordering},
};

use yggdrasil_abi::error::Error;

/// Wrapper struct to ensure a value can only be initialized once and used only after that
#[repr(C)]
pub struct OneTimeInit<T> {
    value: UnsafeCell<MaybeUninit<T>>,
    state: AtomicUsize,
}

unsafe impl<T> Sync for OneTimeInit<T> {}
unsafe impl<T> Send for OneTimeInit<T> {}

impl<T> OneTimeInit<T> {
    const STATE_UNINITIALIZED: usize = 0;
    const STATE_INITIALIZING: usize = 1;
    const STATE_INITIALIZED: usize = 2;

    /// Wraps the value in an [OneTimeInit]
    pub const fn new() -> Self {
        Self {
            value: UnsafeCell::new(MaybeUninit::uninit()),
            state: AtomicUsize::new(Self::STATE_UNINITIALIZED),
        }
    }

    /// Wraps the value in an [OneTimeInit], providing an already initialized value.
    pub const fn new_init(value: T) -> Self {
        Self {
            value: UnsafeCell::new(MaybeUninit::new(value)),
            state: AtomicUsize::new(Self::STATE_INITIALIZED),
        }
    }

    /// Returns `true` if the value has already been initialized
    #[inline]
    pub fn is_initialized(&self) -> bool {
        self.state.load(Ordering::Acquire) == Self::STATE_INITIALIZED
    }

    pub fn try_init_with_opt<F: FnOnce() -> Result<T, Error>>(&self, f: F) -> Result<&T, Error> {
        if !self.try_begin_init() {
            // Already initialized
            return Err(Error::AlreadyExists);
        }

        match f() {
            Ok(value) => {
                let value = unsafe { (*self.value.get()).write(value) };
                self.finish_init();
                Ok(value)
            }
            Err(error) => {
                self.fail_init();
                Err(error)
            }
        }
    }

    pub fn or_init_with<F: FnOnce() -> T>(&self, f: F) -> &T {
        if !self.try_begin_init() {
            return self.wait_for_init();
        }
        let value = unsafe { (*self.value.get()).write(f()) };
        self.finish_init();
        value
    }

    pub fn or_init_with_opt<F: FnOnce() -> Option<T>>(&self, f: F) -> Option<&T> {
        if !self.try_begin_init() {
            return Some(self.wait_for_init());
        }
        match f() {
            Some(value) => {
                let value = unsafe { (*self.value.get()).write(value) };
                self.finish_init();
                Some(value)
            }
            None => {
                self.fail_init();
                None
            }
        }
    }

    pub fn or_try_init_with<F: FnOnce() -> Result<T, Error>>(&self, f: F) -> Result<&T, Error> {
        if !self.try_begin_init() {
            return Ok(self.wait_for_init());
        }
        match f() {
            Ok(value) => {
                let value = unsafe { (*self.value.get()).write(value) };
                self.finish_init();
                Ok(value)
            }
            Err(error) => {
                // Init failed
                self.fail_init();
                Err(error)
            }
        }
    }

    fn try_begin_init(&self) -> bool {
        self.state
            .compare_exchange(
                Self::STATE_UNINITIALIZED,
                Self::STATE_INITIALIZING,
                Ordering::Acquire,
                Ordering::Relaxed,
            )
            .is_ok()
    }

    fn finish_init(&self) {
        self.state
            .compare_exchange(
                Self::STATE_INITIALIZING,
                Self::STATE_INITIALIZED,
                Ordering::Release,
                Ordering::Relaxed,
            )
            .unwrap();
    }

    fn fail_init(&self) {
        self.state
            .compare_exchange(
                Self::STATE_INITIALIZING,
                Self::STATE_UNINITIALIZED,
                Ordering::Release,
                Ordering::Relaxed,
            )
            .unwrap();
    }

    fn wait_for_init(&self) -> &T {
        while self.state.load(Ordering::Acquire) != Self::STATE_INITIALIZED {
            core::hint::spin_loop();
        }

        unsafe { (*self.value.get()).assume_init_ref() }
    }

    pub fn try_init_with<F: FnOnce() -> T>(&self, f: F) -> Result<&T, Error> {
        if !self.try_begin_init() {
            // Already initialized
            return Err(Error::AlreadyExists);
        }

        let value = unsafe { (*self.value.get()).write(f()) };

        self.finish_init();

        Ok(value)
    }

    /// Sets the underlying value of the [OneTimeInit]. If already initialized, panics.
    #[track_caller]
    pub fn init(&self, value: T) -> &T {
        // Transition to "initializing" state
        if self
            .state
            .compare_exchange(
                Self::STATE_UNINITIALIZED,
                Self::STATE_INITIALIZING,
                Ordering::Release,
                Ordering::Relaxed,
            )
            .is_err()
        {
            panic!(
                "{:?}: Double initialization of OneTimeInit<T>",
                panic::Location::caller()
            );
        }

        let value = unsafe { (*self.value.get()).write(value) };

        // Transition to "initialized" state. This must not fail
        self.state
            .compare_exchange(
                Self::STATE_INITIALIZING,
                Self::STATE_INITIALIZED,
                Ordering::Release,
                Ordering::Relaxed,
            )
            .unwrap();

        value
    }

    /// Returns an immutable reference to the underlying value and panics if it hasn't yet been
    /// initialized
    #[track_caller]
    pub fn get(&self) -> &T {
        // TODO check for INITIALIZING state and wait until it becomes INITIALIZED?
        if !self.is_initialized() {
            panic!(
                "{:?}: Attempt to dereference an uninitialized value",
                panic::Location::caller()
            );
        }

        unsafe { (*self.value.get()).assume_init_ref() }
    }

    /// Returns an immutable reference to the underlying value and [None] if the value hasn't yet
    /// been initialized
    pub fn try_get(&self) -> Option<&T> {
        if self.is_initialized() {
            Some(self.get())
        } else {
            None
        }
    }
}
