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

use anyhow::{Context, Result, anyhow, bail};
use log::{error, info, warn};
use tici_common::{
    codec::{table_end_key, table_start_key},
    index::IndexKey,
};
use tici_grpc::{
    AppendBaseFragMetaRequest, AppendFragMetaResponse, CreateIndexRequest, CreateIndexResponse,
    DropIndexRequest, DropIndexResponse, ErrorCode, IndexInfo, KeyRange, PartitionTableInfo,
    TableInfo, util::create_partition_table_and_index_from_def,
};
use tici_shard::writer::{cdc_file_poller::CDCFilePoller, cdc_file_prefix::CDCPrefix};

use crate::{
    HandleRange, RangeBound, ShardManifest,
    meta::writer::{
        AppendBaseFragTask, FTSIndex, ManualCompactionTask, MaybeSplitOrCompactReason,
        MaybeSplitOrCompactTask, NeedRescheduleTask, RescheduleShardReason, ShardBuildProgress,
        ShardTaskCtx, ShardTaskHeader, ShardTaskManager, shard_task::ShardTask,
    },
};

impl ShardTaskManager {
    async fn initial_shard_progress(&self, index: &FTSIndex) -> Result<ShardBuildProgress> {
        // Correctness of starting from the latest CDC file visible to this
        // create_index call:
        //
        // 1. TiDB's add-index scan uses an MVCC snapshot whose timestamp is at least
        //    scan_snapshot_ts. Therefore, it covers all committed versions with
        //    commit_ts <= scan_snapshot_ts, and may also cover some versions with
        //    commit_ts > scan_snapshot_ts.
        // 2. TiDB calls TiCI create_index before entering reorg/backfill and allocating
        //    job.SnapshotVer as scan_snapshot_ts.
        // 3. TiCDC only writes committed transactions to cloud storage. A CDC file
        //    already visible through CDC.index during this create_index call contains
        //    transactions whose commit_ts were allocated before TiDB later obtains
        //    scan_snapshot_ts. PD TSO is globally monotonic, so those commit_ts are <
        //    scan_snapshot_ts.
        //
        // Therefore, initializing the first shard progress to the latest
        // indexed CDC file is safe: skipped CDC rows are covered by TiDB's
        // TiKV scan. Files indexed after this listing are still consumed, even
        // if some of their rows are also covered by the base snapshot.
        let poller = CDCFilePoller::new(self.backend_storage.clone(), index.get_cdc_s3_bucket());
        let prefix = CDCPrefix::new(index.get_cdc_s3_prefix().to_string());
        let latest_file = poller.latest_file_for_table(&prefix).await.with_context(|| {
            format!(
                "create_index: failed to initialize shard progress from latest CDC file, index_key={:?}, bucket={}, prefix={}",
                index.index_key(),
                index.get_cdc_s3_bucket(),
                index.get_cdc_s3_prefix()
            )
        })?;

        let mut progress = ShardBuildProgress::default();
        if let Some(latest_file) = latest_file {
            progress.cdc_s3_last_file = latest_file.to_string();
            info!(
                "create_index: initialized shard progress from latest CDC file, index_key={:?}, cdc_s3_last_file={}",
                index.index_key(),
                progress.cdc_s3_last_file
            );
        }
        Ok(progress)
    }

    async fn ensure_initial_shard(&self, index_key: &IndexKey, worker_node: &str) -> Result<bool> {
        let index = self.get_index(index_key).await?;
        if index.shards_len().await > 0 {
            return Ok(false);
        }
        let initial_progress = self.initial_shard_progress(index.as_ref()).await?;

        let range = HandleRange::new(
            table_start_key(index_key).to_vec(),
            table_end_key(index_key).to_vec(),
        );

        info!(
            "create_index: add new shard on {}, {:?}, index_info={:?}, table_info={:?}, range={:?}, initial_progress={:?}",
            worker_node, index_key, index.index_info, index.table_info, range, initial_progress,
        );

        match self
            .bootstrap_shard(
                index_key,
                &range,
                1,
                ShardManifest::new_for_first(),
                worker_node.to_string(),
                initial_progress,
            )
            .await
        {
            Ok(_) => Ok(true),
            Err(e) => {
                if self.get_index(index_key).await?.shards_len().await > 0 {
                    info!(
                        "create_index: initial shard was created concurrently {:?}",
                        index_key
                    );
                    Ok(false)
                } else {
                    Err(e)
                }
            }
        }
    }

    pub async fn on_msg_create_index(
        &self,
        table_info: TableInfo,
        index_info: IndexInfo,
        partition_table_info: PartitionTableInfo,
        request: &CreateIndexRequest,
    ) -> Result<CreateIndexResponse> {
        let index_key = IndexKey::from(&index_info);
        info!(
            "create_index: try start create index, {:?}, index_info={:?}, table_info={:?}, keyspace_id={}, is_partition={}",
            index_key, index_info, table_info, request.keyspace_id, table_info.is_partition,
        );

        // Build array of (table_info, index_info, index_key) pairs
        let mut table_index_pairs: Vec<(TableInfo, IndexInfo, IndexKey)> = Vec::new();

        if table_info.is_partition {
            for def in partition_table_info.definitions.iter() {
                let (partition_table_info, partition_index_info) =
                    create_partition_table_and_index_from_def(
                        &table_info,
                        &index_info,
                        def.clone(),
                    );
                let partition_index_key = IndexKey::from(&partition_index_info);
                table_index_pairs.push((
                    partition_table_info,
                    partition_index_info,
                    partition_index_key,
                ));
            }
        } else {
            table_index_pairs.push((table_info, index_info, index_key.clone()));
        }

        // Bootstrap all indexes, but keep going so retries can fill missing
        // initial shards even if the index metadata already exists.
        let mut needs_initial_shard = false;
        for (t_info, i_info, i_key) in table_index_pairs.iter() {
            if !self.bootstrap_index(t_info.clone(), i_info.clone()).await? {
                info!(
                    "create_index: the index is already bootstrapped {:?}",
                    i_key
                );
            }
            if self.get_index(i_key).await?.shards_len().await == 0 {
                needs_initial_shard = true;
            }
        }

        if !needs_initial_shard {
            return Ok(CreateIndexResponse {
                status: ErrorCode::Success.into(),
                error_message: "".to_string(),
                index_id: "".to_string(),
            });
        }

        // Select worker node
        let worker_node = match self.worker_node_mgr.select_worker_node_addr(None).await {
            Some(worker_node) => worker_node,
            None => {
                return Ok(CreateIndexResponse {
                    status: ErrorCode::UnknownError.into(),
                    error_message: format!(
                        "create_index: failed to add shard, no available worker nodes={:?}",
                        &request,
                    ),
                    index_id: "".to_string(),
                });
            }
        };

        // Create initial shards for all indexes
        for (_, _, i_key) in table_index_pairs.iter() {
            let _ = self
                .ensure_initial_shard(i_key, worker_node.as_str())
                .await?;
        }

        Ok(CreateIndexResponse {
            status: ErrorCode::Success.into(),
            error_message: "".to_string(),
            index_id: "".to_string(),
        })
    }

    pub async fn on_msg_drop_index(&self, request: &DropIndexRequest) -> Result<DropIndexResponse> {
        // For adhoc tasks like compaction, we just drop them when return.
        // TODO(P2) remove index builder tasks on these indexes, like import into
        let keyspace_id = request.keyspace_id;
        let index_key = IndexKey::new(keyspace_id, request.table_id, request.index_id);
        let mut target_keys = Vec::new();
        let mut should_search_partitions = false;

        match self.get_index(&index_key).await {
            // This index is normal table index
            Ok(index) => {
                target_keys.push(index_key.clone());
                if index.table_info.is_partition {
                    // Non-partition table index marked as partition table index
                    return Ok(DropIndexResponse {
                        status: ErrorCode::InvalidArguments.into(),
                        error_message: format!(
                            "Index meta invalid: Non-partition table index marked as partition table index, index={:?}",
                            index_key).to_string(),
                    });
                }
            }
            // This index is partition table index
            Err(_) => {
                should_search_partitions = true;
            }
        }

        if should_search_partitions {
            let mut partition_keys = self
                .collect_partition_index_keys(keyspace_id, request.table_id, request.index_id)
                .await;
            target_keys.append(&mut partition_keys);
        }

        if target_keys.is_empty() {
            warn!("drop index failed: can't find index {:?}", index_key);
        } else {
            for key in target_keys {
                self.remove_index(&key).await?;
            }
        }

        Ok(DropIndexResponse {
            status: ErrorCode::Success.into(),
            error_message: "".to_string(),
        })
    }

    pub async fn on_msg_append_base_frag_meta(
        &self,
        request: &AppendBaseFragMetaRequest,
    ) -> Result<AppendFragMetaResponse> {
        let index = self
            .get_index(&IndexKey::new(
                request.keyspace_id,
                request.table_id,
                request.index_id,
            ))
            .await?;
        if index.is_finished_append_base() {
            return Ok(AppendFragMetaResponse {
                status: ErrorCode::InvalidOp.into(),
            });
        }
        let mut shard_frags_map = HashMap::new();
        let mut shard_map = HashMap::new();
        for frag in request.frag_metas.iter() {
            let key_range = match &frag.property {
                Some(prop) => KeyRange {
                    start_key: prop.min_handle.clone(),
                    end_key: prop.max_handle.clone(),
                },
                None => {
                    // invalid request without frag_meta property
                    let msg = format!(
                        "append_base_frag_meta: reject fragments without property {:?}",
                        frag,
                    );
                    error!("{}", msg);
                    bail!("{}", msg);
                }
            };

            // According to the logic of fragment generation, the `key_range` here means
            // the fragment with keys in [start_key, end_key]. So find_overlapping_shards
            // should be called with `include_end_key=true`.
            let shards = index
                .find_overlapping_shards(&[key_range.clone()], true)
                .await;
            if shards.is_empty() {
                warn!(
                    "append_base_frag_meta: can't find any overlapped shard, maybe the index is not bootstrap or already dropped, key_range={:?} request {:?}",
                    &key_range, request
                );
                return Ok(AppendFragMetaResponse {
                    status: tici_grpc::ErrorCode::ShardNotFound.into(),
                });
            }
            // Insert the frag into all shards that overlap with the frag's range.
            for shard in shards.iter() {
                let shard_id = shard.read().await.shard_id;
                shard_frags_map
                    .entry(shard_id)
                    .or_insert_with(Vec::new)
                    .push(frag.clone());
                shard_map.entry(shard_id).or_insert_with(|| shard.clone());
            }
        }
        let mut rxs = vec![];
        for (shard_id, frags) in shard_frags_map.iter() {
            let shard = if let Some(shard) = shard_map.get(shard_id) {
                shard
            } else {
                error!(
                    "append_base_frag_meta: not found shard_id, shard={}",
                    shard_id
                );
                bail!(
                    "append_base_frag_meta: not found shard_id, shard={}",
                    shard_id
                );
            };
            let (tx, rx) = tokio::sync::oneshot::channel();
            if let Err(e) = self
                .shard_task_tx
                .send(ShardTask::AppendBaseFrag(AppendBaseFragTask {
                    ctx: ShardTaskCtx::new(Some("append_base".to_string()), tx),
                    header: ShardTaskHeader::from(&*shard.read().await),
                    index: index.clone(),
                    frags: frags.clone(),
                }))
            {
                error!(
                    "append_base_frag_meta: send shard task failed, shard={:?}, err={}",
                    shard_id, e
                );
                bail!(
                    "append_base_frag_meta: send shard task failed, shard={:?}, err={}",
                    shard_id,
                    e
                );
            } else {
                rxs.push(rx);
            }
        }
        let mut has_error = false;
        for rx in rxs {
            match rx.await? {
                Ok((rsp, _update)) => {
                    if rsp.status != ErrorCode::Success as i32 {
                        has_error = true;
                        error!(
                            "append_base_frag_meta: non-success status, status={}",
                            rsp.status
                        );
                    }
                }
                Err(e) => {
                    has_error = true;
                    error!("append_base_frag_meta: error, err={}", e);
                }
            }
        }
        if has_error {
            Ok(AppendFragMetaResponse {
                status: tici_grpc::ErrorCode::UnknownError.into(),
            })
        } else {
            Ok(AppendFragMetaResponse {
                status: tici_grpc::ErrorCode::Success.into(),
            })
        }
    }

    #[allow(unused)]
    pub async fn debug_get_shard_manifest(
        &self,
        request: &tici_grpc::DebugGetShardManifestRequest,
    ) -> Result<tici_grpc::DebugGetShardManifestResponse> {
        #[cfg(any(test, feature = "get_manifest_from_writer", feature = "testexport"))]
        {
            match self
                .get_index(&IndexKey::new(
                    request.keyspace_id,
                    request.table_id,
                    request.index_id,
                ))
                .await
            {
                Ok(fts_index) => {
                    if request.key_range.is_none() {
                        error!("no keyrange provided, req={:?}", request);
                    }
                    let key_range = if let Some(kr) = request.key_range.as_ref() {
                        kr
                    } else {
                        return Ok(tici_grpc::DebugGetShardManifestResponse {
                            status: ErrorCode::InvalidArguments.into(),
                            ..Default::default()
                        });
                    };
                    match fts_index
                        .get_shard(&RangeBound::make_start(key_range.start_key.clone()))
                        .await
                    {
                        None => Ok(tici_grpc::DebugGetShardManifestResponse {
                            status: ErrorCode::ShardNotFound.into(),
                            ..Default::default()
                        }),
                        Some(x) => {
                            let s = x.read().await;
                            Ok(tici_grpc::DebugGetShardManifestResponse {
                                status: ErrorCode::Success.into(),
                                frag_metas: s.manifest.cloned_fragments(),
                                shard_id: s.shard_id,
                                seq: s.seq,
                            })
                        }
                    }
                }
                Err(_) => {
                    // Index not found, return empty result with success status
                    Ok(tici_grpc::DebugGetShardManifestResponse {
                        status: ErrorCode::IndexNotFound.into(),
                        ..Default::default()
                    })
                }
            }
        }
        #[cfg(not(any(test, feature = "get_manifest_from_writer", feature = "testexport")))]
        {
            Ok(tici_grpc::DebugGetShardManifestResponse {
                status: ErrorCode::Deprecated.into(),
                ..Default::default()
            })
        }
    }

    pub async fn debug_get_shard_info(
        &self,
        request: &tici_grpc::DebugGetShardInfoRequest,
    ) -> Result<tici_grpc::DebugGetShardInfoResponse> {
        let to_find_shard_ids: HashSet<u64> = request.shard_ids.iter().cloned().collect();
        let mut resp: tici_grpc::DebugGetShardInfoResponse =
            tici_grpc::DebugGetShardInfoResponse::default();
        let all_indexes: Vec<Arc<FTSIndex>> = self
            .indices
            .read()
            .await
            .iter()
            .map(|x| x.1.clone())
            .collect();
        let mut founded = HashSet::<u64>::new();
        for index in all_indexes {
            for shard in index.clone_shards().await {
                if to_find_shard_ids.contains(&shard.read().await.shard_id) {
                    resp.shards.push(tici_grpc::GrpcInternalShardInfo {
                        table_info: Some(index.table_info.clone()),
                        index_info: Some(index.index_info.clone()),
                        shard_info: Some(shard.read().await.get_header()),
                        frag_metas: shard.read().await.manifest.cloned_fragments(),
                        s3_bucket: shard.read().await.frag_upload_location.bucket.clone(),
                    });
                    founded.insert(shard.read().await.shard_id);
                }
            }
        }
        let diff = to_find_shard_ids.difference(&founded).collect::<Vec<_>>();
        if !diff.is_empty() {
            info!(
                "failed to find shards={}",
                diff.iter()
                    .map(|id| id.to_string())
                    .collect::<Vec<_>>()
                    .join(";")
            );
        }
        Ok(resp)
    }

    pub async fn on_msg_manual_schedule_compaction_by_shard(
        &self,
        index_key: &IndexKey,
        start_key: Vec<u8>,
        end_key: Vec<u8>,
        range: &str,
    ) -> Result<(ErrorCode, String)> {
        let shards = self
            .get_index(index_key)
            .await?
            .find_overlapping_shards(
                &[KeyRange {
                    start_key: start_key.clone(),
                    end_key: end_key.clone(),
                }],
                false,
            )
            .await;

        let mut to_wait = vec![];
        for shard in shards {
            let (tx, rx) = tokio::sync::oneshot::channel();
            let shard_id = shard.read().await.shard_id;
            if let Err(e) =
                self.shard_task_tx
                    .send(ShardTask::ManualCompaction(ManualCompactionTask {
                        ctx: ShardTaskCtx::new(None, tx),
                        header: ShardTaskHeader::new(
                            index_key.clone(),
                            shard.read().await.range.sortable_start(),
                            shard_id,
                        ),
                        range: range.to_string(),
                    }))
            {
                return Err(anyhow!(
                    "manual_compaction: failed to send shard task, shard_id={}, {:?}, start_key={:?}, end_key={:?}, range={:?}, e={:?}",
                    shard_id,
                    index_key,
                    start_key,
                    end_key,
                    range,
                    e
                ));
            } else {
                to_wait.push((shard_id, rx));
            }
        }
        let scheduled = to_wait.len();
        info!(
            "manual_compaction: scheduled {} shards, {:?}, start_key={:?}, end_key={:?}, range={:?}",
            scheduled, index_key, start_key, end_key, range
        );
        let mut success = 0;
        let mut error_shards = vec![];
        for (shard_id, w) in to_wait {
            match w.await? {
                Ok(code) => {
                    if code == ErrorCode::Success {
                        success += 1;
                    }
                }
                Err(e) => {
                    error_shards.push((shard_id, e));
                }
            }
        }
        if success == scheduled {
            Ok((ErrorCode::Success, "".to_string()))
        } else {
            let error_msg = format!(
                "{:?}, start_key={:?}, end_key={:?}, range={:?}, errors={:?}",
                index_key,
                start_key,
                end_key,
                range,
                error_shards
                    .iter()
                    .map(|x| { format!("{}->{}", x.0, x.1.root_cause()) })
            );
            Ok((ErrorCode::TryAgain, error_msg))
        }
    }

    pub async fn on_msg_manual_reschedule_by_shard(
        &self,
        index_key: &IndexKey,
        start_key: Vec<u8>,
        end_key: Vec<u8>,
        writer: Option<String>,
    ) -> Result<(ErrorCode, String)> {
        let index = self.get_index(index_key).await?;
        let shards = index
            .find_overlapping_shards(
                &[KeyRange {
                    start_key: start_key.clone(),
                    end_key: end_key.clone(),
                }],
                false,
            )
            .await;
        let shards_len = shards.len();

        let mut to_wait = vec![];
        for shard in shards {
            let (tx, rx) = tokio::sync::oneshot::channel();
            let shard_id = shard.read().await.shard_id;
            if let Err(e) = self
                .shard_task_tx
                .send(ShardTask::NeedReschedule(NeedRescheduleTask {
                    ctx: ShardTaskCtx::new(None, tx),
                    header: ShardTaskHeader::new(
                        index_key.clone(),
                        shard.read().await.range.sortable_start(),
                        shard_id,
                    ),
                    designated_writer: writer.clone(),
                    reason: RescheduleShardReason::Manual,
                }))
            {
                return Err(anyhow!(
                    "manual_reschedule: failed to send shard task, shard_id={}, {:?}, start_key={:?}, end_key={:?}, writer={:?}, e={:?}",
                    shard_id,
                    index_key,
                    start_key,
                    end_key,
                    writer,
                    e,
                ));
            } else {
                to_wait.push((shard_id, rx));
            }
        }
        let scheduled = to_wait.len();
        let sk = hex::encode(start_key.clone());
        let ek = hex::encode(end_key.clone());
        info!(
            "manual_reschedule: scheduled {} shards, {:?}, start_key={:?}|{}, end_key={:?}|{}, writer={:?}, shards {}/{}",
            scheduled,
            index_key,
            start_key,
            sk,
            end_key,
            ek,
            writer,
            shards_len,
            index.shards_len().await,
        );
        let mut success = 0;
        let mut error_shards = vec![];
        for (shard_id, w) in to_wait {
            match w.await? {
                Ok(_) => {
                    success += 1;
                }
                Err(e) => {
                    error_shards.push((shard_id, e.root_cause().to_string()));
                }
            }
        }
        if success == scheduled {
            Ok((ErrorCode::Success, "".to_string()))
        } else {
            let error_msg = format!(
                "{:?}, start_key={:?}, end_key={:?}, level={:?}, errors={:?}",
                index_key,
                start_key,
                end_key,
                writer,
                error_shards.iter().map(|x| { format!("{}->{}", x.0, x.1) })
            );
            Ok((ErrorCode::TryAgain, error_msg))
        }
    }

    pub async fn on_msg_manual_split_by_shard(
        &self,
        index_key: &IndexKey,
        start_key: Vec<u8>,
    ) -> Result<(ErrorCode, String)> {
        let shard = if let Some(s) = self
            .get_index(index_key)
            .await?
            .get_shard(&RangeBound::make_start(start_key.clone()))
            .await
        {
            s
        } else {
            return Err(anyhow!(
                "manual_split: fshard not found, {:?}, start_key={:?}",
                index_key,
                start_key,
            ));
        };

        {
            let (tx, rx) = tokio::sync::oneshot::channel();
            let shard_id = shard.read().await.shard_id;
            if let Err(e) =
                self.shard_task_tx
                    .send(ShardTask::MaybeSplitOrCompact(MaybeSplitOrCompactTask {
                        ctx: ShardTaskCtx::new(None, tx),
                        header: ShardTaskHeader::new(
                            index_key.clone(),
                            shard.read().await.range.sortable_start(),
                            shard_id,
                        ),
                        reason: MaybeSplitOrCompactReason::ManualSplit,
                    }))
            {
                Err(anyhow!(
                    "manual_split: failed to send shard task, shard_id={}, {:?}, start_key={:?}, e={:?}",
                    shard_id,
                    index_key,
                    start_key,
                    e
                ))
            } else {
                info!(
                    "manual_split: scheduled, wait, {:?}, start_key={:?}",
                    index_key, start_key
                );
                let res = rx.await?;
                match res {
                    Ok(()) => Ok((ErrorCode::Success, "".to_string())),
                    Err(e) => {
                        let error_msg =
                            format!("{:?}, start_key={:?}, error={:?}", index_key, start_key, e);
                        Ok((ErrorCode::TryAgain, error_msg))
                    }
                }
            }
        }
    }
}