use std::{
    collections::{
        HashMap,
        hash_map::Entry::{Occupied as HOccupied, Vacant as HVacant},
    },
    sync::Arc,
    time::Instant,
};

use anyhow::{Context, Result, anyhow, bail};
use log::{debug, error, info, warn};
use tici_common::{cdc_file_name::CDCFileName, index::IndexKey};
use tici_grpc::{
    AddShardRequest, CompactFragRequest, GrpcWorkerServiceClient, ImportWriteStatus,
    WorkerNodeHeartbeatRequest, WorkerNodeHeartbeatResponse, WorkerNodeShardStatus,
    WorkerNodeStatus,
};
use tici_import::{ImportTaskID, metrics::ImportIntoInflightGuard, writer::Writer};
use tici_shard::{
    CDCCheckpoint, FragLocation, ShardWriterMeta, ShardWriterMetaVersion,
    compactor::shard_compactor,
    config::WorkerNodeConfig,
    metadata::metadata_store::{MetadataStore, index_meta_key},
    utils::{
        frag_path_generator::FragPathGenerator, meta_client_pool::MetaClientPool,
        metrics::observe_compaction, schema::IndexMeta,
    },
    writer::{frag_batch_uploader::FragBatchUploader, shard_writer_manager::AddShardOutcome},
};
use tici_system::readable::ReadableDuration;
use tokio::{sync::RwLock, task::JoinHandle};
use tonic::Status;

use crate::{meta::writer::CompactionStateKey, worker::server::DispatchShardWriter};

const IMPORT_FINISH_RETRY_MAX_ATTEMPTS: u32 = 50;
const IMPORT_FINISH_RETRY_MAX_TASKS_PER_TICK: usize = 8;

type WorkerBackgroundTask =
    Arc<parking_lot::Mutex<Option<(JoinHandle<()>, tokio::sync::oneshot::Sender<()>)>>>;

#[derive(Clone, Debug)]
struct PendingImportFinish {
    err_msg: Option<String>,
    rows_processed: u64,
    bytes_processed: u64,
    input_file_size: u64,
    retry_attempts: u32,
}

#[derive(Clone, Debug)]
struct ImportWriteTaskState {
    heartbeat_status: ImportWriteStatus,
    pending_finish: Option<PendingImportFinish>,
}

#[derive(Debug)]
struct PendingImportFinishRetry {
    task_id: ImportTaskID,
    heartbeat_status: ImportWriteStatus,
    pending_finish: PendingImportFinish,
}

#[derive(Debug)]
pub enum AddShardSource {
    Rpc,
    Heartbeat,
    Restore,
}

pub struct WorkerNodeServiceCore {
    /// Meta service heartbeat
    meta_service_heartbeat_thread: WorkerBackgroundTask,
    import_finish_retry_thread: WorkerBackgroundTask,
    pub compaction_status: Arc<RwLock<HashMap<CompactionStateKey, CompactFragRequest>>>,
    import_write_status: Arc<RwLock<HashMap<ImportTaskID, ImportWriteTaskState>>>,
    pub config: WorkerNodeConfig,
    shard_writer_manager: DispatchShardWriter,
    heartbeat_interval: ReadableDuration,
    pub meta_client_pool: Arc<dyn MetaClientPool>,
    pub metadata_store: Arc<MetadataStore>,
    pub frag_path_gen: Arc<FragPathGenerator>,
    frag_batch_uploader: Arc<FragBatchUploader>,
}

impl WorkerNodeServiceCore {
    fn finish_retry_cap_reached(pending_finish: &PendingImportFinish) -> bool {
        pending_finish.retry_attempts >= IMPORT_FINISH_RETRY_MAX_ATTEMPTS
    }

    fn should_stop_reporting_import_task(pending_finish: &PendingImportFinish) -> bool {
        pending_finish.err_msg.is_some() && Self::finish_retry_cap_reached(pending_finish)
    }

    async fn finish_partition_build_once(
        meta_client: &mut dyn tici_shard::utils::meta_client::MetaClient,
        index_key: IndexKey,
        task_id: ImportTaskID,
        owner_addr: String,
        owner_seq: u64,
        pending_finish: &mut PendingImportFinish,
    ) -> bool {
        if Self::should_stop_reporting_import_task(pending_finish) {
            error!(
                "import_write: finish partition build already exhausted retries, task_id={}, retries={}",
                task_id, pending_finish.retry_attempts
            );
            return false;
        }

        match meta_client
            .finish_partition_build(
                index_key,
                task_id,
                owner_addr,
                owner_seq,
                pending_finish.err_msg.clone(),
                pending_finish.rows_processed,
                pending_finish.bytes_processed,
                pending_finish.input_file_size,
            )
            .await
        {
            Ok(()) => {
                if pending_finish.retry_attempts > 0 {
                    info!(
                        "import_write: finish partition build recovered after retries, task_id={}, retries={}",
                        task_id, pending_finish.retry_attempts
                    );
                } else {
                    info!(
                        "import_write: finish partition build success, task_id={}",
                        task_id
                    );
                }
                true
            }
            Err(e) => {
                pending_finish.retry_attempts = pending_finish.retry_attempts.saturating_add(1);
                if Self::should_stop_reporting_import_task(pending_finish) {
                    error!(
                        "import_write: finish partition build exhausted retries, task_id={}, retries={}, err={:?}",
                        task_id, pending_finish.retry_attempts, e
                    );
                } else if Self::finish_retry_cap_reached(pending_finish) {
                    warn!(
                        "import_write: finish partition build still pending after retry cap, task_id={}, retries={}, err={:?}",
                        task_id, pending_finish.retry_attempts, e
                    );
                } else {
                    warn!(
                        "import_write: finish partition build failed, task_id={}, retry={}/{}, err={:?}",
                        task_id, pending_finish.retry_attempts, IMPORT_FINISH_RETRY_MAX_ATTEMPTS, e
                    );
                }
                false
            }
        }
    }

    pub async fn new(
        config: WorkerNodeConfig,
        shard_writer_manager: DispatchShardWriter,
        heartbeat_interval: ReadableDuration,
        frag_batch_uploader: Arc<FragBatchUploader>,
        meta_client_pool: Arc<dyn MetaClientPool>,
        metadata_store: Arc<MetadataStore>,
        frag_path_gen: Arc<FragPathGenerator>,
    ) -> Result<Self> {
        let service = WorkerNodeServiceCore {
            meta_service_heartbeat_thread: Arc::new(parking_lot::Mutex::new(None)),
            import_finish_retry_thread: Arc::new(parking_lot::Mutex::new(None)),
            compaction_status: Arc::new(RwLock::new(HashMap::new())),
            import_write_status: Arc::new(RwLock::new(HashMap::new())),
            config,
            shard_writer_manager,
            heartbeat_interval,
            meta_client_pool,
            metadata_store,
            frag_path_gen,
            frag_batch_uploader,
        };

        Ok(service)
    }

    async fn retry_pending_import_finishes(&self) {
        let owner_addr = self.config.server.get_advertise_addr().to_string();
        let retry_candidates = {
            let mut guard = self.import_write_status.write().await;
            Self::take_pending_import_finish_retries(
                &mut guard,
                IMPORT_FINISH_RETRY_MAX_TASKS_PER_TICK,
            )
        };

        for PendingImportFinishRetry {
            task_id,
            heartbeat_status,
            mut pending_finish,
        } in retry_candidates
        {
            let mut meta_client = match self.meta_client_pool.get_meta_client().await {
                Ok(client) => client,
                Err(e) => {
                    warn!(
                        "import_write: failed to get meta client for pending finish retry, task_id={}, err={:?}",
                        task_id, e
                    );
                    let mut guard = self.import_write_status.write().await;
                    if let Some(state) = guard.get_mut(&task_id) {
                        state.pending_finish = Some(pending_finish);
                    } else {
                        warn!(
                            "import_write: pending finish retry task state disappeared while restoring after meta client failure, task_id={}",
                            task_id
                        );
                    }
                    continue;
                }
            };
            let success = Self::finish_partition_build_once(
                meta_client.as_mut(),
                IndexKey::new(
                    heartbeat_status.keyspace_id,
                    heartbeat_status.table_id,
                    heartbeat_status.index_id,
                ),
                task_id,
                owner_addr.clone(),
                heartbeat_status.owner_seq,
                &mut pending_finish,
            )
            .await;
            {
                let mut guard = self.import_write_status.write().await;
                if success {
                    guard.remove(&task_id);
                } else if WorkerNodeServiceCore::should_stop_reporting_import_task(&pending_finish)
                {
                    error!(
                        "import_write: stop reporting task after retry exhaustion, task_id={}, retries={}",
                        task_id, pending_finish.retry_attempts
                    );
                    guard.remove(&task_id);
                } else if let Some(state) = guard.get_mut(&task_id) {
                    state.pending_finish = Some(pending_finish);
                } else {
                    warn!(
                        "import_write: pending finish retry task state disappeared while restoring after finish failure, task_id={}",
                        task_id
                    );
                }
            }
        }
    }

    fn take_pending_import_finish_retries(
        states: &mut HashMap<ImportTaskID, ImportWriteTaskState>,
        max_tasks: usize,
    ) -> Vec<PendingImportFinishRetry> {
        if max_tasks == 0 {
            return Vec::new();
        }

        let mut candidates = Vec::new();
        for (task_id, state) in states.iter_mut() {
            let Some(pending_finish) = state.pending_finish.take() else {
                continue;
            };
            candidates.push(PendingImportFinishRetry {
                task_id: *task_id,
                heartbeat_status: state.heartbeat_status.clone(),
                pending_finish,
            });
        }

        candidates
            .sort_by_key(|candidate| (candidate.pending_finish.retry_attempts, candidate.task_id));
        let skipped = candidates.split_off(candidates.len().min(max_tasks));
        for skipped in skipped {
            if let Some(state) = states.get_mut(&skipped.task_id) {
                state.pending_finish = Some(skipped.pending_finish);
            }
        }
        candidates
    }

    #[cfg(any(test, feature = "testexport"))]
    pub fn shard_writer_manager(&self) -> DispatchShardWriter {
        self.shard_writer_manager.clone()
    }

    pub fn parse_add_shard_req(req: &AddShardRequest) -> Result<(ShardWriterMeta, IndexMeta)> {
        let table_info = req
            .table_info
            .as_ref()
            .ok_or_else(|| Status::invalid_argument("Missing table info"))?;
        let index_info = req
            .index_info
            .as_ref()
            .ok_or_else(|| Status::invalid_argument("Missing index info"))?;
        let shard_header = req
            .shard
            .as_ref()
            .ok_or_else(|| Status::invalid_argument("Missing shard header"))?;
        let cdc_location = req
            .cdc_location
            .as_ref()
            .ok_or_else(|| Status::invalid_argument("Missing CDC location"))?;
        let upload_location = req
            .upload_location
            .as_ref()
            .ok_or_else(|| Status::invalid_argument("Missing upload location"))?;

        let frag_location = FragLocation {
            bucket: upload_location.bucket.clone(),
            prefix: upload_location.prefix.clone(),
        };
        let cdc_checkpoint = CDCCheckpoint {
            bucket: cdc_location.bucket.clone(),
            table_prefix: cdc_location.prefix.clone(),
            last_file: CDCFileName::parse(&req.cdc_s3_last_file).with_context(|| {
                format!(
                    "Invalid AddShardRequest.cdc_s3_last_file={}",
                    req.cdc_s3_last_file
                )
            })?,
        };
        let shard_meta = tici_shard::ShardWriterMeta {
            version: ShardWriterMetaVersion::LATEST,
            index_key: IndexKey::from(index_info),
            shard_header: shard_header.clone(),
            frag_location,
            cdc_checkpoint,
        };
        let index_meta = tici_shard::utils::schema::IndexMeta::new(table_info, index_info)?;
        Ok((shard_meta, index_meta))
    }

    pub fn shard_statuses(&self) -> Vec<WorkerNodeShardStatus> {
        self.shard_writer_manager.shard_statuses()
    }

    #[cfg(any(test, feature = "testexport"))]
    pub fn shard_count(&self) -> usize {
        self.shard_statuses().len()
    }

    #[cfg(any(test, feature = "testexport"))]
    pub fn contains_shard(&self, shard_id: u64) -> bool {
        self.shard_statuses()
            .iter()
            .any(|status| status.shard_id == shard_id)
    }

    pub async fn add_shard(
        &self,
        shard_meta: ShardWriterMeta,
        index_meta: IndexMeta,
        from: AddShardSource,
    ) -> Result<()> {
        match &self.shard_writer_manager {
            DispatchShardWriter::Real(shard_writer_manager) => {
                let add_outcome = shard_writer_manager
                    .add_shard(&shard_meta, &index_meta)
                    .await
                    .with_context(|| {
                        format!(
                            "{:?}: Failed to add header={} to shard_writer_manager",
                            from,
                            shard_meta.shard_header.to_long_ident()
                        )
                    })?;
                match (from, add_outcome) {
                    (AddShardSource::Rpc | AddShardSource::Heartbeat, AddShardOutcome::Added) => {
                        // Try to persist index meta and shard writer meta asynchronously.
                        let _ = self
                            .metadata_store
                            .persist_worker_local_shard_metadata(&shard_meta, &index_meta)
                            .await;
                    }
                    (_, AddShardOutcome::AlreadyExists) | (AddShardSource::Restore, _) => {}
                }
            }
            DispatchShardWriter::Mock(mock) => {
                mock.add_shard(&shard_meta).await;
            }
        }
        Ok(())
    }

    pub async fn remove_shard(&self, shard_id: u64) -> Result<()> {
        match &self.shard_writer_manager {
            DispatchShardWriter::Real(shard_writer_manager) => {
                shard_writer_manager.remove_shard(shard_id).await?;
                // Try to remove shard writer meta asynchronously.
                let _ = self.metadata_store.remove_shard_writer_meta(shard_id).await;
            }
            DispatchShardWriter::Mock(mock) => {
                mock.remove_shard(&shard_id).await;
            }
        }
        Ok(())
    }

    pub async fn compact_fragments(
        &self,
        compaction_key: CompactionStateKey,
        req: CompactFragRequest,
        output_frag_path: String,
    ) -> Result<()> {
        let config = self.config.clone();
        let limits = config.compaction.max_concurrency;

        fail::fail_point!("worker_service_mock_drop_compact_shard", |_| { Ok(()) });

        let mut meta_client = self.meta_client_pool.get_meta_client().await?;
        {
            let mut status = self.compaction_status.write().await;
            let compactings = status.len();
            match status.entry(compaction_key.clone()) {
                HOccupied(_) => {
                    info!(
                        "compact_fragments: duplicated compaction task: {:?}",
                        compaction_key
                    );
                    return Ok(());
                }

                HVacant(v) => {
                    if compactings >= limits {
                        bail!(
                            "limits reached, compactings={}, limits={}",
                            compactings,
                            limits
                        );
                    }
                    v.insert(req.clone());
                }
            };
        }
        observe_compaction(&req, "schedule");

        let compaction_status = self.compaction_status.clone();
        match &self.shard_writer_manager {
            DispatchShardWriter::Real(_) => {
                let frag_batch_uploader = self.frag_batch_uploader.clone();
                tokio::task::spawn_blocking(move || {
                    // TODO: Refine the internal synchronous and asynchronous logic in
                    // compact_frags and avoid block on.
                    let handle = tokio::runtime::Handle::current();
                    handle.block_on(async {
                        let finish_req = match shard_compactor::compact_frags(
                            req.clone(),
                            config,
                            output_frag_path,
                            frag_batch_uploader,
                        )
                        .await
                        {
                            Err(e) => {
                                warn!("compact_fragments: err={:?}, req={:?}", e, req);
                                // Later, meta will abort the compaction after received the
                                // heartbeat
                                compaction_status.write().await.remove(&compaction_key);
                                return;
                            }
                            Ok(r) => r,
                        };
                        observe_compaction(&finish_req, "execute");
                        if let Err(e) = meta_client.finish_compact_frags(finish_req.clone()).await {
                            warn!(
                                "compact_fragments: finish_compact_frags failed finish_req={:?}, err={:?}",
                                finish_req, e
                            );
                        } else {
                            info!(
                                "compact_fragments: finish_compact_frags success finish_req={:?}",
                                finish_req
                            );
                        }
                        // We must firstly send the notification, otherwise the peer may think
                        // the compaction is aborted.
                        // No matter success or failed, we must remove the compaction status.
                        compaction_status.write().await.remove(&compaction_key);
                    });
                });
            }
            DispatchShardWriter::Mock(m) => {
                let m = m.clone();
                tokio::task::spawn(async move {
                    if m.mock_compact_process
                        .load(std::sync::atomic::Ordering::Relaxed)
                    {
                        let finish_req = m.compact_fragments(req.clone()).await.unwrap();
                        if let Err(e) = meta_client.finish_compact_frags(finish_req.clone()).await {
                            info!(
                                "compact_fragments: send finish req err={:?}, req={:?}, finish_req={:?}",
                                e, &req, &finish_req
                            );
                        } else {
                            info!(
                                "compact_fragments: removed compaction={:?}",
                                &compaction_key
                            );
                            // We must firstly send the notification, otherwise the peer may think
                            // the compaction is aborted.
                            compaction_status.write().await.remove(&compaction_key);
                        }
                    }
                });
            }
        }
        Ok(())
    }

    pub async fn import_write(&self, req: tici_grpc::ImportWriteRequest) -> Result<()> {
        let file_location = req.file_location;

        // TODO: refactor the FragUploader to accept the backend storage directly.
        let (input_backend_bucket, input_backend_prefix) = {
            match storage_backend::backend::Backend::parse(&file_location).inspect_err(|e| {
                warn!("import_write: parse backend err={:?}", e);
            })? {
                storage_backend::backend::Backend::S3(s3) => (s3.bucket, s3.prefix),
                storage_backend::backend::Backend::GCS(gcs) => (gcs.bucket, gcs.prefix),
            }
        };
        let output_backend_prefix = {
            match storage_backend::backend::Backend::parse(&req.frag_prefix).inspect_err(|e| {
                warn!("import_write: parse backend err={:?}", e);
            })? {
                storage_backend::backend::Backend::S3(s3) => s3.prefix,
                storage_backend::backend::Backend::GCS(gcs) => gcs.prefix,
            }
        };

        let task_id = req.task_id;
        info!("import_write: start, task_id={}", task_id);
        match self.import_write_status.write().await.entry(task_id) {
            HVacant(v) => {
                v.insert(ImportWriteTaskState {
                    heartbeat_status: ImportWriteStatus {
                        keyspace_id: req.keyspace_id,
                        table_id: req.table_id,
                        index_id: req.index_id,
                        file_prefix: req.frag_prefix.clone(),
                        task_id,
                        owner_seq: req.owner_seq,
                    },
                    pending_finish: None,
                });
            }
            HOccupied(_) => {
                return Err(anyhow!(
                    "import_write: duplicated task, task_id={}",
                    task_id
                ));
            }
        };
        let import_write_status = self.import_write_status.clone();
        let inflight_guard = ImportIntoInflightGuard::new();
        let metadata_store = self.metadata_store.clone();
        let meta_client_pool = self.meta_client_pool.clone();
        let frag_path_gen = self.frag_path_gen.clone();
        let mut meta_client = meta_client_pool.get_meta_client().await?;
        match &self.shard_writer_manager {
            DispatchShardWriter::Real(_) => {
                let worker_config = self.config.clone();
                let frag_batch_uploader = self.frag_batch_uploader.clone();
                // run the import write task in a blocking thread to avoid blocking the async
                // runtime. But we do not await the thread, so that this function can return
                // immediately.
                tokio::task::spawn_blocking(move || {
                    let _inflight_guard = inflight_guard;
                    tokio::runtime::Handle::current().block_on(async {
                        let owner_addr = worker_config.server.get_advertise_addr().to_string();
                        let writer = Writer {
                            worker_config,
                            s3_bucket: input_backend_bucket,
                            import_file: input_backend_prefix,
                            frag_upload_prefix: output_backend_prefix,
                            index_key: IndexKey::new(req.keyspace_id, req.table_id, req.index_id),
                            frag_path_gen,
                            task_id,
                            metadata_store,
                            frag_batch_uploader,
                        };
                        let pending_finish = match writer.write(meta_client_pool.clone()).await {
                            Ok(summary) => PendingImportFinish {
                                err_msg: None,
                                rows_processed: summary.rows_processed,
                                bytes_processed: summary.bytes_processed,
                                input_file_size: summary.input_file_size,
                                retry_attempts: 0,
                            },
                            Err(e) => PendingImportFinish {
                                err_msg: Some(format!("{:#}", e)),
                                rows_processed: 0,
                                bytes_processed: 0,
                                input_file_size: 0,
                                retry_attempts: 0,
                            },
                        };
                        info!(
                            "import_write: writer.write finished, task_id={}, err_msg={:?}",
                            task_id, pending_finish.err_msg
                        );
                        // TODO: Return the frag meta to meta service.
                        let mut pending_finish = pending_finish;
                        let success = Self::finish_partition_build_once(
                            meta_client.as_mut(),
                            IndexKey::new(req.keyspace_id, req.table_id, req.index_id),
                            req.task_id,
                            owner_addr,
                            req.owner_seq,
                            &mut pending_finish,
                        )
                        .await;
                        if success {
                            import_write_status.write().await.remove(&task_id);
                        } else if Self::should_stop_reporting_import_task(&pending_finish) {
                            error!(
                                "import_write: stop reporting task after retry exhaustion, task_id={}, retries={}",
                                task_id, pending_finish.retry_attempts
                            );
                            import_write_status.write().await.remove(&task_id);
                        } else {
                            let mut guard = import_write_status.write().await;
                            if let Some(state) = guard.get_mut(&task_id) {
                                state.pending_finish = Some(pending_finish);
                            }
                        }
                    });
                });
                Ok(())
            }
            DispatchShardWriter::Mock(_) => {
                unimplemented!();
            }
        }
    }

    pub fn collect_worker_node_status(
        proc_stats: &mut tici_system::cpu_time::ProcessStat,
    ) -> WorkerNodeStatus {
        WorkerNodeStatus {
            cpu_usage: match proc_stats.cpu_usage() {
                Ok(f) => f,
                Err(e) => {
                    error!("heartbeat_loop: failed to get cpu usage {:?}", e);
                    0.0
                }
            },
            memory_usage: {
                tici_system::record_global_memory_usage();
                tici_system::get_global_memory_usage()
            },
        }
    }

    async fn run_one(
        advertise_addr: &str,
        worker_node: &WorkerNodeServiceCore,
        proc_stats: &mut tici_system::cpu_time::ProcessStat,
    ) {
        let compaction_status = worker_node.compaction_status.clone();
        let import_write_status = worker_node.import_write_status.clone();

        let mut client = match worker_node.meta_client_pool.get_meta_client().await {
            Ok(mc) => mc,
            Err(e) => {
                warn!(
                    "heartbeat_loop: failed to get meta client for heartbeat {}, err={:?}",
                    advertise_addr, e
                );
                return;
            }
        };

        let mock_writer = match &worker_node.shard_writer_manager {
            DispatchShardWriter::Real(_) => None,
            DispatchShardWriter::Mock(mock) => {
                if mock.mock_skip_heartbeat() {
                    info!(
                        "heartbeat_loop: mock_skip_heartbeat called, {}",
                        advertise_addr
                    );
                    return;
                }
                Some(mock.clone())
            }
        };
        let node_stat = if let Some(mock) = mock_writer.as_ref() {
            mock.mock_collect_worker_node_status(proc_stats)
        } else {
            WorkerNodeServiceCore::collect_worker_node_status(proc_stats)
        };

        let mut req = WorkerNodeHeartbeatRequest {
            addr: advertise_addr.to_owned(),
            status: Some(node_stat),
            ..Default::default()
        };

        for s in worker_node.shard_statuses() {
            if mock_writer
                .as_ref()
                .is_some_and(|mock| mock.mock_should_skip_shard_heartbeat(&s.shard_id))
            {
                continue;
            }
            let shards: &mut Vec<tici_grpc::WorkerNodeShardStatus> = &mut req.shards;
            shards.push(s);
        }
        for (_, c) in compaction_status.read().await.iter() {
            let compactions: &mut Vec<tici_grpc::WorkerNodeCompactionStatus> = &mut req.compactions;
            compactions.push(tici_grpc::WorkerNodeCompactionStatus {
                req: Some(c.clone()),
            });
        }
        for state in import_write_status.read().await.values() {
            let import_writes: &mut Vec<tici_grpc::ImportWriteStatus> = &mut req.import_writes;
            import_writes.push(state.heartbeat_status.clone());
        }

        debug!(
            "heartbeat send, advertise_addr={}, shards_ids={:?}, req={:?}",
            advertise_addr,
            req.shards
                .iter()
                .map(|s| s.shard_id.to_string())
                .collect::<Vec<_>>()
                .join(","),
            req
        );

        match client.worker_node_heartbeat(req.clone()).await {
            Err(e) => {
                warn!(
                    "heartbeat_loop: meta service heartbeat err={:?}, shard_len={}, compaction_len={}, import_task_len={}",
                    e,
                    req.shards.len(),
                    req.compactions.len(),
                    req.import_writes.len()
                );
            }
            Ok(responses) => {
                for resp in responses {
                    worker_node.handle_heartbeat_response(resp).await;
                }
            }
        };
    }

    pub async fn heartbeat_loop(
        worker_node: Arc<WorkerNodeServiceCore>,
        interval: u64,
        advertise_addr: String,
        mut rx: tokio::sync::oneshot::Receiver<()>,
    ) -> JoinHandle<()> {
        let mut proc_stats = tici_system::cpu_time::ProcessStat::cur_proc_stat().unwrap();
        tokio::spawn(async move {
            loop {
                Self::run_one(
                    advertise_addr.as_str(),
                    worker_node.as_ref(),
                    &mut proc_stats,
                )
                .await;
                tokio::select! {
                    _ = &mut rx => {
                        info!("Stop heartbeat to meta service");
                        break;
                    },
                    _ = tokio::time::sleep(std::time::Duration::from_millis(interval)) => {
                    }
                }
            }
        })
    }

    pub async fn import_finish_retry_loop(
        worker_node: Arc<WorkerNodeServiceCore>,
        interval: u64,
        mut rx: tokio::sync::oneshot::Receiver<()>,
    ) -> JoinHandle<()> {
        tokio::spawn(async move {
            loop {
                worker_node.retry_pending_import_finishes().await;
                tokio::select! {
                    _ = &mut rx => {
                        info!("Stop import finish retry loop");
                        break;
                    },
                    _ = tokio::time::sleep(std::time::Duration::from_millis(interval)) => {
                    }
                }
            }
        })
    }

    pub async fn handle_heartbeat_response(&self, resp: WorkerNodeHeartbeatResponse) {
        if resp.status != 0 {
            error!("heartbeat: meta service heartbeat err={}", resp.status);
            return;
        }

        // Add shards if needed.
        if !resp.to_add_shards.is_empty() {
            if let DispatchShardWriter::Mock(m) = &self.shard_writer_manager {
                m.maybe_check_to_add_shards(resp.to_add_shards.as_slice())
                    .await;
            }
            for req in resp.to_add_shards.iter() {
                let (shard_meta, index_meta) = match Self::parse_add_shard_req(req) {
                    Ok((shard_meta, index_meta)) => (shard_meta, index_meta),
                    Err(e) => {
                        error!("heartbeat: parse_add_shard_req err={:?}", e);
                        continue;
                    }
                };
                let _ = self
                    .add_shard(shard_meta, index_meta, AddShardSource::Heartbeat)
                    .await
                    .inspect_err(|e| error!("heartbeat: add_shard err={:?}", e));
            }
        }

        // Remove shards if needed.
        for remove in resp.to_remove_shards.iter() {
            info!("heartbeat: remove shard={}", remove);
            let _ = self
                .remove_shard(*remove)
                .await
                .inspect_err(|e| error!("heartbeat: remove_shard err={:?}", e));
        }
    }

    pub async fn start_heartbeat_tasks(self: &Arc<Self>) -> Result<()> {
        let (heartbeat_tx, heartbeat_rx) = tokio::sync::oneshot::channel();
        let (retry_tx, retry_rx) = tokio::sync::oneshot::channel();
        let interval = self.heartbeat_interval.as_millis();
        let advertise_addr = self.config.server.get_advertise_addr().to_string();
        info!(
            "Started sending heartbeats to meta service, advertise_addr={}",
            advertise_addr
        );
        let heartbeat_handle = WorkerNodeServiceCore::heartbeat_loop(
            self.clone(),
            interval,
            advertise_addr,
            heartbeat_rx,
        )
        .await;
        let retry_handle =
            WorkerNodeServiceCore::import_finish_retry_loop(self.clone(), interval, retry_rx).await;
        *self.meta_service_heartbeat_thread.lock() = Some((heartbeat_handle, heartbeat_tx));
        *self.import_finish_retry_thread.lock() = Some((retry_handle, retry_tx));
        Ok(())
    }

    async fn stop_heartbeat_tasks(&self) -> Result<()> {
        Self::stop_background_task("heartbeat", &self.meta_service_heartbeat_thread).await?;
        Self::stop_background_task("import_finish_retry", &self.import_finish_retry_thread).await?;
        Ok(())
    }

    async fn stop_background_task(name: &str, task: &WorkerBackgroundTask) -> Result<()> {
        let res = task.lock().take();
        if let Some((j, tx)) = res {
            tx.send(())
                .map_err(|e| anyhow!("stop_background_task({}): send error={:?}", name, e))?;
            j.await?;
        }
        Ok(())
    }

    pub async fn get_client(
        addr: String,
    ) -> Result<GrpcWorkerServiceClient<tonic::transport::Channel>> {
        let addr = if addr.starts_with("http://") || addr.starts_with("https://") {
            addr.to_string()
        } else {
            format!("http://{}", addr)
        };
        let mut attempts = 0;
        let max_attempts = (|| {
            fail::fail_point!("worker_node_client_fail_fast", |t| {
                t.unwrap().parse::<u64>().unwrap()
            });
            30
        })();
        loop {
            match GrpcWorkerServiceClient::connect(addr.clone()).await {
                std::result::Result::Ok(c) => {
                    info!("Connected to WorkerServiceCore gRPC Server at {}", addr);
                    break Ok(c);
                }
                std::result::Result::Err(e) => {
                    info!(
                        "Failed to connect to worker node server at {} (attempt {}/{})={:?}",
                        addr,
                        attempts + 1,
                        max_attempts,
                        e
                    );
                    std::thread::sleep(std::time::Duration::from_secs(1));
                    attempts += 1;
                    if attempts >= max_attempts {
                        break Err(anyhow!("exceed max attempts"));
                    }
                }
            }
        }
    }

    pub async fn restore(&self) -> Result<()> {
        let start = Instant::now();
        let index_metas = self
            .metadata_store
            .scan_index_metas_skip_invalid()?
            .into_iter()
            .collect::<HashMap<String, IndexMeta>>();
        let shard_writer_metas = self.metadata_store.scan_shard_writer_metas_skip_invalid()?;
        info!(
            "restore: scan meta data: {} index_metas, {} shard_writer_metas, use {:?}",
            index_metas.len(),
            shard_writer_metas.len(),
            start.elapsed()
        );

        for (_, shard_writer_meta) in shard_writer_metas {
            let index_key = index_meta_key(&shard_writer_meta.index_key);
            if let Some(index_meta) = index_metas.get(&index_key) {
                let _ = self
                    .add_shard(
                        shard_writer_meta,
                        index_meta.clone(),
                        AddShardSource::Restore,
                    )
                    .await
                    .inspect_err(|e| error!("restore: add_shard err={:?}", e));
            } else {
                warn!(
                    "restore: index_meta not found for shard_writer_meta={:?}",
                    shard_writer_meta.shard_tracing_id()
                );
            }
        }

        info!("restore: done, use {:?}", start.elapsed());

        Ok(())
    }

    pub async fn shutdown(&self) -> Result<()> {
        // Heartbeat responses can add or remove shards. Stop the heartbeat loop before
        // draining shard writers so shutdown cannot race with a new shard operation.
        let stop_heartbeat_result = self.stop_heartbeat_tasks().await;
        match &self.shard_writer_manager {
            DispatchShardWriter::Real(shard_writer_manager) => {
                shard_writer_manager.shutdown().await;
            }
            DispatchShardWriter::Mock(_) => {}
        }
        stop_heartbeat_result
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use tici_common::index::IndexKey;
    use tici_shard::utils::meta_client::MockMetaClient;

    use super::*;

    fn pending_import_write_state(
        task_id: ImportTaskID,
        retry_attempts: u32,
    ) -> ImportWriteTaskState {
        ImportWriteTaskState {
            heartbeat_status: ImportWriteStatus {
                task_id,
                owner_seq: 1,
                ..Default::default()
            },
            pending_finish: Some(PendingImportFinish {
                err_msg: None,
                rows_processed: 11,
                bytes_processed: 12,
                input_file_size: 13,
                retry_attempts,
            }),
        }
    }

    #[test]
    fn pending_finish_retry_candidates_are_limited_and_preserve_unselected() {
        let mut states = HashMap::new();
        for task_id in 1..=10 {
            states.insert(
                task_id,
                pending_import_write_state(task_id, 10 - task_id as u32),
            );
        }

        let candidates = WorkerNodeServiceCore::take_pending_import_finish_retries(&mut states, 3);
        let candidate_ids = candidates
            .iter()
            .map(|candidate| candidate.task_id)
            .collect::<Vec<_>>();
        assert_eq!(candidate_ids, vec![10, 9, 8]);

        for task_id in [10, 9, 8] {
            assert!(states.get(&task_id).unwrap().pending_finish.is_none());
        }
        for task_id in 1..=7 {
            assert!(states.get(&task_id).unwrap().pending_finish.is_some());
        }
    }

    #[tokio::test]
    async fn finish_partition_build_once_records_pending_failure() {
        let mut client = MockMetaClient::new();
        client.debug_set_finish_partition_build_failures_before_success(1);
        let mut pending_finish = PendingImportFinish {
            err_msg: None,
            rows_processed: 11,
            bytes_processed: 12,
            input_file_size: 13,
            retry_attempts: 0,
        };
        let success = WorkerNodeServiceCore::finish_partition_build_once(
            &mut client,
            IndexKey::new(0, 1, 2),
            7,
            "worker-a".to_string(),
            1,
            &mut pending_finish,
        )
        .await;
        assert!(!success);
        assert_eq!(client.debug_finish_partition_build_call_count(), 1);
        assert_eq!(pending_finish.retry_attempts, 1);
    }

    #[tokio::test]
    async fn finish_partition_build_once_succeeds_after_pending_failure() {
        let mut client = MockMetaClient::new();
        client.debug_set_finish_partition_build_failures_before_success(1);
        let mut pending_finish = PendingImportFinish {
            err_msg: None,
            rows_processed: 11,
            bytes_processed: 12,
            input_file_size: 13,
            retry_attempts: 0,
        };
        let success = WorkerNodeServiceCore::finish_partition_build_once(
            &mut client,
            IndexKey::new(0, 1, 2),
            7,
            "worker-a".to_string(),
            1,
            &mut pending_finish,
        )
        .await;
        assert!(!success);
        assert_eq!(pending_finish.retry_attempts, 1);

        let success = WorkerNodeServiceCore::finish_partition_build_once(
            &mut client,
            IndexKey::new(0, 1, 2),
            7,
            "worker-a".to_string(),
            1,
            &mut pending_finish,
        )
        .await;
        assert!(success);
        assert_eq!(client.debug_finish_partition_build_call_count(), 2);
        assert_eq!(pending_finish.retry_attempts, 1);
    }

    #[tokio::test]
    async fn failed_finish_retry_cap_skips_rpc() {
        let mut client = MockMetaClient::new();
        client.debug_set_finish_partition_build_failures_before_success(100);
        let mut pending_finish = PendingImportFinish {
            err_msg: Some("build failed".to_string()),
            rows_processed: 11,
            bytes_processed: 12,
            input_file_size: 13,
            retry_attempts: IMPORT_FINISH_RETRY_MAX_ATTEMPTS,
        };
        let success = WorkerNodeServiceCore::finish_partition_build_once(
            &mut client,
            IndexKey::new(0, 1, 2),
            7,
            "worker-a".to_string(),
            1,
            &mut pending_finish,
        )
        .await;
        assert!(!success);
        assert_eq!(
            pending_finish.retry_attempts,
            IMPORT_FINISH_RETRY_MAX_ATTEMPTS
        );
        assert_eq!(client.debug_finish_partition_build_call_count(), 0);
    }

    #[tokio::test]
    async fn successful_finish_retry_cap_keeps_reporting() {
        let mut client = MockMetaClient::new();
        client.debug_set_finish_partition_build_failures_before_success(100);
        let mut pending_finish = PendingImportFinish {
            err_msg: None,
            rows_processed: 11,
            bytes_processed: 12,
            input_file_size: 13,
            retry_attempts: IMPORT_FINISH_RETRY_MAX_ATTEMPTS,
        };
        let success = WorkerNodeServiceCore::finish_partition_build_once(
            &mut client,
            IndexKey::new(0, 1, 2),
            7,
            "worker-a".to_string(),
            1,
            &mut pending_finish,
        )
        .await;
        assert!(!success);
        assert_eq!(
            pending_finish.retry_attempts,
            IMPORT_FINISH_RETRY_MAX_ATTEMPTS + 1
        );
        assert_eq!(client.debug_finish_partition_build_call_count(), 1);
    }
}