//! Test-only crash and fault control for syscall-coupled storage checkpoints.

use std::{
    collections::BTreeMap,
    sync::{
        Condvar, Mutex, OnceLock,
        atomic::{AtomicUsize, Ordering},
    },
    time::Duration,
};

static TRAVERSAL_CALLS: AtomicUsize = AtomicUsize::new(0);

pub fn reset_traversal_calls() {
    TRAVERSAL_CALLS.store(0, Ordering::SeqCst);
}

pub fn traversal_calls() -> usize {
    TRAVERSAL_CALLS.load(Ordering::SeqCst)
}

pub(crate) fn record_traversal_call() {
    TRAVERSAL_CALLS.fetch_add(1, Ordering::SeqCst);
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum FaultAction {
    Exit,
    Fail,
    Hold,
}

struct State {
    checkpoint: Option<(String, FaultAction, usize)>,
    hits: BTreeMap<String, usize>,
    held: bool,
}
fn slot() -> &'static (Mutex<State>, Condvar) {
    static SLOT: OnceLock<(Mutex<State>, Condvar)> = OnceLock::new();
    SLOT.get_or_init(|| {
        (
            Mutex::new(State {
                checkpoint: None,
                hits: BTreeMap::new(),
                held: false,
            }),
            Condvar::new(),
        )
    })
}

pub fn set_checkpoint(stage: impl Into<String>, action: FaultAction) {
    set_checkpoint_on_hit(stage, action, 1);
}

pub fn set_checkpoint_on_hit(stage: impl Into<String>, action: FaultAction, hit: usize) {
    assert!(hit > 0, "checkpoint hit is one-based");
    let (lock, _) = slot();
    let mut state = lock.lock().expect("artifact test checkpoint lock poisoned");
    state.checkpoint = Some((stage.into(), action, hit));
    state.hits.clear();
}

pub fn checkpoint_hits(stage: &str) -> usize {
    let (lock, _) = slot();
    lock.lock()
        .expect("artifact test checkpoint lock poisoned")
        .hits
        .get(stage)
        .copied()
        .unwrap_or(0)
}

pub fn clear_checkpoint() {
    let (lock, wake) = slot();
    let mut state = lock.lock().expect("artifact test checkpoint lock poisoned");
    state.checkpoint = None;
    state.held = false;
    wake.notify_all();
}
pub fn wait_until_held(timeout: Duration) -> bool {
    let (lock, wake) = slot();
    let state = lock.lock().expect("artifact test checkpoint lock poisoned");
    let (state, _) = wake
        .wait_timeout_while(state, timeout, |state| !state.held)
        .expect("artifact test checkpoint lock poisoned");
    state.held
}

pub(crate) fn checkpoint(stage: &str) -> Option<FaultAction> {
    let (lock, wake) = slot();
    let mut state = lock.lock().expect("artifact test checkpoint lock poisoned");
    let configured = state.checkpoint.clone();
    let action = configured.and_then(|(expected, action, target_hit)| {
        if expected != stage {
            return None;
        }
        let hit = state.hits.entry(stage.to_owned()).or_insert(0);
        if *hit == 0 {
            *hit = 1;
        }
        (*hit == target_hit).then_some(action)
    });
    if action == Some(FaultAction::Hold) {
        state.held = true;
        wake.notify_all();
        while state.checkpoint.is_some() {
            state = wake
                .wait(state)
                .expect("artifact test checkpoint lock poisoned");
        }
    }
    action
}

pub(crate) fn action(stage: &str) -> Option<FaultAction> {
    let (lock, _) = slot();
    let mut state = lock.lock().expect("artifact test checkpoint lock poisoned");
    *state.hits.entry(stage.to_owned()).or_insert(0) += 1;
    let hit = state.hits[stage];
    state
        .checkpoint
        .as_ref()
        .and_then(|(expected, action, target_hit)| {
            (expected == stage && hit == *target_hit).then_some(*action)
        })
}