use std::sync::Arc;

use anyhow::{Result, anyhow};
use either::Either;
use log::info;
use tici_common::index::IndexKey;
use tokio::sync::{RwLock, RwLockReadGuard, oneshot};

use crate::{
    HandleRange, ShardInfo, ShardManifest, SplittableFrag,
    meta::{
        metrics::{
            SHARD_EVENT_COUNTER_VEC, SHARD_SPLIT_IMBALANCE_HISTOGRAM,
            SHARD_TASK_TRIGGER_REASON_COUNTER_VEC,
        },
        writer::{
            ExplicitSplitReason, ExplicitSplitResult, ExplicitSplitTask, FTSIndex,
            MergeShardReason, RescheduleShardReason, ShardBuildProgress, ShardTask, ShardTaskCtx,
            ShardTaskHeader, ShardTaskManager, shard_task_runner::*,
        },
    },
};

#[derive(Clone, Debug)]
pub enum ShardSplitReason {
    // Normal size check.
    NewAppend,
    // Trigger then append speed is higher than comapction speed.
    WritePressure,
    Manual,
}

pub const FRAG_FILE_PREFIX: &str = "frags";

impl ShardTaskManager {
    // Create a shard which is the first of the index.
    // Don't use this function if the shard is created by split or migration.
    pub async fn bootstrap_shard(
        &self,
        index_key: &IndexKey,
        range: &HandleRange,
        epoch: u64,
        manifest: ShardManifest,
        writer: String,
        progress: ShardBuildProgress,
    ) -> Result<Arc<RwLock<ShardInfo>>, anyhow::Error> {
        let index = self.get_index(index_key).await?;
        let shard_info = index
            .bootstrap_shard(
                range,
                epoch,
                manifest,
                writer.clone(),
                self.config.load().s3.bucket.clone(),
                format!("{}/{}", &self.config.load().s3.prefix, FRAG_FILE_PREFIX),
                progress,
            )
            .await?;

        let shard = shard_info.read().await;
        self.helper_after_add_shard(index, &shard).await?;
        SHARD_EVENT_COUNTER_VEC
            .with_label_values(&["bootstrap_shard"])
            .inc();
        Ok(shard_info.clone())
    }

    // NOTE Do not directly call this, you have to schedule a ShardTask.
    pub(super) async fn maybe_split_shards(
        &self,
        shard: Arc<RwLock<ShardInfo>>,
        reason: ShardSplitReason,
        batch: usize,
    ) -> Result<ShardSplitResult> {
        let shard = shard.read().await;
        // TODO(P2) let split and compaction parallel here
        if !shard.progress.compaction_state.is_empty() {
            let mut v = vec![];
            for (_, cs) in shard.progress.compaction_state.iter() {
                v.push((cs.get_compaction_state_key(), cs.get_elapsed_secs()));
            }
            return Ok(ShardSplitResult::Yield(v));
        }
        // TODO(P2) perform split only if there is at least one work node.
        if !shard.should_split(&self.config.load().shard, reason.clone()) {
            return Ok(ShardSplitResult::NoNeed(
                shard.get_total_estimated_bytes(),
                shard.get_total_estimated_rows(),
            ));
        }
        let split_points = shard.manifest.compute_split_points(
            &shard.range,
            self.config.load().shard.split.split_point_scan_size_ratio,
        )?;
        if !split_points.is_empty() {
            let mut new_ranges = Vec::new();
            let mut prev = shard.range.sortable_start();
            for split_point in &split_points {
                new_ranges.push(HandleRange::new_from_sorted(
                    prev.clone(),
                    split_point.clone(),
                ));
                prev = split_point.clone();
            }
            new_ranges.push(HandleRange::new_from_sorted(
                prev,
                shard.range.sortable_end(),
            ));
            let new_shards = self
                .split_shards(shard, new_ranges, batch, "normal")
                .await?;
            match reason {
                ShardSplitReason::NewAppend => {
                    SHARD_TASK_TRIGGER_REASON_COUNTER_VEC
                        .with_label_values(&["split_size"])
                        .inc();
                }
                ShardSplitReason::WritePressure => {
                    SHARD_TASK_TRIGGER_REASON_COUNTER_VEC
                        .with_label_values(&["split_write_pressure"])
                        .inc();
                }
                ShardSplitReason::Manual => {
                    SHARD_TASK_TRIGGER_REASON_COUNTER_VEC
                        .with_label_values(&["split_manual"])
                        .inc();
                }
            }
            Ok(ShardSplitResult::Ok(new_shards))
        } else {
            info!(
                "maybe_split_shards: need to split, but can't get a split point. shard={}, size={}",
                shard.to_long_ident(),
                shard.get_total_estimated_bytes(),
            );
            Ok(ShardSplitResult::CannotSplit)
        }
    }

    #[cfg(any(test, feature = "testexport"))]
    pub async fn debug_maybe_merge_shards(
        &self,
        v: Vec<Arc<RwLock<ShardInfo>>>,
        reason: &MergeShardReason,
    ) -> Result<ShardMergeResult> {
        self.maybe_merge_shards(v, reason).await
    }

    // NOTE Do not directly call this, you have to schedule a ShardTask.
    pub(super) async fn maybe_merge_shards(
        &self,
        v: Vec<Arc<RwLock<ShardInfo>>>,
        reason: &MergeShardReason,
    ) -> Result<ShardMergeResult> {
        let mut guards = Vec::with_capacity(v.len());
        for shard in &v {
            guards.push(shard.read().await);
        }
        let v = guards;
        // TODO(P2) Maybe we need to check if the other part is much frequently written.

        if v.iter().map(|x| x.get_total_estimated_bytes()).sum::<u64>()
            > self.config.load().shard.get_split_size_with_propotion(0.8)
        {
            // If the merged shard could be splitted immediately, we reject this merge.
            return Ok(ShardMergeResult::NoNeed);
        }

        SHARD_TASK_TRIGGER_REASON_COUNTER_VEC
            .with_label_values(&[reason.to_metrics_label()])
            .inc();

        let writer = v
            .first()
            .map(|shard| shard.writer.clone())
            .ok_or_else(|| anyhow!("maybe_merge_shards: empty shards"))?;

        match self.merge_shards(writer, v).await? {
            Either::Left(merged) => Ok(ShardMergeResult::Merged(merged)),
            Either::Right((left, right)) => Ok(ShardMergeResult::NeedMergeOverlapCompaction(
                Box::new((left, right)),
            )),
        }
    }

    // NOTE Do not directly call this, you have to schedule a ShardTask.
    // TODO(P1) Rewrite test framework to remove this.
    #[cfg(any(test, feature = "testexport"))]
    pub async fn debug_split_shards(
        &self,
        original_shard: Arc<RwLock<ShardInfo>>,
        new_ranges: Vec<HandleRange>,
        batch: usize,
    ) -> Result<Vec<ShardInfo>, anyhow::Error> {
        let guard = original_shard.read().await;
        self.split_shards(guard, new_ranges, batch, "debug_normal")
            .await
    }

    pub(super) async fn explicit_split_shards(
        &self,
        original_shard: RwLockReadGuard<'_, ShardInfo>,
        new_ranges: Vec<HandleRange>,
        batch: usize,
    ) -> Result<Vec<ShardInfo>, anyhow::Error> {
        SHARD_TASK_TRIGGER_REASON_COUNTER_VEC
            .with_label_values(&["presplit"])
            .inc();
        self.split_shards(original_shard, new_ranges, batch, "explicit")
            .await
    }

    // NOTE Do not directly call this, you have to schedule a ShardTask.
    async fn split_shards(
        &self,
        original_shard: RwLockReadGuard<'_, ShardInfo>,
        new_ranges: Vec<HandleRange>,
        batch: usize,
        tag: &str,
    ) -> Result<Vec<ShardInfo>, anyhow::Error> {
        info!(
            "try split shards for {}, from={:?}, into={:?}",
            tag, original_shard, new_ranges
        );
        if let Some(first) = new_ranges.first() {
            if first.sortable_start_ref() != original_shard.range.sortable_start_ref() {
                anyhow::bail!(
                    "split_shards: new ranges doesn't consistent with original shard, new_ranges={:?}, original_shard={:?}",
                    new_ranges,
                    original_shard
                );
            }
        } else {
            anyhow::bail!("split_shards: new ranges is empty");
        }
        let mut writers = vec![];
        for _ in 0..new_ranges.len() {
            writers.push(
                self.worker_node_mgr
                    .select_worker_node_addr(None)
                    .await
                    .unwrap_or("".to_string()),
            );
        }
        let shard_infos =
            ShardInfo::new_from_split(&original_shard, new_ranges.as_slice(), writers);
        let index = self.get_index(&original_shard.index_key).await?;
        let topology_permit = index.acquire_topology_mutation_permit().await?;
        let original_shard_for_helper = original_shard.clone();
        let mut new_shards = self
            .meta_store
            .split_shards(&original_shard, shard_infos, batch)
            .await?;
        for shard in &mut new_shards {
            shard.progress.mark_cdc_last_file_persisted();
        }
        let topology_rx = index
            .submit_topology_split_shards(original_shard.range.sortable_start(), new_shards.clone())
            .await?;
        drop(original_shard);
        topology_rx
            .await
            .map_err(|e| anyhow!("split_shards: topology actor canceled, err={:?}", e))??;
        drop(topology_permit);
        self.helper_after_split(index, new_shards.clone(), &original_shard_for_helper)
            .await?;

        let (min_shard_size, max_shard_size) = new_shards
            .iter()
            .map(ShardInfo::get_total_estimated_bytes)
            .fold((u64::MAX, 0_u64), |(min_size, max_size), size| {
                (min_size.min(size), max_size.max(size))
            });
        let split_size_ratio = if min_shard_size == u64::MAX {
            0.0
        } else if min_shard_size == 0 {
            f64::INFINITY
        } else {
            max_shard_size as f64 / min_shard_size as f64
        };

        SHARD_SPLIT_IMBALANCE_HISTOGRAM
            .with_label_values(&["size"])
            .observe(split_size_ratio);

        info!(
            "split_shards: splitted. {:?} ===> {}, thres={}, new_range={:?}, new_sizes={}, split_size_ratio={:.4}",
            original_shard_for_helper,
            new_shards
                .iter()
                .map(|x| format!("shard_id={}", x.shard_id))
                .collect::<Vec<_>>()
                .join(", "),
            self.config.load().shard.get_split_size(),
            new_ranges,
            new_shards
                .iter()
                .map(|x| format!("size={}", x.get_total_estimated_bytes()))
                .collect::<Vec<_>>()
                .join(","),
            split_size_ratio,
        );

        Ok(new_shards)
    }

    pub fn schedule_explicit_split(
        &self,
        header: ShardTaskHeader,
        index: Arc<FTSIndex>,
        split_keys: Vec<Vec<u8>>,
        reason: ExplicitSplitReason,
    ) -> Result<oneshot::Receiver<Result<ExplicitSplitResult>>> {
        let (ctx, rx) = ShardTaskCtx::create(Some(reason.to_metrics_label().to_string()));
        self.shard_task_tx
            .send(ShardTask::ExplicitSplit(ExplicitSplitTask {
                ctx,
                header,
                index,
                split_keys,
                reason,
            }))
            .map_err(|e| {
                anyhow!(
                    "schedule_explicit_split: send shard task failed, err={:?}",
                    e
                )
            })?;
        Ok(rx)
    }

    // NOTE Do not directly call this, you have to schedule a ShardTask.
    pub(super) async fn resche_shards(
        &self,
        shard_info: &mut ShardInfo,
        index_key: &IndexKey,
        designated_writer: Option<String>,
        reason: RescheduleShardReason,
    ) -> Result<()> {
        let index = if let Ok(index) = self.get_index(index_key).await {
            index
        } else {
            return Err(anyhow!(
                "resche_shards: failed to find index, {:?}, shard={}, reason={:?}",
                index_key,
                shard_info.to_long_ident(),
                reason,
            ));
        };
        let prev_worker_node = shard_info.writer.clone();
        let worker_node = if let Some(w) = designated_writer {
            if w.is_empty()
                || self
                    .worker_node_mgr
                    .has_schedulable_worker_node(w.as_str())
                    .await
            {
                w
            } else {
                return Err(anyhow!(
                    "resche_shards: do not have designated_writer {}, index={:?}, shard={}, reason={:?}",
                    w,
                    index_key,
                    shard_info.to_long_ident(),
                    reason,
                ));
            }
        } else {
            self.worker_node_mgr
                .select_worker_node_addr(Some(prev_worker_node.clone()))
                .await
                .unwrap_or("".to_string())
        };
        SHARD_TASK_TRIGGER_REASON_COUNTER_VEC
            .with_label_values(&[reason.to_metrics_label()])
            .inc();
        let mut update = shard_info.process_resche(worker_node.clone())?;
        match self
            .meta_store
            .update_shard("resche", shard_info, update.clone())
            .await
        {
            Err(e) => return Err(e),
            Ok(_) => {
                update.mark_progress_cdc_persisted();
                shard_info.apply_update(update);
            }
        };
        if worker_node.is_empty() {
            info!(
                "resche_shards: can't resche shard, {:?}, shard={}, new_writer={}, reason={:?}",
                index_key,
                shard_info.to_long_ident(),
                shard_info.writer,
                reason,
            );
        } else {
            info!(
                "resche_shards: rescheduled shard, {:?}, shard={}, node: {}->{}, reason={:?}",
                index_key,
                shard_info.to_long_ident(),
                prev_worker_node,
                &worker_node,
                reason,
            );
        }
        self.helper_after_move(shard_info, prev_worker_node, index)
            .await?;
        Ok(())
    }

    // NOTE Do not directly call this, you have to schedule a ShardTask.
    async fn merge_shards(
        &self,
        writer: String,
        input: Vec<RwLockReadGuard<'_, ShardInfo>>,
    ) -> Result<either::Either<Arc<RwLock<ShardInfo>>, (SplittableFrag, SplittableFrag)>> {
        info!(
            "merge_shards: begin, {}",
            input
                .iter()
                .map(|x| x.to_long_ident())
                .collect::<Vec<_>>()
                .join(",")
        );
        if input.len() < 2 {
            return Err(anyhow!(
                "merge_shards: can't merge, less than 2 shards, original_shards={:?}",
                input
            ));
        }
        let mut input = input;
        input.sort_by(|x, y| x.range.sortable_start().cmp(&y.range.sortable_start()));
        let input = input;
        for w in input.windows(2) {
            if w[0].range.sortable_end() != w[1].range.sortable_start() {
                return Err(anyhow!(
                    "merge_shards: can't merge, not adjacent, original_shards={:?}, left={:?} right={:?}",
                    input,
                    w[0].range.sortable_end(),
                    w[1].range.sortable_start(),
                ));
            }
        }
        let notified_shards: Vec<ShardInfo> = input.iter().map(|shard| (**shard).clone()).collect();
        let input_for_helper: Vec<ShardInfo> =
            input.iter().map(|shard| (**shard).clone()).collect();
        let input_refs = input.iter().map(|shard| &**shard).collect::<Vec<_>>();

        let new_shard = match ShardInfo::new_from_merge(writer, input_refs.as_slice())? {
            Either::Left(shard) => shard,
            Either::Right((left, right)) => {
                return Ok(Either::Right((left, right)));
            }
        };
        let index = self.get_index(&input[0].index_key).await?;
        let topology_permit = index.acquire_topology_mutation_permit().await?;
        let mut new_shard = match self.meta_store.merge_shards(input_refs, new_shard).await {
            Ok(ok) => ok,
            Err(e) => {
                SHARD_EVENT_COUNTER_VEC
                    .with_label_values(&["merge_failed_backend"])
                    .inc();
                return Err(e);
            }
        };
        new_shard.progress.mark_cdc_last_file_persisted();

        info!("merge_shards: merged, {:?} =====> {:?}", input, new_shard);

        let input_starts = input
            .iter()
            .map(|shard| shard.range.sortable_start())
            .collect::<Vec<_>>();
        let topology_rx = index
            .submit_topology_merge_shards(input_starts, new_shard)
            .await?;
        drop(input);
        let new_shard = topology_rx
            .await
            .map_err(|e| anyhow!("merge_shards: topology actor canceled, err={:?}", e))??;
        drop(topology_permit);
        let new_shard_guard = new_shard.read().await;
        let input_refs = input_for_helper.iter().collect::<Vec<_>>();
        self.helper_after_merge(
            input_refs.as_slice(),
            &new_shard_guard,
            index,
            notified_shards,
        )
        .await?;
        drop(new_shard_guard);
        Ok(Either::Left(new_shard))
    }
}