use core::mem::ManuallyDrop;
use core::ops::{Deref, DerefMut};

pub const LOCK_ORDER_UNTRACKED: u8 = u8::MAX;

mod order_check {
    use crate::arch::syscall::{LOCK_ORDER_MAX_DEPTH, this_cpu};

    pub fn push(order: u8) {
        if order == super::LOCK_ORDER_UNTRACKED {
            return;
        }
        let pc = this_cpu();
        let d = unsafe { core::ptr::addr_of!((*pc).lock_order_depth).read_volatile() } as usize;
        if d > 0 {
            let top = unsafe {
                core::ptr::addr_of!((*pc).lock_order_stack)
                    .cast::<u8>()
                    .add(d - 1)
                    .read_volatile()
            };
            assert!(
                order > top,
                "lock order violation: acquiring order {} while holding order {}",
                order,
                top
            );
        }
        assert!(d < LOCK_ORDER_MAX_DEPTH, "lock nesting exceeds max depth");
        unsafe {
            core::ptr::addr_of_mut!((*pc).lock_order_stack)
                .cast::<u8>()
                .add(d)
                .write_volatile(order);
            core::ptr::addr_of_mut!((*pc).lock_order_depth).write_volatile((d + 1) as u8);
        }
    }

    pub fn pop(order: u8) {
        if order == super::LOCK_ORDER_UNTRACKED {
            return;
        }
        let pc = this_cpu();
        let d = unsafe { core::ptr::addr_of!((*pc).lock_order_depth).read_volatile() } as usize;
        assert!(d > 0, "lock order stack underflow");
        let top = unsafe {
            core::ptr::addr_of!((*pc).lock_order_stack)
                .cast::<u8>()
                .add(d - 1)
                .read_volatile()
        };
        assert!(
            top == order,
            "lock order pop mismatch: expected {}, got {}",
            order,
            top
        );
        unsafe {
            core::ptr::addr_of_mut!((*pc).lock_order_depth).write_volatile((d - 1) as u8);
        }
    }
}

pub struct InterruptsDisabledToken(());

impl InterruptsDisabledToken {
    pub fn new_checked() -> Option<Self> {
        if x86_64::instructions::interrupts::are_enabled() {
            None
        } else {
            Some(Self(()))
        }
    }
}

#[repr(transparent)]
pub struct SyncUnsafe<T>(core::cell::UnsafeCell<T>);

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

impl<T> SyncUnsafe<T> {
    pub const fn new(val: T) -> Self {
        Self(core::cell::UnsafeCell::new(val))
    }

    pub const fn get(&self) -> *mut T {
        self.0.get()
    }

    #[allow(clippy::mut_from_ref)]
    pub unsafe fn as_mut_unchecked(&self) -> &mut T {
        unsafe { &mut *self.0.get() }
    }

    pub unsafe fn as_ref_unchecked(&self) -> &T {
        unsafe { &*self.0.get() }
    }
}

pub struct IrqMutex<T, const ORDER: u8> {
    inner: spin::Mutex<T>,
}

#[must_use = "dropping the guard immediately re-enables interrupts and releases the lock"]
pub struct IrqMutexGuard<'a, T, const ORDER: u8> {
    guard: ManuallyDrop<spin::MutexGuard<'a, T>>,
    irq_was_enabled: bool,
}

impl<T, const ORDER: u8> IrqMutex<T, ORDER> {
    pub const fn new(val: T) -> Self {
        Self {
            inner: spin::Mutex::new(val),
        }
    }

    pub fn lock(&self) -> IrqMutexGuard<'_, T, ORDER> {
        let irq_was_enabled = x86_64::instructions::interrupts::are_enabled();
        x86_64::instructions::interrupts::disable();
        order_check::push(ORDER);
        let guard = self.inner.lock();
        IrqMutexGuard {
            guard: ManuallyDrop::new(guard),
            irq_was_enabled,
        }
    }

    pub fn lock_after<const HELD: u8, U>(
        &self,
        _held: &IrqMutexGuard<'_, U, HELD>,
    ) -> IrqMutexGuard<'_, T, ORDER> {
        const { assert!(ORDER > HELD) };
        self.lock()
    }

    #[allow(dead_code)]
    pub fn try_lock(&self) -> Option<IrqMutexGuard<'_, T, ORDER>> {
        let irq_was_enabled = x86_64::instructions::interrupts::are_enabled();
        x86_64::instructions::interrupts::disable();
        order_check::push(ORDER);
        match self.inner.try_lock() {
            Some(guard) => Some(IrqMutexGuard {
                guard: ManuallyDrop::new(guard),
                irq_was_enabled,
            }),
            None => {
                order_check::pop(ORDER);
                if irq_was_enabled {
                    x86_64::instructions::interrupts::enable();
                }
                None
            }
        }
    }
}

unsafe impl<T: Send, const ORDER: u8> Send for IrqMutex<T, ORDER> {}
unsafe impl<T: Send, const ORDER: u8> Sync for IrqMutex<T, ORDER> {}

impl<'a, T, const ORDER: u8> Deref for IrqMutexGuard<'a, T, ORDER> {
    type Target = T;
    fn deref(&self) -> &T {
        &self.guard
    }
}

impl<'a, T, const ORDER: u8> DerefMut for IrqMutexGuard<'a, T, ORDER> {
    fn deref_mut(&mut self) -> &mut T {
        &mut self.guard
    }
}

impl<'a, T, const ORDER: u8> Drop for IrqMutexGuard<'a, T, ORDER> {
    fn drop(&mut self) {
        unsafe { ManuallyDrop::drop(&mut self.guard) };
        order_check::pop(ORDER);
        if self.irq_was_enabled {
            x86_64::instructions::interrupts::enable();
        }
    }
}