// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2026 PingCAP, Inc.

//! Runtime context for the task currently dispatched by a YATP worker.
//!
//! The context is installed immediately before `Runner::handle` and remains
//! available through thread-local storage until that dispatch finishes. A task
//! can therefore have multiple contexts during its lifetime, one for each
//! dispatch.

use std::{
    cell::RefCell,
    time::{Duration, Instant},
};

use yatp::{queue::TaskCell as _, task::future::TaskCell};

thread_local! {
    static YATP_DISPATCH_CONTEXT: RefCell<Option<YatpDispatchTlsState>> = const { RefCell::new(None) };
}

/// Intermediate and installed states of the dispatch context.
///
/// The optional prepare state lets code attach supplementary information, such
/// as fair-queue observations, before the complete context is initialized in
/// `Runner::handle`. A context can still be installed without this state; when
/// present, it makes the resulting dispatch information more complete.
enum YatpDispatchTlsState {
    /// Optional supplementary information captured before context
    /// initialization.
    Prepare(YatpDispatchPrepareContext),
    /// Complete context available while the worker handles the task.
    Context(YatpDispatchContext),
}

/// Optional dispatch information available before `Runner::handle` starts.
pub(super) struct YatpDispatchPrepareContext {
    /// Number of task slices dispatched ahead of this task while it waited in
    /// the fair queue.
    pub fair_queue_waited_task_slices: Option<u64>,
}

/// Information associated with the current task dispatch on a YATP worker.
pub struct YatpDispatchContext {
    /// When the task was scheduled into the YATP queue for this dispatch.
    ///
    /// This has the same semantics as `Extras::schedule_time`: it records the
    /// enqueue time, not when the task is popped from the queue.
    schedule_time: Option<Instant>,
    /// When the YATP worker started handling the task after it was popped from
    /// the queue.
    ///
    /// This is captured at the beginning of `Runner::handle`, immediately
    /// before the task is dispatched to the wrapped runner.
    dispatch_start_time: Instant,
    /// Number of task slices dispatched ahead of this task while it waited in
    /// the fair queue.
    fair_queue_waited_task_slices: Option<u64>,
}

impl YatpDispatchContext {
    /// Stores optional information collected before the worker installs the
    /// complete dispatch context.
    pub(super) fn prepare(ctx: YatpDispatchPrepareContext) {
        YATP_DISPATCH_CONTEXT.with(|current| {
            let previous = current
                .borrow_mut()
                .replace(YatpDispatchTlsState::Prepare(ctx));
            if let Some(YatpDispatchTlsState::Context(_)) = previous {
                debug_assert!(false, "YATP task dispatch must not be nested");
            }
        });
    }

    /// Installs the context for one task dispatch on the current worker.
    ///
    /// This reads the current scheduling time from the task's YATP extras and
    /// merges any supplementary information previously stored by `prepare`.
    pub(super) fn install(task_cell: &mut TaskCell) -> YatpDispatchContextGuard {
        YATP_DISPATCH_CONTEXT.with(|current| {
            let extras = task_cell.mut_extras();
            let schedule_time = extras.schedule_time();
            let mut ctx = YatpDispatchContext {
                schedule_time,
                dispatch_start_time: Instant::now(),
                fair_queue_waited_task_slices: None,
            };
            match current.borrow_mut().take() {
                Some(YatpDispatchTlsState::Prepare(p)) => {
                    ctx.fair_queue_waited_task_slices = p.fair_queue_waited_task_slices;
                }
                Some(YatpDispatchTlsState::Context(_)) => {
                    debug_assert!(false, "YATP task dispatch must not be nested");
                }
                None => {}
            }
            _ = current
                .borrow_mut()
                .replace(YatpDispatchTlsState::Context(ctx));
        });
        YatpDispatchContextGuard
    }

    /// Runs `f` with the task dispatch context installed on the current thread.
    ///
    /// The argument is `None` before installation, after the dispatch guard is
    /// dropped, and while only a prepare context is installed.
    pub fn with_tls<R>(f: impl FnOnce(Option<&YatpDispatchContext>) -> R) -> R {
        YATP_DISPATCH_CONTEXT.with(|current| {
            let current = current.borrow();
            let ctx = match current.as_ref() {
                Some(YatpDispatchTlsState::Context(ctx)) => Some(ctx),
                Some(YatpDispatchTlsState::Prepare(_)) | None => None,
            };
            f(ctx)
        })
    }

    /// Returns when the task was scheduled into the YATP queue for this
    /// dispatch, or `None` if the queue did not record the enqueue time.
    ///
    /// This does not represent when the task is popped from the queue.
    #[inline]
    pub fn schedule_time(&self) -> Option<Instant> {
        self.schedule_time
    }

    /// Returns the identifier for this task dispatch.
    ///
    /// The dispatch start time currently serves as the identifier. Callers
    /// should use this method instead of depending on that representation so
    /// the identifier can change independently in the future.
    #[inline]
    pub fn dispatch_id(&self) -> Instant {
        self.dispatch_start_time
    }

    /// Returns when the YATP worker started handling the task after it was
    /// popped from the queue.
    #[inline]
    pub fn dispatch_start_time(&self) -> Instant {
        self.dispatch_start_time
    }

    /// Returns the time from entering the YATP queue until the worker starts
    /// handling this dispatch.
    ///
    /// Returns `None` when this dispatch has no scheduling timestamp.
    #[inline]
    pub fn ready_to_dispatch_duration(&self) -> Option<Duration> {
        self.schedule_time.map(|schedule_time| {
            self.dispatch_start_time
                .saturating_duration_since(schedule_time)
        })
    }

    /// Returns the number of task slices dispatched ahead of this task while
    /// it waited in the fair queue.
    ///
    /// `Some(0)` means the task was dispatched directly from the fair queue
    /// without waiting for another slice. `None` means fair-queue information
    /// is unavailable for this dispatch.
    #[inline]
    pub fn fair_queue_waited_task_slices(&self) -> Option<u64> {
        self.fair_queue_waited_task_slices
    }
}

/// Uninstalls the current task dispatch context when the dispatch finishes.
///
/// Keeping this guard alive for the whole `Runner::handle` call ensures the
/// context is also removed on early returns or unwinding.
#[must_use]
pub(super) struct YatpDispatchContextGuard;

impl Drop for YatpDispatchContextGuard {
    fn drop(&mut self) {
        YATP_DISPATCH_CONTEXT.with(|current| {
            let installed = current.borrow_mut().take();
            debug_assert!(installed.is_some());
        });
    }
}

#[cfg(test)]
mod tests {
    use yatp::{queue::Extras, task::future::TaskCell};

    use super::*;

    fn new_test_task_cell(task_id: u64) -> TaskCell {
        TaskCell::new(
            std::future::pending::<()>(),
            Extras::new_multilevel(task_id, None),
        )
    }

    #[test]
    fn test_yatp_dispatch_context_lifecycle_and_prepare_merge() {
        YatpDispatchContext::with_tls(|ctx| assert!(ctx.is_none()));

        YatpDispatchContext::prepare(YatpDispatchPrepareContext {
            fair_queue_waited_task_slices: Some(0),
        });
        // The prepare state is supplementary and must not be exposed as a
        // complete dispatch context.
        YatpDispatchContext::with_tls(|ctx| assert!(ctx.is_none()));

        let mut task_cell = new_test_task_cell(1);
        let before_install = Instant::now();
        let guard = YatpDispatchContext::install(&mut task_cell);
        let after_install = Instant::now();
        YatpDispatchContext::with_tls(|ctx| {
            let ctx = ctx.unwrap();
            assert!(ctx.schedule_time().is_none());
            assert!(ctx.ready_to_dispatch_duration().is_none());
            assert!(ctx.dispatch_start_time() >= before_install);
            assert!(ctx.dispatch_start_time() <= after_install);
            assert_eq!(ctx.fair_queue_waited_task_slices(), Some(0));
        });
        drop(guard);
        YatpDispatchContext::with_tls(|ctx| assert!(ctx.is_none()));

        let mut task_cell = new_test_task_cell(2);
        let guard = YatpDispatchContext::install(&mut task_cell);
        YatpDispatchContext::with_tls(|ctx| {
            assert_eq!(ctx.unwrap().fair_queue_waited_task_slices(), None);
        });
        drop(guard);
        YatpDispatchContext::with_tls(|ctx| assert!(ctx.is_none()));
    }
}