use core::{
    future::poll_fn,
    ptr,
    sync::atomic::{AtomicU32, Ordering},
    task::Poll,
};

use alloc::sync::Arc;
use libk_mm::process::ProcessAddressSpace;
use libk_util::waker::QueueWaker;
use yggdrasil_abi::error::Error;

use crate::task::mem::ForeignAtomic;

/// User-space mutex (like BSD/Linux's futex) data structure
pub struct UserspaceMutex {
    queue: QueueWaker,
    space: Arc<ProcessAddressSpace>,
    atomic: *const AtomicU32,
}

impl UserspaceMutex {
    /// Creates a new [UserspaceMutex] associated with given `address`
    pub fn new(space: &Arc<ProcessAddressSpace>, address: usize) -> Result<Self, Error> {
        Ok(Self {
            queue: QueueWaker::new(),
            space: space.clone(),
            atomic: ptr::with_exposed_provenance(address),
        })
    }

    fn load(&self) -> Result<u32, Error> {
        self.atomic
            .atomic_load_foreign(&self.space, Ordering::Acquire)
    }

    pub async fn wait_until<P: Fn(u32) -> bool>(&self, predicate: P) -> Result<(), Error> {
        poll_fn(|cx| {
            let result = self.load();
            match result {
                Err(err) => {
                    self.queue.remove(cx.waker());
                    Poll::Ready(Err(err))
                }
                Ok(val) if predicate(val) => {
                    self.queue.remove(cx.waker());
                    Poll::Ready(Ok(()))
                }
                Ok(_) => {
                    self.queue.register(cx.waker());
                    Poll::Pending
                }
            }
        })
        .await
    }

    // pub async fn wait_until(self: Arc<Self>, compare_value: u32) -> Result<(), Error> {
    //     self.wait_predicate(|value| value == compare_value).await
    // }

    // /// Blocks until the value at the mutex's address becomes different from `compare_value`
    // pub async fn wait(self: Arc<Self>, compare_value: u32) -> Result<(), Error> {
    //     self.wait_predicate(|value| value != compare_value).await
    // }

    /// Wakes up at most `limit` threads waiting on this mutex
    pub fn wake(&self, limit: usize) {
        self.queue.wake_some(limit);
    }
}

unsafe impl Send for UserspaceMutex {}
unsafe impl Sync for UserspaceMutex {}
