// Copyright 2026 PingCAP, Inc.

use std::{
    path::{Path, PathBuf},
    sync::{
        Arc, OnceLock, RwLock, Weak,
        atomic::{AtomicU64, Ordering},
    },
};

use tikv_util::time::monotonic_raw_now_millis;

#[derive(Debug)]
struct RegisteredLocalDir {
    dir: PathBuf,
    last_sync_success_at_millis: Weak<AtomicU64>,
}

fn registry() -> &'static RwLock<Vec<RegisteredLocalDir>> {
    static REGISTRY: OnceLock<RwLock<Vec<RegisteredLocalDir>>> = OnceLock::new();
    REGISTRY.get_or_init(|| RwLock::new(Vec::new()))
}

/// Keeps one per-local-dir sync-progress timestamp registered for fail-fast.
#[derive(Debug)]
pub struct LocalDiskSyncProgressRegistration {
    dir: PathBuf,
    last_sync_success_at_millis: Arc<AtomicU64>,
}

impl LocalDiskSyncProgressRegistration {
    pub fn new(dir: PathBuf) -> Self {
        let last_sync_success_at_millis = Arc::new(AtomicU64::new(0));
        registry().write().unwrap().push(RegisteredLocalDir {
            dir: dir.clone(),
            last_sync_success_at_millis: Arc::downgrade(&last_sync_success_at_millis),
        });
        Self {
            dir,
            last_sync_success_at_millis,
        }
    }

    pub fn dir(&self) -> &Path {
        &self.dir
    }

    pub fn last_sync_success_at_millis(&self) -> Arc<AtomicU64> {
        self.last_sync_success_at_millis.clone()
    }
}

impl Drop for LocalDiskSyncProgressRegistration {
    fn drop(&mut self) {
        let last_sync_success_at_ptr = Arc::as_ptr(&self.last_sync_success_at_millis);
        registry().write().unwrap().retain(|entry| {
            entry.dir != self.dir
                || entry.last_sync_success_at_millis.upgrade().is_some_and(
                    |last_sync_success_at_millis| {
                        Arc::as_ptr(&last_sync_success_at_millis) != last_sync_success_at_ptr
                    },
                )
        });
    }
}

/// Records one successful local temp-file publish for the given registered dir.
pub fn record_local_dir_sync_success(dir: &Path) {
    let now = monotonic_raw_now_millis();
    let registrations = registry().read().unwrap();
    // Map subdirectory publishes back to the most specific registered local
    // dir, so progress under `<local_dir>/ia/...` only vetoes that local dir.
    let mut best_match_len = None;
    let mut best_matches = Vec::new();
    for entry in registrations
        .iter()
        .filter(|entry| dir.starts_with(&entry.dir))
    {
        if let Some(last_sync_success_at_millis) = entry.last_sync_success_at_millis.upgrade() {
            let component_count = entry.dir.components().count();
            match best_match_len {
                Some(best_len) if component_count < best_len => {}
                Some(best_len) if component_count == best_len => {
                    best_matches.push(last_sync_success_at_millis);
                }
                _ => {
                    best_match_len = Some(component_count);
                    best_matches.clear();
                    best_matches.push(last_sync_success_at_millis);
                }
            }
        }
    }
    for last_sync_success_at_millis in best_matches {
        last_sync_success_at_millis.fetch_max(now, Ordering::Relaxed);
    }
}

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

    use tempfile::tempdir;

    use super::*;

    #[test]
    fn test_record_local_dir_sync_success_updates_matching_dir_only() {
        let dir1 = tempdir().unwrap();
        let dir2 = tempdir().unwrap();
        let reg1 = LocalDiskSyncProgressRegistration::new(dir1.path().to_path_buf());
        let reg2 = LocalDiskSyncProgressRegistration::new(dir2.path().to_path_buf());

        record_local_dir_sync_success(dir1.path());

        assert_ne!(
            reg1.last_sync_success_at_millis().load(Ordering::Relaxed),
            0
        );
        assert_eq!(
            reg2.last_sync_success_at_millis().load(Ordering::Relaxed),
            0
        );
    }

    #[test]
    fn test_record_local_dir_sync_success_uses_longest_registered_ancestor() {
        let root = tempdir().unwrap();
        let nested_dir = root.path().join("ia");
        std::fs::create_dir(&nested_dir).unwrap();
        let reg_root = LocalDiskSyncProgressRegistration::new(root.path().to_path_buf());
        let reg_nested = LocalDiskSyncProgressRegistration::new(nested_dir.clone());

        record_local_dir_sync_success(&nested_dir.join("meta"));

        assert_eq!(
            reg_root
                .last_sync_success_at_millis()
                .load(Ordering::Relaxed),
            0
        );
        assert_ne!(
            reg_nested
                .last_sync_success_at_millis()
                .load(Ordering::Relaxed),
            0
        );
    }
}