use std::{
    marker::PhantomData,
    sync::{
        atomic::{AtomicPtr, Ordering},
        Arc,
    },
};

pub struct AtomicArc<T> {
    inner: AtomicPtr<T>,
    _phantom: PhantomData<Arc<T>>,
}

impl<T> Drop for AtomicArc<T> {
    fn drop(&mut self) {
        unsafe { Arc::decrement_strong_count(self.load()) }
    }
}

impl<T> AtomicArc<T> {
    fn load(&self) -> *const T {
        self.inner.load(Ordering::Relaxed)
    }

    pub fn new(inner: Arc<T>) -> Self {
        Self {
            inner: AtomicPtr::new(Arc::into_raw(inner) as *mut T),
            _phantom: PhantomData,
        }
    }

    pub fn get(&self) -> Arc<T> {
        unsafe {
            let ptr = self.load();
            Arc::increment_strong_count(ptr);
            Arc::from_raw(ptr)
        }
    }

    #[must_use]
    pub fn swap(&self, new: Arc<T>) -> Arc<T> {
        unsafe {
            let old = self
                .inner
                .swap(Arc::into_raw(new) as *mut T, Ordering::Relaxed);

            Arc::from_raw(old)
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use super::AtomicArc;

    #[test]
    fn test_basic() {
        let atomic_arc = AtomicArc::new(Arc::new(1));

        let first = atomic_arc.get();
        assert_eq!(*first, 1);

        _ = atomic_arc.swap(Arc::new(2));

        assert_eq!(*first, 1);

        Arc::into_inner(first).unwrap();

        let second = atomic_arc.get();

        std::mem::drop(atomic_arc);

        Arc::into_inner(second).unwrap();
    }
}