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

use crate::{guard::IrqGuard, Architecture};

pub struct Spinlock<A: Architecture, T> {
    value: UnsafeCell<T>,
    state: AtomicBool,
    _pd: PhantomData<A>,
}

pub struct SpinlockGuard<'a, A: Architecture, T> {
    lock: &'a Spinlock<A, T>,
}

/// Spinlock implementation which prevents interrupts to avoid deadlocks when an interrupt handler
/// tries to acquire a lock taken before the IRQ fired.
pub struct IrqSafeSpinlock<A: Architecture, T> {
    inner: Spinlock<A, T>,
}

/// Token type allowing safe access to the underlying data of the [IrqSafeSpinlock]. Resumes normal
/// IRQ operation (if enabled before acquiring) when the lifetime is over.
pub struct IrqSafeSpinlockGuard<'a, A: Architecture, T> {
    // Must come first to ensure the lock is dropped first and only then IRQs are re-enabled
    inner: SpinlockGuard<'a, A, T>,
    _irq: IrqGuard<A>,
}

// Spinlock impls
impl<A: Architecture, T> Spinlock<A, T> {
    pub const fn new(value: T) -> Self {
        Self {
            value: UnsafeCell::new(value),
            state: AtomicBool::new(false),
            _pd: PhantomData,
        }
    }

    pub fn lock(&self) -> SpinlockGuard<'_, A, T> {
        // Loop until the lock can be acquired
        if LOCK_HACK.load(Ordering::Acquire) {
            return SpinlockGuard { lock: self };
        }
        while self
            .state
            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
            .is_err()
        {
            core::hint::spin_loop();
        }

        SpinlockGuard { lock: self }
    }
}

impl<A: Architecture, T> Deref for SpinlockGuard<'_, A, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        unsafe { &*self.lock.value.get() }
    }
}

impl<A: Architecture, T> DerefMut for SpinlockGuard<'_, A, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut *self.lock.value.get() }
    }
}

impl<A: Architecture, T> Drop for SpinlockGuard<'_, A, T> {
    fn drop(&mut self) {
        if !LOCK_HACK.load(Ordering::Acquire) {
            self.lock
                .state
                .compare_exchange(true, false, Ordering::Release, Ordering::Relaxed)
                .unwrap();
        }
    }
}

unsafe impl<A: Architecture, T: Send> Sync for Spinlock<A, T> {}
unsafe impl<A: Architecture, T: Send> Send for Spinlock<A, T> {}

// IrqSafeSpinlock impls
impl<A: Architecture, T> IrqSafeSpinlock<A, T> {
    /// Wraps the value in a spinlock primitive
    pub const fn new(value: T) -> Self {
        Self {
            inner: Spinlock::new(value),
        }
    }

    #[inline]
    pub fn replace(&self, value: T) -> T {
        let mut lock = self.lock();
        mem::replace(&mut lock, value)
    }

    /// Attempts to acquire a lock. IRQs will be disabled until the lock is released.
    pub fn lock(&self) -> IrqSafeSpinlockGuard<'_, A, T> {
        // Disable IRQs to avoid IRQ handler trying to acquire the same lock
        let irq_guard = IrqGuard::acquire();

        // Acquire the inner lock
        let inner = self.inner.lock();

        IrqSafeSpinlockGuard {
            inner,
            _irq: irq_guard,
        }
    }

    /// Returns an unsafe reference to the inner value.
    ///
    /// # Safety
    ///
    /// Unsafe: explicitly ignores proper access sharing.
    #[allow(clippy::mut_from_ref)]
    pub unsafe fn grab(&self) -> &mut T {
        unsafe { &mut *self.inner.value.get() }
    }
}

impl<A: Architecture, T: Clone> IrqSafeSpinlock<A, T> {
    pub fn get_cloned(&self) -> T {
        self.lock().clone()
    }
}

impl<A: Architecture, T: Clone> Clone for IrqSafeSpinlock<A, T> {
    fn clone(&self) -> Self {
        let inner = self.lock();
        IrqSafeSpinlock::new(inner.clone())
    }
}

impl<A: Architecture, T> Deref for IrqSafeSpinlockGuard<'_, A, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.inner.deref()
    }
}

impl<A: Architecture, T> DerefMut for IrqSafeSpinlockGuard<'_, A, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.inner.deref_mut()
    }
}

static LOCK_HACK: AtomicBool = AtomicBool::new(false);

/// "Hacks" all the locks in the kernel to make them function as "NULL"-locks instead of spinlocks.
///
/// # Safety
///
/// Only meant to be called from panic handler when the caller is sure other CPUs are halted.
pub unsafe fn hack_locks() {
    LOCK_HACK.store(true, Ordering::Release);
}
