// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0.

use std::{
    sync::{
        Arc,
        atomic::{AtomicU8, Ordering},
        mpsc::Sender,
    },
    time::{Duration, Instant},
};

use tikv::read_pool::{TaskDispatchProfile, TaskPollProfile, YatpPollEventTracker};

use crate::metrics::{SlowPollReport, SlowTaskReport, TaskPollMetrics};

const TASK_RESULT_UNKNOWN: u8 = 0;
const TASK_RESULT_CANCELLED: u8 = 1;
const TASK_RESULT_COMPLETED: u8 = 2;

/// Completion delivered to the workload controller after the task's final poll
/// has been observed.
#[derive(Debug)]
pub struct TaskCompletion {
    pub runtime_id: usize,
    pub completed: bool,
    pub latency: Duration,
    pub slow_task: Option<SlowTaskReport>,
}

/// Handle used by the task body to publish whether its simulated work
/// completed or was cancelled.
#[derive(Clone, Debug)]
pub struct TaskResultHandle {
    result: Arc<AtomicU8>,
}

impl TaskResultHandle {
    pub fn record(&self, completed: bool) {
        let result = if completed {
            TASK_RESULT_COMPLETED
        } else {
            TASK_RESULT_CANCELLED
        };
        self.result.store(result, Ordering::Release);
    }
}

/// The independent observers installed on one simulated task.
pub struct TaskObservers {
    pub poll_tracker: Box<TaskPollObserver>,
    pub result_handle: TaskResultHandle,
}

impl TaskObservers {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        task_id: u64,
        runtime_id: usize,
        offer_sequence: u64,
        simulation_started_at: Instant,
        submitted_at: Instant,
        slow_threshold: Duration,
        completion_tx: Sender<TaskCompletion>,
        task_poll_metrics: TaskPollMetrics,
    ) -> Self {
        let task_result = Arc::new(AtomicU8::new(TASK_RESULT_UNKNOWN));
        let poll_tracker = Box::new(TaskPollObserver {
            task_id,
            runtime_id,
            offer_sequence,
            simulation_started_at,
            submitted_at,
            slow_threshold,
            completion_tx: Some(completion_tx),
            task_result: task_result.clone(),
            task_poll_metrics,
            pending_poll: None,
            last_poll_end: None,
            last_dispatch_id: None,
            dispatch_index: 0,
            slow_trace: (slow_threshold > Duration::ZERO).then(SlowTraceState::new),
        });
        Self {
            poll_tracker,
            result_handle: TaskResultHandle {
                result: task_result,
            },
        }
    }
}

#[derive(Clone, Copy, Debug)]
struct PendingPoll {
    poll_start: Instant,
    interval_start: Instant,
    wait: Duration,
    queue_wait: Duration,
    other_wait: Duration,
    dispatch_index: usize,
    immediate_repoll: bool,
    fair_queue_waited_task_slices: Option<u64>,
}

#[derive(Debug)]
struct SlowTraceState {
    polls: Vec<SlowPollReport>,
}

impl SlowTraceState {
    fn new() -> Self {
        Self { polls: Vec::new() }
    }
}

/// Observes every poll for task-runtime CPU and completion. When a non-zero
/// slow threshold is configured, it additionally retains the per-poll trace
/// until the task completes.
pub struct TaskPollObserver {
    task_id: u64,
    runtime_id: usize,
    offer_sequence: u64,
    simulation_started_at: Instant,
    submitted_at: Instant,
    slow_threshold: Duration,
    completion_tx: Option<Sender<TaskCompletion>>,
    task_result: Arc<AtomicU8>,
    task_poll_metrics: TaskPollMetrics,
    pending_poll: Option<PendingPoll>,
    last_poll_end: Option<Instant>,
    last_dispatch_id: Option<Instant>,
    dispatch_index: usize,
    slow_trace: Option<SlowTraceState>,
}

impl TaskPollObserver {
    fn record_poll_start(&mut self, profile: &TaskDispatchProfile) {
        debug_assert!(self.pending_poll.is_none());
        let interval_start = self.last_poll_end.unwrap_or(self.submitted_at);
        let dispatch_id = profile.dispatch_id();
        let immediate_repoll = self.last_dispatch_id == Some(dispatch_id);
        if !immediate_repoll {
            self.dispatch_index += 1;
            self.last_dispatch_id = Some(dispatch_id);
        }

        let poll_start = profile.poll_start_time;
        let wait = poll_start.saturating_duration_since(interval_start);
        let queue_wait = if immediate_repoll {
            Duration::ZERO
        } else {
            profile.queue_wait.unwrap_or_default().min(wait)
        };
        if !immediate_repoll {
            if profile.queue_wait.is_some() {
                self.task_poll_metrics
                    .slice_wait_seconds
                    .observe(queue_wait.as_secs_f64());
            }
            if let Some(waits) = profile.fair_queue_waited_task_slices {
                self.task_poll_metrics.slice_waits.observe(waits as f64);
            }
        }
        self.pending_poll = Some(PendingPoll {
            poll_start,
            interval_start,
            wait,
            queue_wait,
            other_wait: wait.saturating_sub(queue_wait),
            dispatch_index: self.dispatch_index,
            immediate_repoll,
            fair_queue_waited_task_slices: if immediate_repoll {
                None
            } else {
                profile.fair_queue_waited_task_slices
            },
        });
    }

    fn record_poll_end(
        &mut self,
        task_finished: bool,
        profile: &TaskPollProfile,
        observed_at: Instant,
    ) {
        self.task_poll_metrics
            .cpu_seconds
            .inc_by(profile.cpu_time.as_secs_f64());
        self.task_poll_metrics
            .wall_seconds
            .inc_by(profile.wall_time.as_secs_f64());

        let pending = self.pending_poll.take().unwrap_or_else(|| {
            let poll_start = observed_at
                .checked_sub(profile.wall_time)
                .unwrap_or(observed_at);
            let interval_start = self.last_poll_end.unwrap_or(self.submitted_at);
            let wait = poll_start.saturating_duration_since(interval_start);
            PendingPoll {
                poll_start,
                interval_start,
                wait,
                queue_wait: Duration::ZERO,
                other_wait: wait,
                dispatch_index: self.dispatch_index,
                immediate_repoll: false,
                fair_queue_waited_task_slices: None,
            }
        });
        let poll_end = pending
            .poll_start
            .checked_add(profile.wall_time)
            .unwrap_or(observed_at)
            .min(observed_at);
        self.last_poll_end = Some(poll_end);
        if let Some(trace) = self.slow_trace.as_mut() {
            trace.polls.push(SlowPollReport {
                poll_index: trace.polls.len() + 1,
                dispatch_index: pending.dispatch_index,
                immediate_repoll: pending.immediate_repoll,
                interval_start_seconds: seconds_since(
                    pending.interval_start,
                    self.simulation_started_at,
                ),
                poll_start_seconds: seconds_since(pending.poll_start, self.simulation_started_at),
                poll_end_seconds: seconds_since(poll_end, self.simulation_started_at),
                wait_seconds: pending.wait.as_secs_f64(),
                queue_wait_seconds: pending.queue_wait.as_secs_f64(),
                other_wait_seconds: pending.other_wait.as_secs_f64(),
                poll_wall_seconds: profile.wall_time.as_secs_f64(),
                poll_cpu_seconds: profile.cpu_time.as_secs_f64(),
                fair_queue_waited_task_slices: pending.fair_queue_waited_task_slices,
            });
        }

        if task_finished {
            self.send_completion(observed_at);
        }
    }

    fn send_completion(&mut self, observed_at: Instant) {
        let Some(completion_tx) = self.completion_tx.take() else {
            return;
        };
        let completed = self.task_result.load(Ordering::Acquire) == TASK_RESULT_COMPLETED;
        let latency = observed_at.saturating_duration_since(self.submitted_at);
        let slow_task = (completed
            && self.slow_threshold > Duration::ZERO
            && latency > self.slow_threshold)
            .then(|| SlowTaskReport {
                task_id: self.task_id,
                runtime_id: self.runtime_id,
                offer_sequence: self.offer_sequence,
                submitted_at_seconds: seconds_since(self.submitted_at, self.simulation_started_at),
                completed_at_seconds: seconds_since(observed_at, self.simulation_started_at),
                latency_seconds: latency.as_secs_f64(),
                slow_threshold_seconds: self.slow_threshold.as_secs_f64(),
                polls: self
                    .slow_trace
                    .take()
                    .map(|trace| trace.polls)
                    .unwrap_or_default(),
            });
        let _ = completion_tx.send(TaskCompletion {
            runtime_id: self.runtime_id,
            completed,
            latency,
            slow_task,
        });
    }
}

impl YatpPollEventTracker for TaskPollObserver {
    #[inline]
    fn on_poll_start(&mut self, profile: &TaskDispatchProfile) {
        self.record_poll_start(profile);
    }

    #[inline]
    fn on_poll_end(&mut self, task_finished: bool, profile: &TaskPollProfile) {
        self.record_poll_end(task_finished, profile, Instant::now());
    }
}

impl Drop for TaskPollObserver {
    fn drop(&mut self) {
        // Normal tasks send from their final on_poll_end callback. This only
        // prevents the workload controller from hanging if a task is dropped
        // before reaching that callback.
        self.send_completion(Instant::now());
    }
}

fn seconds_since(instant: Instant, start: Instant) -> f64 {
    instant.saturating_duration_since(start).as_secs_f64()
}

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

    use cpu_time::ThreadTime;

    use super::*;

    fn test_poll_metrics(runtime_id: usize) -> TaskPollMetrics {
        TaskPollMetrics::new(runtime_id).unwrap()
    }

    #[test]
    fn test_poll_observer_reports_completion_with_slow_poll_breakdown() {
        let simulation_start = Instant::now();
        let submitted_at = simulation_start + Duration::from_millis(1);
        let (completion_tx, completion_rx) = mpsc::channel();
        let poll_metrics = test_poll_metrics(100);
        let observers = TaskObservers::new(
            9,
            2,
            17,
            simulation_start,
            submitted_at,
            Duration::from_millis(5),
            completion_tx,
            poll_metrics.clone(),
        );
        let mut tracker = observers.poll_tracker;

        let first_dispatch = submitted_at + Duration::from_millis(3);
        tracker.on_poll_start(&TaskDispatchProfile {
            schedule_time: Some(submitted_at + Duration::from_millis(1)),
            dispatch_start_time: first_dispatch,
            poll_start_thread_time: ThreadTime::now(),
            poll_start_time: first_dispatch + Duration::from_millis(1),
            queue_wait: Some(Duration::from_millis(2)),
            fair_queue_waited_task_slices: Some(3),
        });
        tracker.record_poll_end(
            false,
            &TaskPollProfile {
                wall_time: Duration::from_millis(2),
                cpu_time: Duration::from_millis(1),
            },
            first_dispatch + Duration::from_millis(3),
        );

        // Immediate repolls stay in the same dispatch and therefore do not
        // repeat queue wait or fair-queue wait counts.
        tracker.on_poll_start(&TaskDispatchProfile {
            schedule_time: Some(submitted_at + Duration::from_millis(1)),
            dispatch_start_time: first_dispatch,
            poll_start_thread_time: ThreadTime::now(),
            poll_start_time: first_dispatch + Duration::from_millis(4),
            queue_wait: Some(Duration::from_millis(2)),
            fair_queue_waited_task_slices: Some(3),
        });
        observers.result_handle.record(true);
        tracker.record_poll_end(
            true,
            &TaskPollProfile {
                wall_time: Duration::from_millis(2),
                cpu_time: Duration::from_millis(1),
            },
            first_dispatch + Duration::from_millis(6),
        );

        // The completion is emitted directly by the final poll callback.
        let completion = completion_rx.recv().unwrap();
        assert!(completion.completed);
        assert_eq!(completion.runtime_id, 2);
        assert_eq!(completion.latency, Duration::from_millis(9));
        assert!((poll_metrics.cpu_seconds.get() - 0.002).abs() < 1e-12);
        assert!((poll_metrics.wall_seconds.get() - 0.004).abs() < 1e-12);
        assert_eq!(poll_metrics.slice_wait_seconds.get_sample_count(), 1);
        assert!((poll_metrics.slice_wait_seconds.get_sample_sum() - 0.002).abs() < 1e-12);
        assert_eq!(poll_metrics.slice_waits.get_sample_count(), 1);
        assert!((poll_metrics.slice_waits.get_sample_sum() - 3.0).abs() < 1e-12);
        let slow_task = completion.slow_task.unwrap();
        assert_eq!(slow_task.task_id, 9);
        assert_eq!(slow_task.offer_sequence, 17);
        assert_eq!(slow_task.polls.len(), 2);
        assert_eq!(slow_task.polls[0].dispatch_index, 1);
        assert!((slow_task.polls[0].queue_wait_seconds - 0.002).abs() < 1e-12);
        assert!(
            (slow_task.polls[0].poll_end_seconds
                - slow_task.polls[0].poll_start_seconds
                - slow_task.polls[0].poll_wall_seconds)
                .abs()
                < 1e-12
        );
        assert_eq!(slow_task.polls[0].fair_queue_waited_task_slices, Some(3));
        assert!(slow_task.polls[1].immediate_repoll);
        assert_eq!(slow_task.polls[1].queue_wait_seconds, 0.0);
        assert!(
            (slow_task.polls[1].poll_end_seconds
                - slow_task.polls[1].poll_start_seconds
                - slow_task.polls[1].poll_wall_seconds)
                .abs()
                < 1e-12
        );
        assert_eq!(slow_task.polls[1].fair_queue_waited_task_slices, None);

        drop(tracker);
        completion_rx.try_recv().unwrap_err();
    }

    #[test]
    fn test_poll_observer_discards_fast_and_cancelled_traces() {
        for (completed, threshold) in [
            (true, Duration::from_secs(1)),
            (false, Duration::from_nanos(1)),
        ] {
            let now = Instant::now();
            let (completion_tx, completion_rx) = mpsc::channel();
            let poll_metrics = test_poll_metrics(101 + completed as usize);
            let observers =
                TaskObservers::new(1, 0, 0, now, now, threshold, completion_tx, poll_metrics);
            let mut tracker = observers.poll_tracker;
            tracker.record_poll_start(&TaskDispatchProfile {
                schedule_time: Some(now),
                dispatch_start_time: now,
                poll_start_thread_time: ThreadTime::now(),
                poll_start_time: now,
                queue_wait: Some(Duration::ZERO),
                fair_queue_waited_task_slices: None,
            });
            observers.result_handle.record(completed);
            tracker.record_poll_end(
                true,
                &TaskPollProfile {
                    wall_time: Duration::from_millis(1),
                    cpu_time: Duration::from_millis(1),
                },
                now + Duration::from_millis(1),
            );
            assert!(completion_rx.recv().unwrap().slow_task.is_none());
        }
    }

    #[test]
    fn test_zero_threshold_still_reports_completion_without_trace() {
        let now = Instant::now();
        let (completion_tx, completion_rx) = mpsc::channel();
        let poll_metrics = test_poll_metrics(103);
        let observers = TaskObservers::new(
            1,
            4,
            0,
            now,
            now,
            Duration::ZERO,
            completion_tx,
            poll_metrics.clone(),
        );
        let mut tracker = observers.poll_tracker;
        tracker.record_poll_start(&TaskDispatchProfile {
            schedule_time: None,
            dispatch_start_time: now,
            poll_start_thread_time: ThreadTime::now(),
            poll_start_time: now,
            queue_wait: None,
            fair_queue_waited_task_slices: None,
        });
        observers.result_handle.record(true);
        tracker.record_poll_end(
            true,
            &TaskPollProfile {
                wall_time: Duration::from_millis(1),
                cpu_time: Duration::from_millis(1),
            },
            now + Duration::from_millis(1),
        );

        let completion = completion_rx.recv().unwrap();
        assert!(completion.completed);
        assert!(completion.slow_task.is_none());
        assert_eq!(poll_metrics.slice_wait_seconds.get_sample_count(), 0);
    }

    #[test]
    fn test_poll_observer_drop_reports_cancelled_completion() {
        let now = Instant::now();
        let (completion_tx, completion_rx) = mpsc::channel();
        let observers = TaskObservers::new(
            1,
            4,
            0,
            now,
            now,
            Duration::ZERO,
            completion_tx,
            test_poll_metrics(104),
        );

        drop(observers.poll_tracker);
        assert!(!completion_rx.recv().unwrap().completed);
        completion_rx.try_recv().unwrap_err();
    }
}