//! Instrumented wrapper for [`std::sync::Mutex`].

use std::sync::Mutex as StdMutex;

use crate::instant::Instant;
use crate::mutexes::{
    cancel_wait_stamp, elapsed_nanos, register_mutex, send_mutex_event, wait_stamp,
    InstrumentMutex, MutexEvent,
};

/// Instrumented drop-in replacement for [`std::sync::Mutex`].
///
/// Not constructed directly - use the [`mutex!`](crate::mutex) macro.
pub struct Mutex<T> {
    inner: StdMutex<T>,
    id: u32,
}

impl<T> Mutex<T> {
    /// Drop-in constructor for the `hotpath_meta::wrap` prefix migration. Captures the
    /// caller location as the registered source.
    #[track_caller]
    #[deprecated(note = "construct via the hotpath_meta::mutex! macro instead of new()")]
    pub fn new(value: T) -> Self {
        let loc = std::panic::Location::caller();
        let source: &'static str =
            Box::leak(format!("{}:{}", loc.file(), loc.line()).into_boxed_str());
        Self::__new_instrumented(StdMutex::new(value), source, None)
    }

    #[doc(hidden)]
    pub fn __new_instrumented(
        inner: StdMutex<T>,
        source: &'static str,
        label: Option<String>,
    ) -> Self {
        let id = register_mutex::<T>(source, label);
        Self { inner, id }
    }

    pub fn lock(&self) -> std::sync::LockResult<MutexGuard<'_, T>> {
        // Stamp before acquisition to measure wait time; the guard then measures acquire time.
        let wait_start = wait_stamp();
        match self.inner.lock() {
            Ok(inner) => Ok(self.guard(inner, wait_start.map(elapsed_nanos))),
            Err(poison) => Err(std::sync::PoisonError::new(
                self.guard(poison.into_inner(), wait_start.map(elapsed_nanos)),
            )),
        }
    }

    pub fn try_lock(&self) -> std::sync::TryLockResult<MutexGuard<'_, T>> {
        let wait_start = wait_stamp();
        match self.inner.try_lock() {
            Ok(inner) => Ok(self.guard(inner, wait_start.map(elapsed_nanos))),
            Err(std::sync::TryLockError::Poisoned(poison)) => Err(
                std::sync::TryLockError::Poisoned(std::sync::PoisonError::new(
                    self.guard(poison.into_inner(), wait_start.map(elapsed_nanos)),
                )),
            ),
            Err(std::sync::TryLockError::WouldBlock) => {
                cancel_wait_stamp();
                Err(std::sync::TryLockError::WouldBlock)
            }
        }
    }

    pub fn into_inner(self) -> std::sync::LockResult<T> {
        self.inner.into_inner()
    }

    pub fn get_mut(&mut self) -> std::sync::LockResult<&mut T> {
        self.inner.get_mut()
    }

    fn guard<'a>(
        &'a self,
        inner: std::sync::MutexGuard<'a, T>,
        wait_nanos: Option<u64>,
    ) -> MutexGuard<'a, T> {
        MutexGuard {
            inner: Some(inner),
            start: wait_nanos.map(|_| Instant::now()),
            wait_nanos,
            id: self.id,
        }
    }
}

/// Guard returned by [`Mutex::lock`]. Emits wait and acquire durations on drop.
#[must_use = "if unused the Mutex will immediately unlock"]
pub struct MutexGuard<'a, T> {
    inner: Option<std::sync::MutexGuard<'a, T>>,
    start: Option<Instant>,
    wait_nanos: Option<u64>,
    id: u32,
}

impl<T> std::ops::Deref for MutexGuard<'_, T> {
    type Target = T;
    fn deref(&self) -> &T {
        self.inner.as_ref().expect("guard held until drop")
    }
}

impl<T> std::ops::DerefMut for MutexGuard<'_, T> {
    fn deref_mut(&mut self) -> &mut T {
        self.inner.as_mut().expect("guard held until drop")
    }
}

impl<T> Drop for MutexGuard<'_, T> {
    fn drop(&mut self) {
        // Release the real lock before stamping/sending so the held duration
        // excludes the event-send cost and the lock frees as early as possible.
        drop(self.inner.take());
        let acquire_nanos = self
            .start
            .map(|start| Instant::now().duration_since(start).as_nanos() as u64);
        send_mutex_event(MutexEvent::Released {
            id: self.id,
            wait_nanos: self.wait_nanos,
            acquire_nanos,
        });
    }
}

impl<T> InstrumentMutex for StdMutex<T> {
    type Output = Mutex<T>;
    fn instrument(self, source: &'static str, label: Option<String>) -> Self::Output {
        Mutex::__new_instrumented(self, source, label)
    }
}
