// Copyright 2026 PingCAP, Inc.

use std::{
    fs,
    path::PathBuf,
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, AtomicU64, Ordering},
        mpsc::{Receiver, channel},
    },
    thread::JoinHandle,
    time::Duration,
};

use tikv_util::{
    config::VersionTrack, crit, error, logger, metrics::CRITICAL_ERROR,
    sys::thread::StdThreadBuildWrapper, time::monotonic_raw_now_millis, warn, worker::Worker,
};

use super::{Config, disk_probe::ProbeRunner, metrics::DISK_PROBE_DURATION_HISTOGRAM};

const PROBE_INTERVAL: Duration = Duration::from_secs(1);
const FAIL_FAST_LOG_FLUSH_TIMEOUT: Duration = Duration::from_secs(2);
const PROBE_THREAD_STOP_TIMEOUT: Duration = Duration::from_secs(2);
const PROBE_PAYLOAD: &[u8] = b"rfstore disk fail fast probe";

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum ProbeKind {
    RaftWal,
    RaftData,
    Kv,
}

impl ProbeKind {
    fn name(self) -> &'static str {
        match self {
            Self::RaftWal => "raft_wal",
            Self::RaftData => "raft_data",
            Self::Kv => "kv",
        }
    }

    fn probe_filename(self) -> &'static str {
        match self {
            Self::RaftWal => ".raft_wal_fail_fast_probe.tmp",
            Self::RaftData => ".raft_data_fail_fast_probe.tmp",
            Self::Kv => ".kv_fail_fast_probe.tmp",
        }
    }

    fn stuck_critical_error_kind(self) -> &'static str {
        match self {
            Self::RaftWal => "disk_probe_stuck_raft_wal",
            Self::RaftData => "disk_probe_stuck_raft_data",
            Self::Kv => "disk_probe_stuck_kv",
        }
    }

    fn fail_fast_critical_error_kind(self) -> &'static str {
        match self {
            Self::RaftWal => "disk_probe_fail_fast_raft_wal",
            Self::RaftData => "disk_probe_fail_fast_raft_data",
            Self::Kv => "disk_probe_fail_fast_kv",
        }
    }
}

#[derive(Clone, Debug)]
pub(crate) struct FailFastTarget {
    // Probe classification controls both synthetic probe semantics and metric/log labels.
    kind: ProbeKind,
    dir: PathBuf,
    // The monitor consumes one symmetric target shape, but the veto source remains
    // target-specific: raft WAL uses successful `rfengine.persist`, kv uses same-dir
    // local temp-file publish progress, and raft data uses successful async WAL sync.
    last_progress_at_millis: Option<Arc<AtomicU64>>,
}

impl FailFastTarget {
    /// One synchronous WAL target. Recent successful `rfengine.persist` may
    /// veto fail-fast for either WAL target in a double-writer setup.
    pub(crate) fn raft_wal(dir: PathBuf, last_progress_at_millis: Arc<AtomicU64>) -> Self {
        Self {
            kind: ProbeKind::RaftWal,
            dir,
            last_progress_at_millis: Some(last_progress_at_millis),
        }
    }

    /// The secondary raft target used only when raft data lives on a different
    /// mount. Recent successful async WAL sync may veto fail-fast here.
    pub(crate) fn raft_data(dir: PathBuf, last_progress_at_millis: Arc<AtomicU64>) -> Self {
        Self {
            kind: ProbeKind::RaftData,
            dir,
            last_progress_at_millis: Some(last_progress_at_millis),
        }
    }

    /// One independent kv local-dir target. Only progress on this same dir may
    /// veto fail-fast for this target.
    pub(crate) fn kv(dir: PathBuf, last_progress_at_millis: Arc<AtomicU64>) -> Self {
        Self {
            kind: ProbeKind::Kv,
            dir,
            last_progress_at_millis: Some(last_progress_at_millis),
        }
    }
}

#[derive(Clone, Debug)]
struct ProbeTarget {
    kind: ProbeKind,
    probe: ProbeRunner,
    last_progress_at_millis: Option<Arc<AtomicU64>>,
}

impl ProbeTarget {
    fn new(target: FailFastTarget) -> Self {
        let FailFastTarget {
            kind,
            dir,
            last_progress_at_millis,
        } = target;
        Self {
            kind,
            // All targets use a simple write+sync probe. Target-specific forward
            // progress signals can still veto probe-only failures.
            probe: ProbeRunner::new(dir.join(kind.probe_filename()), PROBE_PAYLOAD),
            last_progress_at_millis,
        }
    }
}

fn record_probe_duration(disk: &str, outcome: &str, duration: Duration) {
    DISK_PROBE_DURATION_HISTOGRAM
        .with_label_values(&[disk, outcome])
        .observe(duration.as_secs_f64());
}

fn fail_fast_timeout(cfg: &Config) -> Option<Duration> {
    cfg.disk_hang_timeout
        .and_then(|timeout| (!timeout.0.is_zero()).then_some(timeout.0))
}

fn run_probe(stop: &AtomicBool, target: &ProbeTarget) -> bool {
    let started = target.probe.try_start_probe();
    debug_assert!(started, "fail-fast probe must not overlap");
    if !started {
        return true;
    }

    match target.probe.probe_once() {
        Ok(duration) => {
            target.probe.finish_probe_success();
            if stop.load(Ordering::Acquire) {
                return false;
            }
            record_probe_duration(target.kind.name(), "success", duration);
        }
        Err(e) => {
            let duration = target.probe.current_probe_elapsed().unwrap_or_default();
            target.probe.finish_probe_failure();
            if stop.load(Ordering::Acquire) {
                return false;
            }
            record_probe_duration(target.kind.name(), "failure", duration);
            warn!(
                "fail-fast probe failed";
                "disk" => target.kind.name(),
                "path" => %target.probe.path().display(),
                "elapsed" => ?duration,
                "err" => ?e,
            );
        }
    }

    true
}

fn spawn_probe_thread(
    name: &'static str,
    cfg: Arc<VersionTrack<Config>>,
    stop: Arc<AtomicBool>,
    targets: Vec<ProbeTarget>,
) -> std::io::Result<(JoinHandle<()>, Receiver<()>)> {
    let (exit_tx, exit_rx) = channel();
    let handle = std::thread::Builder::new()
        .name(name.to_owned())
        .spawn_wrapper(move || {
            while !stop.load(Ordering::Acquire) {
                if fail_fast_timeout(&cfg.value()).is_none() {
                    std::thread::sleep(PROBE_INTERVAL);
                    continue;
                }

                for target in &targets {
                    if !run_probe(&stop, target) {
                        break;
                    }
                }
                if stop.load(Ordering::Acquire) {
                    break;
                }

                std::thread::sleep(PROBE_INTERVAL);
            }
            let _ = exit_tx.send(());
        })?;
    Ok((handle, exit_rx))
}

type ProbeThread = (JoinHandle<()>, Receiver<()>);

fn probe_thread_groups(targets: &[ProbeTarget]) -> Vec<(&'static str, Vec<ProbeTarget>)> {
    let wal_target_count = targets
        .iter()
        .filter(|target| target.kind == ProbeKind::RaftWal)
        .count();
    if wal_target_count < 2 {
        return vec![("rfstore-fail-fast-probe", targets.to_vec())];
    }

    let (wal_targets, other_targets): (Vec<_>, Vec<_>) = targets
        .iter()
        .cloned()
        .partition(|target| target.kind == ProbeKind::RaftWal);
    let mut groups = vec![("rfstore-fail-fast-wal-probe", wal_targets)];
    if !other_targets.is_empty() {
        groups.push(("rfstore-fail-fast-disk-probe", other_targets));
    }
    groups
}

fn should_fail_fast(target: &ProbeTarget, timeout: Duration) -> bool {
    let current_probe_elapsed = target.probe.current_probe_elapsed();
    let stuck_log_threshold = timeout / 2;
    if let Some(elapsed) = current_probe_elapsed.filter(|elapsed| *elapsed >= stuck_log_threshold) {
        CRITICAL_ERROR
            .with_label_values(&[target.kind.stuck_critical_error_kind()])
            .inc();
        error!(
            "fail-fast probe: disk still not responsive after elapsed";
            "disk" => target.kind.name(),
            "path" => %target.probe.path().display(),
            "elapsed" => ?elapsed,
        );
    }

    // A probe must prove that it is configured correctly before it can make
    // the process fail fast.
    if target.probe.time_since_last_success().is_none() {
        return false;
    }

    // The checker only cares whether the target has observed recent forward
    // progress. It deliberately does not care whether that came from raft
    // append success or kv local sync progress.
    let last_progress_elapsed = last_progress_elapsed(target.last_progress_at_millis.as_deref());
    let base_signal = current_probe_elapsed.is_some_and(|elapsed| elapsed >= timeout)
        || should_fail_fast_on_repeated_failures(&target.probe, timeout);
    if !base_signal {
        return false;
    }

    if last_progress_elapsed.is_some_and(|elapsed| elapsed < timeout) {
        warn!(
            "fail-fast vetoed by recent progress";
            "disk" => target.kind.name(),
            "timeout" => ?timeout,
            "current_probe_elapsed" => ?current_probe_elapsed,
            "last_progress_elapsed" => ?last_progress_elapsed,
        );
        return false;
    }

    CRITICAL_ERROR
        .with_label_values(&[target.kind.fail_fast_critical_error_kind()])
        .inc();
    crit!(
        "fail-fast: disk probe timed out, shutting down to avoid serving on an unhealthy disk";
        "disk" => target.kind.name(),
        "timeout" => ?timeout,
        "current_probe_elapsed" => ?current_probe_elapsed,
        "last_progress_elapsed" => ?last_progress_elapsed,
    );
    true
}

pub struct FailFastMonitor {
    stop: Arc<AtomicBool>,
    targets: Vec<ProbeTarget>,
    probe_threads: Mutex<Vec<ProbeThread>>,
}

impl FailFastMonitor {
    pub(crate) fn new(
        cfg: Arc<VersionTrack<Config>>,
        targets: Vec<FailFastTarget>,
        check_worker: Worker,
    ) -> std::io::Result<Self> {
        let stop = Arc::new(AtomicBool::new(false));
        // Normalize the external target description once so the monitor only
        // reasons about one internal probe shape afterward.
        let targets = targets
            .into_iter()
            .map(ProbeTarget::new)
            .collect::<Vec<_>>();

        let mut probe_threads = Vec::new();
        for (name, thread_targets) in probe_thread_groups(&targets) {
            match spawn_probe_thread(name, cfg.clone(), stop.clone(), thread_targets) {
                Ok(thread) => probe_threads.push(thread),
                Err(e) => {
                    stop.store(true, Ordering::Release);
                    return Err(e);
                }
            }
        }

        let stop_for_check = stop.clone();
        let targets_for_check = targets.clone();
        let mut fail_fast_enabled = fail_fast_timeout(&cfg.value()).is_some();
        check_worker.spawn_interval_task(PROBE_INTERVAL, move || {
            if stop_for_check.load(Ordering::Acquire) {
                return;
            }

            let Some(timeout) = fail_fast_timeout(&cfg.value()) else {
                for target in &targets_for_check {
                    target.probe.reset_failure_window();
                }
                fail_fast_enabled = false;
                return;
            };
            if !fail_fast_enabled {
                for target in &targets_for_check {
                    target.probe.reset_failure_window();
                }
                fail_fast_enabled = true;
                return;
            }

            if targets_for_check
                .iter()
                .any(|target| should_fail_fast(target, timeout))
            {
                logger::panic_after_best_effort_flush(
                    FAIL_FAST_LOG_FLUSH_TIMEOUT,
                    "fail-fast: rfstore self-killed due to storage disk unavailability",
                );
            }
        });

        Ok(Self {
            stop,
            targets,
            probe_threads: Mutex::new(probe_threads),
        })
    }

    pub fn stop(&self) {
        self.stop.store(true, Ordering::Release);
        self.try_join_probe_threads();
        self.cleanup_probe_files();
    }

    fn try_join_probe_threads(&self) {
        for (probe_thread_handle, exit_rx) in self.probe_threads.lock().unwrap().drain(..) {
            if exit_rx.recv_timeout(PROBE_THREAD_STOP_TIMEOUT).is_ok() {
                if let Err(e) = probe_thread_handle.join() {
                    warn!("fail-fast probe thread panicked during shutdown"; "err" => ?e);
                }
            }
        }
    }

    fn cleanup_probe_files(&self) {
        for target in &self.targets {
            match fs::remove_file(target.probe.path()) {
                Ok(()) => {}
                Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
                Err(e) => {
                    warn!(
                        "failed to remove fail-fast probe file";
                        "err" => ?e,
                        "path" => %target.probe.path().display(),
                    );
                }
            }
        }
    }
}

fn should_fail_fast_on_repeated_failures(probe: &ProbeRunner, timeout: Duration) -> bool {
    probe.current_probe_elapsed().is_none()
        && probe.failure_count_since_last_success() > 0
        && probe
            .time_since_last_success()
            .is_some_and(|elapsed| elapsed >= timeout)
}

fn last_progress_elapsed(last_progress_at_millis: Option<&AtomicU64>) -> Option<Duration> {
    let last = last_progress_at_millis?.load(Ordering::Relaxed);
    if last == 0 {
        return None;
    }
    Some(Duration::from_millis(
        monotonic_raw_now_millis().saturating_sub(last),
    ))
}

impl Drop for FailFastMonitor {
    fn drop(&mut self) {
        self.cleanup_probe_files();
    }
}

#[cfg(test)]
mod tests {
    use tempfile::tempdir;

    use super::*;

    fn new_target(
        kind: ProbeKind,
        dir: PathBuf,
        last_progress_at_millis: Option<Arc<AtomicU64>>,
    ) -> ProbeTarget {
        let target = match kind {
            ProbeKind::RaftWal => FailFastTarget {
                kind,
                dir,
                last_progress_at_millis,
            },
            ProbeKind::RaftData => FailFastTarget {
                kind,
                dir,
                last_progress_at_millis,
            },
            ProbeKind::Kv => FailFastTarget {
                kind,
                dir,
                last_progress_at_millis,
            },
        };
        ProbeTarget::new(target)
    }

    #[test]
    fn test_fail_fast_timeout_none_disables() {
        let mut cfg = Config::default();
        cfg.disk_hang_timeout = None;
        assert_eq!(fail_fast_timeout(&cfg), None);
    }

    #[test]
    fn test_should_fail_fast_on_timeout_only_after_probe_success() {
        let dir = tempdir().unwrap();
        let target = new_target(ProbeKind::RaftWal, dir.path().to_path_buf(), None);
        let timeout = Duration::from_millis(10);

        assert!(target.probe.try_start_probe());
        std::thread::sleep(timeout + Duration::from_millis(5));
        assert!(!should_fail_fast(&target, timeout));

        target.probe.finish_probe_success();
        assert!(target.probe.try_start_probe());
        std::thread::sleep(timeout + Duration::from_millis(5));
        assert!(should_fail_fast(&target, timeout));
    }

    #[test]
    fn test_should_fail_fast_on_repeated_failures_after_prior_success() {
        let dir = tempdir().unwrap();
        let target = new_target(ProbeKind::RaftWal, dir.path().to_path_buf(), None);
        let timeout = Duration::from_millis(10);

        assert!(!should_fail_fast_on_repeated_failures(
            &target.probe,
            timeout
        ));
        target.probe.finish_probe_success();
        assert!(target.probe.try_start_probe());
        target.probe.finish_probe_failure();
        std::thread::sleep(timeout + Duration::from_millis(5));
        assert!(should_fail_fast_on_repeated_failures(
            &target.probe,
            timeout
        ));

        target.probe.reset_failure_window();
        assert!(target.probe.try_start_probe());
        target.probe.finish_probe_failure();
        std::thread::sleep(timeout + Duration::from_millis(5));
        assert!(!should_fail_fast_on_repeated_failures(
            &target.probe,
            timeout
        ));
    }

    #[test]
    fn test_should_fail_fast_repeated_failures_respect_progress_veto() {
        let dir = tempdir().unwrap();
        let progress = Arc::new(AtomicU64::new(0));
        let target = new_target(
            ProbeKind::RaftWal,
            dir.path().to_path_buf(),
            Some(progress.clone()),
        );
        let timeout = Duration::from_millis(10);

        target.probe.finish_probe_success();
        assert!(target.probe.try_start_probe());
        target.probe.finish_probe_failure();
        std::thread::sleep(timeout + Duration::from_millis(5));

        progress.store(monotonic_raw_now_millis(), Ordering::Relaxed);
        assert!(!should_fail_fast(&target, timeout));

        progress.store(
            monotonic_raw_now_millis().saturating_sub(timeout.as_millis() as u64 + 1),
            Ordering::Relaxed,
        );
        assert!(should_fail_fast(&target, timeout));
    }

    #[test]
    fn test_raft_data_respects_async_wal_progress_veto() {
        let dir = tempdir().unwrap();
        let progress = Arc::new(AtomicU64::new(0));
        let target = new_target(
            ProbeKind::RaftData,
            dir.path().to_path_buf(),
            Some(progress.clone()),
        );
        let timeout = Duration::from_millis(10);

        target.probe.finish_probe_success();
        assert!(target.probe.try_start_probe());
        std::thread::sleep(timeout + Duration::from_millis(5));
        progress.store(monotonic_raw_now_millis(), Ordering::Relaxed);

        assert!(!should_fail_fast(&target, timeout));
    }

    #[test]
    fn test_probe_threads_split_only_for_redundant_wal() {
        let dir = tempdir().unwrap();
        let progress = Arc::new(AtomicU64::new(0));
        let primary = new_target(
            ProbeKind::RaftWal,
            dir.path().join("primary"),
            Some(progress.clone()),
        );
        let raft_data = new_target(ProbeKind::RaftData, dir.path().join("raft-data"), None);
        let kv = new_target(ProbeKind::Kv, dir.path().join("kv"), None);

        let groups = probe_thread_groups(&[primary.clone(), raft_data.clone(), kv.clone()]);
        assert_eq!(groups.len(), 1);
        assert_eq!(groups[0].1.len(), 3);

        let secondary = new_target(
            ProbeKind::RaftWal,
            dir.path().join("secondary"),
            Some(progress),
        );
        let groups = probe_thread_groups(&[primary, secondary, raft_data, kv]);
        assert_eq!(groups.len(), 2);
        assert!(
            groups[0]
                .1
                .iter()
                .all(|target| target.kind == ProbeKind::RaftWal)
        );
        assert!(
            groups[1]
                .1
                .iter()
                .all(|target| target.kind != ProbeKind::RaftWal)
        );
    }
}