use super::consumer::SPSCConsumer;
use super::producer::SPSCProducer;
use crate::util::{CACHE_LINE_SIZE, CachePadded, assert_valid_capacity};
use core::cell::UnsafeCell;
use core::mem::MaybeUninit;
use core::sync::atomic::{AtomicUsize, Ordering};
use std::alloc::{Layout, alloc, handle_alloc_error};
use std::ptr;

/// Lock-free single-producer / single-consumer ring buffer storage.
///
/// `S` must be a power of two and ≥ 2. Usable capacity is `S - 1` elements.
/// Call [`split`](Self::split) to obtain the producer and consumer handles.
///
/// The two atomic indices are cache-padded onto separate lines to prevent
/// false-sharing between producer and consumer cores.
pub struct SPSCRBuffer<T, const S: usize>
where
    T: Send,
{
    /// Written exclusively by the producer.
    pub(super) real_write_index: CachePadded<AtomicUsize>,
    /// Written exclusively by the consumer.
    pub(super) real_read_index: CachePadded<AtomicUsize>,
    /// Backing storage.
    pub(super) ring: [UnsafeCell<MaybeUninit<T>>; S],
}

unsafe impl<T: Send, const S: usize> Sync for SPSCRBuffer<T, S> {}

impl<T, const S: usize> Drop for SPSCRBuffer<T, S>
where
    T: Send,
{
    fn drop(&mut self) {
        let write = self.real_write_index.0.load(Ordering::Acquire);
        let mut read = self.real_read_index.0.load(Ordering::Acquire);

        let mask = S - 1;
        while read != write {
            // SAFETY: Slots in [read, write) are producer-initialised and not yet consumed.
            unsafe { (*self.ring[read & mask].get()).assume_init_drop() };
            read = read.wrapping_add(1);
        }
    }
}

impl<T, const S: usize> Default for SPSCRBuffer<T, S>
where
    T: Send,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<T, const S: usize> SPSCRBuffer<T, S>
where
    T: Send,
{
    /// Creates an empty ring buffer on the stack.
    ///
    /// **Note:** Use [`new_boxed`](Self::new_boxed) for large `S`.
    ///
    /// # Panics
    /// If `S` is not a power of two, `S < 2`, or `T` is zero-sized.
    pub fn new() -> Self {
        assert_valid_capacity::<T, S>();
        Self {
            real_write_index: CachePadded(AtomicUsize::new(0)),
            real_read_index: CachePadded(AtomicUsize::new(0)),
            ring: core::array::from_fn(|_| UnsafeCell::new(MaybeUninit::uninit())),
        }
    }

    /// Heap allocation. Bypasses stack to prevent overflows.
    ///
    /// # Panics
    /// Same conditions as [`new`](Self::new).
    pub fn new_boxed() -> Box<Self> {
        assert_valid_capacity::<T, S>();

        let layout = Layout::new::<Self>();

        unsafe {
            let ptr = alloc(layout).cast::<Self>();
            if ptr.is_null() {
                handle_alloc_error(layout);
            }

            // Write each field directly to avoid reading uninitialised memory.
            ptr::write(
                ptr::addr_of_mut!((*ptr).real_write_index),
                CachePadded(AtomicUsize::new(0)),
            );
            ptr::write(
                ptr::addr_of_mut!((*ptr).real_read_index),
                CachePadded(AtomicUsize::new(0)),
            );

            let ring_base = ptr::addr_of_mut!((*ptr).ring).cast::<UnsafeCell<MaybeUninit<T>>>();
            for i in 0..S {
                ptr::write(ring_base.add(i), UnsafeCell::new(MaybeUninit::uninit()));
            }

            Box::from_raw(ptr)
        }
    }

    /// Pre-faults OS pages and warms the CPU cache for the ring storage.
    ///
    /// Iterates through the buffer at page and cache-line boundaries to
    /// eliminate cold-start latency and page faults on the hot path.
    #[cfg(feature = "perf")]
    pub fn warm_up(&mut self) {
        use crate::util::page_size::get_page_size;

        let page_size = get_page_size();
        let slot_size = core::mem::size_of::<UnsafeCell<MaybeUninit<T>>>();
        let page_stride = (page_size / slot_size).max(1);
        let cache_stride = (CACHE_LINE_SIZE / slot_size).max(1);

        let mut i = 0usize;
        while i < S {
            // SAFETY: i < S; writing one byte into MaybeUninit<T> is always
            // sound — no validity invariant exists on uninitialised storage.
            unsafe { ptr::write_volatile(self.ring[i].get().cast::<u8>(), 0u8) };

            let page_end = (i / page_stride + 1) * page_stride;
            let next_cache = i + cache_stride;
            i = if next_cache < page_end.min(S) {
                next_cache
            } else {
                page_end
            };
        }
    }

    /// Locks the buffer's memory in RAM using `mlock`, preventing swapping.
    ///
    /// Requies the `perf` feature and a Unix-like target.
    ///
    /// # Errors
    /// - `EPERM` — insufficient privileges (`CAP_IPC_LOCK` or `RLIMIT_MEMLOCK`).
    /// - `ENOMEM` — per-process lock limit exceeded.
    #[cfg(all(feature = "perf", target_family = "unix"))]
    pub fn lock_memory(&self) -> Result<(), std::io::Error> {
        let addr = (self as *const Self).cast::<libc::c_void>();
        let size = core::mem::size_of::<Self>();
        // SAFETY: addr/size describe a live region; mlock only advises the kernel.
        let rc = unsafe { libc::mlock(addr, size) };
        if rc == 0 {
            Ok(())
        } else {
            Err(std::io::Error::last_os_error())
        }
    }

    /// Splits the buffer into producer and consumer handles.
    ///
    /// Both handles borrow `self` for their lifetime, ensuring only one pair
    /// exists at a time and that the buffer outlives both. Re-splitting resumes
    /// from the current indices.
    pub fn split(&mut self) -> (SPSCProducer<'_, T, S>, SPSCConsumer<'_, T, S>) {
        // &mut self guarantees exclusivity; Relaxed ordering is sufficient.
        let write = self.real_write_index.0.load(Ordering::Relaxed);
        let read = self.real_read_index.0.load(Ordering::Relaxed);
        (
            SPSCProducer {
                buffer: self,
                write_index: write,
                cached_read_index: read,
            },
            SPSCConsumer {
                buffer: self,
                cached_write_index: write,
                read_index: read,
            },
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use core::sync::atomic::{AtomicUsize, Ordering};
    use std::panic::{AssertUnwindSafe, catch_unwind};

    #[derive(Debug)]
    struct DropTracker<'a> {
        id: u32,
        drops: &'a AtomicUsize,
        clones: &'a AtomicUsize,
    }

    impl<'a> DropTracker<'a> {
        fn new(id: u32, drops: &'a AtomicUsize, clones: &'a AtomicUsize) -> Self {
            Self { id, drops, clones }
        }
    }

    impl Clone for DropTracker<'_> {
        fn clone(&self) -> Self {
            self.clones.fetch_add(1, Ordering::Relaxed);
            Self {
                id: self.id,
                drops: self.drops,
                clones: self.clones,
            }
        }
    }

    impl Drop for DropTracker<'_> {
        fn drop(&mut self) {
            self.drops.fetch_add(1, Ordering::Relaxed);
        }
    }

    fn occupied_len<T: Send, const S: usize>(buffer: &SPSCRBuffer<T, S>) -> usize {
        let write = buffer.real_write_index.0.load(Ordering::Acquire);
        let read = buffer.real_read_index.0.load(Ordering::Acquire);
        write.wrapping_sub(read) & (S - 1)
    }

    fn assert_buffer_state<T: Send, const S: usize>(
        buffer: &SPSCRBuffer<T, S>,
        expected_len: usize,
    ) {
        let write = buffer.real_write_index.0.load(Ordering::Acquire);
        let read = buffer.real_read_index.0.load(Ordering::Acquire);
        assert!(write < S, "write index {write} out of bounds (S={S})");
        assert!(read < S, "read index {read} out of bounds (S={S})");
        let len = occupied_len(buffer);
        assert_eq!(len, expected_len, "unexpected occupancy");
        assert!(len < S, "occupancy {len} exceeds max capacity {}", S - 1);
    }

    #[test]
    #[should_panic]
    fn rejects_non_power_of_two_capacity() {
        let _ = SPSCRBuffer::<u8, 3>::new();
    }

    #[test]
    #[should_panic]
    fn rejects_capacity_of_one() {
        let _ = SPSCRBuffer::<u8, 1>::new();
    }

    #[test]
    fn unread_items_are_dropped_exactly_once() {
        let drops = AtomicUsize::new(0);
        let clones = AtomicUsize::new(0);
        {
            let mut buffer = SPSCRBuffer::<DropTracker<'_>, 8>::new();
            let (mut producer, _consumer) = buffer.split();
            for id in 0..5 {
                producer
                    .try_push(DropTracker::new(id, &drops, &clones))
                    .expect("ring should have space");
            }
            assert_eq!(drops.load(Ordering::Acquire), 0);
            assert_eq!(clones.load(Ordering::Acquire), 0);
        }
        assert_eq!(
            drops.load(Ordering::Acquire),
            5,
            "all 5 items must be dropped"
        );
        assert_eq!(clones.load(Ordering::Acquire), 0, "no clones expected");
    }

    #[test]
    fn partially_consumed_items_are_not_double_dropped() {
        let drops = AtomicUsize::new(0);
        let clones = AtomicUsize::new(0);
        {
            let mut buffer = SPSCRBuffer::<DropTracker<'_>, 8>::new();
            let (mut producer, mut consumer) = buffer.split();
            for id in 0..4 {
                producer
                    .try_push(DropTracker::new(id, &drops, &clones))
                    .expect("ring should have space");
            }
            assert_buffer_state(producer.buffer, 4);

            let v0 = consumer.try_pop().expect("should pop item 0");
            let v1 = consumer.try_pop().expect("should pop item 1");
            assert_eq!(v0.id, 0);
            assert_eq!(v1.id, 1);
            drop(v0);
            drop(v1);
            assert_eq!(drops.load(Ordering::Acquire), 2);
            assert_buffer_state(consumer.buffer, 2);
        }
        assert_eq!(drops.load(Ordering::Acquire), 4);
        assert_eq!(clones.load(Ordering::Acquire), 0);
    }

    #[test]
    fn panic_unwind_drops_all_buffered_items() {
        let drops = AtomicUsize::new(0);
        let clones = AtomicUsize::new(0);
        let result = catch_unwind(AssertUnwindSafe(|| {
            let mut buffer = SPSCRBuffer::<DropTracker<'_>, 8>::new();
            let (mut producer, _consumer) = buffer.split();
            for id in 0..3 {
                producer
                    .try_push(DropTracker::new(id, &drops, &clones))
                    .expect("ring should have space");
            }
            panic!("forced unwind");
        }));
        assert!(result.is_err(), "panic should be captured by catch_unwind");
        assert_eq!(drops.load(Ordering::Acquire), 3);
        assert_eq!(clones.load(Ordering::Acquire), 0);
    }

    #[test]
    fn resplit_resumes_from_current_indices() {
        let mut buffer = SPSCRBuffer::<u32, 8>::new();
        {
            let (mut producer, mut consumer) = buffer.split();
            assert_eq!(producer.try_push(1), Ok(()));
            assert_eq!(producer.try_push(2), Ok(()));
            assert_eq!(producer.try_push(3), Ok(()));
            assert_eq!(consumer.try_pop(), Some(1));
            assert_eq!(consumer.try_pop(), Some(2));
            assert_eq!(consumer.try_pop(), Some(3));
            assert_buffer_state(consumer.buffer, 0);
        }
        let expected_write = buffer.real_write_index.0.load(Ordering::Acquire);
        let expected_read = buffer.real_read_index.0.load(Ordering::Acquire);
        assert_eq!(expected_write, expected_read);
        assert_ne!(expected_write, 0, "indices must have advanced");

        let (producer, consumer) = buffer.split();
        assert_eq!(producer.write_index, expected_write);
        assert_eq!(producer.cached_read_index, expected_read);
        assert_eq!(consumer.read_index, expected_read);
        assert_eq!(consumer.cached_write_index, expected_write);
    }

    #[test]
    fn sequential_resplit_stress_preserves_invariant() {
        let mut buffer = SPSCRBuffer::<u32, 8>::new();
        for cycle in 0..128_u32 {
            let (mut producer, mut consumer) = buffer.split();
            assert_eq!(producer.try_push(cycle), Ok(()));
            assert_eq!(consumer.try_pop(), Some(cycle));
            assert_buffer_state(consumer.buffer, 0);
            drop(producer);
            drop(consumer);

            let exp_write = buffer.real_write_index.0.load(Ordering::Acquire);
            let exp_read = buffer.real_read_index.0.load(Ordering::Acquire);
            assert_eq!(exp_write, exp_read, "cycle={cycle}");

            let (p2, c2) = buffer.split();
            assert_eq!(p2.write_index, exp_write, "cycle={cycle}");
            assert_eq!(p2.cached_read_index, exp_read, "cycle={cycle}");
            assert_eq!(c2.read_index, exp_read, "cycle={cycle}");
            assert_eq!(c2.cached_write_index, exp_write, "cycle={cycle}");
        }
    }

    #[test]
    fn boxed_buffer_behaves_identically_to_stack_buffer() {
        let drops = AtomicUsize::new(0);
        let clones = AtomicUsize::new(0);
        {
            let mut buffer = SPSCRBuffer::<DropTracker<'_>, 16>::new_boxed();
            let (mut producer, mut consumer) = buffer.split();
            for id in 0..7 {
                producer
                    .try_push(DropTracker::new(id, &drops, &clones))
                    .expect("ring should have space");
            }
            for expected_id in 0..4_u32 {
                let item = consumer.try_pop().expect("should have item");
                assert_eq!(item.id, expected_id);
            }
            assert_buffer_state(consumer.buffer, 3);
        }
        assert_eq!(drops.load(Ordering::Acquire), 7);
        assert_eq!(clones.load(Ordering::Acquire), 0);
    }

    #[test]
    fn warm_up_does_not_corrupt_subsequent_data() {
        let mut buffer = SPSCRBuffer::<u64, 16>::new();
        buffer.warm_up();

        let (mut producer, mut consumer) = buffer.split();
        for i in 0..15_u64 {
            assert_eq!(producer.try_push(i * 3), Ok(()));
        }
        for i in 0..15_u64 {
            assert_eq!(consumer.try_pop(), Some(i * 3));
        }
        assert_buffer_state(consumer.buffer, 0);
    }
}
