use core::sync::atomic::{AtomicU32, Ordering};

use super::platform;
use super::raw_sys;

const UNLOCKED: u32 = 0;
const LOCKED: u32 = 1;
const CONTENDED: u32 = 2;

#[derive(Debug)]
pub struct FutexMutex {
    state: AtomicU32,
}

impl FutexMutex {
    pub const fn new() -> Self {
        Self {
            state: AtomicU32::new(UNLOCKED),
        }
    }

    pub fn lock(&self) -> FutexMutexGuard<'_> {
        if self
            .state
            .compare_exchange(UNLOCKED, LOCKED, Ordering::Acquire, Ordering::Relaxed)
            .is_err()
        {
            self.lock_contended();
        }

        FutexMutexGuard { mutex: self }
    }

    fn lock_contended(&self) {
        loop {
            if self.state.swap(CONTENDED, Ordering::Acquire) == UNLOCKED {
                return;
            }

            match unsafe {
                platform::futex_wait(self.state_ptr(), raw_sys::FUTEX_WAIT_PRIVATE, CONTENDED)
            } {
                Ok(()) => {}
                Err(error) if matches!(error.errno(), raw_sys::EAGAIN | raw_sys::EINTR) => {}
                Err(_) => {}
            }
        }
    }

    fn unlock(&self) {
        if self.state.fetch_sub(1, Ordering::Release) != LOCKED {
            self.state.store(UNLOCKED, Ordering::Release);
            let _ =
                unsafe { platform::futex_wake(self.state_ptr(), raw_sys::FUTEX_WAKE_PRIVATE, 1) };
        }
    }

    #[inline(always)]
    fn state_ptr(&self) -> *const u32 {
        (&self.state as *const AtomicU32).cast::<u32>()
    }
}

impl Default for FutexMutex {
    fn default() -> Self {
        Self::new()
    }
}

pub struct FutexMutexGuard<'a> {
    mutex: &'a FutexMutex,
}

impl Drop for FutexMutexGuard<'_> {
    fn drop(&mut self) {
        self.mutex.unlock();
    }
}

pub(crate) fn futex_wait_for_value(state: &AtomicU32, expected: u32) {
    match unsafe {
        platform::futex_wait(
            (state as *const AtomicU32).cast::<u32>(),
            raw_sys::FUTEX_WAIT_PRIVATE,
            expected,
        )
    } {
        Ok(()) => {}
        Err(error) if matches!(error.errno(), raw_sys::EAGAIN | raw_sys::EINTR) => {}
        Err(_) => {}
    }
}

pub(crate) fn futex_wake_all(state: &AtomicU32) {
    let _ = unsafe {
        platform::futex_wake(
            (state as *const AtomicU32).cast::<u32>(),
            raw_sys::FUTEX_WAKE_PRIVATE,
            i32::MAX as u32,
        )
    };
}

#[cfg(test)]
mod tests {
    use core::sync::atomic::{AtomicUsize, Ordering};

    use std::sync::Arc;
    use std::thread;
    use std::vec::Vec;

    use super::FutexMutex;

    #[test]
    fn futex_mutex_serializes_multiple_threads() {
        let mutex = Arc::new(FutexMutex::new());
        let counter = Arc::new(AtomicUsize::new(0));
        let mut threads = Vec::new();

        for _ in 0..4 {
            let mutex = Arc::clone(&mutex);
            let counter = Arc::clone(&counter);
            threads.push(thread::spawn(move || {
                for _ in 0..5000 {
                    let _guard = mutex.lock();
                    counter.fetch_add(1, Ordering::Relaxed);
                }
            }));
        }

        for thread in threads {
            thread.join().unwrap();
        }

        assert_eq!(counter.load(Ordering::Acquire), 20_000);
    }
}
