// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::{Error, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tokio::sync::RwLock;
use tracing::{debug, warn};
use uuid::Uuid;

use super::{BUCKET_META_PREFIX, DiskError, DiskStore, HealDiskExt as _, RUSTFS_META_BUCKET};

const LOG_COMPONENT_HEAL: &str = "heal";
const LOG_SUBSYSTEM_RESUME: &str = "resume";
const EVENT_HEAL_RESUME_STATE: &str = "heal_resume_state";
const EVENT_HEAL_CHECKPOINT_STATE: &str = "heal_checkpoint_state";

/// resume state file constants
const RESUME_STATE_FILE: &str = "ahm_resume_state.json";
const RESUME_PROGRESS_FILE: &str = "ahm_progress.json";
const RESUME_CHECKPOINT_FILE: &str = "ahm_checkpoint.json";

/// Current on-disk schema version for `ResumeState`. Snapshots written by an
/// older schema (which tracked latest-only object names and a positional
/// cursor) are incompatible with the per-version resume cursor, so they are
/// discarded on load and the scan restarts from the beginning.
const CURRENT_RESUME_SCHEMA: u32 = 2;
/// Current on-disk schema version for `ResumeCheckpoint`. Same rationale as
/// `CURRENT_RESUME_SCHEMA`: pre-per-version dedup identities are not comparable
/// to the new `compose_key` identities, so a stale checkpoint is discarded.
const CURRENT_CHECKPOINT_SCHEMA: u32 = 2;

/// Build the canonical, provably-injective dedup identity for an object
/// version. Length-prefixing the object key makes the encoding injective: no
/// two distinct `(object, version_id)` pairs can collide, even for adversarial
/// keys containing `:` or embedded null bytes. This is the single source of
/// truth for per-version dedup across the heal loop and the checkpoint sets.
pub fn compose_key(object: &str, version_id: Option<&str>) -> String {
    format!("{}:{}{}", object.len(), object, version_id.unwrap_or(""))
}

/// Persistence throttle for per-object bookkeeping: flush after this many
/// buffered mutations or once the interval elapses, whichever comes first.
/// Object heal is idempotent, so a crash re-heals at most one throttle window.
const PERSIST_EVERY_MUTATIONS: usize = 1000;
const PERSIST_INTERVAL: Duration = Duration::from_secs(5);

/// Tracks buffered mutations between persisted snapshots.
#[derive(Debug)]
struct PersistThrottle {
    pending: usize,
    last_save: Instant,
}

impl PersistThrottle {
    fn new() -> Self {
        Self {
            pending: 0,
            last_save: Instant::now(),
        }
    }

    /// Record one mutation; returns true when the batch should be flushed.
    fn record(&mut self) -> bool {
        self.pending += 1;
        self.pending >= PERSIST_EVERY_MUTATIONS || self.last_save.elapsed() >= PERSIST_INTERVAL
    }

    fn mark_saved(&mut self) {
        self.pending = 0;
        self.last_save = Instant::now();
    }
}

/// Helper function to convert Path to &str, returning an error if conversion fails
fn path_to_str(path: &Path) -> Result<&str> {
    path.to_str()
        .ok_or_else(|| Error::other(format!("Invalid UTF-8 path: {path:?}")))
}

/// resume state
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResumeState {
    /// on-disk schema version; absent in legacy snapshots (defaults to 0)
    #[serde(default)]
    pub schema_version: u32,
    /// authoritative opaque `(marker, version_marker)` continuation token for
    /// the version listing. `None` means "start from the beginning".
    #[serde(default)]
    pub resume_cursor: Option<String>,
    /// task id
    pub task_id: String,
    /// task type
    pub task_type: String,
    /// set disk identifier (for erasure set tasks)
    #[serde(default)]
    pub set_disk_id: String,
    /// start time
    pub start_time: u64,
    /// last update time
    pub last_update: u64,
    /// completed
    pub completed: bool,
    /// total objects
    pub total_objects: u64,
    /// processed objects
    pub processed_objects: u64,
    /// successful objects
    pub successful_objects: u64,
    /// failed objects
    pub failed_objects: u64,
    /// skipped objects
    pub skipped_objects: u64,
    /// current bucket
    pub current_bucket: Option<String>,
    /// current object
    pub current_object: Option<String>,
    /// completed buckets
    pub completed_buckets: Vec<String>,
    /// pending buckets
    pub pending_buckets: Vec<String>,
    /// error message
    pub error_message: Option<String>,
    /// retry count
    pub retry_count: u32,
    /// max retries
    pub max_retries: u32,
}

impl ResumeState {
    pub fn new(task_id: String, task_type: String, set_disk_id: String, buckets: Vec<String>) -> Self {
        Self {
            schema_version: CURRENT_RESUME_SCHEMA,
            resume_cursor: None,
            task_id,
            task_type,
            set_disk_id,
            start_time: SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs(),
            last_update: SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs(),
            completed: false,
            total_objects: 0,
            processed_objects: 0,
            successful_objects: 0,
            failed_objects: 0,
            skipped_objects: 0,
            current_bucket: None,
            current_object: None,
            completed_buckets: Vec::new(),
            pending_buckets: buckets,
            error_message: None,
            retry_count: 0,
            max_retries: 3,
        }
    }

    pub fn update_progress(&mut self, processed: u64, successful: u64, failed: u64, skipped: u64) {
        self.processed_objects = processed;
        self.successful_objects = successful;
        self.failed_objects = failed;
        self.skipped_objects = skipped;
        self.last_update = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
    }

    pub fn set_current_item(&mut self, bucket: Option<String>, object: Option<String>) {
        self.current_bucket = bucket;
        self.current_object = object;
        self.last_update = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
    }

    /// Read the authoritative version-listing continuation cursor.
    pub fn resume_cursor(&self) -> Option<String> {
        self.resume_cursor.clone()
    }

    /// Persist the authoritative version-listing continuation cursor. `None`
    /// resets the scan to the beginning of the current bucket.
    pub fn set_resume_cursor(&mut self, cursor: Option<String>) {
        self.resume_cursor = cursor;
        self.last_update = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
    }

    pub fn complete_bucket(&mut self, bucket: &str) {
        if !self.completed_buckets.contains(&bucket.to_string()) {
            self.completed_buckets.push(bucket.to_string());
        }
        if let Some(pos) = self.pending_buckets.iter().position(|b| b == bucket) {
            self.pending_buckets.remove(pos);
        }
        self.last_update = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
    }

    pub fn mark_completed(&mut self) {
        self.completed = true;
        self.last_update = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
    }

    /// Reset per-pass progress so a retry re-scans the whole set from the
    /// start. `retry_count`/`max_retries` are intentionally preserved so
    /// retries stay bounded.
    pub fn reset_for_retry(&mut self) {
        self.completed_buckets.clear();
        self.processed_objects = 0;
        self.successful_objects = 0;
        self.failed_objects = 0;
        self.skipped_objects = 0;
        self.completed = false;
        // A retry re-scans every bucket from the beginning, so the version
        // cursor must be cleared too — otherwise the retry would resume mid-scan.
        self.resume_cursor = None;
        self.last_update = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
    }

    pub fn set_error(&mut self, error: String) {
        self.error_message = Some(error);
        self.last_update = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
    }

    pub fn increment_retry(&mut self) {
        self.retry_count += 1;
        self.last_update = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
    }

    pub fn can_retry(&self) -> bool {
        self.retry_count < self.max_retries
    }

    pub fn get_progress_percentage(&self) -> f64 {
        if self.total_objects == 0 {
            return 0.0;
        }
        (self.processed_objects as f64 / self.total_objects as f64) * 100.0
    }

    pub fn get_success_rate(&self) -> f64 {
        let total = self.successful_objects + self.failed_objects;
        if total == 0 {
            return 0.0;
        }
        (self.successful_objects as f64 / total as f64) * 100.0
    }
}

/// resume manager
pub struct ResumeManager {
    disk: DiskStore,
    state: Arc<RwLock<ResumeState>>,
    throttle: Mutex<PersistThrottle>,
}

impl ResumeManager {
    /// create new resume manager
    pub async fn new(
        disk: DiskStore,
        task_id: String,
        task_type: String,
        set_disk_id: String,
        buckets: Vec<String>,
    ) -> Result<Self> {
        let state = ResumeState::new(task_id, task_type, set_disk_id, buckets);
        let manager = Self {
            disk,
            state: Arc::new(RwLock::new(state)),
            throttle: Mutex::new(PersistThrottle::new()),
        };

        // save initial state
        if let Err(e) = manager.save_state().await {
            warn!(
                target: "rustfs::heal::resume",
                event = EVENT_HEAL_RESUME_STATE,
                component = LOG_COMPONENT_HEAL,
                subsystem = LOG_SUBSYSTEM_RESUME,
                state = "initial_save_failed",
                error = %e,
                "Heal resume state persistence failed"
            );
        }
        Ok(manager)
    }

    /// load resume state from disk
    pub async fn load_from_disk(disk: DiskStore, task_id: &str) -> Result<Self> {
        let state_data = Self::read_state_file(&disk, task_id).await?;
        let mut state: ResumeState = serde_json::from_slice(&state_data).map_err(|e| Error::TaskExecutionFailed {
            message: format!("Failed to deserialize resume state: {e}"),
        })?;

        // A snapshot written by an older schema tracked a latest-only positional
        // cursor that is meaningless under per-version resume. Discard the stale
        // progress so the scan restarts cleanly, then stamp the current schema.
        if state.schema_version < CURRENT_RESUME_SCHEMA {
            warn!(
                target: "rustfs::heal::resume",
                event = EVENT_HEAL_RESUME_STATE,
                component = LOG_COMPONENT_HEAL,
                subsystem = LOG_SUBSYSTEM_RESUME,
                task_id,
                found_schema = state.schema_version,
                current_schema = CURRENT_RESUME_SCHEMA,
                state = "schema_discarded",
                "Heal resume state schema is stale; discarding cursor and progress"
            );
            state.resume_cursor = None;
            state.processed_objects = 0;
            state.successful_objects = 0;
            state.failed_objects = 0;
            state.skipped_objects = 0;
            state.completed = false;
            state.schema_version = CURRENT_RESUME_SCHEMA;
        }

        Ok(Self {
            disk,
            state: Arc::new(RwLock::new(state)),
            throttle: Mutex::new(PersistThrottle::new()),
        })
    }

    /// check if resume state exists
    pub async fn has_resume_state(disk: &DiskStore, task_id: &str) -> bool {
        let file_path = Path::new(BUCKET_META_PREFIX).join(format!("{task_id}_{RESUME_STATE_FILE}"));
        match path_to_str(&file_path) {
            Ok(path_str) => match disk.read_all(RUSTFS_META_BUCKET, path_str).await {
                Ok(data) => !data.is_empty(),
                Err(_) => false,
            },
            Err(_) => false,
        }
    }

    /// get current state
    pub async fn get_state(&self) -> ResumeState {
        self.state.read().await.clone()
    }

    /// update progress
    pub async fn update_progress(&self, processed: u64, successful: u64, failed: u64, skipped: u64) -> Result<()> {
        let mut state = self.state.write().await;
        state.update_progress(processed, successful, failed, skipped);
        drop(state);
        self.save_state_throttled().await
    }

    /// Set current item. Called once per healed object, so persistence is
    /// throttled: the in-memory state always updates, but the snapshot is only
    /// written every `PERSIST_EVERY_MUTATIONS` calls or `PERSIST_INTERVAL`.
    pub async fn set_current_item(&self, bucket: Option<String>, object: Option<String>) -> Result<()> {
        let mut state = self.state.write().await;
        state.set_current_item(bucket, object);
        drop(state);
        let should_save = self.throttle.lock().map(|mut throttle| throttle.record()).unwrap_or(true);
        if !should_save {
            return Ok(());
        }
        self.save_state_throttled().await
    }

    async fn save_state_throttled(&self) -> Result<()> {
        let result = self.save_state().await;
        if result.is_ok()
            && let Ok(mut throttle) = self.throttle.lock()
        {
            throttle.mark_saved();
        }
        result
    }

    /// Read the authoritative version-listing continuation cursor.
    pub async fn resume_cursor(&self) -> Option<String> {
        self.state.read().await.resume_cursor()
    }

    /// Persist the authoritative version-listing continuation cursor. This is
    /// written unthrottled (once per completed page) so a crash always resumes
    /// from a real page boundary.
    pub async fn set_resume_cursor(&self, cursor: Option<String>) -> Result<()> {
        let mut state = self.state.write().await;
        state.set_resume_cursor(cursor);
        drop(state);
        self.save_state().await
    }

    /// complete bucket
    pub async fn complete_bucket(&self, bucket: &str) -> Result<()> {
        let mut state = self.state.write().await;
        state.complete_bucket(bucket);
        drop(state);
        self.save_state_throttled().await
    }

    /// mark task completed
    pub async fn mark_completed(&self) -> Result<()> {
        let mut state = self.state.write().await;
        state.mark_completed();
        drop(state);
        self.save_state_throttled().await
    }

    /// set error message
    pub async fn set_error(&self, error: String) -> Result<()> {
        let mut state = self.state.write().await;
        state.set_error(error);
        drop(state);
        self.save_state().await
    }

    /// increment retry count
    pub async fn increment_retry(&self) -> Result<()> {
        let mut state = self.state.write().await;
        state.increment_retry();
        drop(state);
        self.save_state().await
    }

    /// Arm a bounded retry: if the retry budget remains, bump the retry
    /// counter, reset per-pass progress (so the next resume re-scans the whole
    /// set), persist, and return `true`. Returns `false` when retries are
    /// exhausted, leaving the state untouched.
    pub async fn schedule_retry(&self) -> Result<bool> {
        let mut state = self.state.write().await;
        if !state.can_retry() {
            return Ok(false);
        }
        state.increment_retry();
        state.reset_for_retry();
        drop(state);
        self.save_state().await?;
        Ok(true)
    }

    /// cleanup resume state
    pub async fn cleanup(&self) -> Result<()> {
        let state = self.state.read().await;
        let task_id = &state.task_id;

        // delete state files
        let state_file = Path::new(BUCKET_META_PREFIX).join(format!("{task_id}_{RESUME_STATE_FILE}"));
        let progress_file = Path::new(BUCKET_META_PREFIX).join(format!("{task_id}_{RESUME_PROGRESS_FILE}"));
        let checkpoint_file = Path::new(BUCKET_META_PREFIX).join(format!("{task_id}_{RESUME_CHECKPOINT_FILE}"));

        // ignore delete errors, files may not exist
        if let Ok(path_str) = path_to_str(&state_file) {
            let _ = self.disk.delete(RUSTFS_META_BUCKET, path_str, Default::default()).await;
        }
        if let Ok(path_str) = path_to_str(&progress_file) {
            let _ = self.disk.delete(RUSTFS_META_BUCKET, path_str, Default::default()).await;
        }
        if let Ok(path_str) = path_to_str(&checkpoint_file) {
            let _ = self.disk.delete(RUSTFS_META_BUCKET, path_str, Default::default()).await;
        }

        debug!(
            target: "rustfs::heal::resume",
            event = EVENT_HEAL_RESUME_STATE,
            component = LOG_COMPONENT_HEAL,
            subsystem = LOG_SUBSYSTEM_RESUME,
            task_id,
            state = "cleaned",
            "Heal resume state cleaned"
        );
        Ok(())
    }

    /// save state to disk
    async fn save_state(&self) -> Result<()> {
        let state = self.state.read().await;
        let state_data = serde_json::to_vec(&*state).map_err(|e| Error::TaskExecutionFailed {
            message: format!("Failed to serialize resume state: {e}"),
        })?;

        let file_path = Path::new(BUCKET_META_PREFIX).join(format!("{}_{}", state.task_id, RESUME_STATE_FILE));

        let path_str = path_to_str(&file_path)?;
        if let Err(e) = self.disk.write_all(RUSTFS_META_BUCKET, path_str, state_data.into()).await {
            if matches!(e, DiskError::UnformattedDisk) {
                warn!(
                    target: "rustfs::heal::resume",
                    event = EVENT_HEAL_RESUME_STATE,
                    component = LOG_COMPONENT_HEAL,
                    subsystem = LOG_SUBSYSTEM_RESUME,
                    task_id = %state.task_id,
                    state = "skipped_unformatted_disk",
                    "Heal resume state persistence skipped on unformatted disk"
                );
                return Ok(());
            }
            return Err(Error::TaskExecutionFailed {
                message: format!("Failed to save resume state: {e}"),
            });
        }

        debug!(
            target: "rustfs::heal::resume",
            event = EVENT_HEAL_RESUME_STATE,
            component = LOG_COMPONENT_HEAL,
            subsystem = LOG_SUBSYSTEM_RESUME,
            task_id = %state.task_id,
            state = "saved",
            "Heal resume state persisted"
        );
        Ok(())
    }

    /// read state file from disk
    async fn read_state_file(disk: &DiskStore, task_id: &str) -> Result<Vec<u8>> {
        let file_path = Path::new(BUCKET_META_PREFIX).join(format!("{task_id}_{RESUME_STATE_FILE}"));

        let path_str = path_to_str(&file_path)?;
        disk.read_all(RUSTFS_META_BUCKET, path_str)
            .await
            .map(|bytes| bytes.to_vec())
            .map_err(|e| Error::TaskExecutionFailed {
                message: format!("Failed to read resume state file: {e}"),
            })
    }
}

/// resume checkpoint
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResumeCheckpoint {
    /// on-disk schema version; absent in legacy snapshots (defaults to 0)
    #[serde(default)]
    pub schema_version: u32,
    /// task id
    pub task_id: String,
    /// checkpoint time
    pub checkpoint_time: u64,
    /// current bucket index
    pub current_bucket_index: usize,
    /// current object index
    pub current_object_index: usize,
    /// Objects healed since the last completed page. HashSet: with the
    /// previous Vec the per-object `contains` was O(n) and made large-bucket
    /// heals O(N²). Only spans the in-flight page — completed pages are
    /// covered by `current_object_index`, so `complete_page` prunes the sets.
    pub processed_objects: HashSet<String>,
    /// failed objects
    pub failed_objects: HashSet<String>,
    /// skipped objects
    pub skipped_objects: HashSet<String>,
}

impl ResumeCheckpoint {
    pub fn new(task_id: String) -> Self {
        Self {
            schema_version: CURRENT_CHECKPOINT_SCHEMA,
            task_id,
            checkpoint_time: SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs(),
            current_bucket_index: 0,
            current_object_index: 0,
            processed_objects: HashSet::new(),
            failed_objects: HashSet::new(),
            skipped_objects: HashSet::new(),
        }
    }

    pub fn update_position(&mut self, bucket_index: usize, object_index: usize) {
        self.current_bucket_index = bucket_index;
        self.current_object_index = object_index;
        self.checkpoint_time = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
    }

    pub fn add_processed_object(&mut self, object: String) {
        self.processed_objects.insert(object);
    }

    pub fn add_failed_object(&mut self, object: String) {
        self.failed_objects.insert(object);
    }

    pub fn add_skipped_object(&mut self, object: String) {
        self.skipped_objects.insert(object);
    }

    /// Advance past a fully-processed page: objects below `object_index` are
    /// skipped by position on resume, so the per-object sets no longer need
    /// their entries and would otherwise grow with the whole bucket.
    pub fn complete_page(&mut self, bucket_index: usize, object_index: usize) {
        self.update_position(bucket_index, object_index);
        self.processed_objects.clear();
        self.skipped_objects.clear();
        self.failed_objects.clear();
    }

    /// Reset the scan to the start and clear the per-object sets so a retry
    /// re-scans the whole set.
    pub fn reset_for_retry(&mut self) {
        self.update_position(0, 0);
        self.processed_objects.clear();
        self.skipped_objects.clear();
        self.failed_objects.clear();
    }
}

/// resume checkpoint manager
pub struct CheckpointManager {
    disk: DiskStore,
    checkpoint: Arc<RwLock<ResumeCheckpoint>>,
    throttle: Mutex<PersistThrottle>,
}

impl CheckpointManager {
    /// create new checkpoint manager
    pub async fn new(disk: DiskStore, task_id: String) -> Result<Self> {
        let checkpoint = ResumeCheckpoint::new(task_id);
        let manager = Self {
            disk,
            checkpoint: Arc::new(RwLock::new(checkpoint)),
            throttle: Mutex::new(PersistThrottle::new()),
        };

        // save initial checkpoint
        if let Err(e) = manager.save_checkpoint().await {
            warn!(
                target: "rustfs::heal::resume",
                event = EVENT_HEAL_CHECKPOINT_STATE,
                component = LOG_COMPONENT_HEAL,
                subsystem = LOG_SUBSYSTEM_RESUME,
                state = "initial_save_failed",
                error = %e,
                "Heal checkpoint persistence failed"
            );
        }
        Ok(manager)
    }

    /// load checkpoint from disk
    pub async fn load_from_disk(disk: DiskStore, task_id: &str) -> Result<Self> {
        let checkpoint_data = Self::read_checkpoint_file(&disk, task_id).await?;
        let mut checkpoint: ResumeCheckpoint =
            serde_json::from_slice(&checkpoint_data).map_err(|e| Error::TaskExecutionFailed {
                message: format!("Failed to deserialize checkpoint: {e}"),
            })?;

        // A checkpoint from an older schema stored latest-only dedup identities
        // that are not comparable to the new per-version `compose_key`
        // identities. Discard the stale sets and position, then stamp the
        // current schema so the scan restarts cleanly.
        if checkpoint.schema_version < CURRENT_CHECKPOINT_SCHEMA {
            warn!(
                target: "rustfs::heal::resume",
                event = EVENT_HEAL_CHECKPOINT_STATE,
                component = LOG_COMPONENT_HEAL,
                subsystem = LOG_SUBSYSTEM_RESUME,
                task_id,
                found_schema = checkpoint.schema_version,
                current_schema = CURRENT_CHECKPOINT_SCHEMA,
                state = "schema_discarded",
                "Heal checkpoint schema is stale; discarding dedup sets and position"
            );
            checkpoint.processed_objects.clear();
            checkpoint.failed_objects.clear();
            checkpoint.skipped_objects.clear();
            checkpoint.current_object_index = 0;
            checkpoint.schema_version = CURRENT_CHECKPOINT_SCHEMA;
        }

        Ok(Self {
            disk,
            checkpoint: Arc::new(RwLock::new(checkpoint)),
            throttle: Mutex::new(PersistThrottle::new()),
        })
    }

    /// check if checkpoint exists
    pub async fn has_checkpoint(disk: &DiskStore, task_id: &str) -> bool {
        let file_path = Path::new(BUCKET_META_PREFIX).join(format!("{task_id}_{RESUME_CHECKPOINT_FILE}"));
        match path_to_str(&file_path) {
            Ok(path_str) => match disk.read_all(RUSTFS_META_BUCKET, path_str).await {
                Ok(data) => !data.is_empty(),
                Err(_) => false,
            },
            Err(_) => false,
        }
    }

    /// get current checkpoint
    pub async fn get_checkpoint(&self) -> ResumeCheckpoint {
        self.checkpoint.read().await.clone()
    }

    /// update position
    pub async fn update_position(&self, bucket_index: usize, object_index: usize) -> Result<()> {
        let mut checkpoint = self.checkpoint.write().await;
        checkpoint.update_position(bucket_index, object_index);
        drop(checkpoint);
        self.save_checkpoint_throttled().await
    }

    /// Advance past a completed page and prune the per-object sets, then persist.
    pub async fn complete_page(&self, bucket_index: usize, object_index: usize) -> Result<()> {
        let mut checkpoint = self.checkpoint.write().await;
        checkpoint.complete_page(bucket_index, object_index);
        drop(checkpoint);
        self.save_checkpoint_throttled().await
    }

    /// Reset the checkpoint to the start of the scan for a retry, then persist.
    pub async fn reset_for_retry(&self) -> Result<()> {
        let mut checkpoint = self.checkpoint.write().await;
        checkpoint.reset_for_retry();
        drop(checkpoint);
        self.save_checkpoint_throttled().await
    }

    /// Add a processed object. Called once per healed object, so persistence
    /// is batched (`PERSIST_EVERY_MUTATIONS` / `PERSIST_INTERVAL`); positions
    /// and page boundaries still persist unconditionally.
    pub async fn add_processed_object(&self, object: String) -> Result<()> {
        let mut checkpoint = self.checkpoint.write().await;
        checkpoint.add_processed_object(object);
        drop(checkpoint);
        self.save_checkpoint_if_due().await
    }

    /// add failed object (batched, see `add_processed_object`)
    pub async fn add_failed_object(&self, object: String) -> Result<()> {
        let mut checkpoint = self.checkpoint.write().await;
        checkpoint.add_failed_object(object);
        drop(checkpoint);
        self.save_checkpoint_if_due().await
    }

    /// add skipped object (batched, see `add_processed_object`)
    pub async fn add_skipped_object(&self, object: String) -> Result<()> {
        let mut checkpoint = self.checkpoint.write().await;
        checkpoint.add_skipped_object(object);
        drop(checkpoint);
        self.save_checkpoint_if_due().await
    }

    async fn save_checkpoint_if_due(&self) -> Result<()> {
        let should_save = self.throttle.lock().map(|mut throttle| throttle.record()).unwrap_or(true);
        if !should_save {
            return Ok(());
        }
        self.save_checkpoint_throttled().await
    }

    async fn save_checkpoint_throttled(&self) -> Result<()> {
        let result = self.save_checkpoint().await;
        if result.is_ok()
            && let Ok(mut throttle) = self.throttle.lock()
        {
            throttle.mark_saved();
        }
        result
    }

    /// cleanup checkpoint
    pub async fn cleanup(&self) -> Result<()> {
        let checkpoint = self.checkpoint.read().await;
        let task_id = &checkpoint.task_id;

        let checkpoint_file = Path::new(BUCKET_META_PREFIX).join(format!("{task_id}_{RESUME_CHECKPOINT_FILE}"));
        if let Ok(path_str) = path_to_str(&checkpoint_file) {
            let _ = self.disk.delete(RUSTFS_META_BUCKET, path_str, Default::default()).await;
        }

        debug!(
            target: "rustfs::heal::resume",
            event = EVENT_HEAL_CHECKPOINT_STATE,
            component = LOG_COMPONENT_HEAL,
            subsystem = LOG_SUBSYSTEM_RESUME,
            task_id,
            state = "cleaned",
            "Heal checkpoint cleaned"
        );
        Ok(())
    }

    /// save checkpoint to disk
    async fn save_checkpoint(&self) -> Result<()> {
        let checkpoint = self.checkpoint.read().await;
        let checkpoint_data = serde_json::to_vec(&*checkpoint).map_err(|e| Error::TaskExecutionFailed {
            message: format!("Failed to serialize checkpoint: {e}"),
        })?;

        let file_path = Path::new(BUCKET_META_PREFIX).join(format!("{}_{}", checkpoint.task_id, RESUME_CHECKPOINT_FILE));

        let path_str = path_to_str(&file_path)?;
        self.disk
            .write_all(RUSTFS_META_BUCKET, path_str, checkpoint_data.into())
            .await
            .map_err(|e| Error::TaskExecutionFailed {
                message: format!("Failed to save checkpoint: {e}"),
            })?;

        debug!(
            target: "rustfs::heal::resume",
            event = EVENT_HEAL_CHECKPOINT_STATE,
            component = LOG_COMPONENT_HEAL,
            subsystem = LOG_SUBSYSTEM_RESUME,
            task_id = %checkpoint.task_id,
            state = "saved",
            "Heal checkpoint persisted"
        );
        Ok(())
    }

    /// read checkpoint file from disk
    async fn read_checkpoint_file(disk: &DiskStore, task_id: &str) -> Result<Vec<u8>> {
        let file_path = Path::new(BUCKET_META_PREFIX).join(format!("{task_id}_{RESUME_CHECKPOINT_FILE}"));

        let path_str = path_to_str(&file_path)?;
        disk.read_all(RUSTFS_META_BUCKET, path_str)
            .await
            .map(|bytes| bytes.to_vec())
            .map_err(|e| Error::TaskExecutionFailed {
                message: format!("Failed to read checkpoint file: {e}"),
            })
    }
}

/// resume utils
pub struct ResumeUtils;

impl ResumeUtils {
    /// generate unique task id
    pub fn generate_task_id() -> String {
        Uuid::new_v4().to_string()
    }

    /// check if task can be resumed
    pub async fn can_resume_task(disk: &DiskStore, task_id: &str) -> bool {
        ResumeManager::has_resume_state(disk, task_id).await
    }

    /// get all resumable task ids
    pub async fn get_resumable_tasks(disk: &DiskStore) -> Result<Vec<String>> {
        // List all files in the buckets metadata directory
        let entries = match disk.list_dir("", RUSTFS_META_BUCKET, BUCKET_META_PREFIX, -1).await {
            Ok(entries) => entries,
            Err(e) => {
                debug!(
                    target: "rustfs::heal::resume",
                    event = EVENT_HEAL_RESUME_STATE,
                    component = LOG_COMPONENT_HEAL,
                    subsystem = LOG_SUBSYSTEM_RESUME,
                    state = "list_failed",
                    error = %e,
                    "Heal resume state listing failed"
                );
                return Ok(Vec::new());
            }
        };

        let mut task_ids = Vec::new();

        // Filter files that end with ahm_resume_state.json and extract task IDs
        for entry in entries {
            if entry.ends_with(&format!("_{RESUME_STATE_FILE}")) {
                // Extract task ID from filename: {task_id}_ahm_resume_state.json
                if let Some(task_id) = entry.strip_suffix(&format!("_{RESUME_STATE_FILE}"))
                    && !task_id.is_empty()
                {
                    task_ids.push(task_id.to_string());
                }
            }
        }

        debug!(
            target: "rustfs::heal::resume",
            event = EVENT_HEAL_RESUME_STATE,
            component = LOG_COMPONENT_HEAL,
            subsystem = LOG_SUBSYSTEM_RESUME,
            task_count = task_ids.len(),
            state = "listed",
            "Heal resume states listed"
        );
        Ok(task_ids)
    }

    /// cleanup expired resume states
    pub async fn cleanup_expired_states(disk: &DiskStore, max_age_hours: u64) -> Result<()> {
        let task_ids = Self::get_resumable_tasks(disk).await?;
        let current_time = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs();

        for task_id in task_ids {
            if let Ok(resume_manager) = ResumeManager::load_from_disk(disk.clone(), &task_id).await {
                let state = resume_manager.get_state().await;
                let age_hours = (current_time - state.last_update) / 3600;

                if age_hours > max_age_hours {
                    debug!(
                        target: "rustfs::heal::resume",
                        event = EVENT_HEAL_RESUME_STATE,
                        component = LOG_COMPONENT_HEAL,
                        subsystem = LOG_SUBSYSTEM_RESUME,
                        task_id,
                        age_hours,
                        state = "expired_cleanup_started",
                        "Heal resume cleanup started"
                    );
                    if let Err(e) = resume_manager.cleanup().await {
                        warn!(
                            target: "rustfs::heal::resume",
                            event = EVENT_HEAL_RESUME_STATE,
                            component = LOG_COMPONENT_HEAL,
                            subsystem = LOG_SUBSYSTEM_RESUME,
                            task_id,
                            age_hours,
                            state = "expired_cleanup_failed",
                            error = %e,
                            "Heal resume state cleanup failed"
                        );
                    }
                }
            }
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_resume_state_creation() {
        let task_id = ResumeUtils::generate_task_id();
        let buckets = vec!["bucket1".to_string(), "bucket2".to_string()];
        let state = ResumeState::new(task_id.clone(), "erasure_set".to_string(), "pool_0_set_0".to_string(), buckets);

        assert_eq!(state.task_id, task_id);
        assert_eq!(state.task_type, "erasure_set");
        assert!(!state.completed);
        assert_eq!(state.processed_objects, 0);
        assert_eq!(state.pending_buckets.len(), 2);
    }

    #[tokio::test]
    async fn test_resume_state_progress() {
        let task_id = ResumeUtils::generate_task_id();
        let buckets = vec!["bucket1".to_string()];
        let mut state = ResumeState::new(task_id, "erasure_set".to_string(), "pool_0_set_0".to_string(), buckets);

        state.update_progress(10, 8, 1, 1);
        assert_eq!(state.processed_objects, 10);
        assert_eq!(state.successful_objects, 8);
        assert_eq!(state.failed_objects, 1);
        assert_eq!(state.skipped_objects, 1);

        let progress = state.get_progress_percentage();
        assert_eq!(progress, 0.0); // total_objects is 0

        state.total_objects = 100;
        let progress = state.get_progress_percentage();
        assert_eq!(progress, 10.0);
    }

    #[test]
    fn reset_for_retry_clears_progress_but_keeps_retry_budget() {
        // backlog#855 / #799 B6: a retry must re-scan from the start without
        // spending the retry budget's identity.
        let buckets = vec!["bucket1".to_string(), "bucket2".to_string()];
        let mut state = ResumeState::new("t".to_string(), "erasure_set".to_string(), "pool_0_set_0".to_string(), buckets);
        state.update_progress(10, 8, 2, 0);
        state.complete_bucket("bucket1");
        state.increment_retry();
        state.mark_completed();

        state.reset_for_retry();

        assert!(!state.completed, "retry must un-complete the task");
        assert_eq!(state.completed_buckets.len(), 0, "all buckets must be re-scanned");
        assert_eq!(state.processed_objects, 0);
        assert_eq!(state.successful_objects, 0);
        assert_eq!(state.failed_objects, 0);
        assert_eq!(state.skipped_objects, 0);
        assert_eq!(state.retry_count, 1, "retry budget must be preserved");
    }

    #[test]
    fn can_retry_is_bounded_by_max_retries() {
        let mut state = ResumeState::new("t".to_string(), "erasure_set".to_string(), "pool_0_set_0".to_string(), vec![]);
        assert!(state.can_retry());
        for _ in 0..state.max_retries {
            assert!(state.can_retry());
            state.increment_retry();
        }
        assert!(!state.can_retry(), "retries must stop after max_retries");
    }

    #[test]
    fn checkpoint_reset_for_retry_rewinds_position_and_clears_sets() {
        let mut checkpoint = ResumeCheckpoint::new("task".to_string());
        checkpoint.update_position(3, 42);
        checkpoint.add_processed_object("bucket/a".to_string());
        checkpoint.add_failed_object("bucket/b".to_string());
        checkpoint.add_skipped_object("bucket/c".to_string());

        checkpoint.reset_for_retry();

        assert_eq!(checkpoint.current_bucket_index, 0);
        assert_eq!(checkpoint.current_object_index, 0);
        assert!(checkpoint.processed_objects.is_empty());
        assert!(checkpoint.failed_objects.is_empty());
        assert!(checkpoint.skipped_objects.is_empty());
    }

    #[tokio::test]
    async fn test_resume_state_bucket_completion() {
        let task_id = ResumeUtils::generate_task_id();
        let buckets = vec!["bucket1".to_string(), "bucket2".to_string()];
        let mut state = ResumeState::new(task_id, "erasure_set".to_string(), "pool_0_set_0".to_string(), buckets);

        assert_eq!(state.pending_buckets.len(), 2);
        assert_eq!(state.completed_buckets.len(), 0);

        state.complete_bucket("bucket1");
        assert_eq!(state.pending_buckets.len(), 1);
        assert_eq!(state.completed_buckets.len(), 1);
        assert!(state.completed_buckets.contains(&"bucket1".to_string()));
    }

    #[test]
    fn test_checkpoint_object_sets_dedupe_and_prune() {
        let mut checkpoint = ResumeCheckpoint::new("task".to_string());
        checkpoint.add_processed_object("bucket/a".to_string());
        checkpoint.add_processed_object("bucket/a".to_string());
        checkpoint.add_skipped_object("bucket/b".to_string());
        checkpoint.add_failed_object("bucket/c".to_string());
        assert_eq!(checkpoint.processed_objects.len(), 1);
        assert!(checkpoint.processed_objects.contains("bucket/a"));

        checkpoint.complete_page(2, 2000);
        assert_eq!(checkpoint.current_bucket_index, 2);
        assert_eq!(checkpoint.current_object_index, 2000);
        assert!(checkpoint.processed_objects.is_empty());
        assert!(checkpoint.skipped_objects.is_empty());
        assert!(checkpoint.failed_objects.is_empty());
    }

    #[test]
    fn test_checkpoint_loads_legacy_vec_format() {
        // Checkpoints written before the HashSet migration stored the object
        // lists as JSON arrays (possibly with duplicates); they must still load.
        let legacy = r#"{
            "task_id": "t1",
            "checkpoint_time": 1700000000,
            "current_bucket_index": 1,
            "current_object_index": 42,
            "processed_objects": ["a", "b", "a"],
            "failed_objects": [],
            "skipped_objects": ["c"]
        }"#;
        let checkpoint: ResumeCheckpoint = serde_json::from_str(legacy).unwrap();
        assert_eq!(checkpoint.current_object_index, 42);
        assert_eq!(checkpoint.processed_objects.len(), 2);
        assert!(checkpoint.processed_objects.contains("a"));
        assert!(checkpoint.skipped_objects.contains("c"));
    }

    #[test]
    fn test_compose_key_injective_with_adversarial_keys() {
        // Length-prefixing must keep the encoding injective even when keys
        // contain the delimiter, embedded nulls, or look like a composed key.
        assert_ne!(compose_key("a\0b", None), compose_key("a", Some("b")));
        assert_ne!(compose_key("3:xy", None), compose_key("x", Some("y")));
        assert_ne!(compose_key("a:b", None), compose_key("a", Some("b")));
        assert_ne!(compose_key("", Some("x")), compose_key("x", None));
        // Identical inputs must produce identical keys (stable identity).
        assert_eq!(compose_key("obj", Some("v1")), compose_key("obj", Some("v1")));
    }

    #[test]
    fn test_composite_key_dedup_distinguishes_versions() {
        // Two versions of the same object must be distinct dedup identities, and
        // the delete-marker/nil (None) version must not collide with a real one.
        let mut checkpoint = ResumeCheckpoint::new("task".to_string());
        checkpoint.add_processed_object(compose_key("obj", Some("v1")));
        checkpoint.add_processed_object(compose_key("obj", Some("v2")));
        checkpoint.add_processed_object(compose_key("obj", None));
        assert_eq!(checkpoint.processed_objects.len(), 3);
        assert!(checkpoint.processed_objects.contains(&compose_key("obj", Some("v1"))));
        assert!(checkpoint.processed_objects.contains(&compose_key("obj", Some("v2"))));
        assert!(checkpoint.processed_objects.contains(&compose_key("obj", None)));
        // A different object with the same version id is still distinct.
        assert!(!checkpoint.processed_objects.contains(&compose_key("other", Some("v1"))));
    }

    #[tokio::test]
    async fn test_resumestate_schema_v0_discarded_on_load() {
        use super::super::{DiskOption, Endpoint, new_disk};
        use tempfile::TempDir;

        let temp_dir = TempDir::new().unwrap();
        let disk_path = temp_dir.path().join("test_disk");
        std::fs::create_dir_all(&disk_path).unwrap();
        let endpoint = Endpoint::try_from(disk_path.to_string_lossy().as_ref()).unwrap();
        let disk = new_disk(
            &endpoint,
            &DiskOption {
                cleanup: false,
                health_check: false,
            },
        )
        .await
        .unwrap();
        let _ = disk.make_volume(RUSTFS_META_BUCKET).await;
        let _ = disk.make_volume(&format!("{RUSTFS_META_BUCKET}/{BUCKET_META_PREFIX}")).await;

        // Legacy snapshot: no schema_version, a stale positional cursor and progress.
        let legacy = r#"{
            "task_id": "old-task",
            "task_type": "erasure_set",
            "set_disk_id": "pool_0_set_0",
            "start_time": 1700000000,
            "last_update": 1700000000,
            "completed": true,
            "total_objects": 100,
            "processed_objects": 50,
            "successful_objects": 40,
            "failed_objects": 10,
            "skipped_objects": 0,
            "current_bucket": null,
            "current_object": null,
            "completed_buckets": ["b1"],
            "pending_buckets": [],
            "error_message": null,
            "retry_count": 1,
            "max_retries": 3,
            "resume_cursor": "v1:stale-token"
        }"#;
        let file_path = format!("{BUCKET_META_PREFIX}/old-task_{RESUME_STATE_FILE}");
        disk.write_all(RUSTFS_META_BUCKET, &file_path, legacy.as_bytes().to_vec().into())
            .await
            .unwrap();

        let manager = ResumeManager::load_from_disk(disk.clone(), "old-task").await.unwrap();
        let state = manager.get_state().await;
        assert_eq!(state.schema_version, CURRENT_RESUME_SCHEMA, "schema must be stamped current");
        assert_eq!(state.resume_cursor, None, "stale cursor must be cleared");
        assert_eq!(state.processed_objects, 0);
        assert_eq!(state.successful_objects, 0);
        assert_eq!(state.failed_objects, 0);
        assert!(!state.completed);
        temp_dir.close().unwrap();
    }

    #[tokio::test]
    async fn test_checkpoint_schema_v0_discarded_on_load() {
        use super::super::{DiskOption, Endpoint, new_disk};
        use tempfile::TempDir;

        let temp_dir = TempDir::new().unwrap();
        let disk_path = temp_dir.path().join("test_disk");
        std::fs::create_dir_all(&disk_path).unwrap();
        let endpoint = Endpoint::try_from(disk_path.to_string_lossy().as_ref()).unwrap();
        let disk = new_disk(
            &endpoint,
            &DiskOption {
                cleanup: false,
                health_check: false,
            },
        )
        .await
        .unwrap();
        let _ = disk.make_volume(RUSTFS_META_BUCKET).await;
        let _ = disk.make_volume(&format!("{RUSTFS_META_BUCKET}/{BUCKET_META_PREFIX}")).await;

        // Legacy checkpoint: no schema_version, stale position and dedup sets.
        let legacy = r#"{
            "task_id": "old-task",
            "checkpoint_time": 1700000000,
            "current_bucket_index": 2,
            "current_object_index": 500,
            "processed_objects": ["a", "b"],
            "failed_objects": ["c"],
            "skipped_objects": ["d"]
        }"#;
        let file_path = format!("{BUCKET_META_PREFIX}/old-task_{RESUME_CHECKPOINT_FILE}");
        disk.write_all(RUSTFS_META_BUCKET, &file_path, legacy.as_bytes().to_vec().into())
            .await
            .unwrap();

        let manager = CheckpointManager::load_from_disk(disk.clone(), "old-task").await.unwrap();
        let checkpoint = manager.get_checkpoint().await;
        assert_eq!(checkpoint.schema_version, CURRENT_CHECKPOINT_SCHEMA, "schema must be stamped current");
        assert_eq!(checkpoint.current_object_index, 0, "stale position must be reset");
        assert!(checkpoint.processed_objects.is_empty());
        assert!(checkpoint.failed_objects.is_empty());
        assert!(checkpoint.skipped_objects.is_empty());
        temp_dir.close().unwrap();
    }

    #[test]
    fn test_persist_throttle_batches_until_threshold() {
        let mut throttle = PersistThrottle::new();
        for _ in 0..PERSIST_EVERY_MUTATIONS - 1 {
            assert!(!throttle.record(), "must not flush below the mutation threshold");
        }
        assert!(throttle.record(), "must flush at the mutation threshold");
        throttle.mark_saved();
        assert!(!throttle.record(), "counter must reset after a save");
    }

    #[tokio::test]
    async fn test_resume_utils() {
        let task_id1 = ResumeUtils::generate_task_id();
        let task_id2 = ResumeUtils::generate_task_id();

        assert_ne!(task_id1, task_id2);
        assert_eq!(task_id1.len(), 36); // UUID length
        assert_eq!(task_id2.len(), 36);
    }

    #[tokio::test]
    async fn test_get_resumable_tasks_integration() {
        use super::super::{DiskOption, Endpoint, new_disk};
        use tempfile::TempDir;

        // Create a temporary directory for testing
        let temp_dir = TempDir::new().unwrap();
        let disk_path = temp_dir.path().join("test_disk");
        std::fs::create_dir_all(&disk_path).unwrap();

        // Create a local disk for testing
        let endpoint = Endpoint::try_from(disk_path.to_string_lossy().as_ref()).unwrap();
        let disk_option = DiskOption {
            cleanup: false,
            health_check: false,
        };
        let disk = new_disk(&endpoint, &disk_option).await.unwrap();

        // Create necessary directories first (ignore if already exist)
        let _ = disk.make_volume(RUSTFS_META_BUCKET).await;
        let _ = disk.make_volume(&format!("{RUSTFS_META_BUCKET}/{BUCKET_META_PREFIX}")).await;

        // Create some test resume state files
        let task_ids = vec![
            "test-task-1".to_string(),
            "test-task-2".to_string(),
            "test-task-3".to_string(),
        ];

        // Save resume state files for each task
        for task_id in &task_ids {
            let state = ResumeState::new(
                task_id.clone(),
                "erasure_set".to_string(),
                "pool_0_set_0".to_string(),
                vec!["bucket1".to_string(), "bucket2".to_string()],
            );

            let state_data = serde_json::to_vec(&state).unwrap();
            let file_path = format!("{BUCKET_META_PREFIX}/{task_id}_{RESUME_STATE_FILE}");

            disk.write_all(RUSTFS_META_BUCKET, &file_path, state_data.into())
                .await
                .unwrap();
        }

        // Also create some non-resume state files to test filtering
        let non_resume_files = vec![
            "other_file.txt",
            "task4_ahm_checkpoint.json",
            "task5_ahm_progress.json",
            "_ahm_resume_state.json", // Invalid: empty task ID
        ];

        for file_name in non_resume_files {
            let file_path = format!("{BUCKET_META_PREFIX}/{file_name}");
            disk.write_all(RUSTFS_META_BUCKET, &file_path, b"test data".to_vec().into())
                .await
                .unwrap();
        }

        // Now call get_resumable_tasks to see if it finds the correct files
        let found_task_ids = ResumeUtils::get_resumable_tasks(&disk).await.unwrap();

        // Verify that only the valid resume state files are found
        assert_eq!(found_task_ids.len(), 3);
        for task_id in &task_ids {
            assert!(found_task_ids.contains(task_id), "Task ID {task_id} not found");
        }

        // Verify that invalid files are not included
        assert!(!found_task_ids.contains(&"".to_string()));
        assert!(!found_task_ids.contains(&"task4".to_string()));
        assert!(!found_task_ids.contains(&"task5".to_string()));

        // Clean up
        temp_dir.close().unwrap();
    }
}
