//! Runtime implementations for synchronized mutable-only-once values.
use core::{
    cell::UnsafeCell,
    mem::MaybeUninit,
    sync::atomic::{AtomicU32, Ordering},
};

use abi::process::MutexOperation;

/// Describes a value which can only be initialized once at runtime and provides only immutable
/// access to `T`.
pub struct Once<T> {
    value: UnsafeCell<MaybeUninit<T>>,
    state: AtomicU32,
}

impl<T> Once<T> {
    const UNINIT: u32 = 0;
    const PENDING: u32 = 1;
    const INIT: u32 = 2;

    /// Constructs an uninitialized [Once].
    pub const fn uninit() -> Self {
        Self {
            value: UnsafeCell::new(MaybeUninit::uninit()),
            state: AtomicU32::new(Self::UNINIT),
        }
    }

    /// Attempts to initialize this [Once], if not already.
    ///
    /// Returns `true` if the value was initialized from the `producer` function, `false` if
    /// the [Once] is already initialized.
    pub fn try_init_with<F: FnOnce() -> T>(&self, producer: F) -> bool {
        // Try state -> PENDING
        if self
            .state
            .compare_exchange(
                Self::UNINIT,
                Self::PENDING,
                Ordering::Acquire,
                Ordering::Relaxed,
            )
            .is_err()
        {
            return false;
        }

        // Write the value while state == PENDING
        unsafe { (*self.value.get()).write(producer()) };

        // state -> INIT
        self.state.store(Self::INIT, Ordering::Release);

        // Notify anyone waiting
        self.notify();

        true
    }

    /// Initializes this [Once] with the `producer` given.
    ///
    /// # Panics
    ///
    /// Panics if the value has already been initialized.
    #[track_caller]
    pub fn init_with<F: FnOnce() -> T>(&self, producer: F) {
        if !self.try_init_with(producer) {
            panic!("Once<T> has already been initialized");
        }
    }

    /// Waits until this [Once] becomes initialized and returns a reference to the value.
    pub fn wait(&self) -> &'static T {
        loop {
            if let Some(value) = self.try_get() {
                return value;
            }

            self.wait_initialized();
        }
    }

    /// Returns a reference to the value, if it was initialized previously, or [None] if
    /// the value is uninitialized at the moment of this call.
    pub fn try_get(&self) -> Option<&'static T> {
        if self.state.load(Ordering::Acquire) == Self::INIT {
            Some(unsafe { (*self.value.get()).assume_init_ref() })
        } else {
            None
        }
    }

    /// Returns a reference to the value.
    ///
    /// # Panics
    ///
    /// Panics if the value has not yet been initialized.
    #[track_caller]
    pub fn get(&self) -> &'static T {
        match self.try_get() {
            Some(value) => value,
            None => {
                panic!("Attempt to use an uninitialized Once<T>")
            }
        }
    }

    fn notify(&self) {
        unsafe { crate::sys::mutex(&self.state, &MutexOperation::Wake(u32::MAX)) }.ok();
    }

    fn wait_initialized(&self) {
        unsafe {
            crate::sys::mutex(
                &self.state,
                &MutexOperation::WaitUntilEqual(Self::INIT, None),
            )
        }
        .expect("Kernel mutex wait failed");
    }
}

unsafe impl<T: Sync> Sync for Once<T> {}
