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

use anyhow::{Result, anyhow};
use arc_swap::ArcSwap;
use log::{info, warn};
use tokio::sync::mpsc::UnboundedSender;

use crate::{
    RangeBound, ShardInfo,
    config::{MetaServiceConfig, ShardConfig},
    meta::{
        metrics::*,
        util::BgRunner,
        writer::{
            FTSIndex, MergeShardReason, MergeTask, NeedRescheduleTask, RescheduleShardReason,
            ShardStats, ShardTaskCtx, ShardTaskHeader, shard_task::ShardTask,
        },
    },
};

// Collects shard stats, report to metrics, and schedule merge if it is needed.
pub struct ShardStatsPatrol {
    tx: UnboundedSender<ShardStatsCollectorTask>,
    // BgRunner should be held, otherwise the inner loop is stopped.
    _prom_runner: BgRunner,
}

pub struct AddShardStatsTask {
    shard_id: u64,
    index: Arc<FTSIndex>,
    start: RangeBound,
    stats: ShardStats,
}

pub enum ShardStatsCollectorTask {
    Add(AddShardStatsTask),
    Remove(u64),
    Debug((u64, tokio::sync::oneshot::Sender<ShardPatrolResult>)),
    GetConfig(tokio::sync::oneshot::Sender<Arc<MetaServiceConfig>>),
}

#[derive(Default, Debug)]
pub struct ShardPatrolResult {
    pub sched_try_merge_count: usize,
    pub sched_no_heartbeat_count: usize,
}

impl ShardStatsPatrol {
    #[cfg(any(test, feature = "testexport"))]
    pub fn get_tx(&self) -> UnboundedSender<ShardStatsCollectorTask> {
        self.tx.clone()
    }

    pub async fn debug_run_one(&self, tick: u64) -> Result<ShardPatrolResult> {
        let (tx, rx) = tokio::sync::oneshot::channel();
        if let Err(e) = self.tx.send(ShardStatsCollectorTask::Debug((tick, tx))) {
            warn!("send error={:?}", e);
        }
        Ok(rx.await?)
    }

    pub fn insert(
        &self,
        shard_id: u64,
        index: Arc<FTSIndex>,
        start: RangeBound,
        st: ShardStats,
    ) -> Result<()> {
        let task = AddShardStatsTask {
            shard_id,
            index,
            start,
            stats: st,
        };
        if let Err(e) = self.tx.send(ShardStatsCollectorTask::Add(task)) {
            return Err(anyhow!("ShardStatsPatrol::insert err={}", e));
        }
        Ok(())
    }

    pub fn restore(&self, index: Arc<FTSIndex>, shard: &ShardInfo) -> Result<()> {
        self.insert(
            shard.shard_id,
            index,
            shard.range.sortable_start(),
            shard.stats.clone(),
        )
    }

    pub fn remove(&self, shard_id: u64) -> Result<()> {
        if let Err(e) = self.tx.send(ShardStatsCollectorTask::Remove(shard_id)) {
            return Err(anyhow!("ShardStatsPatrol::remove err={}", e));
        }
        Ok(())
    }

    async fn run_one(
        shard_stats: &HashMap<u64, (Arc<FTSIndex>, RangeBound, ShardStats)>,
        shard_config: &ShardConfig,
        shard_task_tx: &mut UnboundedSender<ShardTask>,
        topology_change_enabled: &AtomicBool,
        tick: u64,
    ) -> ShardPatrolResult {
        let mut potential_merge_shard = HashMap::new();
        let mut no_heartbeat_shard = vec![];
        let stats = shard_stats
            .iter()
            .map(|x| (*x.0, x.1.clone()))
            .collect::<Vec<_>>();
        // Use a dedicated loop for prom metrics
        for (_, s) in stats.iter() {
            let size = s.2.get_size();
            SHARD_SIZE_HISTOGRAM.observe(size as f64);
            SHARD_FRAGMENTS_HISTOGRAM.observe(s.2.get_fragment_count() as f64);
        }

        let sched_no_heartbeat_count = if tick % shard_config.check_dead_shard_tick == 0 {
            for (sid, (index, start_key, stats)) in stats.iter() {
                if shard_config.is_abnormal_shard_no_heartbeat(
                    stats
                        .last_heartbeat_ts
                        .load(std::sync::atomic::Ordering::Relaxed),
                ) {
                    // Merge can also reassign worker, so just do merge.
                    no_heartbeat_shard.push((index.clone(), start_key.clone(), sid));
                }
            }

            let no_heartbeat_count = no_heartbeat_shard.len();
            if no_heartbeat_count != 0 {
                let no_heartbeat_shard = no_heartbeat_shard
                    .drain(..)
                    .map(|(index, start_key, shard_id)| ShardTaskHeader {
                        index_key: index.index_key(),
                        start_key,
                        shard_id: *shard_id,
                    })
                    .collect::<Vec<_>>();
                for header in no_heartbeat_shard {
                    if let Err(e) =
                        shard_task_tx.send(ShardTask::NeedReschedule(NeedRescheduleTask {
                            ctx: ShardTaskCtx::new_no_cb(Some("resche_patrol".to_string())),
                            header,
                            designated_writer: None,
                            reason: RescheduleShardReason::Dead,
                        }))
                    {
                        warn!("can not send NeedReschedule, e={:?}", e);
                    }
                }
            }
            no_heartbeat_count
        } else {
            0
        };

        let sched_try_merge_count = if topology_change_enabled
            .load(std::sync::atomic::Ordering::Relaxed)
            && shard_config.enable_merge
            && stats.len() > 1
            && tick % shard_config.merge_strategy.merge_check_tick == 0
        {
            for (sid, (index, start_key, stats)) in stats.iter() {
                let size = stats.get_size();
                if size < shard_config.get_merge_size()
                    && !index.is_topology_frozen()
                    && stats.get_write_qps().await < shard_config.merge_strategy.write_qps
                {
                    potential_merge_shard.insert(sid, (index.clone(), start_key.clone()));
                }
            }
            let try_merge_count = potential_merge_shard.len();
            if try_merge_count != 0 {
                let potential_merge_shard = potential_merge_shard
                    .drain()
                    .map(|(k, v)| MergeTask {
                        ctx: ShardTaskCtx::new_no_cb(Some("merge".to_string())),
                        index: v.0,
                        start_key: v.1,
                        shard_id: *k,
                        reason: MergeShardReason::Auto,
                    })
                    .collect::<Vec<_>>();
                // TODO(P2) introduce cold-down to scheduled and failed shards.
                for shard in potential_merge_shard {
                    if let Err(e) = shard_task_tx.send(ShardTask::TryMerge(shard)) {
                        warn!("send merge task error, e={:?}", e);
                    }
                }
            }
            try_merge_count
        } else {
            0
        };

        SHARD_COUNT_GAUGE_VEC
            .with_label_values(&["dead"])
            .set(sched_no_heartbeat_count as i64);

        if sched_try_merge_count != 0 || sched_no_heartbeat_count != 0 {
            info!(
                "ShardStatsPatrol: sched_try_merge_count={}, sched_no_heartbeat_count={}",
                sched_try_merge_count, sched_no_heartbeat_count
            );
        }

        ShardPatrolResult {
            sched_try_merge_count,
            sched_no_heartbeat_count,
        }
    }

    pub fn new(
        config: Arc<ArcSwap<MetaServiceConfig>>,
        shard_task_tx: UnboundedSender<ShardTask>,
        topology_change_enabled: Arc<AtomicBool>,
    ) -> Self {
        let mut shard_task_tx = shard_task_tx;
        let duration = config.load().shard.patrol_interval.0;
        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();

        let prom_runner = BgRunner::spawn(move |mut rf| async move {
            let mut shard_stats: HashMap<u64, (Arc<FTSIndex>, RangeBound, ShardStats)> =
                Default::default();
            use tokio::time::interval;
            let mut interval = interval(duration);
            let mut tick = 0;
            loop {
                tokio::select! {
                    _ = interval.tick() => {
                        Self::run_one(
                            &shard_stats,
                            &config.load().shard,
                            &mut shard_task_tx,
                            topology_change_enabled.as_ref(),
                            tick,
                        ).await;
                        tick += 1;
                    },
                    msg = rx.recv() => {
                        match msg {
                            Some(task) => {
                                match task {
                                    ShardStatsCollectorTask::Add(AddShardStatsTask {
                                        shard_id,
                                        index,
                                        start,
                                        stats,
                                    }) => {
                                        shard_stats.insert(shard_id, (index, start, stats));
                                    }
                                    ShardStatsCollectorTask::Remove(shard_id) => {
                                        shard_stats.remove(&shard_id);
                                    }
                                    ShardStatsCollectorTask::Debug((tick, k)) => {
                                        let r = ShardStatsPatrol::run_one(
                                            &shard_stats,
                                            &config.load().shard,
                                            &mut shard_task_tx,
                                            topology_change_enabled.as_ref(),
                                            tick,
                                        )
                                        .await;
                                        if let Err(e) = k.send(r) {
                                            warn!("send error, e={:?}", e);
                                        }
                                    }
                                    ShardStatsCollectorTask::GetConfig(tx) => {
                                        if let Err(e) = tx.send(config.load_full()) {
                                            warn!("send error, e={:?}", e);
                                        }
                                    }
                                }
                            }
                            None => {
                                break;
                            }
                        }
                    }
                    _ = &mut rf => {
                        return;
                    }
                }
            }
        });
        Self {
            tx,
            _prom_runner: prom_runner,
        }
    }
}