use std::{
    collections::{
        BTreeMap, HashSet,
        btree_map::Entry::{Occupied as BOccupied, Vacant as BVacant},
    },
    fmt::Debug,
    sync::{
        Arc,
        atomic::{AtomicBool, AtomicUsize},
    },
};

use anyhow::{Result, anyhow, bail};
use arc_swap::ArcSwap;
use log::{debug, error, info, warn};
use tici_grpc::{
    AppendDeltaFragMetaRequest, AppendFragMetaResponse, FinishCompactFragRequest,
    FinishCompactFragResponse, WorkerNodeHeartbeatRequest, WorkerNodeHeartbeatResponse,
    WorkerNodeStatus,
};
use tokio::sync::{
    Notify, RwLock,
    mpsc::{UnboundedReceiver, UnboundedSender},
};

use crate::{
    config::MetaServiceConfig,
    meta::{
        MetaStore,
        util::BgRunner,
        writer::{
            AppendDeltaFragTask, CompactionAbortTask, CompactionFinishTask, CompactionStateKey,
            HealthCheckMgr, HealthCheckerCommands, NeedRescheduleTask, RescheduleShardReason,
            ShardTask, ShardTaskCtx, ShardTaskHeader, WorkerNodeMeta, task::*,
        },
    },
};

const IMPORT_TASK_NOTIFY_RETRY_DELAY_SECS: u64 = 10;

pub struct WorkerNodeMgr {
    worker_nodes: RwLock<BTreeMap<String, Arc<WorkerNodeMeta>>>,
    blocked_worker_addrs: RwLock<HashSet<String>>,
    meta_store: Arc<dyn MetaStore>,
    shard_task_tx: UnboundedSender<ShardTask>,
    wn_task_tx: UnboundedSender<WorkerNodeTask>,
    current_wn_idx: AtomicUsize,
    pub(super) config: Arc<ArcSwap<MetaServiceConfig>>,
    stopped: AtomicBool,
    _runner: BgRunner,
    _health_runner: BgRunner,
    health_tx: UnboundedSender<HealthCheckerCommands>,
    import_dispatch_progress: Notify,
}

impl Debug for WorkerNodeMgr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("WorkerNodeMgr")
            .field(
                "stopped",
                &self.stopped.load(std::sync::atomic::Ordering::Relaxed),
            )
            .finish()
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum WorkerNodeScheduleStrategy {
    RoundRobin,
    BalancedCount,
}

struct WorkerNodeCandidate {
    index: usize,
    addr: String,
    worker: Arc<WorkerNodeMeta>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) enum ImportWorkerNodeSelection {
    Selected(String),
    BlockedByPendingDispatch,
    NoAvailableWorker,
}

impl WorkerNodeScheduleStrategy {
    fn from_config(value: &str) -> Self {
        if value.eq_ignore_ascii_case("balanced-count")
            || value.eq_ignore_ascii_case("balanced_count")
            || value.eq_ignore_ascii_case("balancedcount")
        {
            return Self::BalancedCount;
        }
        Self::RoundRobin
    }
}

impl WorkerNodeMgr {
    pub fn new(
        shard_task_tx: UnboundedSender<ShardTask>,
        wn_task_tx: UnboundedSender<WorkerNodeTask>,
        wn_task_rx: UnboundedReceiver<WorkerNodeTask>,
        config: Arc<ArcSwap<MetaServiceConfig>>,
        meta_store: Arc<dyn MetaStore>,
    ) -> Arc<Self> {
        let (runner, stx) = BgRunner::wrapper(
            "WorkerNodeMgr".to_string(),
            wn_task_rx,
            |mgr: Arc<WorkerNodeMgr>, x: WorkerNodeTask| async move {
                if Self::handle_task(mgr.as_ref(), x).await {
                    mgr.stopped
                        .store(true, std::sync::atomic::Ordering::Relaxed);
                    return true;
                }
                false
            },
            |mgr: Arc<WorkerNodeMgr>| {
                mgr.stopped
                    .store(true, std::sync::atomic::Ordering::Relaxed);
            },
        );

        let (health_runner, health_tx, stx2) = HealthCheckMgr::create_runner();

        let s = Arc::new(Self {
            worker_nodes: RwLock::new(BTreeMap::new()),
            blocked_worker_addrs: RwLock::new(HashSet::new()),
            meta_store,
            shard_task_tx,
            wn_task_tx,
            // So the first selection starts from index 0.
            current_wn_idx: AtomicUsize::new(usize::MAX),
            config,
            stopped: AtomicBool::new(false),
            _runner: runner,
            _health_runner: health_runner,
            health_tx,
            import_dispatch_progress: Notify::new(),
        });
        if stx.send(s.clone()).is_err() {
            warn!("error send");
        }
        if stx2.send(s.clone()).is_err() {
            warn!("error send 2");
        }
        s
    }

    pub fn get_tx(&self) -> UnboundedSender<WorkerNodeTask> {
        self.wn_task_tx.clone()
    }

    pub fn stop(&self) {
        if let Err(e) = self.wn_task_tx.send(WorkerNodeTask::Hup) {
            warn!("send hup to wn_task_tx error, e={:?}", e);
        }
        if let Err(e) = self.health_tx.send(HealthCheckerCommands::Hup) {
            warn!("send hup to health_tx error, e={:?}", e);
        }
    }

    pub fn is_stopped(&self) -> bool {
        self.stopped.load(std::sync::atomic::Ordering::Relaxed)
    }

    pub(crate) async fn wait_import_dispatch_progress(&self) {
        self.import_dispatch_progress.notified().await;
    }

    fn notify_import_dispatch_progress(&self) {
        self.import_dispatch_progress.notify_one();
    }

    pub async fn block_workers(&self, addrs: &[String]) {
        let mut g = self.blocked_worker_addrs.write().await;
        for addr in addrs {
            g.insert(addr.clone());
        }
    }

    pub async fn unblock_workers(&self, addrs: &[String]) {
        let mut g = self.blocked_worker_addrs.write().await;
        for addr in addrs {
            g.remove(addr);
        }
    }

    pub async fn list_blocked_workers(&self) -> Vec<String> {
        let mut out: Vec<String> = self
            .blocked_worker_addrs
            .read()
            .await
            .iter()
            .cloned()
            .collect();
        out.sort();
        out
    }

    pub async fn is_worker_blocked(&self, addr: &str) -> bool {
        self.blocked_worker_addrs.read().await.contains(addr)
    }

    #[cfg(any(test, feature = "testexport"))]
    pub fn cloned_config(&self) -> Arc<MetaServiceConfig> {
        self.config.load_full()
    }
}

impl WorkerNodeMgr {
    async fn handle_task(mgr: &WorkerNodeMgr, task: WorkerNodeTask) -> bool {
        match task {
            WorkerNodeTask::Hup => {
                info!("WorkerNodeMgr: hup");
                return true;
            }
            WorkerNodeTask::CreateWorkerNode(CreateWorkerNodeTask {
                ctx,
                status,
                meta_store,
            }) => {
                mgr.create_worker_node(meta_store, ctx.addr.clone(), status)
                    .await;
                ctx.report_ok(());
            }
            WorkerNodeTask::ModifyShard(task) => match mgr.get_worker_node(&task.ctx.addr).await {
                Some(wn) => {
                    wn.send_wn_task(WorkerNodeTask::ModifyShard(task));
                }
                None => {
                    let addr = task.ctx.addr.clone();
                    if !task.to_add.is_empty() {
                        warn!(
                            "handle_task(ModifyShard): add shard to empty addr, task={:?}",
                            task
                        );
                    } else {
                        info!(
                            "remove shard from a non-existing work node, maybe already removed={}, task={:?}",
                            addr, task,
                        );
                    }
                    task.ctx.report_ok(());
                }
            },
            WorkerNodeTask::AddCompaction(task) => {
                match mgr.get_worker_node(&task.ctx.addr).await {
                    Some(wn) => {
                        wn.send_wn_task(WorkerNodeTask::AddCompaction(task));
                    }
                    None => {
                        task.ctx.report_error(anyhow!(
                            "handle_task(AddCompaction): failed to get worker, shard={:?}",
                            task.header,
                        ));
                    }
                }
            }
            WorkerNodeTask::AddImportFragBuildTask(import_task) => {
                let wn_addr = import_task.owner_addr.clone();
                match mgr.get_worker_node(wn_addr.as_str()).await {
                    Some(wn) => {
                        if !wn
                            .is_pending_import_task_assignment(
                                import_task.task_id,
                                &import_task.owner_addr,
                                import_task.owner_seq,
                            )
                            .await
                        {
                            info!(
                                "handle_task(AddImportFragBuildTask): skip stale or non-pending import task dispatch, import_task={:?} addr={:?}",
                                &import_task, wn_addr
                            );
                            return false;
                        }

                        // If error happens, the import task will be reclaimed by heartbeat
                        // reconcile or worker removal.
                        if let Err(e) = wn
                            .notify_import_frag_build(&import_task, &mgr.config.load().s3)
                            .await
                        {
                            error!(
                                "handle_task(AddImportFragBuildTask): failed to notify worker, import_task={:?} addr={:?} err={}",
                                &import_task, wn_addr, e
                            );
                            let retry_tx = mgr.wn_task_tx.clone();
                            tokio::spawn(async move {
                                tokio::time::sleep(tokio::time::Duration::from_secs(
                                    IMPORT_TASK_NOTIFY_RETRY_DELAY_SECS,
                                ))
                                .await;
                                if let Err(e) = retry_tx
                                    .send(WorkerNodeTask::AddImportFragBuildTask(import_task))
                                {
                                    warn!(
                                        "handle_task(AddImportFragBuildTask): failed to enqueue retry, task={:?}",
                                        e.0
                                    );
                                }
                            });
                        } else {
                            let marked = wn
                                .mark_import_task_running(
                                    import_task.task_id,
                                    &import_task.owner_addr,
                                    import_task.owner_seq,
                                )
                                .await;
                            if !marked {
                                warn!(
                                    "handle_task(AddImportFragBuildTask): notified stale import task, import_task={:?} addr={:?}",
                                    &import_task, wn_addr
                                );
                            } else {
                                mgr.notify_import_dispatch_progress();
                            }
                            info!(
                                "handle_task(AddImportFragBuildTask): notified the worker node, import_task={:?} addr={:?}",
                                &import_task, wn_addr
                            )
                        }
                    }
                    None => {
                        if let Err(e) = mgr
                            .meta_store
                            .reclaim_import_task(import_task.task_id, &import_task.owner_addr)
                            .await
                        {
                            error!(
                                "handle_task(AddImportFragBuildTask): failed to reclaim orphaned import task, import_task={:?}, err={:?}",
                                import_task, e
                            );
                        }
                        error!(
                            "handle_task(AddImportFragBuildTask): failed to get worker, import_task={:?} addr={:?}",
                            import_task, wn_addr
                        );
                    }
                }
            }
            _ => {}
        }
        false
    }
}

impl WorkerNodeMgr {
    pub async fn remove_worker_node(&self, addr: &str, reason: &str) {
        info!("remove_worker_node={}, reason={}", addr, reason);
        let removed = {
            let mut g = self.worker_nodes.write().await;
            g.remove(addr)
        };
        if let Some(wn) = removed {
            match self.meta_store.reclaim_import_tasks_by_worker(addr).await {
                Ok(task_ids) => {
                    if !task_ids.is_empty() {
                        info!(
                            "remove_worker_node: reclaimed import tasks, addr={}, task_ids={:?}",
                            addr, task_ids
                        );
                    }
                }
                Err(e) => {
                    warn!(
                        "remove_worker_node: failed to reclaim import tasks, addr={}, err={:?}",
                        addr, e
                    );
                }
            }
            // Notify the worker node to remove all shards.
            let (vs, vc) = wn.destroy().await;
            for header in vs {
                if let Err(e) =
                    self.shard_task_tx
                        .send(ShardTask::NeedReschedule(NeedRescheduleTask {
                            ctx: ShardTaskCtx::new_no_cb(Some("resche_removeworker".to_string())),
                            header,
                            designated_writer: None,
                            reason: RescheduleShardReason::ScaleIn,
                        }))
                {
                    warn!(
                        "remove_worker_node: send NeedRescheduleTask, error={:?}, task={:?}",
                        e, e.0
                    );
                }
            }
            for state in vc {
                info!("remove_worker_node: abort compaction, state={:?}", state);
                if let Err(e) =
                    self.shard_task_tx
                        .send(ShardTask::CompactionAborted(CompactionAbortTask {
                            ctx: ShardTaskCtx::new_no_cb(Some("compaction_abort".to_string())),
                            header: ShardTaskHeader::from(&state),
                            key: state.get_compaction_state_key(),
                        }))
                {
                    warn!(
                        "remove_worker_node: send CompactionAbortTask, error={:?}, task={:?}",
                        e, e.0
                    );
                }
            }
        }
    }

    pub async fn create_worker_node(
        &self,
        meta_store: Arc<dyn MetaStore>,
        addr: String,
        status: Option<WorkerNodeStatus>,
    ) -> Arc<WorkerNodeMeta> {
        match self.worker_nodes.write().await.entry(addr.clone()) {
            BOccupied(o) => o.get().clone(),
            BVacant(v) => {
                info!("create_worker_node: worker node registering {}", addr);
                let w = WorkerNodeMeta::new(
                    addr.clone(),
                    status.unwrap_or_default(),
                    self.health_tx.clone(),
                    meta_store,
                    self.config.load().server.grpc_client.clone(),
                );
                if let Err(e) = self
                    .health_tx
                    .send(HealthCheckerCommands::Add(addr.clone()))
                {
                    warn!(
                        "create_worker_node send HealthCheckerCommands error, e={:?}",
                        e
                    );
                }
                v.insert(w.clone());
                w
            }
        }
    }

    pub async fn has_worker_node(&self, addr: &str) -> bool {
        self.worker_nodes.read().await.contains_key(addr)
    }

    pub async fn has_schedulable_worker_node(&self, addr: &str) -> bool {
        if self.is_worker_blocked(addr).await {
            return false;
        }
        self.has_worker_node(addr).await
    }

    pub async fn get_worker_node(&self, addr: &str) -> Option<Arc<WorkerNodeMeta>> {
        self.worker_nodes.read().await.get(addr).cloned()
    }

    pub async fn collect_worker_nodes(&self) -> Vec<(String, Arc<WorkerNodeMeta>)> {
        let m_w: Vec<(String, Arc<WorkerNodeMeta>)> = self
            .worker_nodes
            .read()
            .await
            .iter()
            .map(|(a, b)| (a.clone(), b.clone()))
            .collect::<Vec<_>>();
        m_w
    }

    pub async fn get_worker_nodes_len(&self) -> usize {
        self.worker_nodes.read().await.len()
    }

    fn schedule_strategy(&self) -> WorkerNodeScheduleStrategy {
        let config = self.config.load();
        WorkerNodeScheduleStrategy::from_config(&config.shard.writer_scheduling_strategy)
    }

    fn select_round_robin_addr(
        &self,
        worker_nodes: &BTreeMap<String, Arc<WorkerNodeMeta>>,
        blocked: &HashSet<String>,
        other_than: Option<&str>,
    ) -> Option<String> {
        let len = worker_nodes.len();
        let start = self
            .current_wn_idx
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
            .wrapping_add(1)
            % len;
        let candidates = Self::collect_schedulable_workers(worker_nodes, blocked, other_than);
        if candidates.is_empty() {
            return None;
        }

        if let Some(candidate) = candidates.iter().find(|candidate| candidate.index >= start) {
            return Some(candidate.addr.clone());
        }

        candidates.first().map(|candidate| candidate.addr.clone())
    }

    fn collect_schedulable_workers(
        worker_nodes: &BTreeMap<String, Arc<WorkerNodeMeta>>,
        blocked: &HashSet<String>,
        other_than: Option<&str>,
    ) -> Vec<WorkerNodeCandidate> {
        let mut preferred = Vec::new();
        let mut other_than_candidate: Option<WorkerNodeCandidate> = None;

        for (index, (addr, worker)) in worker_nodes.iter().enumerate() {
            if blocked.contains(addr) {
                continue;
            }
            if let Some(o) = other_than
                && addr.as_str() == o
            {
                other_than_candidate.get_or_insert(WorkerNodeCandidate {
                    index,
                    addr: addr.clone(),
                    worker: worker.clone(),
                });
                continue;
            }
            preferred.push(WorkerNodeCandidate {
                index,
                addr: addr.clone(),
                worker: worker.clone(),
            });
        }

        if !preferred.is_empty() {
            return preferred;
        }

        if let Some(candidate) = other_than_candidate {
            return vec![candidate];
        }

        vec![]
    }

    async fn select_balanced_count_addr(
        &self,
        candidates: Vec<WorkerNodeCandidate>,
    ) -> Option<String> {
        if candidates.is_empty() {
            return None;
        }

        let mut min_count = usize::MAX;
        let mut min_candidates: Vec<String> = Vec::new();

        for candidate in candidates {
            let count = candidate.worker.shards_len().await;
            match count.cmp(&min_count) {
                std::cmp::Ordering::Less => {
                    min_count = count;
                    min_candidates.clear();
                    min_candidates.push(candidate.addr);
                }
                std::cmp::Ordering::Equal => {
                    min_candidates.push(candidate.addr);
                }
                std::cmp::Ordering::Greater => {}
            }
        }

        if min_candidates.is_empty() {
            return None;
        }

        let idx = self
            .current_wn_idx
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
            % min_candidates.len();
        Some(min_candidates[idx].clone())
    }

    async fn select_balanced_import_count_addr(
        &self,
        candidates: Vec<WorkerNodeCandidate>,
    ) -> ImportWorkerNodeSelection {
        if candidates.is_empty() {
            return ImportWorkerNodeSelection::NoAvailableWorker;
        }

        let mut min_count = usize::MAX;
        let mut min_candidates: Vec<String> = Vec::new();
        let mut pending_blocked_count = 0;

        for candidate in candidates {
            let untracked_count = candidate.worker.untracked_import_tasks_len().await;
            if untracked_count > 0 {
                debug!(
                    "select_balanced_import_count_addr: skip worker with untracked import tasks, addr={}, untracked_count={}",
                    candidate.addr, untracked_count
                );
                continue;
            }
            let pending_count = candidate.worker.pending_import_tasks_len().await;
            if pending_count > 0 {
                pending_blocked_count += 1;
                debug!(
                    "select_balanced_import_count_addr: skip worker with pending import dispatch, addr={}, pending_count={}",
                    candidate.addr, pending_count
                );
                continue;
            }
            let count = candidate.worker.import_tasks_len().await;
            match count.cmp(&min_count) {
                std::cmp::Ordering::Less => {
                    min_count = count;
                    min_candidates.clear();
                    min_candidates.push(candidate.addr);
                }
                std::cmp::Ordering::Equal => {
                    min_candidates.push(candidate.addr);
                }
                std::cmp::Ordering::Greater => {}
            }
        }

        if min_candidates.is_empty() {
            if pending_blocked_count > 0 {
                return ImportWorkerNodeSelection::BlockedByPendingDispatch;
            }
            return ImportWorkerNodeSelection::NoAvailableWorker;
        }

        let idx = self
            .current_wn_idx
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
            % min_candidates.len();
        ImportWorkerNodeSelection::Selected(min_candidates[idx].clone())
    }

    // Select a worker node address
    pub async fn select_worker_node_addr(&self, other_than: Option<String>) -> Option<String> {
        let blocked = self.blocked_worker_addrs.read().await;
        let g = self.worker_nodes.read().await;

        if g.is_empty() {
            return None;
        }

        let other_than = other_than.as_deref();
        match self.schedule_strategy() {
            WorkerNodeScheduleStrategy::RoundRobin => {
                self.select_round_robin_addr(&g, &blocked, other_than)
            }
            WorkerNodeScheduleStrategy::BalancedCount => {
                let candidates = Self::collect_schedulable_workers(&g, &blocked, other_than);
                drop(g);
                drop(blocked);
                self.select_balanced_count_addr(candidates).await
            }
        }
    }

    pub(crate) async fn select_import_worker_node(&self) -> ImportWorkerNodeSelection {
        let blocked = self.blocked_worker_addrs.read().await;
        let g = self.worker_nodes.read().await;

        if g.is_empty() {
            return ImportWorkerNodeSelection::NoAvailableWorker;
        }

        let candidates = Self::collect_schedulable_workers(&g, &blocked, None);
        drop(g);
        drop(blocked);
        self.select_balanced_import_count_addr(candidates).await
    }

    pub async fn select_import_worker_node_addr(&self) -> Option<String> {
        match self.select_import_worker_node().await {
            ImportWorkerNodeSelection::Selected(addr) => Some(addr),
            ImportWorkerNodeSelection::BlockedByPendingDispatch
            | ImportWorkerNodeSelection::NoAvailableWorker => None,
        }
    }

    pub async fn on_import_task_finished(&self, owner_addr: &str, task_id: u64) -> bool {
        if let Some(worker) = self.get_worker_node(owner_addr).await {
            worker.remove_import_task(task_id).await.is_some()
        } else {
            false
        }
    }

    pub async fn report_worker_node_status(
        &self,
        req: WorkerNodeHeartbeatRequest,
    ) -> Result<Vec<WorkerNodeHeartbeatResponse>> {
        let status = if let Some(ref s) = req.status {
            s
        } else {
            warn!(
                "heartbeat: worker node heartbeat status is None, addr={}",
                req.addr
            );
            return Ok(vec![WorkerNodeHeartbeatResponse {
                status: tici_grpc::ErrorCode::UnknownError.into(),
                ..Default::default()
            }]);
        };

        let current_worker = self
            .create_worker_node(
                self.meta_store.clone(),
                req.addr.clone(),
                Some(status.clone()),
            )
            .await;
        let pagination = {
            let config = self.config.load();
            HeartbeatTaskPagination {
                response_size_limit: Self::worker_heartbeat_response_size_limit(config.as_ref()),
            }
        };

        let (tx, rx) = tokio::sync::oneshot::channel();
        current_worker.send_wn_task(WorkerNodeTask::Heartbeat(HeartbeatTask {
            ctx: WorkerNodeTaskCtx::new(req.addr.clone(), tx),
            heartbeat: req,
            shard_task_tx: self.shard_task_tx.clone(),
            pagination,
        }));
        rx.await?
    }

    pub async fn on_msg_append_delta_frag_meta(
        &self,
        request: &AppendDeltaFragMetaRequest,
    ) -> Result<AppendFragMetaResponse> {
        let worker_node = match self
            .get_worker_node(request.worker_node_addr.as_str())
            .await
        {
            Some(s) => s,
            None => {
                info!(
                    "append_delta_frag_meta: failed to get worker_node={:?}, request={:?}",
                    request.worker_node_addr, &request
                );
                return Ok(AppendFragMetaResponse {
                    status: tici_grpc::ErrorCode::WorkerNotFound.into(),
                });
            }
        };

        if let Some(shard_schedule) = worker_node.get_shard_schedule(&request.shard_id).await {
            let (tx, rx) = tokio::sync::oneshot::channel();
            if let Err(e) =
                self.shard_task_tx
                    .send(ShardTask::AppendDeltaFrag(AppendDeltaFragTask {
                        ctx: ShardTaskCtx::new(Some("append_delta".to_string()), tx),
                        header: ShardTaskHeader::from(&shard_schedule),
                        index: shard_schedule.index_cloned(),
                        req: request.clone(),
                    }))
            {
                warn!(
                    "append_delta_frag_meta: send shard task failed, err={:?}",
                    e
                );
                bail!("append_delta_frag_meta: send shard task failed, err={}", e);
            }

            match rx.await? {
                Ok((rsp, update)) => {
                    let (ctx, rx) = WorkerNodeTaskCtx::create(worker_node.addr.clone());
                    if let Some(update) = update {
                        worker_node
                            .send_wn_task(WorkerNodeTask::Update(UpdateTask { ctx, update }));
                        let _ = rx.await?;
                    }
                    Ok(rsp)
                }
                Err(e) => {
                    error!(
                        "append_delta_frag_meta: error, request={:?}, err={:?}",
                        request, e
                    );
                    Ok(AppendFragMetaResponse {
                        status: tici_grpc::ErrorCode::UnknownError.into(),
                    })
                }
            }
        } else {
            info!(
                "append_frag_meta: can't find shard, maybe not assigned to this worker node, request={:?}",
                request
            );
            Ok(AppendFragMetaResponse {
                status: tici_grpc::ErrorCode::ShardNotScheduled.into(),
            })
        }
    }

    pub async fn on_compaction_finish(
        &self,
        request: &FinishCompactFragRequest,
    ) -> Result<FinishCompactFragResponse, anyhow::Error> {
        let compaction_key = CompactionStateKey::try_from(request)?;
        let output_frags = request.output_frags.clone();
        let worker_node = match self.get_worker_node(&compaction_key.get_worker()).await {
            Some(s) => s,
            None => {
                info!(
                    "on_compaction_finish: failed to get worker node, request={:?}",
                    &request
                );
                return Ok(FinishCompactFragResponse {
                    status: tici_grpc::ErrorCode::WorkerNotFound.into(),
                });
            }
        };

        let compaction = if let Some(compaction) = worker_node.get_compaction(&compaction_key).await
        {
            compaction.clone()
        } else {
            // Maybe splitted.
            error!(
                "on_compaction_finish: can't find compaction={:?}",
                &compaction_key
            );
            return Ok(FinishCompactFragResponse {
                status: tici_grpc::ErrorCode::CompactionNotFound.into(),
            });
        };

        let header = compaction.state.get_header();
        let index = compaction.state.index.clone();

        let (tx, rx) = tokio::sync::oneshot::channel();
        if let Err(e) = self
            .shard_task_tx
            .send(ShardTask::CompactionFinish(CompactionFinishTask {
                ctx: ShardTaskCtx::new(None, tx),
                header: header.clone(),
                index,
                compaction_key: compaction_key.clone(),
                compaction_state: compaction.state.clone(),
                output_frags,
            }))
        {
            warn!(
                "on_compaction_finish: send message, shard={:?}, err={:?}",
                header, e
            );
            return Err(anyhow!(
                "on_compaction_finish: send message, shard={:?}, err={}",
                header,
                e
            ));
        }

        let rsp = match rx.await? {
            Ok(rsp) => {
                // `rsp` could be ShardNotFound or Success
                let x = worker_node.remove_compaction(&compaction_key).await;
                if x.is_none() {
                    warn!(
                        "on_compaction_finish: can't remove compaction key in worker node, key={:?}",
                        &compaction_key
                    );
                } else {
                    info!(
                        "on_compaction_finish: removed compaction from worker, compaction_key={:?}, after_len={}",
                        &compaction_key,
                        worker_node.compaction_len().await
                    );
                }
                rsp
            }
            Err(e) => {
                error!(
                    "on_compaction_finish: error, request={:?}, err={:?}",
                    request, e
                );
                FinishCompactFragResponse {
                    status: tici_grpc::ErrorCode::UnknownError.into(),
                }
            }
        };
        Ok(rsp)
    }
}

impl WorkerNodeMgr {
    fn worker_heartbeat_response_size_limit(config: &MetaServiceConfig) -> usize {
        let server_encoding_limit = config.server.grpc.max_encoding_message_size.0 as usize;
        // The remote worker's decoding limit is not directly available on meta, so use
        // the locally configured client decoding limit as a conservative
        // counterpart.
        let local_client_decoding_limit =
            config.server.grpc_client.max_decoding_message_size.0 as usize;
        server_encoding_limit.min(local_client_decoding_limit)
    }
}

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

    use arc_swap::ArcSwap;
    use tici_import::{ImportTask, ImportTaskParams};

    use super::*;
    use crate::meta::store::test_utils::{DummyMetaStore, RecordingMetaStore};

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn select_worker_node_skip_blocked() {
        let (shard_task_tx, _shard_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let (wn_task_tx, wn_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let config = Arc::new(ArcSwap::from_pointee(MetaServiceConfig::default()));
        let store: Arc<dyn MetaStore> = Arc::new(DummyMetaStore::default());
        let mgr = WorkerNodeMgr::new(shard_task_tx, wn_task_tx, wn_task_rx, config, store.clone());

        mgr.create_worker_node(store.clone(), "a".to_string(), None)
            .await;
        mgr.create_worker_node(store.clone(), "b".to_string(), None)
            .await;

        mgr.block_workers(&["a".to_string()]).await;

        for _ in 0..10 {
            assert_eq!(
                mgr.select_worker_node_addr(None).await.as_deref(),
                Some("b")
            );
        }
        assert!(!mgr.has_schedulable_worker_node("a").await);
        assert!(mgr.has_schedulable_worker_node("b").await);

        // Blocked addrs should stay blocked even if the worker registers later.
        mgr.block_workers(&["c".to_string()]).await;
        mgr.create_worker_node(store.clone(), "c".to_string(), None)
            .await;
        for _ in 0..10 {
            assert_eq!(
                mgr.select_worker_node_addr(None).await.as_deref(),
                Some("b")
            );
        }
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn select_worker_node_skip_other_than_when_possible() {
        let (shard_task_tx, _shard_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let (wn_task_tx, wn_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let config = Arc::new(ArcSwap::from_pointee(MetaServiceConfig::default()));
        let store: Arc<dyn MetaStore> = Arc::new(DummyMetaStore::default());
        let mgr = WorkerNodeMgr::new(shard_task_tx, wn_task_tx, wn_task_rx, config, store.clone());

        mgr.create_worker_node(store.clone(), "a".to_string(), None)
            .await;
        mgr.create_worker_node(store.clone(), "b".to_string(), None)
            .await;

        // `a` is a valid worker, but it should be skipped when there are other choices.
        for _ in 0..10 {
            assert_eq!(
                mgr.select_worker_node_addr(Some("a".to_string()))
                    .await
                    .as_deref(),
                Some("b")
            );
        }

        // If `other_than` is the only available worker, return it.
        mgr.block_workers(&["b".to_string()]).await;
        for _ in 0..10 {
            assert_eq!(
                mgr.select_worker_node_addr(Some("a".to_string()))
                    .await
                    .as_deref(),
                Some("a")
            );
        }
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn select_worker_node_all_blocked_returns_empty() {
        let (shard_task_tx, _shard_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let (wn_task_tx, wn_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let config = Arc::new(ArcSwap::from_pointee(MetaServiceConfig::default()));
        let store: Arc<dyn MetaStore> = Arc::new(DummyMetaStore::default());
        let mgr = WorkerNodeMgr::new(shard_task_tx, wn_task_tx, wn_task_rx, config, store.clone());

        mgr.create_worker_node(store.clone(), "a".to_string(), None)
            .await;
        mgr.block_workers(&["a".to_string()]).await;

        assert_eq!(mgr.select_worker_node_addr(None).await.as_deref(), None);

        // Even if the blocked worker is removed from the active set, the selection
        // should still return empty (unassigned).
        mgr.remove_worker_node("a", "test").await;
        assert_eq!(mgr.select_worker_node_addr(None).await.as_deref(), None);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 8)]
    async fn select_worker_node_high_concurrency_is_nearly_even() {
        let (shard_task_tx, _shard_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let (wn_task_tx, wn_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let config = Arc::new(ArcSwap::from_pointee(MetaServiceConfig::default()));
        let store: Arc<dyn MetaStore> = Arc::new(DummyMetaStore::default());
        let mgr = WorkerNodeMgr::new(shard_task_tx, wn_task_tx, wn_task_rx, config, store.clone());

        let worker_count = 32_usize;
        let worker_addrs: Vec<String> = (0..worker_count).map(|i| format!("w{i:02}")).collect();
        for addr in &worker_addrs {
            mgr.create_worker_node(store.clone(), addr.clone(), None)
                .await;
        }

        let iterations_per_task = 128_usize;
        let task_count = worker_count;
        let request_count = task_count * iterations_per_task;
        let start_barrier = Arc::new(tokio::sync::Barrier::new(task_count + 1));
        let mut handles = Vec::with_capacity(task_count);

        for _ in 0..task_count {
            let mgr = mgr.clone();
            let start_barrier = start_barrier.clone();
            handles.push(tokio::spawn(async move {
                start_barrier.wait().await;
                let mut selected = Vec::with_capacity(iterations_per_task);
                for _ in 0..iterations_per_task {
                    selected.push(mgr.select_worker_node_addr(None).await);
                }
                selected
            }));
        }

        start_barrier.wait().await;

        let mut counts = std::collections::BTreeMap::new();
        for addr in &worker_addrs {
            counts.insert(addr.clone(), 0_usize);
        }

        for handle in handles {
            let selected_addrs = handle
                .await
                .expect("select_worker_node_addr task should not panic");
            for selected in selected_addrs {
                let selected = selected.expect("worker should be selected");
                *counts
                    .get_mut(&selected)
                    .expect("selected worker must be known") += 1;
            }
        }
        let total_selected: usize = counts.values().sum();
        assert_eq!(total_selected, request_count);

        let max_count = *counts
            .values()
            .max()
            .expect("count map should contain workers");
        let min_count = *counts
            .values()
            .min()
            .expect("count map should contain workers");
        assert!(
            max_count - min_count <= 1,
            "expected near-even distribution under high concurrency, counts={:?}",
            counts
        );
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn remove_worker_node_reclaims_owned_import_tasks() {
        let (shard_task_tx, _shard_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let (wn_task_tx, wn_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let config = Arc::new(ArcSwap::from_pointee(MetaServiceConfig::default()));
        let store = Arc::new(RecordingMetaStore::default());
        let mgr = WorkerNodeMgr::new(shard_task_tx, wn_task_tx, wn_task_rx, config, store.clone());

        mgr.create_worker_node(store.clone(), "a".to_string(), None)
            .await;
        mgr.remove_worker_node("a", "timeout").await;

        let reclaimed = match store.reclaimed_import_workers.lock() {
            Ok(guard) => guard.clone(),
            Err(poisoned) => poisoned.into_inner().clone(),
        };
        assert_eq!(reclaimed, vec!["a".to_string()]);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn stale_import_dispatch_is_skipped_without_reinserting_task() {
        let (shard_task_tx, _shard_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let (wn_task_tx, wn_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let config = Arc::new(ArcSwap::from_pointee(MetaServiceConfig::default()));
        let store = Arc::new(RecordingMetaStore::default());
        let mgr = WorkerNodeMgr::new(shard_task_tx, wn_task_tx, wn_task_rx, config, store.clone());

        mgr.create_worker_node(store, "worker-a".to_string(), None)
            .await;

        let task = ImportTask {
            task_id: 42,
            job_id: 7,
            owner_addr: "worker-a".to_string(),
            owner_seq: 3,
            parameters: ImportTaskParams::new(
                "a".to_string(),
                "z".to_string(),
                "mock://import/stale-dispatch".to_string(),
            ),
        };

        assert!(
            !WorkerNodeMgr::handle_task(
                mgr.as_ref(),
                WorkerNodeTask::AddImportFragBuildTask(task.clone()),
            )
            .await
        );

        let worker = mgr.get_worker_node(&task.owner_addr).await.unwrap();
        assert!(!worker.has_import_task(task.task_id).await);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn running_import_dispatch_retry_is_skipped_without_notify_retry() {
        let (shard_task_tx, _shard_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let (retry_tx, mut retry_rx) = tokio::sync::mpsc::unbounded_channel();
        let (_runner_tx, runner_rx) = tokio::sync::mpsc::unbounded_channel();
        let config = Arc::new(ArcSwap::from_pointee(MetaServiceConfig::default()));
        let store = Arc::new(RecordingMetaStore::default());
        let store_dyn: Arc<dyn MetaStore> = store.clone();
        let mgr = WorkerNodeMgr::new(
            shard_task_tx,
            retry_tx,
            runner_rx,
            config,
            store_dyn.clone(),
        );

        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
            .await
            .expect("test listener should bind");
        let worker_addr = format!(
            "http://{}",
            listener
                .local_addr()
                .expect("test listener should have a local addr")
        );
        let accept_task = tokio::spawn(async move {
            if let Ok((_socket, _peer)) = listener.accept().await {
                tokio::time::sleep(Duration::from_secs(5)).await;
            }
        });

        let task = ImportTask {
            task_id: 43,
            job_id: 7,
            owner_addr: worker_addr,
            owner_seq: 4,
            parameters: ImportTaskParams::new(
                "a".to_string(),
                "z".to_string(),
                "mock://import/running-retry".to_string(),
            ),
        };
        let worker = mgr
            .create_worker_node(store_dyn, task.owner_addr.clone(), None)
            .await;
        assert!(worker.reserve_import_task(task.clone()).await);
        assert!(
            worker
                .mark_import_task_running(task.task_id, &task.owner_addr, task.owner_seq)
                .await
        );

        assert!(
            !worker
                .is_pending_import_task_assignment(task.task_id, &task.owner_addr, task.owner_seq)
                .await
        );
        let handled = tokio::time::timeout(
            Duration::from_millis(100),
            WorkerNodeMgr::handle_task(
                mgr.as_ref(),
                WorkerNodeTask::AddImportFragBuildTask(task.clone()),
            ),
        )
        .await;
        match handled {
            Ok(stop) => assert!(!stop),
            Err(_) => panic!("running import retry should be skipped before grpc notify"),
        }
        assert!(
            retry_rx.try_recv().is_err(),
            "running import retry should not enqueue another retry"
        );
        assert!(worker.has_import_task(task.task_id).await);
        assert_eq!(worker.pending_import_tasks_len().await, 0);
        accept_task.abort();
    }

    #[test]
    fn worker_heartbeat_response_size_limit_uses_smaller_encode_decode_limit() {
        let mut config = MetaServiceConfig::default();
        config.server.grpc.max_encoding_message_size = tici_system::readable::ReadableSize::mb(32);
        config.server.grpc_client.max_decoding_message_size =
            tici_system::readable::ReadableSize::mb(16);

        assert_eq!(
            WorkerNodeMgr::worker_heartbeat_response_size_limit(&config),
            tici_system::readable::ReadableSize::mb(16).0 as usize
        );
    }
}