// Copyright 2026 PingCAP, Inc.

use std::time::Duration;

use pd_client::PdClient;

use super::*;
use crate::{Error, Result, UnsafeRecoverClearTarget, wal::SkippedStores};

impl ReplicationWorker {
    pub(super) fn handle_unsafe_recover_clear(
        &mut self,
        target: UnsafeRecoverClearTarget,
        dry_run: bool,
    ) -> Result<Vec<u64>> {
        info!(
            "handle_unsafe_recover_clear: start";
            "target" => ?target,
            "dry_run" => dry_run,
        );

        let requested_regions = self.resolve_unsafe_recover_clear_regions(target)?;
        info!(
            "handle_unsafe_recover_clear: resolved regions";
            "target" => ?target,
            "dry_run" => dry_run,
            "requested_regions" => ?requested_regions,
        );

        self.precheck_unsafe_recover_clear_upstream_stores(&requested_regions)?;
        info!(
            "handle_unsafe_recover_clear: upstream precheck success";
            "target" => ?target,
            "dry_run" => dry_run,
            "requested_regions" => ?requested_regions,
        );

        self.merged_engine
            .unsafe_recover_check_clear_regions(&requested_regions)
            .map_err(|err| Error::UnsafeRecoverClearRejected(err.to_string()))?;
        info!(
            "handle_unsafe_recover_clear: merged-engine dependency check success";
            "target" => ?target,
            "dry_run" => dry_run,
            "requested_regions" => ?requested_regions,
        );

        if dry_run {
            info!(
                "handle_unsafe_recover_clear: dry-run success";
                "target" => ?target,
                "dry_run" => dry_run,
                "requested_regions" => ?requested_regions,
            );
            return Ok(requested_regions);
        }

        for &region_id in &requested_regions {
            self.apply_ctx.flush_observer_region(region_id);
        }
        info!(
            "handle_unsafe_recover_clear: observer flush success";
            "target" => ?target,
            "dry_run" => dry_run,
            "requested_regions" => ?requested_regions,
        );

        let destroyed_regions = self
            .merged_engine
            .unsafe_recover_execute_clear_regions(&requested_regions)?;
        info!(
            "handle_unsafe_recover_clear: execute clear success";
            "target" => ?target,
            "dry_run" => dry_run,
            "requested_regions" => ?requested_regions,
            "destroyed_regions" => ?destroyed_regions,
        );

        let requested_set: HashSet<u64> = requested_regions.iter().copied().collect();
        let extra_destroyed_regions = destroyed_regions
            .iter()
            .copied()
            .filter(|region_id| !requested_set.contains(region_id))
            .collect::<Vec<_>>();
        if !extra_destroyed_regions.is_empty() {
            for region_id in &extra_destroyed_regions {
                self.apply_ctx.flush_observer_region(*region_id);
            }
        }

        for region_id in &destroyed_regions {
            // Unsafe recover clear bypasses the normal `CdcMsg::RegionDestroyed`
            // path, so it must explicitly reuse the same CDC destroy cleanup.
            self.handle_region_destroyed(*region_id, 0);
        }
        info!(
            "handle_unsafe_recover_clear: completed";
            "target" => ?target,
            "dry_run" => dry_run,
            "requested_regions" => ?requested_regions,
            "destroyed_regions" => ?destroyed_regions,
        );
        Ok(requested_regions)
    }

    fn resolve_unsafe_recover_clear_regions(
        &mut self,
        target: UnsafeRecoverClearTarget,
    ) -> Result<Vec<u64>> {
        let raft = self.merged_engine.get_raft();
        let mut requested_regions = match target {
            UnsafeRecoverClearTarget::Region(region_id) => {
                let region_is_live = self.merged_engine.get_region_progress(region_id).is_some()
                    && Self::get_region_for_rep(&raft, region_id).is_some();
                if !region_is_live {
                    return Err(Error::UnsafeRecoverClearNotFound(format!(
                        "region {region_id} not found"
                    )));
                }
                vec![region_id]
            }
            UnsafeRecoverClearTarget::Keyspace(keyspace_id) => {
                let mut regions = self.merged_engine.get_keyspace_regions(keyspace_id);
                regions.retain(|region_id| Self::get_region_for_rep(&raft, *region_id).is_some());
                if regions.is_empty() {
                    return Err(Error::UnsafeRecoverClearNotFound(format!(
                        "no live regions in keyspace {keyspace_id}"
                    )));
                }
                regions
            }
        };
        requested_regions.sort_unstable();
        requested_regions.dedup();
        Ok(requested_regions)
    }

    fn precheck_unsafe_recover_clear_upstream_stores(
        &self,
        requested_regions: &[u64],
    ) -> Result<()> {
        let pd = self.ctx.pd.clone();
        let skip_store_addr_keywords = self.config.skip_store_addr_keywords.clone();
        let include_only_store_addr_keywords = self.config.include_only_store_addr_keywords.clone();
        let refresh_interval = self.config.wal.min_target_time_span.0;
        let timeout = self.config.unsafe_recover_clear_pd_precheck_timeout.0;
        self.runtime
            .block_on(precheck_unsafe_recover_clear_upstream_stores_with_timeout(
                pd,
                skip_store_addr_keywords,
                include_only_store_addr_keywords,
                refresh_interval,
                requested_regions,
                timeout,
            ))
    }
}

async fn precheck_unsafe_recover_clear_upstream_stores_impl(
    pd: std::sync::Arc<dyn PdClient>,
    skip_store_addr_keywords: Vec<String>,
    include_only_store_addr_keywords: Vec<String>,
    refresh_interval: Duration,
    requested_regions: &[u64],
) -> Result<()> {
    let mut skipped_stores = SkippedStores::new(
        pd.clone(),
        skip_store_addr_keywords,
        include_only_store_addr_keywords,
        refresh_interval,
    );
    skipped_stores.refresh().await?;

    for &region_id in requested_regions {
        let Some(region) = pd.get_region_by_id(region_id).await? else {
            continue;
        };
        let mut offending_store_ids = region
            .get_peers()
            .iter()
            .map(|peer| peer.get_store_id())
            .filter(|store_id| !skipped_stores.contains(*store_id))
            .collect::<Vec<_>>();
        if !offending_store_ids.is_empty() {
            offending_store_ids.sort_unstable();
            offending_store_ids.dedup();
            return Err(Error::UnsafeRecoverClearRejected(format!(
                "region {region_id} still has an upstream peer on non-skipped stores {:?}",
                offending_store_ids,
            )));
        }
    }
    Ok(())
}

async fn precheck_unsafe_recover_clear_upstream_stores_with_timeout(
    pd: std::sync::Arc<dyn PdClient>,
    skip_store_addr_keywords: Vec<String>,
    include_only_store_addr_keywords: Vec<String>,
    refresh_interval: Duration,
    requested_regions: &[u64],
    timeout: Duration,
) -> Result<()> {
    tokio::time::timeout(
        timeout,
        precheck_unsafe_recover_clear_upstream_stores_impl(
            pd,
            skip_store_addr_keywords,
            include_only_store_addr_keywords,
            refresh_interval,
            requested_regions,
        ),
    )
    .await
    .map_err(|_| {
        Error::PdTimeout(format!(
            "unsafe recover clear upstream PD precheck timed out after {timeout:?}",
        ))
    })?
}

#[cfg(test)]
mod tests {
    use std::{sync::Arc, time::Duration};

    use async_trait::async_trait;
    use futures::{FutureExt, future::pending};
    use kvproto::metapb;
    use pd_client::{PdClient, PdFuture};
    use security::GetSecurityManager;

    use super::*;

    fn new_store(store_id: u64, addr: &str) -> metapb::Store {
        let mut store = metapb::Store::default();
        store.set_id(store_id);
        store.set_address(addr.to_string());
        store
    }

    struct PendingStoreListPdClient;

    impl GetSecurityManager for PendingStoreListPdClient {}

    #[async_trait]
    impl PdClient for PendingStoreListPdClient {
        fn get_all_stores_async(&self, _: bool) -> PdFuture<Vec<metapb::Store>> {
            pending::<pd_client::Result<Vec<metapb::Store>>>().boxed()
        }
    }

    struct PendingRegionLookupPdClient;

    impl GetSecurityManager for PendingRegionLookupPdClient {}

    #[async_trait]
    impl PdClient for PendingRegionLookupPdClient {
        fn get_all_stores_async(&self, _: bool) -> PdFuture<Vec<metapb::Store>> {
            futures::future::ready(Ok(vec![new_store(1, "included-store-1")])).boxed()
        }

        fn get_region_by_id(&self, _: u64) -> PdFuture<Option<metapb::Region>> {
            pending::<pd_client::Result<Option<metapb::Region>>>().boxed()
        }
    }

    #[tokio::test]
    async fn test_unsafe_recover_clear_precheck_timeout() {
        let cases: Vec<(&str, Arc<dyn PdClient>)> = vec![
            (
                "refreshing skipped stores",
                Arc::new(PendingStoreListPdClient),
            ),
            ("loading regions", Arc::new(PendingRegionLookupPdClient)),
        ];

        for (case_name, pd_client) in cases {
            let err = precheck_unsafe_recover_clear_upstream_stores_with_timeout(
                pd_client,
                vec![],
                vec!["included-store".to_string()],
                Duration::from_secs(1),
                &[1],
                Duration::from_millis(10),
            )
            .await
            .unwrap_err();

            assert!(
                matches!(err, Error::PdTimeout(ref message) if message.contains("unsafe recover clear upstream PD precheck timed out")),
                "case {case_name} returned unexpected error: {err:?}",
            );
        }
    }
}