use std::{collections::HashMap, sync::Arc};

use anyhow::{Context, Result};
use log::info;
use tokio::sync::mpsc::UnboundedSender;

use crate::{
    ShardInfo, ShardInfoAppendUpdate,
    meta::{
        metrics::*,
        writer::{
            FTSIndex, MaybeSplitOrCompactReason, MaybeSplitOrCompactTask, ModifyShardTaskAdd,
            ModifyShardTaskRemove, ShardEventIndexInfo, ShardTask, ShardTaskCtx, ShardTaskHeader,
            ShardTaskManager, WorkerNodeMgr, WorkerNodeTask,
        },
    },
};

#[derive(Default)]
pub struct ModifyCollector {
    workers: HashMap<String, (Vec<ModifyShardTaskAdd>, Vec<ModifyShardTaskRemove>)>,
}

impl ModifyCollector {
    pub fn set_add_shard(&mut self, shard: &ShardInfo, index: Arc<FTSIndex>, notified: bool) {
        let entry = self
            .workers
            .entry(shard.writer.clone())
            .or_insert_with(|| (Vec::new(), Vec::new()));
        entry.0.push(ModifyShardTaskAdd {
            header: shard.get_header(),
            progress: shard.progress.clone(),
            upload_location: shard.get_upload_location(),
            shard_stats: shard.stats.clone(),
            index,
            notified,
        });
    }

    pub fn set_remove_shard(&mut self, addr: String, shard_id: u64) {
        let entry = self
            .workers
            .entry(addr)
            .or_insert_with(|| (Vec::new(), Vec::new()));
        entry.1.push(ModifyShardTaskRemove { shard_id });
    }

    pub async fn send_all(self, wn_task_tx: UnboundedSender<WorkerNodeTask>) -> Result<()> {
        for (addr, (to_add, to_remove)) in self.workers.into_iter() {
            WorkerNodeMgr::send_modify_shard(wn_task_tx.clone(), addr.clone(), to_add, to_remove)
                .await
                .with_context(|| format!("addr={}", addr))?;
        }
        Ok(())
    }

    pub fn modifies_count(&self) -> usize {
        let mut count = 0;
        for (_, v) in self.workers.iter() {
            count += v.0.len();
            count += v.1.len();
        }
        count
    }
}

impl ShardTaskManager {
    pub(super) async fn helper_after_append(
        &self,
        index: Arc<FTSIndex>,
        shard: &ShardInfo,
        update_append: ShardInfoAppendUpdate,
        size_added: u64,
    ) {
        if shard.get_total_estimated_bytes() > self.config.load().shard.max_size.0 {
            index
                .stats
                .big_shard_after_write
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        }
        if let Err(e) =
            self.shard_task_tx
                .send(ShardTask::MaybeSplitOrCompact(MaybeSplitOrCompactTask {
                    ctx: ShardTaskCtx::new_no_cb(Some("after_append_suc".to_string())),
                    header: ShardTaskHeader::from(shard),
                    reason: MaybeSplitOrCompactReason::AfterAppend,
                }))
        {
            info!(
                "append_frag_meta: register MaybeSplitOrCompact failure, error={:?}, shard={}",
                e,
                shard.to_long_ident()
            );
        }
        self.event_pubsub
            .on_shard_append(ShardEventIndexInfo::from(index.as_ref()), update_append);
        shard.stats.on_append(shard, size_added).await;
    }

    pub(super) async fn helper_after_compaction(
        &self,
        success: bool,
        shard: &ShardInfo,
        index: Arc<FTSIndex>,
    ) {
        if success {
            shard.stats.observe_compaction_finish();
            index.stats.observe_compaction_finish();
            shard.stats.on_write(shard).await;
            SHARD_EVENT_COUNTER_VEC
                .with_label_values(&["compaction_finish"])
                .inc();
        } else {
            shard.stats.observe_compaction_abort();
            index.stats.observe_compaction_abort();
            SHARD_EVENT_COUNTER_VEC
                .with_label_values(&["compaction_abort"])
                .inc();
        }
        self.event_pubsub
            .on_shard_update(ShardEventIndexInfo::from(index.as_ref()), shard.clone());
    }

    pub(super) async fn helper_after_add_shard(
        &self,
        index: Arc<FTSIndex>,
        shard: &ShardInfo,
    ) -> Result<()> {
        let mut collector = ModifyCollector::default();
        collector.set_add_shard(shard, index.clone(), true);
        collector.send_all(self.wn_task_tx.clone()).await?;
        self.shard_patrol.insert(
            shard.shard_id,
            index.clone(),
            shard.range.sortable_start(),
            shard.stats.clone(),
        )?;
        self.stats
            .total_shards_count
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        if shard.writer.is_empty() {
            self.add_unscheduled_shards(1);
        }
        self.event_pubsub
            .on_shard_add(ShardEventIndexInfo::from(index.as_ref()), shard.clone());
        self.on_index_change();
        Ok(())
    }

    pub(super) async fn helper_after_split(
        &self,
        index: Arc<FTSIndex>,
        new_shards: Vec<ShardInfo>,
        original_shard: &ShardInfo,
    ) -> Result<()> {
        index.stats.observe_split_finish();
        let mut collector = ModifyCollector::default();
        collector.set_remove_shard(original_shard.writer.clone(), original_shard.shard_id);
        self.shard_patrol.remove(original_shard.shard_id)?;

        for shard in new_shards.iter() {
            collector.set_add_shard(shard, index.clone(), true);
            self.shard_patrol.insert(
                shard.shard_id,
                index.clone(),
                shard.range.sortable_start(),
                shard.stats.clone(),
            )?;
        }

        if !new_shards.is_empty() {
            self.stats
                .total_shards_count
                .fetch_add(new_shards.len() - 1, std::sync::atomic::Ordering::Relaxed);
        }
        self.replace_unscheduled_shards(
            usize::from(original_shard.writer.is_empty()),
            new_shards
                .iter()
                .filter(|shard| shard.writer.is_empty())
                .count(),
        );
        SHARD_EVENT_COUNTER_VEC.with_label_values(&["split"]).inc();
        // Notify shard split asynchronously (non-blocking)
        self.event_pubsub.on_shard_split(
            ShardEventIndexInfo::from(index.as_ref()),
            original_shard.clone(),
            new_shards.clone(),
        );
        self.on_index_change();
        collector.send_all(self.wn_task_tx.clone()).await?;
        Ok(())
    }

    pub(super) async fn helper_after_merge(
        &self,
        input: &[&ShardInfo],
        new_shard: &ShardInfo,
        index: Arc<FTSIndex>,
        notified_shards: Vec<ShardInfo>,
    ) -> Result<()> {
        let mut collector = ModifyCollector::default();
        for shard in input.iter() {
            collector.set_remove_shard(shard.writer.clone(), shard.shard_id);
            self.shard_patrol.remove(shard.shard_id)?;
        }
        collector.set_add_shard(new_shard, index.clone(), true);
        self.replace_unscheduled_shards(
            input.iter().filter(|shard| shard.writer.is_empty()).count(),
            usize::from(new_shard.writer.is_empty()),
        );
        collector.send_all(self.wn_task_tx.clone()).await?;

        new_shard.stats.on_write(new_shard).await;
        self.shard_patrol.insert(
            new_shard.shard_id,
            index.clone(),
            new_shard.range.sortable_start(),
            new_shard.stats.clone(),
        )?;
        if !input.is_empty() {
            self.stats
                .total_shards_count
                .fetch_sub(input.len() - 1, std::sync::atomic::Ordering::Relaxed);
        }
        self.event_pubsub.on_shard_merge(
            ShardEventIndexInfo::from(index.as_ref()),
            notified_shards,
            new_shard.clone(),
        );
        index
            .stats
            .merge_finish
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        SHARD_EVENT_COUNTER_VEC
            .with_label_values(&["merge_succeed"])
            .inc();
        self.on_index_change();
        Ok(())
    }

    pub(super) async fn helper_after_move(
        &self,
        shard_info: &ShardInfo,
        prev_worker_node: String,
        index: Arc<FTSIndex>,
    ) -> Result<()> {
        let mut collector = ModifyCollector::default();
        let prev_writer_was_empty = prev_worker_node.is_empty();
        if !prev_worker_node.is_empty() && prev_worker_node != shard_info.writer {
            // If there is only one worker node, then the shard could be failed over again
            // to this worker node. So we can't remove this node.
            collector.set_remove_shard(prev_worker_node, shard_info.shard_id);
        }
        // If we fail over to the same node, we have to send the add shard request
        // again, because the worker node may failover, thus having no shards
        // scheduled.
        if !shard_info.writer.is_empty() {
            collector.set_add_shard(shard_info, index.clone(), true);
        }

        self.shard_patrol.insert(
            shard_info.shard_id,
            index.clone(),
            shard_info.range.sortable_start(),
            shard_info.stats.clone(),
        )?;

        if shard_info.writer.is_empty() {
            SHARD_EVENT_COUNTER_VEC
                .with_label_values(&["resche_failed"])
                .inc();
        } else {
            SHARD_EVENT_COUNTER_VEC
                .with_label_values(&["resche_ok"])
                .inc();
        }
        self.event_pubsub.on_shard_update(
            ShardEventIndexInfo::from(index.as_ref()),
            shard_info.clone(),
        );
        self.replace_unscheduled_shards(
            usize::from(prev_writer_was_empty),
            usize::from(shard_info.writer.is_empty()),
        );
        collector.send_all(self.wn_task_tx.clone()).await?;
        Ok(())
    }

    pub(super) async fn helper_after_remove_shard(
        &self,
        shards: Vec<(String, u64)>,
        index: Arc<FTSIndex>,
    ) -> Result<()> {
        let mut collector = ModifyCollector::default();
        let len = shards.len();
        let unscheduled_len = shards
            .iter()
            .filter(|(writer, _)| writer.is_empty())
            .count();
        let mut shard_ids = vec![];
        for (writer, shard_id) in shards.into_iter() {
            collector.set_remove_shard(writer, shard_id);
            self.shard_patrol.remove(shard_id)?;
            shard_ids.push(shard_id);
        }
        self.event_pubsub
            .on_shard_remove(ShardEventIndexInfo::from(index.as_ref()), shard_ids);
        if let Err(e) = collector.send_all(self.wn_task_tx.clone()).await {
            info!("helper_after_remove_shard error={:?}", e);
        }
        self.stats
            .total_shards_count
            .fetch_sub(len, std::sync::atomic::Ordering::Relaxed);
        self.remove_unscheduled_shards(unscheduled_len);
        let index_count = self.indexes_count().await - 1;
        INDEX_COUNT_GAUGE_VEC
            .with_label_values(&["total"])
            .set(index_count as i64);
        self.on_index_change();
        Ok(())
    }
}