use crate::mpsc_rbuffer::consumer::MPSCConsumer;
use crate::mpsc_rbuffer::producer::MPSCProducer;
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;

/// A single storage cell in the ring buffer.
///
/// Each slot contains a value and a `sequence` number used to coordinate
/// multiple producers and the single consumer.
pub struct Slot<T>
where
    T: Send,
{
    /// The actual data stored in the slot.
    pub value: UnsafeCell<MaybeUninit<T>>,
    /// The monotonic sequence number used for synchronization.
    /// A producer can write when `sequence == position`.
    /// A consumer can read when `sequence == position + 1`.
    pub sequence: AtomicUsize,
}

impl<T: Send> Slot<T> {
    /// Creates an empty slot whose initial sequence equals its ring index `i`.
    fn new(i: usize) -> Self {
        Self {
            value: UnsafeCell::new(MaybeUninit::uninit()),
            sequence: AtomicUsize::new(i),
        }
    }
}

/// Lock-free multi-producer / single-consumer (MPSC) ring buffer storage.
///
/// Producers coordinate via an atomic fetch-and-add on `real_write_index` to
/// claim a slot, then use the per-slot `sequence` to ensure the consumer
/// only reads fully initialized data.
///
/// # Constraints
/// * `S` must be a power of two and ≥ 2.
/// * Usable capacity is `S - 1` to distinguish between full and empty states.
/// * `T` must be `Send` but does not require `Sync`.
///
/// # Usage
/// Call [`split`](Self::split) to obtain a producer/consumer pair. The producer
/// handle can be cloned for use in multiple threads.
pub struct MPSCRBuffer<T, const S: usize>
where
    T: Send,
{
    /// Shared write index claimed atomically by all producers.
    pub(super) real_write_index: CachePadded<AtomicUsize>,
    /// Private read index for the single consumer.
    pub(super) real_read_index: CachePadded<UnsafeCell<usize>>,
    /// Backing storage consisting of sequence-tracked slots.
    pub(super) ring: [UnsafeCell<Slot<T>>; S],
}

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

impl<T, const S: usize> Drop for MPSCRBuffer<T, S>
where
    T: Send,
{
    fn drop(&mut self) {
        let write_position = self.real_write_index.0.load(Ordering::Acquire);
        // SAFETY: exclusive access via &mut self; no concurrent consumer.
        let mut read_position = unsafe { *self.real_read_index.0.get() };

        while read_position != write_position {
            let index = read_position & (S - 1);
            // SAFETY: index is in-bounds by construction.
            let slot = unsafe { &*self.ring.get_unchecked(index).get() };

            let expected = read_position.wrapping_add(1);
            let sequence = slot.sequence.load(Ordering::Acquire);

            if sequence == expected {
                // SAFETY: sequence confirms the producer wrote this slot.
                unsafe { (*slot.value.get()).assume_init_drop() };
                read_position = read_position.wrapping_add(1);
            } else {
                unreachable!(
                    "MPSCRBuffer::drop: slot {index} has unexpected sequence \
                     {sequence} (expected {expected}); a producer may have been \
                     abandoned mid-write. Slot will be leaked."
                );
            }
        }
    }
}

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

impl<T, const S: usize> MPSCRBuffer<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(UnsafeCell::new(0)),
            ring: core::array::from_fn(|i| UnsafeCell::new(Slot::new(i))),
        }
    }

    //// 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);
            }

            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(UnsafeCell::new(0)),
            );

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

            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<Slot<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 {
            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 a producer and a consumer handle.
    ///
    /// The producer handle can be cloned to support multiple producers.
    /// Both handles borrow the buffer, ensuring memory safety for the duration
    /// of the handles' lifetimes.
    pub fn split(&mut self) -> (MPSCProducer<'_, T, S>, MPSCConsumer<'_, T, S>) {
        // SAFETY: &mut self guarantees no concurrent access.
        let read = unsafe { *self.real_read_index.0.get() };
        (
            MPSCProducer { buffer: self },
            MPSCConsumer {
                buffer: self,
                position: read,
            },
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mpsc_rbuffer::producer::MPSCProducer;
    use core::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::{Arc, Barrier};
    use std::thread;
    use std::vec::Vec;

    fn positions<T: Send, const S: usize>(buf: &MPSCRBuffer<T, S>) -> (usize, usize) {
        let write = buf.real_write_index.0.load(Ordering::Acquire);
        // SAFETY: test-only read; no concurrent consumer in these call sites.
        let read = unsafe { *buf.real_read_index.0.get() };
        (write, read)
    }

    fn assert_distance_within_capacity<T: Send, const S: usize>(buf: &MPSCRBuffer<T, S>) {
        let (write, read) = positions(buf);
        assert!(
            write >= read,
            "write index regressed: write={write} read={read}"
        );
        // S - 1 is the maximum number of live elements; S would mean the
        // buffer is simultaneously full and empty (impossible by design).
        assert!(
            write - read < S,
            "buffer distance exceeded capacity: write={write} read={read} S={S}"
        );
    }

    #[derive(Debug)]
    struct DropTracker {
        id: usize,
        per_id: Arc<[AtomicUsize]>,
        total: Arc<AtomicUsize>,
    }

    impl Drop for DropTracker {
        fn drop(&mut self) {
            self.per_id[self.id].fetch_add(1, Ordering::AcqRel);
            self.total.fetch_add(1, Ordering::AcqRel);
        }
    }

    #[derive(Debug)]
    struct CycleDrop<'a> {
        drops: &'a AtomicUsize,
    }

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

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

    #[test]
    #[should_panic]
    fn rejects_zero_capacity() {
        let _ = MPSCRBuffer::<u8, 0>::new();
    }

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

    #[test]
    fn slot_reuse_across_many_cycles_large_capacity() {
        const SIZE: usize = 4096;
        const CYCLES: usize = 16;

        let mut buffer = MPSCRBuffer::<usize, SIZE>::new();
        let (mut producer, mut consumer) = buffer.split();

        for cycle in 0..CYCLES {
            let base = cycle * SIZE;
            for i in 0..SIZE {
                assert!(producer.push(base + i).is_ok());
            }
            for i in 0..SIZE {
                assert_eq!(consumer.pop(), base + i);
            }
            assert_eq!(consumer.try_pop(), None);
            assert_distance_within_capacity(consumer.buffer);
        }
    }

    #[test]
    fn multi_producer_high_contention_no_loss_no_duplicates() {
        const PRODUCERS: usize = 6;
        const ITEMS_PER_PRODUCER: usize = 20_000;
        const TOTAL: usize = PRODUCERS * ITEMS_PER_PRODUCER;

        let mut buffer = MPSCRBuffer::<u64, 1024>::new();
        let (producer0, mut consumer) = buffer.split();
        let shared = producer0.buffer;

        let barrier = Arc::new(Barrier::new(PRODUCERS + 1));
        let mut next_expected = [0usize; PRODUCERS];
        let mut seen = [0usize; PRODUCERS];

        thread::scope(|scope| {
            let mut first = Some(producer0);
            for pid in 0..PRODUCERS {
                let start = barrier.clone();
                let mut producer = if pid == 0 {
                    first.take().expect("missing first producer")
                } else {
                    MPSCProducer { buffer: shared }
                };

                scope.spawn(move || {
                    start.wait();
                    let mut rnd = 0x9E3779B97F4A7C15_u64 ^ pid as u64;
                    for seq in 0..ITEMS_PER_PRODUCER {
                        let value = ((pid as u64) << 32) | (seq as u64);
                        assert!(producer.push(value).is_ok());
                        rnd ^= rnd << 13;
                        rnd ^= rnd >> 7;
                        rnd ^= rnd << 17;
                        if rnd & 0x3F == 0 {
                            thread::yield_now();
                        }
                    }
                });
            }

            barrier.wait();

            for n in 0..TOTAL {
                let value = consumer.pop();
                let pid = (value >> 32) as usize;
                let seq = (value & 0xFFFF_FFFF) as usize;

                assert!(pid < PRODUCERS, "invalid producer id {pid}");
                assert_eq!(
                    seq, next_expected[pid],
                    "out-of-order item from producer {pid}"
                );
                next_expected[pid] += 1;
                seen[pid] += 1;

                if n & 0x1FFF == 0 {
                    thread::yield_now();
                }
            }
        });

        for (pid, s) in seen.iter().enumerate() {
            assert_eq!(*s, ITEMS_PER_PRODUCER, "missing values from pid={pid}");
        }
        assert_eq!(consumer.try_pop(), None);
    }

    #[test]
    fn multi_producer_stress_repeated_runs() {
        const RUNS: usize = 24;
        const PRODUCERS: usize = 4;
        const ITEMS_PER_PRODUCER: usize = 3_000;
        const TOTAL: usize = PRODUCERS * ITEMS_PER_PRODUCER;

        for run in 0..RUNS {
            let mut buffer = MPSCRBuffer::<u64, 256>::new();
            let (producer0, mut consumer) = buffer.split();
            let shared = producer0.buffer;
            let barrier = Arc::new(Barrier::new(PRODUCERS + 1));
            let mut seen = vec![0usize; TOTAL];

            thread::scope(|scope| {
                let mut first = Some(producer0);
                for pid in 0..PRODUCERS {
                    let start = barrier.clone();
                    let mut producer = if pid == 0 {
                        first.take().expect("missing first producer")
                    } else {
                        MPSCProducer { buffer: shared }
                    };

                    scope.spawn(move || {
                        start.wait();
                        for seq in 0..ITEMS_PER_PRODUCER {
                            let id = pid * ITEMS_PER_PRODUCER + seq;
                            assert!(producer.push(id as u64).is_ok());
                            if (seq + pid + run) & 0x7F == 0 {
                                thread::yield_now();
                            }
                        }
                    });
                }

                barrier.wait();
                for _ in 0..TOTAL {
                    let id = consumer.pop() as usize;
                    assert!(id < TOTAL, "out-of-range id {id}");
                    seen[id] += 1;
                }
            });

            for (id, count) in seen.into_iter().enumerate() {
                assert_eq!(count, 1, "duplicate/lost element at id={id}, run={run}");
            }
            assert_eq!(consumer.try_pop(), None);
        }
    }

    #[test]
    fn drop_with_remaining_elements_drops_exactly_once() {
        const N: usize = 96;
        let per_id: Arc<[AtomicUsize]> = (0..N)
            .map(|_| AtomicUsize::new(0))
            .collect::<Vec<_>>()
            .into();
        let total = Arc::new(AtomicUsize::new(0));

        {
            let mut buffer = MPSCRBuffer::<DropTracker, 128>::new();
            let (mut producer, mut consumer) = buffer.split();

            for id in 0..N {
                assert!(
                    producer
                        .push(DropTracker {
                            id,
                            per_id: per_id.clone(),
                            total: total.clone(),
                        })
                        .is_ok()
                );
            }

            let mut consumed = Vec::with_capacity(32);
            for _ in 0..32 {
                consumed.push(consumer.pop());
            }
            drop(consumed);
        }

        assert_eq!(total.load(Ordering::Acquire), N);
        for id in 0..N {
            let drops = per_id[id].load(Ordering::Acquire);
            assert_eq!(drops, 1, "value id={id} dropped {drops} times");
        }
    }

    #[test]
    fn repeated_create_destroy_cycles_drop_all_values_once() {
        const CYCLES: usize = 200;
        let drops = AtomicUsize::new(0);
        let mut created = 0usize;

        for cycle in 0..CYCLES {
            let mut buffer = MPSCRBuffer::<CycleDrop<'_>, 32>::new();
            let (mut producer, mut consumer) = buffer.split();

            let to_push = (cycle % 24) + 1;
            let to_pop = to_push / 3;

            for _ in 0..to_push {
                created += 1;
                assert!(producer.push(CycleDrop { drops: &drops }).is_ok());
            }
            for _ in 0..to_pop {
                drop(consumer.pop());
            }
        }

        assert_eq!(drops.load(Ordering::Acquire), created);
    }

    #[test]
    fn boxed_buffer_sequence_initialised_correctly() {
        let mut buffer = MPSCRBuffer::<u32, 8>::new_boxed();
        let (mut producer, mut consumer) = buffer.split();

        for i in 0..7_u32 {
            assert!(producer.push(i).is_ok(), "push failed at slot {i}");
        }
        for i in 0..7_u32 {
            assert_eq!(consumer.pop(), i);
        }
        assert_eq!(consumer.try_pop(), None);
    }

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

        let (mut producer, mut consumer) = buffer.split();
        for i in 0..15_u64 {
            assert!(producer.push(i * 7).is_ok());
        }
        for i in 0..15_u64 {
            assert_eq!(consumer.pop(), i * 7);
        }
        assert_eq!(consumer.try_pop(), None);
    }
}
