// Copyright 2026 PingCAP, Inc. Licensed under Apache-2.0.

use std::{
    future::Future,
    pin::Pin,
    task::{Context, Poll},
    time::Duration,
};

use cpu_time::ThreadTime;
use pin_project::pin_project;

/// CpuObserveFuture tracks thread CPU time spent while polling its delegate
/// future.
#[pin_project]
pub struct CpuObserveFuture<F> {
    enabled: bool,
    total_duration: Duration,
    #[pin]
    delegate: F,
}

impl<F> CpuObserveFuture<F> {
    /// Create a new CpuObserveFuture that always tracks CPU time.
    pub fn new(delegate: F) -> Self {
        CpuObserveFuture {
            enabled: true,
            total_duration: Duration::ZERO,
            delegate,
        }
    }

    /// Create a new CpuObserveFuture that tracks CPU time only when `enabled`
    /// is true. When false, polling skips `ThreadTime` sampling and the
    /// observed duration stays zero.
    pub fn new_with_enabled(delegate: F, enabled: bool) -> Self {
        CpuObserveFuture {
            enabled,
            total_duration: Duration::ZERO,
            delegate,
        }
    }

    /// Get total CPU time observed so far.
    pub fn total_duration(&self) -> Duration {
        self.total_duration
    }
}

impl<F: Future> Future for CpuObserveFuture<F> {
    type Output = (Duration, F::Output);

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        let timer = (*this.enabled).then(ThreadTime::now);
        let res = this.delegate.poll(cx);
        if let Some(timer) = timer {
            *this.total_duration += timer.elapsed();
        }
        let dur = *this.total_duration;
        res.map(|r| (dur, r))
    }
}