Documentation
use super::mutual::Mutual;
use super::{InitState, Waiter, Waiters, WaitersExt};
use alloc::sync::Arc;
use core::cell::UnsafeCell;
use core::future::Future;
use core::mem::MaybeUninit;
use core::sync::atomic::Ordering;

/// An async lazy initialization primitive
pub struct Lazy<T, F> {
    state: Mutual<InitState>,
    value: UnsafeCell<MaybeUninit<T>>,
    init_fn: UnsafeCell<Option<F>>,
    waiters: Waiters,
}

impl<T, F> Lazy<T, F> {
    pub const fn new(init: F) -> Self {
        Self {
            state: Mutual::new(),
            value: UnsafeCell::new(MaybeUninit::uninit()),
            init_fn: UnsafeCell::new(Some(init)),
            waiters: Waiters::new(),
        }
    }
}

impl<T: Send + Sync, F: Future<Output = T>> Lazy<T, F> {
    /// Get or initialize the value
    pub async fn get(&self) -> &T {
        loop {
            if self.state.is(&InitState::Initialized) {
                return unsafe { &*(*self.value.get()).as_ptr() };
            }

            if self.state.is(&InitState::Initializing) {
                // Wait for initialization to complete
                let waiter = Waiter::default();
                unsafe { self.waiters.enqueue(waiter) };
                // Check again in case it completed
                if self.state.is(&InitState::Initialized) {
                    return unsafe { &*(*self.value.get()).as_ptr() };
                }
                // Otherwise wait
                continue;
            }

            // Try to become the initializer
            match self.state.compare_exchange(
                InitState::Uninitialized,
                InitState::Initializing,
                Ordering::AcqRel,
                Ordering::Acquire,
            ) {
                Ok(_) => {
                    // We won the race to initialize
                    let init_fn = unsafe {
                        (*self.init_fn.get())
                            .take()
                            .expect("init function already taken")
                    };

                    let value = init_fn.await;

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

                    self.state.set(InitState::Initialized);
                    self.waiters.notify_all();

                    return unsafe { &*(*self.value.get()).as_ptr() };
                }
                Err(_) => {
                    // Lost the race, loop back
                    continue;
                }
            }
        }
    }

    /// Try to get without waiting
    pub fn try_get(&self) -> Option<&T> {
        if self.state.is(&InitState::Initialized) {
            unsafe { Some(&*(*self.value.get()).as_ptr()) }
        } else {
            None
        }
    }
}

unsafe impl<T: Send + Sync, F: Send> Send for Lazy<T, F> {}
unsafe impl<T: Send + Sync, F: Send> Sync for Lazy<T, F> {}

/// A simpler async once cell that can be set manually
pub struct Once<T> {
    state: Mutual<InitState>,
    value: UnsafeCell<MaybeUninit<T>>,
    waiters: Waiters,
}

impl<T> Once<T> {
    pub const fn new() -> Self {
        Self {
            state: Mutual::new(),
            value: UnsafeCell::new(MaybeUninit::uninit()),
            waiters: Waiters::new(),
        }
    }

    /// Initialize with a value
    pub fn init(&self, value: T) -> Result<(), T> {
        match self.state.compare_exchange(
            InitState::Uninitialized,
            InitState::Initialized,
            Ordering::AcqRel,
            Ordering::Acquire,
        ) {
            Ok(_) => {
                unsafe {
                    (*self.value.get()).write(value);
                }
                self.waiters.notify_all();
                Ok(())
            }
            Err(_) => Err(value),
        }
    }

    /// Get the value, waiting if not yet initialized
    pub async fn get(&self) -> &T {
        loop {
            if self.state.is(&InitState::Initialized) {
                return unsafe { &*(*self.value.get()).as_ptr() };
            }

            // Wait for initialization
            let waiter = Waiter::default();
            unsafe { self.waiters.enqueue(waiter) };
        }
    }

    /// Try to get without waiting
    pub fn try_get(&self) -> Option<&T> {
        if self.state.is(&InitState::Initialized) {
            unsafe { Some(&*(*self.value.get()).as_ptr()) }
        } else {
            None
        }
    }
}

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

/// Arc-wrapped async lazy for shared ownership
pub type SharedLazy<T, F> = Arc<Lazy<T, F>>;