//! K419 legacy single-CPU line ownership; K468–K469 isolated SMP owner.
//! The legacy lock is not a compare-and-swap lock and must remain single-CPU.
//! The SMP lock requires both CPUs to map its RAM as Normal Inner-Shareable.

use core::sync::atomic::AtomicU64;
use core::sync::atomic::{AtomicUsize, Ordering};

pub struct SerialLineLock {
    owned: AtomicUsize,
}

#[must_use = "dropping the guard releases the serial line"]
pub struct SerialLineGuard<'a> {
    lock: &'a SerialLineLock,
}

impl SerialLineLock {
    pub const fn new() -> Self {
        Self {
            owned: AtomicUsize::new(0),
        }
    }

    pub fn try_lock(&self) -> Option<SerialLineGuard<'_>> {
        if self.owned.load(Ordering::Acquire) != 0 {
            return None;
        }
        self.owned.store(1, Ordering::Release);
        Some(SerialLineGuard { lock: self })
    }

    pub fn lock(&self) -> SerialLineGuard<'_> {
        loop {
            if let Some(guard) = self.try_lock() {
                return guard;
            }
            core::hint::spin_loop();
        }
    }

    pub fn is_locked(&self) -> bool {
        self.owned.load(Ordering::Acquire) != 0
    }
}

impl Default for SerialLineLock {
    fn default() -> Self {
        Self::new()
    }
}

impl Drop for SerialLineGuard<'_> {
    fn drop(&mut self) {
        self.lock.owned.store(0, Ordering::Release);
    }
}

/// K468: independent SMP UART line owner. Only use after both CPUs map this
/// object as writable Normal Inner-Shareable RAM with caches enabled.
pub struct SmpSerialLineLock {
    owner: AtomicUsize,
    contentions: AtomicU64,
    rejected_unlocks: AtomicU64,
    irq_denials: AtomicU64,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IrqLockError {
    InvalidCpu,
    InterruptContext,
}

pub struct SmpSerialLineGuard<'a> {
    lock: &'a SmpSerialLineLock,
    cpu: usize,
}

impl SmpSerialLineLock {
    pub const fn new() -> Self {
        Self {
            owner: AtomicUsize::new(0),
            contentions: AtomicU64::new(0),
            rejected_unlocks: AtomicU64::new(0),
            irq_denials: AtomicU64::new(0),
        }
    }

    pub fn try_lock(&self, cpu: usize) -> Option<SmpSerialLineGuard<'_>> {
        assert!(cpu < 2, "K468 lock accepts only CPU0/CPU1");
        match self
            .owner
            .compare_exchange(0, cpu + 1, Ordering::Acquire, Ordering::Relaxed)
        {
            Ok(_) => Some(SmpSerialLineGuard { lock: self, cpu }),
            Err(_) => {
                self.contentions.fetch_add(1, Ordering::Relaxed);
                None
            }
        }
    }

    pub fn lock(&self, cpu: usize) -> SmpSerialLineGuard<'_> {
        loop {
            if let Some(guard) = self.try_lock(cpu) {
                return guard;
            }
            core::hint::spin_loop();
        }
    }

    pub fn owner(&self) -> usize {
        self.owner.load(Ordering::Acquire)
    }
    pub fn contention_count(&self) -> u64 {
        self.contentions.load(Ordering::Relaxed)
    }

    /// Rejects a release unless the caller-supplied CPU identity owns the line.
    /// The caller must obtain `cpu` from its current hardware CPU, not a task ID.
    pub fn try_unlock(&self, cpu: usize) -> bool {
        if cpu >= 2 {
            self.rejected_unlocks.fetch_add(1, Ordering::Relaxed);
            return false;
        }
        if self
            .owner
            .compare_exchange(cpu + 1, 0, Ordering::Release, Ordering::Relaxed)
            .is_ok()
        {
            true
        } else {
            self.rejected_unlocks.fetch_add(1, Ordering::Relaxed);
            false
        }
    }

    pub fn rejected_unlock_count(&self) -> u64 {
        self.rejected_unlocks.load(Ordering::Relaxed)
    }

    /// IRQ handlers must never spin on a lock held by their interrupted CPU.
    /// This API intentionally cannot return a guard, even when the lock is free.
    pub fn try_lock_from_irq(&self, cpu: usize) -> Result<(), IrqLockError> {
        if cpu >= 2 {
            return Err(IrqLockError::InvalidCpu);
        }
        self.irq_denials.fetch_add(1, Ordering::Relaxed);
        Err(IrqLockError::InterruptContext)
    }

    pub fn irq_denial_count(&self) -> u64 {
        self.irq_denials.load(Ordering::Relaxed)
    }
}

impl Default for SmpSerialLineLock {
    fn default() -> Self {
        Self::new()
    }
}

impl Drop for SmpSerialLineGuard<'_> {
    fn drop(&mut self) {
        assert!(
            self.lock.try_unlock(self.cpu),
            "K469 guard must still own the line"
        );
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    extern crate std;
    use std::sync::{Arc, Mutex};
    use std::thread;
    use std::{format, string::String, vec::Vec};

    #[test]
    fn k419_second_writer_is_busy_until_first_guard_drops() {
        let lock = SerialLineLock::new();
        assert!(!lock.is_locked());
        let first = lock.try_lock().expect("first writer");
        assert!(lock.is_locked());
        assert!(
            lock.try_lock().is_none(),
            "second writer must not enter a partial line"
        );
        drop(first);
        assert!(!lock.is_locked());
        let second = lock.lock();
        assert!(lock.is_locked());
        drop(second);
        assert!(!lock.is_locked());
    }

    #[test]
    fn k468_two_threads_keep_complete_lines() {
        let lock = Arc::new(SmpSerialLineLock::new());
        let output = Arc::new(Mutex::new(String::new()));
        let held = lock.lock(0);
        assert!(lock.try_lock(1).is_none());
        drop(held);
        let threads: Vec<_> = (0..2)
            .map(|cpu| {
                let lock = lock.clone();
                let output = output.clone();
                thread::spawn(move || {
                    for index in 0..32 {
                        let _guard = lock.lock(cpu);
                        let line = format!(
                            "CPU={cpu} INDEX={index} PAYLOAD={}\n",
                            if cpu == 0 {
                                "A".repeat(64)
                            } else {
                                "B".repeat(64)
                            }
                        );
                        for byte in line.chars() {
                            output.lock().unwrap().push(byte);
                            thread::yield_now();
                        }
                    }
                })
            })
            .collect();
        for handle in threads {
            handle.join().unwrap();
        }
        let output = output.lock().unwrap();
        let lines: Vec<_> = output.lines().collect();
        assert_eq!(lines.len(), 64);
        for cpu in 0..2 {
            for index in 0..32 {
                let expected = format!(
                    "CPU={cpu} INDEX={index} PAYLOAD={}",
                    if cpu == 0 {
                        "A".repeat(64)
                    } else {
                        "B".repeat(64)
                    }
                );
                assert_eq!(lines.iter().filter(|line| **line == expected).count(), 1);
            }
        }
        assert_eq!(lock.owner(), 0);
        assert!(lock.contention_count() > 0);
    }

    #[test]
    fn k469_wrong_cpu_unlock_keeps_original_owner() {
        let lock = Arc::new(SmpSerialLineLock::new());
        let held = lock.lock(0);
        let foreign = lock.clone();
        assert!(!thread::spawn(move || foreign.try_unlock(1)).join().unwrap());
        assert_eq!(lock.owner(), 1);
        assert!(lock.try_lock(1).is_none());
        assert_eq!(lock.rejected_unlock_count(), 1);
        drop(held);
        assert_eq!(lock.owner(), 0);
        let second = lock.lock(1);
        drop(second);
        assert_eq!(lock.owner(), 0);
        assert!(!lock.try_unlock(0), "double release must fail");
        assert_eq!(lock.rejected_unlock_count(), 2);
    }

    #[test]
    fn k470_irq_context_never_spins_or_steals_owner() {
        let lock = SmpSerialLineLock::new();
        let held = lock.lock(0);
        assert_eq!(
            lock.try_lock_from_irq(0),
            Err(IrqLockError::InterruptContext)
        );
        assert_eq!(lock.owner(), 1);
        drop(held);
        assert_eq!(
            lock.try_lock_from_irq(0),
            Err(IrqLockError::InterruptContext)
        );
        assert_eq!(lock.try_lock_from_irq(2), Err(IrqLockError::InvalidCpu));
        assert_eq!(lock.irq_denial_count(), 2);
        assert_eq!(lock.owner(), 0);
        assert_eq!(lock.rejected_unlock_count(), 0);
    }
}