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

use super::IrqGuard;

struct RwLockInner {
    value: AtomicUsize,
}

pub struct IrqSafeRwLock<T> {
    value: UnsafeCell<T>,
    inner: RwLockInner,
}

pub struct IrqSafeRwLockReadGuard<'a, T> {
    lock: &'a IrqSafeRwLock<T>,
    _guard: IrqGuard,
}

pub struct IrqSafeRwLockWriteGuard<'a, T> {
    lock: &'a IrqSafeRwLock<T>,
    _guard: IrqGuard,
}

impl RwLockInner {
    const LOCKED_READ: usize = 1 << 2;
    const LOCKED_WRITE: usize = 1;

    const fn new() -> Self {
        Self {
            value: AtomicUsize::new(0),
        }
    }

    #[inline]
    fn downgrade_write(&self) {
        // Acquire read lock
        let value = self.acquire_read_raw();
        assert_eq!(value & Self::LOCKED_WRITE, Self::LOCKED_WRITE);
        // Release write lock
        self.value.fetch_nand(Self::LOCKED_WRITE, Ordering::Release);
    }

    #[inline]
    fn upgrade(&self) {
        // At least one read lock is held by this task.
        // When there's *exactly* one lock (being this task) held, upgrade is possible
        while !self.try_upgrade() {
            core::hint::spin_loop();
        }
    }

    #[inline]
    fn acquire_read_raw(&self) -> usize {
        let value = self.value.fetch_add(Self::LOCKED_READ, Ordering::Acquire);

        if value > usize::MAX / 2 {
            self.value.fetch_sub(Self::LOCKED_READ, Ordering::Relaxed);
            panic!("Too many read locks acquired");
        }

        value
    }

    #[inline]
    fn try_acquire_read(&self) -> bool {
        let value = self.acquire_read_raw();
        let acquired = value & Self::LOCKED_WRITE != Self::LOCKED_WRITE;

        if !acquired {
            unsafe {
                self.release_read();
            }
        }

        acquired
    }

    #[inline]
    fn try_acquire_write(&self) -> bool {
        self.value
            .compare_exchange(0, Self::LOCKED_WRITE, Ordering::Acquire, Ordering::Relaxed)
            .is_ok()
    }

    #[inline]
    fn try_upgrade(&self) -> bool {
        self.value
            .compare_exchange(
                Self::LOCKED_READ,
                Self::LOCKED_WRITE,
                Ordering::Acquire,
                Ordering::Relaxed,
            )
            .is_ok()
    }

    #[inline]
    fn acquire_read(&self) {
        while !self.try_acquire_read() {
            core::hint::spin_loop();
        }
    }

    #[inline]
    fn acquire_write(&self) {
        while !self.try_acquire_write() {
            core::hint::spin_loop();
        }
    }

    #[inline]
    unsafe fn release_read(&self) {
        self.value.fetch_sub(Self::LOCKED_READ, Ordering::Release);
    }

    #[inline]
    unsafe fn release_write(&self) {
        self.value.fetch_and(!Self::LOCKED_WRITE, Ordering::Release);
    }
}

impl<T> IrqSafeRwLock<T> {
    pub const fn new(value: T) -> Self {
        Self {
            value: UnsafeCell::new(value),
            inner: RwLockInner::new(),
        }
    }

    pub fn read(&self) -> IrqSafeRwLockReadGuard<'_, T> {
        let guard = IrqGuard::acquire();
        self.inner.acquire_read();
        IrqSafeRwLockReadGuard {
            lock: self,
            _guard: guard,
        }
    }

    pub fn write(&self) -> IrqSafeRwLockWriteGuard<'_, T> {
        let guard = IrqGuard::acquire();
        self.inner.acquire_write();
        IrqSafeRwLockWriteGuard {
            lock: self,
            _guard: guard,
        }
    }

    #[inline]
    unsafe fn downgrade(&self) {
        self.inner.downgrade_write();
    }

    #[inline]
    unsafe fn upgrade(&self) {
        self.inner.upgrade();
    }

    unsafe fn release_read(&self) {
        self.inner.release_read();
    }

    unsafe fn release_write(&self) {
        self.inner.release_write();
    }
}

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

impl<T> Deref for IrqSafeRwLockReadGuard<'_, T> {
    type Target = T;

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

impl<T> Drop for IrqSafeRwLockReadGuard<'_, T> {
    fn drop(&mut self) {
        unsafe { self.lock.release_read() }
    }
}

impl<'a, T> IrqSafeRwLockReadGuard<'a, T> {
    pub fn get(guard: &Self) -> *const T {
        guard.lock.value.get()
    }

    pub fn upgrade(guard: IrqSafeRwLockReadGuard<'a, T>) -> IrqSafeRwLockWriteGuard<'a, T> {
        let lock = guard.lock;
        let irq_guard = IrqGuard::acquire();
        // Read lock still held
        core::mem::forget(guard);

        unsafe {
            lock.upgrade();
        }

        IrqSafeRwLockWriteGuard {
            lock,
            _guard: irq_guard,
        }
    }
}

impl<'a, T> IrqSafeRwLockWriteGuard<'a, T> {
    pub fn downgrade(guard: IrqSafeRwLockWriteGuard<'a, T>) -> IrqSafeRwLockReadGuard<'a, T> {
        let lock = guard.lock;
        let irq_guard = IrqGuard::acquire();
        // Read lock still held
        core::mem::forget(guard);

        // Read lock acquired, write lock released
        unsafe {
            lock.downgrade();
        }

        IrqSafeRwLockReadGuard {
            lock,
            _guard: irq_guard,
        }
    }
}

impl<T> Deref for IrqSafeRwLockWriteGuard<'_, T> {
    type Target = T;

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

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

impl<T> Drop for IrqSafeRwLockWriteGuard<'_, T> {
    fn drop(&mut self) {
        unsafe { self.lock.release_write() }
    }
}
