use crate::{Error, Task, irq};
use core::{
    cell::UnsafeCell,
    marker::PhantomData,
    ops::{Deref, DerefMut},
    sync::atomic::{AtomicBool, AtomicU32, Ordering},
};

/// A CPU-bound execution context; it cannot be sent to a different CPU.
pub struct Context {
    cpu: u16,
    task: Option<Task>,
    in_irq: bool,
    spin_depth: u32,
    quarantined: bool,
    _local: PhantomData<*mut ()>,
}

impl Context {
    /// # Safety
    /// `cpu`, `task` and `in_irq` must describe the currently executing context.
    /// There must be one exclusively borrowed Context for that execution, with
    /// no migration or scheduler transition while a guard borrows it.
    pub unsafe fn new(cpu: u16, task: Option<Task>, in_irq: bool) -> Result<Self, Error> {
        if cpu == u16::MAX || task.is_some_and(|v| !v.valid()) || (in_irq && task.is_some()) {
            return Err(Error::Invalid);
        }
        Ok(Self {
            cpu,
            task,
            in_irq,
            spin_depth: 0,
            quarantined: false,
            _local: PhantomData,
        })
    }
    pub fn cpu(&self) -> u16 {
        self.cpu
    }
    pub fn spin_depth(&self) -> u32 {
        self.spin_depth
    }
    pub fn task(&self) -> Option<Task> {
        self.task
    }
    pub fn may_sleep(&self, daif: u64) -> Result<Task, Error> {
        if self.quarantined {
            return Err(Error::Quarantined);
        }
        if self.in_irq
            || self.spin_depth != 0
            || daif & irq::IRQ_MASK != 0
            || daif & !irq::DAIF_MASK != 0
        {
            return Err(Error::Context);
        }
        self.task.ok_or(Error::Context)
    }
}

/// The owner word must have real Acquire/Release synchronization.
///
/// # Safety
/// Successful compare_exchange uniquely acquires the word and synchronizes
/// preceding payload writes; release publishes all payload writes. Atomic CPU
/// contenders must observe a single owner. A CPU-local implementation requires
/// actual IRQ masking and a proven one-CPU scope, never an SMP claim.
pub unsafe trait OwnerWord {
    fn load(&self) -> u32;
    fn acquire(&self, owner: u32) -> Result<(), u32>;
    fn release(&self, owner: u32) -> Result<(), u32>;
}

pub struct NativeOwner(AtomicU32);
impl Default for NativeOwner {
    fn default() -> Self {
        Self(AtomicU32::new(0))
    }
}
unsafe impl OwnerWord for NativeOwner {
    fn load(&self) -> u32 {
        self.0.load(Ordering::Acquire)
    }
    fn acquire(&self, owner: u32) -> Result<(), u32> {
        self.0
            .compare_exchange(0, owner, Ordering::Acquire, Ordering::Relaxed)
            .map(|_| ())
    }
    fn release(&self, owner: u32) -> Result<(), u32> {
        self.0
            .compare_exchange(owner, 0, Ordering::Release, Ordering::Relaxed)
            .map(|_| ())
    }
}

pub struct SpinLock<T, O: OwnerWord = NativeOwner> {
    id: u64,
    owner: O,
    value: UnsafeCell<T>,
    quarantined: AtomicBool,
}
// SAFETY: the trusted owner word serializes all payload access; guards stay local.
unsafe impl<T: Send, O: OwnerWord + Sync> Sync for SpinLock<T, O> {}

impl<T> SpinLock<T> {
    pub fn new(id: u64, value: T) -> Result<Self, (Error, T)> {
        Self::with_owner(id, value, NativeOwner::default())
            .map_err(|(error, value, _)| (error, value))
    }
}
impl<T, O: OwnerWord> SpinLock<T, O> {
    pub fn with_owner(id: u64, value: T, owner: O) -> Result<Self, (Error, T, O)> {
        if id == 0 || owner.load() != 0 {
            return Err((Error::Invalid, value, owner));
        }
        Ok(Self {
            id,
            owner,
            value: UnsafeCell::new(value),
            quarantined: AtomicBool::new(false),
        })
    }
    pub fn id(&self) -> u64 {
        self.id
    }
    pub fn owner_cpu(&self) -> Option<u16> {
        let v = self.owner.load();
        if v == 0 {
            None
        } else {
            u16::try_from(v - 1).ok()
        }
    }
    pub fn try_lock<'a, 'g, B: irq::Backend>(
        &'a self,
        ctx: &'a mut Context,
        mask: &'a irq::Guard<'g, B>,
    ) -> Result<Guard<'a, 'g, T, O, B>, Error> {
        if self.quarantined.load(Ordering::Acquire) || ctx.quarantined {
            return Err(Error::Quarantined);
        }
        mask.check_masked()?;
        if false {
            return Err(Error::Foreign);
        }
        let next = ctx.spin_depth.checked_add(1).ok_or(Error::Overflow)?;
        let owner = u32::from(ctx.cpu) + 1;
        match self.owner.acquire(owner) {
            Ok(()) => {
                ctx.spin_depth = next;
                Ok(Guard {
                    lock: self,
                    context: ctx,
                    mask,
                    active: true,
                    _local: PhantomData,
                })
            }
            Err(old) if old == owner => Err(Error::Recursive),
            Err(_) => Err(Error::Contended),
        }
    }
}

pub struct Guard<'a, 'g, T, O: OwnerWord, B: irq::Backend> {
    lock: &'a SpinLock<T, O>,
    context: &'a mut Context,
    mask: &'a irq::Guard<'g, B>,
    active: bool,
    _local: PhantomData<*mut ()>,
}
impl<'g, T, O: OwnerWord, B: irq::Backend> Guard<'_, 'g, T, O, B> {
    pub fn may_sleep(&self) -> Result<Task, Error> {
        self.context.may_sleep(self.mask.observation().1)
    }
    pub fn nested<'b, U, P: OwnerWord>(
        &'b mut self,
        lock: &'b SpinLock<U, P>,
    ) -> Result<Guard<'b, 'g, U, P, B>, Error> {
        lock.try_lock(self.context, self.mask)
    }
    pub fn finish(mut self) -> Result<(), (Error, Self)> {
        match self.unlock() {
            Ok(()) => Ok(()),
            Err(error) => Err((error, self)),
        }
    }
    fn unlock(&mut self) -> Result<(), Error> {
        if !self.active {
            return Ok(());
        }
        if self.mask.check_masked().is_err() || self.context.spin_depth == 0 {
            return Err(Error::Quarantined);
        }
        self.lock
            .owner
            .release(u32::from(self.context.cpu) + 1)
            .map_err(|_| Error::Quarantined)?;
        self.context.spin_depth -= 1;
        self.active = false;
        Ok(())
    }
}
impl<T, O: OwnerWord, B: irq::Backend> Deref for Guard<'_, '_, T, O, B> {
    type Target = T;
    fn deref(&self) -> &T {
        // SAFETY: only the uniquely acquired guard accesses this payload.
        unsafe { &*self.lock.value.get() }
    }
}
impl<T, O: OwnerWord, B: irq::Backend> DerefMut for Guard<'_, '_, T, O, B> {
    fn deref_mut(&mut self) -> &mut T {
        // SAFETY: exclusive guard borrow and owner word exclude other payload access.
        unsafe { &mut *self.lock.value.get() }
    }
}
impl<T, O: OwnerWord, B: irq::Backend> Drop for Guard<'_, '_, T, O, B> {
    fn drop(&mut self) {
        if self.unlock().is_err() {
            self.context.quarantined = true;
            self.lock.quarantined.store(true, Ordering::Release);
        }
    }
}

#[cfg(test)]
mod exhaustion_tests {
    use super::*;
    struct Registers(u64);
    unsafe impl irq::Backend for Registers {
        fn cpu(&self) -> u16 {
            0
        }
        fn read_daif(&self) -> u64 {
            self.0
        }
        fn write_daif(&mut self, v: u64) {
            self.0 = v
        }
    }
    #[test]
    fn spin_depth_exhaustion_does_not_acquire_owner() {
        let mut irq = irq::IrqState::new(Registers(0)).ok().unwrap();
        let mask = irq.enter().unwrap();
        let mut context = unsafe { Context::new(0, None, true).unwrap() };
        context.spin_depth = u32::MAX;
        let lock = SpinLock::new(1, 17).unwrap();
        assert!(matches!(
            lock.try_lock(&mut context, &mask),
            Err(Error::Overflow)
        ));
        assert_eq!(lock.owner_cpu(), None);
        assert_eq!(context.spin_depth, u32::MAX);
    }
    #[test]
    fn retained_spin_depth_forbids_sleep_even_if_irq_state_is_corrupted() {
        let mut context = unsafe {
            Context::new(
                0,
                Some(Task {
                    id: 1,
                    generation: 1,
                }),
                false,
            )
            .unwrap()
        };
        context.spin_depth = 1;
        assert_eq!(context.may_sleep(0), Err(Error::Context));
        assert_eq!(context.spin_depth(), 1);
    }
}