use std::sync::Arc;

use anyhow::anyhow;
use log::warn;
use tici_grpc::{AppendDeltaFragMetaRequest, AppendFragMetaResponse, ErrorCode, FragMeta};

use crate::{
    MaybeShardInfoUpdate, ShardInfo, ShardInfoAppendUpdate, ShardInfoUpdate,
    meta::{
        metrics::SHARD_EVENT_COUNTER_VEC,
        writer::{
            FTSIndex, MaybeSplitOrCompactReason, MaybeSplitOrCompactTask, NeedRescheduleTask,
            RescheduleShardReason, ShardTask, ShardTaskCtx, ShardTaskHeader, ShardTaskManager,
        },
    },
};

impl ShardTaskManager {
    // NOTE Do not directly call this, you have to schedule a ShardTask.
    pub(super) async fn append_delta_frag_meta(
        &self,
        ctx: ShardTaskCtx<(AppendFragMetaResponse, Option<ShardInfoUpdate>)>,
        shard: &mut ShardInfo,
        index: Arc<FTSIndex>,
        request: AppendDeltaFragMetaRequest,
    ) {
        let max_fragments = self.config.load().shard.max_fragments;
        if shard.manifest.len() as u64 > max_fragments {
            SHARD_EVENT_COUNTER_VEC
                .with_label_values(&["append_fail_toomanyfrags"])
                .inc();
            warn!(
                "append_delta_frag_meta: too many fragments, shard={}, frags={}, max={}, returning={:?}",
                shard.to_long_ident(),
                shard.manifest.len(),
                max_fragments,
                ErrorCode::Throttled
            );
            if let Err(e) =
                self.shard_task_tx
                    .send(ShardTask::MaybeSplitOrCompact(MaybeSplitOrCompactTask {
                        ctx: ShardTaskCtx::new_no_cb(Some("after_append_failed".to_string())),
                        header: ShardTaskHeader::from(&*shard),
                        reason: MaybeSplitOrCompactReason::StallTooManyFrags,
                    }))
            {
                log::info!(
                    "append_delta_frag_meta: schedule MaybeSplitOrCompact failure, error={:?}, shard={}",
                    e,
                    shard.to_long_ident()
                );
            }
            ctx.report_ok((
                AppendFragMetaResponse {
                    status: ErrorCode::Throttled.into(),
                },
                None,
            ));
            return;
        }

        let update = match shard.process_append_delta(&request) {
            Err(e) => {
                ctx.report_err(e);
                return;
            }
            Ok(u) => match u {
                MaybeShardInfoUpdate::Update(update) => update,
                MaybeShardInfoUpdate::ErrorCode(code) => {
                    let reschedule_header = if code == ErrorCode::ProgressUnwinded {
                        if shard.writer == request.worker_node_addr && !shard.writer.is_empty() {
                            Some(ShardTaskHeader::new(
                                index.index_key(),
                                shard.range.sortable_start(),
                                shard.shard_id,
                            ))
                        } else {
                            warn!(
                                "append_delta_frag_meta: detected progress unwind from stale worker, shard={}, meta_writer={}, worker={}",
                                shard.to_long_ident(),
                                shard.writer,
                                request.worker_node_addr,
                            );
                            None
                        }
                    } else {
                        None
                    };
                    let worker_node_addr = request.worker_node_addr.clone();
                    let shard_id = shard.shard_id;
                    let shard_ident = shard.to_long_ident();
                    let index_key = index.index_key();
                    let start_key = shard.range.sortable_start();
                    ctx.report_ok((
                        AppendFragMetaResponse {
                            status: code.into(),
                        },
                        None,
                    ));
                    if let Some(header) = reschedule_header {
                        log::info!(
                            "append_delta_frag_meta: schedule NeedReschedule after progress unwind, index={:?}, shard_id={}, start_key={:?}, worker={}, shard={}",
                            index_key,
                            shard_id,
                            start_key,
                            worker_node_addr,
                            shard_ident
                        );
                        if let Err(e) =
                            self.shard_task_tx
                                .send(ShardTask::NeedReschedule(NeedRescheduleTask {
                                    ctx: ShardTaskCtx::new_no_cb(Some(
                                        "resche_progress_unwind".to_string(),
                                    )),
                                    header,
                                    designated_writer: Some(String::new()),
                                    reason: RescheduleShardReason::ProgressUnwinded,
                                }))
                        {
                            warn!(
                                "append_delta_frag_meta: failed to schedule NeedReschedule after progress unwind, index={:?}, shard_id={}, start_key={:?}, worker={}, shard={}, e={:?}",
                                index_key, shard_id, start_key, worker_node_addr, shard_ident, e,
                            );
                        }
                    }
                    return;
                }
            },
        };
        let mut update = *update;
        let should_persist = shard.should_persist_empty_append(
            &update,
            self.config.load().shard.append.append_empty_persist_freq,
        );

        if should_persist {
            if let Err(e) = self
                .meta_store
                .update_shard("commit_append", &*shard, update.clone())
                .await
            {
                SHARD_EVENT_COUNTER_VEC
                    .with_label_values(&["append_failed_backend_delta"])
                    .inc();
                ctx.report_err(anyhow!(
                    "append_frag_meta: can't commit append, shard={:?}, e={:?}",
                    shard,
                    e
                ));
                return;
            }
            update.mark_progress_cdc_persisted();
        } else {
            log::debug!(
                "append_frag_meta: skip empty append persist due to freq limit, shard={}, seq={}, cdc_progress={:?}",
                shard.to_long_ident(),
                shard.seq,
                update
                    .progress
                    .as_ref()
                    .map(|progress| progress.cdc_s3_last_file.as_str())
            );
        }
        let size_added = update.size_added;
        index.append_qps.record(size_added);
        if update.seq.is_some() {
            shard.reset_empty_append_persist_count();
        }

        let committed_seq = update.seq.unwrap_or(shard.seq);
        let new_frags = update.appended_frags.clone().unwrap_or_default();
        shard.apply_update(update);

        let update_append = ShardInfoAppendUpdate {
            shard_id: shard.shard_id,
            start_key: shard.range.sortable_start(),
            frags: new_frags,
            new_seq: committed_seq,
            is_base: false,
            cdc_s3_last_file: shard.progress.cdc_s3_last_file.clone(),
        };

        self.helper_after_append(index, shard, update_append, size_added)
            .await;

        ctx.report_ok((
            AppendFragMetaResponse {
                status: tici_grpc::ErrorCode::Success.into(),
            },
            Some(ShardInfoUpdate::from(&*shard)),
        ));
    }

    // NOTE Do not directly call this, you have to schedule a ShardTask.
    pub(super) async fn append_base_frag_meta(
        &self,
        ctx: ShardTaskCtx<(AppendFragMetaResponse, Option<ShardInfoUpdate>)>,
        shard: &mut ShardInfo,
        index: Arc<FTSIndex>,
        frags: &[FragMeta],
    ) {
        if let Err(code) = validate_base_frags(shard, frags) {
            ctx.report_ok((
                AppendFragMetaResponse {
                    status: code.into(),
                },
                None,
            ));
            return;
        }
        // NOTE: base fragments are appended by import workers, which are not
        // necessarily `self.writer`. NOTE: base fragments are not currently
        // protected by Throttled check.
        let mut update = shard.process_append_base(frags);
        if let Err(e) = self
            .meta_store
            .update_shard("commit_append", &*shard, update.clone())
            .await
        {
            SHARD_EVENT_COUNTER_VEC
                .with_label_values(&["append_failed_backend_base"])
                .inc();
            ctx.report_err(anyhow!(
                "append_frag_meta: can't commit append, shard={:?}, e={:?}",
                shard,
                e
            ));
            return;
        }

        update.mark_progress_cdc_persisted();
        let size_added = frags.iter().map(|x| x.bytes()).sum();
        let update_append = ShardInfoAppendUpdate {
            shard_id: shard.shard_id,
            start_key: shard.range.sortable_start(),
            frags: update.appended_frags.clone().unwrap_or_default(),
            new_seq: update.seq.unwrap_or(shard.seq),
            is_base: true,
            cdc_s3_last_file: shard.progress.cdc_s3_last_file.clone(),
        };
        shard.apply_update(update);

        ctx.report_ok((
            AppendFragMetaResponse {
                status: tici_grpc::ErrorCode::Success.into(),
            },
            Some(ShardInfoUpdate::from(&*shard)),
        ));

        self.helper_after_append(index, shard, update_append, size_added)
            .await;
    }
}

fn validate_base_frags(shard: &ShardInfo, frags: &[FragMeta]) -> Result<(), ErrorCode> {
    for frag in frags {
        let prop = frag.validate().map_err(|e| {
            warn!(
                "append_base_frag_meta: invalid frag, shard={:?}, frag={:?}, err={:#}",
                shard, frag, e
            );
            ErrorCode::InvalidArguments
        })?;
        if !shard
            .range
            .range_overlaps_minmax(&prop.min_handle, &prop.max_handle)
        {
            warn!(
                "append_base_frag_meta: frag range mismatch, shard={:?}, frag_range=[{:?},{:?}]",
                shard, prop.min_handle, prop.max_handle
            );
            return Err(ErrorCode::FragRangeMismatch);
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use arc_swap::ArcSwap;
    use tici_grpc::{
        AppendDeltaFragMetaRequest, ErrorCode, FragMeta, FragProperty, IndexInfo, S3Location,
        ShardProgress, TableInfo, WorkerNodeStatus,
    };
    use tici_shard::test_utils::{MockBackendStorage, test_mock_frag_storage_refs};

    use super::validate_base_frags;
    use crate::{
        HandleRange, ShardInfo,
        common::test_utils::make_sf,
        config::MetaServiceConfig,
        meta::{
            store::{MetaStore, test_utils::RecordingMetaStore},
            writer::{FTSIndex, ModifyShardTaskAdd, ShardTaskCtx, ShardTaskManager, WorkerNodeMgr},
        },
    };

    fn make_index(store: Arc<dyn MetaStore>) -> Arc<FTSIndex> {
        Arc::new(FTSIndex::new(
            store,
            TableInfo::default(),
            IndexInfo::default(),
            "bucket".to_string(),
            "prefix".to_string(),
        ))
    }

    async fn wait_for_stopped<F>(mut is_stopped: F)
    where
        F: FnMut() -> bool,
    {
        let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(3);
        while !is_stopped() && tokio::time::Instant::now() < deadline {
            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
        }
    }

    #[test]
    fn test_validate_base_frags_overlap_ok() {
        let shard = ShardInfo::new_for_test(HandleRange::new(b"b".to_vec(), b"m".to_vec()));
        let frag = make_sf(b"c", b"d");
        assert!(validate_base_frags(&shard, &[frag]).is_ok());
    }

    #[test]
    fn test_validate_base_frags_overlap_mismatch() {
        let shard = ShardInfo::new_for_test(HandleRange::new(b"b".to_vec(), b"m".to_vec()));
        let frag = make_sf(b"x", b"y");
        assert_eq!(
            validate_base_frags(&shard, &[frag]),
            Err(ErrorCode::FragRangeMismatch)
        );
    }

    #[test]
    fn test_validate_base_frags_missing_property() {
        let shard = ShardInfo::new_for_test(HandleRange::new(b"b".to_vec(), b"m".to_vec()));
        let mut frag = FragMeta::new(
            0,
            "f_1",
            FragProperty {
                min_handle: b"b".to_vec(),
                split_handles: vec![b"c".to_vec()],
                max_handle: b"m".to_vec(),
                size: 1,
                ..Default::default()
            },
            test_mock_frag_storage_refs(),
            1,
        )
        .unwrap();
        frag.property = None;
        assert_eq!(
            validate_base_frags(&shard, &[frag]),
            Err(ErrorCode::InvalidArguments)
        );
    }

    #[test]
    fn test_validate_base_frags_rejects_empty_storage_refs() {
        let shard = ShardInfo::new_for_test(HandleRange::new(b"b".to_vec(), b"m".to_vec()));
        let frag = FragMeta {
            frag_path: "f_1".to_string(),
            property: Some(FragProperty {
                min_handle: b"b".to_vec(),
                split_handles: vec![b"c".to_vec()],
                max_handle: b"m".to_vec(),
                ..Default::default()
            }),
            storage_refs: Vec::new(),
            ..Default::default()
        };
        assert_eq!(
            validate_base_frags(&shard, &[frag]),
            Err(ErrorCode::InvalidArguments)
        );
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_append_delta_progress_unwind_removes_current_worker() {
        let meta_store = Arc::new(RecordingMetaStore::default());
        let meta_store_dyn: Arc<dyn MetaStore> = meta_store.clone();
        let config = Arc::new(ArcSwap::from_pointee(MetaServiceConfig::default()));
        let (shard_task_tx, shard_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let (wn_task_tx, wn_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let worker_node_mgr = WorkerNodeMgr::new(
            shard_task_tx.clone(),
            wn_task_tx,
            wn_task_rx,
            config.clone(),
            meta_store_dyn.clone(),
        );
        let mgr = ShardTaskManager::new(
            meta_store_dyn.clone(),
            worker_node_mgr.clone(),
            config,
            shard_task_tx,
            shard_task_rx,
            Arc::new(MockBackendStorage::new()),
        );

        let worker_addr = "worker-a".to_string();
        let worker = worker_node_mgr
            .create_worker_node(
                meta_store_dyn.clone(),
                worker_addr.clone(),
                Some(WorkerNodeStatus::default()),
            )
            .await;
        let index = make_index(meta_store_dyn);
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"a".to_vec(), b"z".to_vec()));
        shard.shard_id = 42;
        shard.writer = worker_addr.clone();
        shard.progress.cdc_s3_last_file =
            "tici_default_prefix/cdc/10313/461869219056451591/CDC00000000000000000002.json"
                .to_string();
        let shard_start = shard.range.sortable_start();
        let manager_shard = Arc::new(tokio::sync::RwLock::new(shard.clone()));
        mgr.indices
            .write()
            .await
            .insert(index.index_key(), index.clone());
        index
            .insert_restored_shard(shard_start.clone(), manager_shard.clone())
            .await;

        WorkerNodeMgr::send_modify_shard(
            worker_node_mgr.get_tx(),
            worker_addr.clone(),
            vec![ModifyShardTaskAdd {
                header: shard.get_header(),
                progress: shard.progress.clone(),
                upload_location: S3Location {
                    bucket: "bucket".to_string(),
                    prefix: "prefix".to_string(),
                },
                shard_stats: shard.stats.clone(),
                index: index.clone(),
                notified: false,
            }],
            vec![],
        )
        .await
        .unwrap();
        assert!(worker.get_shard_schedule(&shard.shard_id).await.is_some());

        let shard_id = shard.shard_id;
        let (ctx, rx) = ShardTaskCtx::create(Some("test_progress_unwind".to_string()));
        mgr.append_delta_frag_meta(
            ctx,
            &mut shard,
            index,
            AppendDeltaFragMetaRequest {
                shard_id,
                worker_node_addr: worker_addr.clone(),
                progress: Some(ShardProgress {
                    last_cdc_file:
                        "tici_default_prefix/cdc/10313/461869219056451591/CDC00000000000000000001.json"
                            .to_string(),
                }),
                ..Default::default()
            },
        )
        .await;

        let (rsp, update) = rx.await.unwrap().unwrap();
        assert_eq!(rsp.status, ErrorCode::ProgressUnwinded as i32);
        assert!(update.is_none());

        let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(3);
        loop {
            let is_removed = worker.get_shard_schedule(&shard.shard_id).await.is_none();
            let writer = manager_shard.read().await.writer.clone();
            let update_count = meta_store.shard_updates.lock().unwrap().len();
            if is_removed && writer.is_empty() && update_count == 1 {
                break;
            }
            if tokio::time::Instant::now() >= deadline {
                break;
            }
            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
        }

        assert!(manager_shard.read().await.writer.is_empty());
        assert!(worker.get_shard_schedule(&shard.shard_id).await.is_none());
        assert_eq!(
            mgr.stats
                .unscheduled_shards_count
                .load(std::sync::atomic::Ordering::Relaxed),
            1
        );

        {
            let updates = meta_store.shard_updates.lock().unwrap();
            assert_eq!(updates.len(), 1);
            assert_eq!(updates[0].0, "resche");
            assert_eq!(updates[0].1.writer.as_deref(), Some(""));
        }

        mgr.stop();
        worker_node_mgr.stop();
        wait_for_stopped(|| mgr.is_stopped()).await;
        wait_for_stopped(|| worker_node_mgr.is_stopped()).await;
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_append_empty_persist_every_two_progress_updates() {
        let meta_store = Arc::new(RecordingMetaStore::default());
        let meta_store_dyn: Arc<dyn MetaStore> = meta_store.clone();
        let mut config = MetaServiceConfig::default();
        config.shard.append.append_empty_persist_freq = 2;
        let config = Arc::new(ArcSwap::from_pointee(config));
        let (shard_task_tx, shard_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let (wn_task_tx, wn_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let worker_node_mgr = WorkerNodeMgr::new(
            shard_task_tx.clone(),
            wn_task_tx,
            wn_task_rx,
            config.clone(),
            meta_store_dyn.clone(),
        );
        let mgr = ShardTaskManager::new(
            meta_store_dyn.clone(),
            worker_node_mgr.clone(),
            config,
            shard_task_tx,
            shard_task_rx,
            Arc::new(MockBackendStorage::new()),
        );

        let worker_addr = "worker-a".to_string();
        let index = make_index(meta_store_dyn);
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"a".to_vec(), b"z".to_vec()));
        shard.shard_id = 7;
        shard.writer = worker_addr.clone();
        shard.progress.cdc_s3_last_file = "cdc-0001".to_string();
        shard.progress.mark_cdc_last_file_persisted();
        let shard_id = shard.shard_id;

        for cdc_file in ["cdc-0002", "cdc-0003"] {
            let (ctx, rx) = ShardTaskCtx::create(Some("test_append_empty_freq".to_string()));
            mgr.append_delta_frag_meta(
                ctx,
                &mut shard,
                index.clone(),
                AppendDeltaFragMetaRequest {
                    shard_id,
                    worker_node_addr: worker_addr.clone(),
                    progress: Some(ShardProgress {
                        last_cdc_file: cdc_file.to_string(),
                    }),
                    ..Default::default()
                },
            )
            .await;

            let (rsp, update) = rx.await.unwrap().unwrap();
            assert_eq!(rsp.status, ErrorCode::Success as i32);
            assert!(update.is_some());
        }

        assert_eq!(meta_store.shard_updates.lock().unwrap().len(), 1);
        assert_eq!(shard.progress.cdc_s3_last_file, "cdc-0003");
        assert_eq!(shard.progress.persisted_cdc_s3_last_file, "cdc-0003");
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_non_empty_append_resets_empty_append_persist_counter() {
        let meta_store = Arc::new(RecordingMetaStore::default());
        let meta_store_dyn: Arc<dyn MetaStore> = meta_store.clone();
        let mut config = MetaServiceConfig::default();
        config.shard.append.append_empty_persist_freq = 2;
        let config = Arc::new(ArcSwap::from_pointee(config));
        let (shard_task_tx, shard_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let (wn_task_tx, wn_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let worker_node_mgr = WorkerNodeMgr::new(
            shard_task_tx.clone(),
            wn_task_tx,
            wn_task_rx,
            config.clone(),
            meta_store_dyn.clone(),
        );
        let mgr = ShardTaskManager::new(
            meta_store_dyn.clone(),
            worker_node_mgr.clone(),
            config,
            shard_task_tx,
            shard_task_rx,
            Arc::new(MockBackendStorage::new()),
        );

        let worker_addr = "worker-a".to_string();
        let index = make_index(meta_store_dyn);
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"a".to_vec(), b"z".to_vec()));
        shard.shard_id = 8;
        shard.writer = worker_addr.clone();
        shard.progress.cdc_s3_last_file = "cdc-0010".to_string();
        shard.progress.mark_cdc_last_file_persisted();
        let shard_id = shard.shard_id;

        let (ctx, rx) = ShardTaskCtx::create(Some("test_append_empty_skip".to_string()));
        mgr.append_delta_frag_meta(
            ctx,
            &mut shard,
            index.clone(),
            AppendDeltaFragMetaRequest {
                shard_id,
                worker_node_addr: worker_addr.clone(),
                progress: Some(ShardProgress {
                    last_cdc_file: "cdc-0011".to_string(),
                }),
                ..Default::default()
            },
        )
        .await;
        let (rsp, _) = rx.await.unwrap().unwrap();
        assert_eq!(rsp.status, ErrorCode::Success as i32);
        assert_eq!(meta_store.shard_updates.lock().unwrap().len(), 0);

        let (ctx, rx) = ShardTaskCtx::create(Some("test_append_non_empty".to_string()));
        mgr.append_delta_frag_meta(
            ctx,
            &mut shard,
            index.clone(),
            AppendDeltaFragMetaRequest {
                shard_id,
                worker_node_addr: worker_addr.clone(),
                frag_metas: vec![make_sf(b"b", b"c")],
                progress: Some(ShardProgress {
                    last_cdc_file: "cdc-0012".to_string(),
                }),
                ..Default::default()
            },
        )
        .await;
        let (rsp, _) = rx.await.unwrap().unwrap();
        assert_eq!(rsp.status, ErrorCode::Success as i32);

        let (ctx, rx) = ShardTaskCtx::create(Some("test_append_empty_skip_again".to_string()));
        mgr.append_delta_frag_meta(
            ctx,
            &mut shard,
            index,
            AppendDeltaFragMetaRequest {
                shard_id,
                worker_node_addr: worker_addr,
                progress: Some(ShardProgress {
                    last_cdc_file: "cdc-0013".to_string(),
                }),
                ..Default::default()
            },
        )
        .await;
        let (rsp, _) = rx.await.unwrap().unwrap();
        assert_eq!(rsp.status, ErrorCode::Success as i32);

        assert_eq!(meta_store.shard_updates.lock().unwrap().len(), 1);
        assert_eq!(shard.progress.cdc_s3_last_file, "cdc-0013");
        assert_eq!(shard.progress.persisted_cdc_s3_last_file, "cdc-0012");
    }
}