use std::cell::UnsafeCell;
use std::sync::atomic::{AtomicU8, Ordering};
use std::sync::Arc;

const DIRTY: u8 = 0b100;
const INDEX_MASK: u8 = 0b011;

struct Shared<T> {
    slots: [UnsafeCell<T>; 3],
    state: AtomicU8,
}

/// SAFETY: `state` is the only cross-thread coordination primitive. Producer only touches
/// `write_idx`, consumer only touches `read_idx`. The swap protocol ensures producer, consumer,
/// and ready slots are always distinct. `T: Send` is required because values cross threads.
unsafe impl<T: Send> Sync for Shared<T> {}

pub struct Producer<T> {
    shared: Arc<Shared<T>>,
    write_idx: u8,
}

pub struct Consumer<T> {
    shared: Arc<Shared<T>>,
    read_idx: u8,
}

/// Create a producer/consumer pair for triple-buffered communication
pub fn triple_buffer<T: Clone>(value: T) -> (Producer<T>, Consumer<T>) {
    let shared = Arc::new(Shared {
        slots: [
            UnsafeCell::new(value.clone()),
            UnsafeCell::new(value.clone()),
            UnsafeCell::new(value),
        ],
        state: AtomicU8::new(1), // ready=1, not dirty
    });
    (
        Producer {
            shared: shared.clone(),
            write_idx: 0,
        },
        Consumer {
            shared,
            read_idx: 2,
        },
    )
}

impl<T> Producer<T> {
    /// Mutable access to the producer-owned slot
    pub fn write(&mut self) -> &mut T {
        unsafe { &mut *self.shared.slots[self.write_idx as usize].get() }
    }

    /// Publish the current write slot
    ///
    /// The old write slot becomes the shared ready slot, and the old ready slot becomes the
    /// new write slot.
    pub fn publish(&mut self) {
        let old = self
            .shared
            .state
            .swap(self.write_idx | DIRTY, Ordering::AcqRel);
        self.write_idx = old & INDEX_MASK;
    }
}

impl<T> Consumer<T> {
    /// Read the latest published value
    ///
    /// If nothing new was published since the last call, returns the same slot as before.
    pub fn read(&mut self) -> &T {
        self.maybe_swap();
        unsafe { &*self.shared.slots[self.read_idx as usize].get() }
    }

    /// Mutable access to the consumer-owned slot
    ///
    /// Sound because the producer cannot reclaim this slot until the consumer swaps it back.
    pub fn read_mut(&mut self) -> &mut T {
        self.maybe_swap();
        unsafe { &mut *self.shared.slots[self.read_idx as usize].get() }
    }

    fn maybe_swap(&mut self) {
        let state = self.shared.state.load(Ordering::Acquire);
        if state & DIRTY != 0 {
            let old = self.shared.state.swap(self.read_idx, Ordering::AcqRel);
            self.read_idx = old & INDEX_MASK;
        }
    }
}
