use std::{
    collections::{
        HashMap,
        hash_map::Entry::{Occupied as HOccupied, Vacant as HVacant},
    },
    sync::{
        Arc,
        atomic::{AtomicBool, AtomicUsize},
    },
};

use anyhow::Result;
use arc_swap::ArcSwap;
use log::{error, info, warn};
use storage_backend::storage::BackendStorage;
use tici_common::index::IndexKey;
use tici_import::ImportJobMeta;
use tokio::sync::{
    RwLock,
    mpsc::{UnboundedReceiver, UnboundedSender},
};

use crate::{
    ShardInfo,
    config::MetaServiceConfig,
    meta::{
        MetaStore,
        metrics::*,
        writer::{
            FTSIndex, GcWorker, ModifyCollector, NeedRescheduleTask, RescheduleShardReason,
            ShardEventIndexInfo, ShardEventPublishMgr, ShardStatsPatrol, ShardTaskCtx,
            ShardTaskHeader, ShardTaskRunner, WorkerNodeMgr, WorkerNodeTask, shard_task::ShardTask,
            shard_task_runner::*,
        },
    },
};

#[derive(Default, Debug)]
pub struct ShardsStats {
    pub(super) total_shards_count: AtomicUsize,
    pub(super) unscheduled_shards_count: AtomicUsize,
}

// NOTE: The methods in `ShardTaskManager` are supposed to provide atomic
// sematics for basic operations like split / merge / reschedule. DO NOT add
// logics which are not necessary to these operations. You can add them before
// scheduling a `ShardTadk`, or in handler methods such as `handle_split` and
// `handle_merge`.
pub struct ShardTaskManager {
    pub(super) meta_store: Arc<dyn MetaStore>,
    pub(super) backend_storage: Arc<dyn BackendStorage>,
    // (table_id, index_id) -> FTSIndex
    #[allow(clippy::type_complexity)]
    pub(super) indices: Arc<RwLock<HashMap<IndexKey, Arc<FTSIndex>>>>,
    pub(super) worker_node_mgr: Arc<WorkerNodeMgr>,
    pub(super) wn_task_tx: UnboundedSender<WorkerNodeTask>,
    pub config: Arc<ArcSwap<MetaServiceConfig>>,
    pub(super) shard_task_tx: UnboundedSender<ShardTask>,
    pub(super) stats: Arc<ShardsStats>,
    pub(super) stopped: AtomicBool,
    pub(super) topology_change_enabled: Arc<AtomicBool>,
    pub event_pubsub: ShardEventPublishMgr,
    pub(super) shard_patrol: ShardStatsPatrol,
    pub gc_worker: GcWorker,
    _global_runner: ShardTaskRunner,
    pub(super) local_scheduler_runners: Vec<ShardTaskRunner>,
}

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

impl ShardTaskManager {
    pub fn new(
        meta_store: Arc<dyn MetaStore>,
        worker_node_mgr: Arc<WorkerNodeMgr>,
        config_arc: Arc<ArcSwap<MetaServiceConfig>>,
        shard_task_tx: UnboundedSender<ShardTask>,
        shard_task_rx: UnboundedReceiver<ShardTask>,
        backend_storage: Arc<dyn BackendStorage>,
    ) -> Arc<Self> {
        let indices = Arc::new(RwLock::new(HashMap::new()));
        let shard_task_tx_clone = shard_task_tx.clone();
        let mut local_scheduler_runners = vec![];
        let mut local_scheduler_runners_inits = vec![];
        let mut local_scheduler_txs = vec![];
        for i in 0..config_arc.load().shard.shard_scheduler_runner_cnt {
            let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
            let (local_runner, local_stx) =
                ShardTaskRunner::new(rx, LocalRunnerRouter::Receiver(i));
            local_scheduler_runners.push(local_runner);
            local_scheduler_runners_inits.push(local_stx);
            local_scheduler_txs.push(tx);
        }
        let (runner, stx) = ShardTaskRunner::new(
            shard_task_rx,
            LocalRunnerRouter::Sender(local_scheduler_txs),
        );
        let router = runner.router_cloned();
        let wn_task_tx = worker_node_mgr.get_tx();
        let topology_change_enabled = Arc::new(AtomicBool::new(false));
        let shard_patrol = ShardStatsPatrol::new(
            config_arc.clone(),
            shard_task_tx_clone,
            topology_change_enabled.clone(),
        );
        let mgr = Arc::new_cyclic(move |weak| {
            let gc_backend_storage = backend_storage.clone();
            let gc_worker = GcWorker::new(
                config_arc.clone(),
                meta_store.clone(),
                gc_backend_storage,
                weak.clone(),
            );
            Self {
                meta_store,
                backend_storage,
                indices,
                worker_node_mgr,
                wn_task_tx,
                config: config_arc,
                shard_task_tx,
                stats: Arc::new(ShardsStats::default()),
                stopped: AtomicBool::new(false),
                topology_change_enabled,
                event_pubsub: ShardEventPublishMgr::default(),
                shard_patrol,
                local_scheduler_runners,
                _global_runner: runner,
                gc_worker,
            }
        });
        let runner_ctx = ShardTaskRunnerCtx::new(mgr.clone(), router);
        if let Err(e) = stx.send(runner_ctx) {
            warn!("error send runner_ctx, e={:?}", e);
        }
        for (stx, runner) in local_scheduler_runners_inits
            .into_iter()
            .zip(mgr.local_scheduler_runners.iter())
        {
            if let Err(e) = stx.send(ShardTaskRunnerCtx::new(mgr.clone(), runner.router_cloned())) {
                warn!(
                    "error send local runner, l={}, e={:?}",
                    runner.router_cloned().ident("err"),
                    e
                );
            }
        }
        mgr
    }

    pub fn stop(&self) {
        info!("shutdown shard scheduler");
        if let Err(e) = self.shard_task_tx.send(ShardTask::Hup) {
            warn!("shutdown shard schedule error, err={:?}", e);
        }
    }

    pub(super) fn add_unscheduled_shards(&self, count: usize) {
        if count == 0 {
            return;
        }
        let current = self
            .stats
            .unscheduled_shards_count
            .fetch_add(count, std::sync::atomic::Ordering::Relaxed)
            + count;
        UNSCHEDULED_SHARDS_GAUGE.set(current as i64);
    }

    pub(super) fn remove_unscheduled_shards(&self, count: usize) {
        if count == 0 {
            return;
        }

        let mut current = self
            .stats
            .unscheduled_shards_count
            .load(std::sync::atomic::Ordering::Relaxed);
        loop {
            let next = current.saturating_sub(count);
            match self.stats.unscheduled_shards_count.compare_exchange_weak(
                current,
                next,
                std::sync::atomic::Ordering::Relaxed,
                std::sync::atomic::Ordering::Relaxed,
            ) {
                Ok(_) => {
                    UNSCHEDULED_SHARDS_GAUGE.set(next as i64);
                    return;
                }
                Err(actual) => current = actual,
            }
        }
    }

    pub(super) fn set_unscheduled_shards(&self, count: usize) {
        self.stats
            .unscheduled_shards_count
            .store(count, std::sync::atomic::Ordering::Relaxed);
        UNSCHEDULED_SHARDS_GAUGE.set(count as i64);
    }

    pub(super) fn replace_unscheduled_shards(&self, removed_count: usize, added_count: usize) {
        match added_count.cmp(&removed_count) {
            std::cmp::Ordering::Greater => {
                self.add_unscheduled_shards(added_count - removed_count);
            }
            std::cmp::Ordering::Less => {
                self.remove_unscheduled_shards(removed_count - added_count);
            }
            std::cmp::Ordering::Equal => {}
        }
    }
}

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

    pub fn get_shard_task_tx(&self) -> UnboundedSender<ShardTask> {
        self.shard_task_tx.clone()
    }

    pub fn get_gc_worker(&self) -> &GcWorker {
        &self.gc_worker
    }

    pub fn get_shard_patrol(&self) -> &ShardStatsPatrol {
        &self.shard_patrol
    }

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

    pub async fn get_all_table_ids(&self) -> Vec<(u32, i64)> {
        let guard = self.indices.read().await;
        let mut set = std::collections::HashSet::new();
        for index in guard.values() {
            set.insert((index.keyspace_id(), index.table_id()));
        }
        let mut ids: Vec<(u32, i64)> = set.into_iter().collect();
        ids.sort_unstable();
        ids
    }

    async fn collect_import_index_keys(
        &self,
        keyspace_id: u32,
        table_id: i64,
        index_ids: &[i64],
    ) -> Vec<IndexKey> {
        let mut index_keys = Vec::new();
        for index_id in index_ids {
            let mut partition_index_keys = self
                .collect_partition_index_keys(keyspace_id, table_id, *index_id)
                .await;
            if partition_index_keys.is_empty() {
                partition_index_keys.push(IndexKey::new(keyspace_id, table_id, *index_id));
            }
            index_keys.extend(partition_index_keys);
        }
        index_keys.sort_by_key(|index_key| {
            (
                index_key.keyspace_id,
                index_key.table_id,
                index_key.index_id,
            )
        });
        index_keys.dedup();
        index_keys
    }

    pub async fn set_import_topology_frozen(
        &self,
        keyspace_id: u32,
        table_id: i64,
        index_ids: &[i64],
        frozen: bool,
    ) {
        for index_key in self
            .collect_import_index_keys(keyspace_id, table_id, index_ids)
            .await
        {
            match self.get_index(&index_key).await {
                Ok(index) => index.set_topology_frozen(frozen),
                Err(e) => {
                    warn!(
                        "set_import_topology_frozen: skip missing index={:?}, frozen={}, e={:?}",
                        index_key, frozen, e
                    );
                }
            }
        }
    }

    pub async fn finish_import_append_base(
        &self,
        keyspace_id: u32,
        table_id: i64,
        index_ids: &[i64],
    ) -> Result<(), anyhow::Error> {
        for index_key in self
            .collect_import_index_keys(keyspace_id, table_id, index_ids)
            .await
        {
            match self.get_index(&index_key).await {
                Ok(index) => {
                    index.finish_append_base().await?;
                }
                Err(e) => {
                    warn!(
                        "finish_import_append_base: skip missing index={:?}, e={:?}",
                        index_key, e
                    );
                }
            }
        }
        Ok(())
    }

    pub async fn set_import_topology_frozen_for_job(&self, job_meta: &ImportJobMeta, frozen: bool) {
        self.set_import_topology_frozen(
            job_meta.keyspace_id,
            job_meta.table_id,
            &job_meta.index_ids,
            frozen,
        )
        .await;
    }

    pub async fn restore_from_start(&self) -> Result<(), anyhow::Error> {
        info!("restore_from_start from meta table");
        self.topology_change_enabled
            .store(false, std::sync::atomic::Ordering::Relaxed);
        let indexes = self.meta_store.restore_indexes().await?;
        let notified_index = indexes
            .iter()
            .map(|x| ShardEventIndexInfo::new(x.1.clone(), x.0.clone()))
            .collect::<Vec<_>>();
        let index_count = indexes.len();
        {
            let mut guard = self.indices.write().await;
            for (index_info, table_info, cdc_s3_bucket, cdc_s3_prefix, properties) in
                indexes.into_iter()
            {
                let index_key = IndexKey::from(&index_info);
                match guard.entry(index_key.clone()) {
                    HOccupied(_o) => {
                        error!("restore: meet duplicate index={:?}", index_key,);
                    }
                    HVacant(v) => {
                        let fts = Arc::new(FTSIndex::new_from_restore(
                            self.meta_store.clone(),
                            table_info,
                            index_info,
                            cdc_s3_bucket,
                            cdc_s3_prefix,
                            properties,
                        ));
                        v.insert(fts);
                    }
                }
            }
        }
        let shards = self.meta_store.restore_shards().await?;

        self.event_pubsub
            .on_full_sync(notified_index, shards.clone());
        info!(
            "restore: restored {} indexes and {} shards",
            index_count,
            shards.len()
        );
        INDEX_COUNT_GAUGE_VEC
            .with_label_values(&["total"])
            .set(index_count as i64);
        let restored_shards = shards
            .into_iter()
            .map(|shard| Arc::new(RwLock::new(shard)))
            .collect::<Vec<_>>();
        let mut collector = ModifyCollector::default();
        let mut unscheduled_shards_count = 0usize;
        for shard in restored_shards.iter() {
            let restored_shard = shard.read().await.clone();
            let index = match self.get_index(&restored_shard.index_key).await {
                Ok(index) => index,
                Err(_) => {
                    panic!(
                        "restore: meet unknown index keyspace_id={}, table_id={}, index_id={}",
                        restored_shard.keyspace_id(),
                        restored_shard.table_id(),
                        restored_shard.index_id()
                    )
                }
            };
            index
                .insert_restored_shard(restored_shard.range.sortable_start(), shard.clone())
                .await;
            self.stats
                .total_shards_count
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            if restored_shard.writer.is_empty() {
                unscheduled_shards_count += 1;
            }
            self.restore_shard_runtime(&restored_shard, index, &mut collector)
                .await?;
            // NOTE We restore compaction state if we want the compaction.
        }
        self.set_unscheduled_shards(unscheduled_shards_count);
        for job_meta in self.meta_store.get_active_import_jobs().await? {
            self.set_import_topology_frozen_for_job(&job_meta, true)
                .await;
        }
        info!("finished set_import_topology_frozen_for_job");
        self.topology_change_enabled
            .store(true, std::sync::atomic::Ordering::Relaxed);
        collector.send_all(self.wn_task_tx.clone()).await?;
        for shard in restored_shards.iter() {
            let restored_shard = shard.read().await.clone();
            if restored_shard.writer.is_empty()
                || !self
                    .worker_node_mgr
                    .is_worker_blocked(restored_shard.writer.as_str())
                    .await
            {
                continue;
            }
            warn!(
                "restore: shard={} is assigned to blocked worker={}, schedule reschedule",
                restored_shard.to_long_ident(),
                restored_shard.writer
            );
            if let Err(e) = self
                .shard_task_tx
                .send(ShardTask::NeedReschedule(NeedRescheduleTask {
                    ctx: ShardTaskCtx::new_no_cb(Some("resche_blocked_worker".to_string())),
                    header: ShardTaskHeader::from(&restored_shard),
                    designated_writer: None,
                    reason: RescheduleShardReason::Blocked,
                }))
            {
                warn!("restore: send NeedRescheduleTask error={:?}", e);
            }
        }
        self.on_index_change();
        Ok(())
    }

    async fn restore_shard_runtime(
        &self,
        shard: &ShardInfo,
        index: Arc<FTSIndex>,
        collector: &mut ModifyCollector,
    ) -> Result<()> {
        self.shard_patrol.restore(index.clone(), shard)?;
        if shard.writer.is_empty() {
            return Ok(());
        }
        if self
            .worker_node_mgr
            .is_worker_blocked(shard.writer.as_str())
            .await
        {
            return Ok(());
        }

        WorkerNodeMgr::send_create_worker_node(
            self.wn_task_tx.clone(),
            shard.writer.clone(),
            None,
            self.meta_store.clone(),
        )
        .await?;
        collector.set_add_shard(shard, index, false);
        Ok(())
    }
}