use core::sync::atomic::{Ordering, fence};

use osom_lib_alloc::traits::Allocator;
use osom_lib_reprc::macros::reprc;

use crate::immutable::{ImmutableString, internal_string::InternalString};

/// A weak reference to the underlying [`ImmutableString`].
///
/// This object cannot inspect the underlying string. But it does track
/// weak references, and each weak reference can build a strong reference,
/// assuming any other strong reference is alive.
///
/// # Examples
///
/// ```rust
/// # cfg_select! {
/// #    feature="std" => {
/// use osom_lib_strings::immutable::ImmutableStringError;
/// use osom_lib_strings::immutable::std::StdImmutableString;
///
/// fn main() -> Result<(), ImmutableStringError> {
///     // The first instance counts as both strong and weak reference.
///     let strong1 = StdImmutableString::from_str_slice("foo")?;
///     assert_eq!(strong1.strong_count(), 1);
///     assert_eq!(strong1.weak_count(), 1);
///     assert_eq!(strong1, "foo");
///
///     let weak = strong1.downgrade();  // Cheap operation
///     assert_eq!(strong1.strong_count(), 1);
///     assert_eq!(strong1.weak_count(), 2);
///     assert_eq!(weak.strong_count(), 1);
///     assert_eq!(weak.weak_count(), 2);
///
///     let strong2 = weak.upgrade().unwrap();  // Cheap operation
///     assert_eq!(strong1.strong_count(), 2);
///     assert_eq!(strong1.weak_count(), 2);
///     assert_eq!(weak.strong_count(), 2);
///     assert_eq!(weak.weak_count(), 2);
///     assert_eq!(strong2, strong1);
///
///     drop(strong2);
///     assert_eq!(strong1.strong_count(), 1);
///     assert_eq!(strong1.weak_count(), 2);
///     assert_eq!(weak.strong_count(), 1);
///     assert_eq!(weak.weak_count(), 2);
///
///     drop(strong1);
///     // Both counters are decreased, since there are no more strong references.
///     assert_eq!(weak.strong_count(), 0);
///     assert_eq!(weak.weak_count(), 1);
///
///     // No more strong reference, .upgrade() won't work.
///     assert!(weak.upgrade().is_none());
///
///     drop(weak);  // The actual deallocation happens here.
///     Ok(())
/// }
/// #   },
/// #   _ => { }
/// # }
/// ```
#[reprc]
#[repr(transparent)]
#[derive(Debug)]
#[must_use]
pub struct WeakImmutableString<TAllocator: Allocator> {
    internal: InternalString<TAllocator>,
}

impl<TAllocator: Allocator> WeakImmutableString<TAllocator> {
    #[inline(always)]
    pub(crate) fn from_internal(internal: InternalString<TAllocator>) -> Self {
        Self { internal }
    }

    #[inline(always)]
    #[must_use]
    pub fn strong_count(&self) -> u32 {
        self.internal.strong().load(Ordering::Relaxed)
    }

    #[inline(always)]
    #[must_use]
    pub fn weak_count(&self) -> u32 {
        self.internal.weak().load(Ordering::Relaxed)
    }

    /// Upgrades current weak reference to the strong [`ImmutableString`].
    ///
    /// Returns `None` if this this cannot be done (because there were no strong
    /// references alive). Otherwise returns `Some()` with a strong reference.
    #[must_use]
    pub fn upgrade(&self) -> Option<ImmutableString<TAllocator>>
    where
        TAllocator: Clone,
    {
        let strong = self.internal.strong();
        let mut current = strong.load(Ordering::Relaxed);
        loop {
            if current == 0 {
                return None;
            }
            match strong.compare_exchange_weak(current, current + 1, Ordering::Acquire, Ordering::Relaxed) {
                Ok(_) => return Some(ImmutableString::from_internal(self.internal.clone())),
                Err(new) => current = new,
            }
        }
    }

    /// Abandons current weak reference.
    ///
    /// If the internal weak counter is positive it returns false.
    ///
    /// Otherwise it deallocates the underlying memory and returns true.
    /// In particular only single (the last) [`WeakImmutableString`] returns true
    /// by calling this.
    #[inline(always)]
    #[must_use]
    pub fn abandon(mut self) -> bool {
        let result = self.internal_abandon();
        core::mem::forget(self);
        result
    }

    fn internal_abandon(&mut self) -> bool {
        let internal = unsafe { core::ptr::read(&raw const self.internal) };
        let prev = internal.weak().fetch_sub(1, Ordering::Release);
        if prev > 1 {
            return false;
        }

        // Synchronize with all prior Release decrements before deallocating.
        fence(Ordering::Acquire);
        internal.deallocate();
        true
    }
}

impl<TAllocator: Allocator> Drop for WeakImmutableString<TAllocator> {
    fn drop(&mut self) {
        let _ = self.internal_abandon();
    }
}

impl<TAllocator: Allocator> Clone for WeakImmutableString<TAllocator> {
    fn clone(&self) -> Self {
        let internal_clone = self.internal.clone();
        internal_clone.weak().fetch_add(1, Ordering::Relaxed);
        Self {
            internal: internal_clone,
        }
    }
}
