use core::cell::UnsafeCell;
use core::ops::{Deref, DerefMut};
use core::sync::atomic::{AtomicBool, Ordering};

use crate::int_state::{pop_off, push_off};

/// Spin-lock that disables S-mode interrupts while held.
pub struct SpinNoIrq<T> {
    locked: AtomicBool,
    data: UnsafeCell<T>,
}

unsafe impl<T: Send> Sync for SpinNoIrq<T> {}
unsafe impl<T: Send> Send for SpinNoIrq<T> {}

/// 
pub struct SpinNoIrqGuard<'a, T> {
    lock: &'a SpinNoIrq<T>,
}

impl<T> SpinNoIrq<T> {
    /// 
    pub const fn new(value: T) -> Self {
        Self {
            locked: AtomicBool::new(false),
            data: UnsafeCell::new(value),
        }
    }

    /// 
    pub fn lock(&self) -> SpinNoIrqGuard<'_, T> {
        // 1. Disable interrupts and track nesting depth on this hart FIRST.
        push_off();

        // 2. Acquire the lock.
        while self.locked
            .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
            .is_err()
        {
            while self.locked.load(Ordering::Relaxed) {
                core::hint::spin_loop();
            }
        }

        SpinNoIrqGuard { lock: self }
    }
}

impl<'a, T> Deref for SpinNoIrqGuard<'a, T> {
    type Target = T;
    #[inline]
    fn deref(&self) -> &Self::Target {
        unsafe { &*self.lock.data.get() }
    }
}

impl<'a, T> DerefMut for SpinNoIrqGuard<'a, T> {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut *self.lock.data.get() }
    }
}

impl<'a, T> Drop for SpinNoIrqGuard<'a, T> {
    fn drop(&mut self) {
        // 1. Release the lock.
        self.lock.locked.store(false, Ordering::Release);
        
        // 2. Decrement nesting depth, re-enable interrupts if it reaches 0.
        pop_off();
    }
}