// 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.
/// scanner cache locks and the cache snapshot persist/publish path.
use super::*;

pub(crate) fn scanner_cache_lock_resource(cache_name: &str, source: DataUsageCacheSource) -> String {
    let lock_name = format!("{SCANNER_CACHE_LOCK_SUFFIX}.pool-{}.set-{}", source.pool_index, source.set_index);
    path_join_buf(&[crate::BUCKET_META_PREFIX, cache_name, &lock_name])
}

pub(crate) fn scanner_cache_lock_timeout() -> Duration {
    Duration::from_secs(rustfs_utils::get_env_u64("RUSTFS_LOCK_ACQUIRE_TIMEOUT", 5))
}

#[derive(Debug)]
pub(crate) struct ScannerCacheLockGuards {
    scoped: NamespaceLockGuard,
}

impl ScannerCacheLockGuards {
    pub(crate) fn is_lock_lost(&self) -> bool {
        self.scoped.is_lock_lost()
    }
}

#[derive(Debug)]
pub(crate) enum ScannerCacheLockError {
    Create { resource: String, source: Error },
    Acquire { resource: String, source: LockError },
}

impl ScannerCacheLockError {
    pub(crate) fn state(&self) -> &'static str {
        match self {
            Self::Create { .. } => "lock_create_failed",
            Self::Acquire { .. } => "lock_acquire_failed",
        }
    }

    pub(crate) fn is_contention(&self) -> bool {
        matches!(
            self,
            Self::Acquire {
                source: LockError::Timeout { .. } | LockError::AlreadyLocked { .. },
                ..
            }
        )
    }
}

impl std::fmt::Display for ScannerCacheLockError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Create { resource, source } => write!(formatter, "create scanner cache lock {resource}: {source}"),
            Self::Acquire { resource, source } => write!(formatter, "acquire scanner cache lock {resource}: {source}"),
        }
    }
}

pub(crate) async fn acquire_scanner_cache_locks(
    store: &SetDisks,
    cache_name: &str,
    source: DataUsageCacheSource,
) -> std::result::Result<ScannerCacheLockGuards, ScannerCacheLockError> {
    let timeout = scanner_cache_lock_timeout();
    let scoped_resource = scanner_cache_lock_resource(cache_name, source);
    let scoped_lock = store
        .new_ns_lock(RUSTFS_META_BUCKET, &scoped_resource)
        .await
        .map_err(|source| ScannerCacheLockError::Create {
            resource: scoped_resource.clone(),
            source,
        })?;
    let scoped = scoped_lock
        .get_write_lock_quiet(timeout)
        .await
        .map_err(|source| ScannerCacheLockError::Acquire {
            resource: scoped_resource,
            source,
        })?;

    Ok(ScannerCacheLockGuards { scoped })
}

pub(super) async fn await_scanner_disk_shutdown<F>(scan: Pin<&mut F>)
where
    F: Future,
{
    let _ = tokio::time::timeout(SCANNER_CACHE_LOCK_LOSS_SHUTDOWN_TIMEOUT, scan).await;
}

pub(crate) fn current_cache_root_entry(
    cache: &DataUsageCache,
    name: &str,
    source: DataUsageCacheSource,
    next_cycle: u64,
    leader_epoch: u64,
    scan_plan_digest: DataUsageScanPlanDigest,
) -> std::result::Result<Option<DataUsageEntryInfo>, ScannerError> {
    let metadata_is_current = cache.info.name == name
        && cache.info.source == Some(source)
        && cache.info.snapshot_complete
        && cache.info.scan_plan_digest == Some(scan_plan_digest)
        && cache.info.last_update.is_some()
        && cache.info.next_cycle == next_cycle
        && cache.info.leader_epoch == leader_epoch
        && cache.info.cache_key_format == DATA_USAGE_CACHE_KEY_FORMAT;
    if !metadata_is_current {
        return Ok(None);
    }

    cache_root_entry_info(cache).map(Some)
}

pub(crate) enum DataUsageCacheScanState {
    Current(Box<DataUsageEntryInfo>),
    Prepared {
        outcome: DataUsageCachePrepareOutcome,
        invalid_current: Option<ScannerError>,
    },
}

pub(crate) fn current_cache_root_or_prepare(
    cache: &mut DataUsageCache,
    name: &str,
    source: DataUsageCacheSource,
    next_cycle: u64,
    leader_epoch: u64,
    scan_plan_digest: DataUsageScanPlanDigest,
    require_source: bool,
) -> DataUsageCacheScanState {
    match current_cache_root_entry(cache, name, source, next_cycle, leader_epoch, scan_plan_digest) {
        Ok(Some(root)) => DataUsageCacheScanState::Current(Box::new(root)),
        current => DataUsageCacheScanState::Prepared {
            invalid_current: current.err(),
            outcome: cache.prepare_for_scan(name, next_cycle, leader_epoch, source, scan_plan_digest, require_source),
        },
    }
}

#[cfg(test)]
pub(super) fn cache_snapshot_is_current(
    cache: &DataUsageCache,
    name: &str,
    source: DataUsageCacheSource,
    next_cycle: u64,
    leader_epoch: u64,
    scan_plan_digest: DataUsageScanPlanDigest,
) -> bool {
    matches!(
        current_cache_root_entry(cache, name, source, next_cycle, leader_epoch, scan_plan_digest),
        Ok(Some(_))
    )
}

pub(super) fn completed_data_usage_info(
    results: &[DataUsageCache],
    expected_sources: &HashSet<DataUsageCacheSource>,
    all_buckets: &[String],
    bucket_plan_complete: bool,
    budget_elapsed: bool,
    cancelled: bool,
) -> Option<(DataUsageInfo, SystemTime)> {
    if !bucket_plan_complete {
        return None;
    }
    let completed_set_count = results.iter().filter(|result| result.info.last_update.is_some()).count();
    if !should_publish_completed_snapshot(completed_set_count, results.len(), budget_elapsed, cancelled) {
        return None;
    }
    if !scanner_results_form_complete_snapshot(results, expected_sources) {
        return None;
    }

    if results.iter().any(|result| result.root().is_none()) {
        return None;
    }

    let mut total = DataUsageEntry::default();
    let mut buckets_usage = HashMap::with_capacity(all_buckets.len());
    for bucket in all_buckets {
        let mut merged = DataUsageEntry::default();
        for result in results {
            let entry = result.checked_flatten(bucket)?;
            if !merged.checked_merge(&entry) {
                return None;
            }
        }
        if !total.checked_merge(&merged) {
            return None;
        }
        buckets_usage.insert(bucket.clone(), checked_bucket_usage_info(&merged)?);
    }

    let merged_last_update = results.iter().filter_map(|result| result.info.last_update).max()?;
    let bucket_sizes = buckets_usage
        .iter()
        .map(|(bucket, usage)| (bucket.clone(), usage.size))
        .collect();
    let data_usage_info = DataUsageInfo {
        last_update: Some(merged_last_update),
        scanner_cycle: Some(results.first()?.info.next_cycle),
        objects_total_count: u64::try_from(total.objects).ok()?,
        versions_total_count: u64::try_from(total.versions).ok()?,
        delete_markers_total_count: u64::try_from(total.delete_markers).ok()?,
        objects_total_size: u64::try_from(total.size).ok()?,
        tier_stats: total.all_tier_stats.filter(|tiers| !tiers.is_empty()),
        buckets_count: u64::try_from(all_buckets.len()).ok()?,
        bucket_sizes,
        buckets_usage,
        usage_snapshot_complete: true,
        ..Default::default()
    };
    Some((data_usage_info, merged_last_update))
}

pub(super) async fn send_cache_root_entry_info(
    bucket_result_tx: &mpsc::Sender<DataUsageEntryInfo>,
    cache: &DataUsageCache,
    pending_maintenance_work: &AtomicBool,
) -> std::result::Result<(), ScannerError> {
    let root = cache_root_entry_info(cache)?;
    send_cache_root_entry(bucket_result_tx, root, cache, pending_maintenance_work).await
}

pub(super) async fn send_cache_root_entry(
    bucket_result_tx: &mpsc::Sender<DataUsageEntryInfo>,
    root: DataUsageEntryInfo,
    cache: &DataUsageCache,
    pending_maintenance_work: &AtomicBool,
) -> std::result::Result<(), ScannerError> {
    record_bucket_pending_maintenance_work(cache, pending_maintenance_work);
    bucket_result_tx
        .send(root)
        .await
        .map_err(|err| ScannerError::Other(format!("scanner cache root channel closed: {err}")))
}

pub(super) async fn persist_and_publish_cache_snapshot(
    store: Arc<SetDisks>,
    updates: &mpsc::Sender<DataUsageCache>,
    mut cache_snapshot: DataUsageCache,
    cache_cycle_floor: &AtomicU64,
) -> Option<SystemTime> {
    let source = cache_snapshot.info.source?;
    let guard = match acquire_scanner_cache_locks(store.as_ref(), DATA_USAGE_CACHE_NAME, source).await {
        Ok(guard) => guard,
        Err(err) => {
            error!(
                target: "rustfs::scanner::io",
                event = EVENT_SCANNER_CACHE_PERSIST_STATE,
                component = LOG_COMPONENT_SCANNER,
                subsystem = LOG_SUBSYSTEM_IO,
                cache_name = DATA_USAGE_CACHE_NAME,
                state = err.state(),
                error = %err,
                "Scanner cache snapshot lock acquisition failed"
            );
            return None;
        }
    };

    let mut persisted = DataUsageCache::default();
    let revisions = match persisted.load_with_revisions(store.clone(), DATA_USAGE_CACHE_NAME).await {
        Ok(revisions) => revisions,
        Err(err) => {
            error!(
                target: "rustfs::scanner::io",
                event = EVENT_SCANNER_CACHE_PERSIST_STATE,
                component = LOG_COMPONENT_SCANNER,
                subsystem = LOG_SUBSYSTEM_IO,
                cache_name = DATA_USAGE_CACHE_NAME,
                state = "load_or_revision_lookup_failed",
                error = %err,
                "Scanner cache snapshot load or revision lookup failed"
            );
            return None;
        }
    };
    let scan_plan_digest = cache_snapshot.info.scan_plan_digest?;
    if persisted.info.next_cycle > cache_snapshot.info.next_cycle {
        cache_cycle_floor.fetch_max(persisted.info.next_cycle, Ordering::AcqRel);
        warn!(
            target: "rustfs::scanner::io",
            event = EVENT_SCANNER_CACHE_PERSIST_STATE,
            component = LOG_COMPONENT_SCANNER,
            subsystem = LOG_SUBSYSTEM_IO,
            cache_name = DATA_USAGE_CACHE_NAME,
            requested_cycle = cache_snapshot.info.next_cycle,
            persisted_cycle = persisted.info.next_cycle,
            state = "stale_cycle_rejected",
            "Scanner rejected a set cache cycle regression"
        );
        return None;
    }
    if persisted.info.leader_epoch > cache_snapshot.info.leader_epoch {
        error!(
            target: "rustfs::scanner::io",
            event = EVENT_SCANNER_CACHE_PERSIST_STATE,
            component = LOG_COMPONENT_SCANNER,
            subsystem = LOG_SUBSYSTEM_IO,
            cache_name = DATA_USAGE_CACHE_NAME,
            requested_epoch = cache_snapshot.info.leader_epoch,
            persisted_epoch = persisted.info.leader_epoch,
            state = "stale_leader_rejected",
            "Scanner rejected a set cache snapshot from an older leader epoch"
        );
        return None;
    }
    if matches!(
        current_cache_root_entry(
            &persisted,
            DATA_USAGE_ROOT,
            source,
            cache_snapshot.info.next_cycle,
            cache_snapshot.info.leader_epoch,
            scan_plan_digest,
        ),
        Ok(Some(_))
    ) {
        cache_snapshot = persisted;
    } else {
        if guard.is_lock_lost() {
            error!(
                target: "rustfs::scanner::io",
                event = EVENT_SCANNER_CACHE_PERSIST_STATE,
                component = LOG_COMPONENT_SCANNER,
                subsystem = LOG_SUBSYSTEM_IO,
                cache_name = DATA_USAGE_CACHE_NAME,
                state = "lock_lost",
                "Scanner cache snapshot save skipped after lock loss"
            );
            return None;
        }

        let done_save = Metrics::time(Metric::SaveUsage);
        if let Err(e) = cache_snapshot
            .save_with_revisions(store, DATA_USAGE_CACHE_NAME, &revisions)
            .await
        {
            error!(
                target: "rustfs::scanner::io",
                event = EVENT_SCANNER_CACHE_PERSIST_STATE,
                component = LOG_COMPONENT_SCANNER,
                subsystem = LOG_SUBSYSTEM_IO,
                cache_name = DATA_USAGE_CACHE_NAME,
                state = "save_failed",
                error = %e,
                "Scanner cache snapshot persistence failed"
            );
            done_save();
            return None;
        }
        done_save();
    }
    if guard.is_lock_lost() {
        error!(
            target: "rustfs::scanner::io",
            event = EVENT_SCANNER_CACHE_PERSIST_STATE,
            component = LOG_COMPONENT_SCANNER,
            subsystem = LOG_SUBSYSTEM_IO,
            cache_name = DATA_USAGE_CACHE_NAME,
            state = "lock_lost_after_save",
            "Scanner cache snapshot publish skipped after lock loss"
        );
        return None;
    }
    drop(guard);
    let last_update = cache_snapshot.info.last_update;

    if let Err(e) = updates.send(cache_snapshot).await {
        error!(
            target: "rustfs::scanner::io",
            event = EVENT_SCANNER_CACHE_PERSIST_STATE,
            component = LOG_COMPONENT_SCANNER,
            subsystem = LOG_SUBSYSTEM_IO,
            cache_name = DATA_USAGE_CACHE_NAME,
            state = "publish_failed",
            error = %e,
            "Scanner cache snapshot publish failed"
        );
    }

    last_update
}

pub(super) async fn send_data_usage_update(updates: &mpsc::Sender<DataUsageInfo>, data_usage_info: DataUsageInfo) -> Result<()> {
    updates.send(data_usage_info).await.map_err(|e| {
        error!(
            target: "rustfs::scanner::io",
            event = EVENT_SCANNER_DATA_USAGE_STREAM,
            component = LOG_COMPONENT_SCANNER,
            subsystem = LOG_SUBSYSTEM_IO,
            state = "send_failed",
            error = %e,
            "Scanner data usage publish failed"
        );
        StorageError::other("scanner data usage receiver closed before update delivery")
    })
}