use std::collections::HashSet;

use anyhow::Result;
use log::info;
use tici::{
    self, SplittableFrag,
    meta::writer::{CompactShardReason, CompactionScheduleError, CompactionState},
    worker::mocks::*,
};
use tici_common::index::IndexKey;
use tici_grpc::{ErrorCode, FinishCompactFragRequest, FinishCompactFragResponse, FragMeta};

use crate::test_utils::{mock_frag_path, test_suite::E2ETestSuite};

impl E2ETestSuite {
    pub async fn do_compact_scheduled(
        &self,
        index_key: &IndexKey,
        shard_id: u64,
        input: Vec<SplittableFrag>,
        do_check: bool,
    ) -> Result<Option<CompactionState>> {
        info!("do_compact_scheduled: shard_id={}", shard_id);
        let (ind, shard) = self.get_shard_and_index(index_key, shard_id).await?;
        let (prev, state) = {
            let mut g = shard.write().await;
            let prev = g.progress.compaction_state.len();
            let mut scheduled = vec![];
            match self
                .get_meta_service()
                .shard_task_manager
                .debug_schedule_compaction_if_not_overlapped(
                    ind,
                    &mut scheduled,
                    &mut g,
                    input,
                    CompactShardReason::Manual,
                    0,
                )
                .await
            {
                Ok(()) => (prev, scheduled.first().cloned()),
                Err(CompactionScheduleError::OverlapDetected { .. }) => {
                    // Overlap is expected in some tests, return None
                    info!("Compaction skipped due to overlap (expected in tests)");
                    (prev, None)
                }
                Err(CompactionScheduleError::Other(e)) => return Err(e),
            }
        };
        let shard = self.get_shard_from_index(index_key, shard_id).await?;
        if do_check {
            assert!(state.is_some());
            assert_eq!(shard.read().await.progress.compaction_state.len(), prev + 1);
        }
        Ok(state)
    }

    pub async fn do_compact_finished_impl(
        &self,
        shard_id: u64,
        state: CompactionState,
        output: Vec<FragMeta>,
    ) -> Result<(Result<FinishCompactFragResponse>, usize)> {
        // Must clear the worker node part. Otherwise, suggest to use
        // `mock_compact_process`.
        let _ = self
            .get_worker_node_handle(&state.worker)
            .worker_node_service
            .compaction_status
            .write()
            .await
            .remove(&state.get_compaction_state_key());
        let shard_info = self.get_shard_from_restore(shard_id).await?;
        let prev = self
            .collect_worker_nodes()
            .await
            .first()
            .unwrap()
            .1
            .compaction_len()
            .await;

        let req = FinishCompactFragRequest {
            addr: state.worker.clone(),
            shard: Some(shard_info.get_header()),
            keyspace_id: shard_info.keyspace_id(),
            table_id: shard_info.table_id(),
            index_id: shard_info.index_id(),
            seq: state.seq,
            timestamp: state.get_start_ts(),
            leader_term: String::new(),
            input_frags: state
                .input_fragments
                .iter()
                .map(|x| x.f.clone())
                .collect::<Vec<_>>(),
            output_frags: output.clone(),
        };
        let resp = self
            .get_meta_service()
            .worker_node_mgr
            .on_compaction_finish(&req)
            .await;
        Ok((resp, prev))
    }

    pub async fn do_compact_finished_fail(
        &self,
        shard_id: u64,
        state: CompactionState,
        output: Vec<FragMeta>,
    ) -> Result<()> {
        let (resp, _) = self
            .do_compact_finished_impl(shard_id, state.clone(), output)
            .await?;
        if let Ok(resp) = resp {
            let suc: i32 = ErrorCode::Success.into();
            if resp.status == suc {
                return Err(anyhow::anyhow!(
                    "do_compact_finished_impl: failed {:?}",
                    resp
                ));
            }
        }
        Ok(())
    }

    pub async fn do_compact_finished(
        &self,
        shard_id: u64,
        state: CompactionState,
        output: Vec<FragMeta>,
        do_check: bool,
    ) -> Result<()> {
        let config = self.get_meta_service().config.load();
        let skip_compaction_queue_check = config.shard.compaction_policy == "uni"
            && state.remaining_iterations >= config.shard.uni_compaction.max_compaction_iterations;
        let (res, prev) = self
            .do_compact_finished_impl(shard_id, state.clone(), output)
            .await?;
        res?;
        {
            let shard = self.get_shard_from_restore(shard_id).await?;
            let set = shard
                .manifest
                .get_fragments()
                .iter()
                .map(|f| f.frag_path.clone())
                .collect::<HashSet<_>>();
            for frag in state.input_fragments.iter() {
                assert!(!set.contains(&frag.frag_path));
            }
            for frag in shard.manifest.get_fragments().iter() {
                assert!(set.contains(&frag.frag_path));
            }
            if !skip_compaction_queue_check {
                assert_eq!(shard.progress.compaction_state.len(), 0);
            }
        }
        if do_check && !skip_compaction_queue_check {
            assert_eq!(
                self.collect_worker_nodes()
                    .await
                    .first()
                    .unwrap()
                    .1
                    .compaction_len()
                    .await
                    + 1,
                prev
            );
        }
        Ok(())
    }

    pub async fn check_compact_aborted(
        &self,
        shard_id: u64,
        input: Vec<SplittableFrag>,
    ) -> Result<()> {
        let shard = self.get_shard_from_restore(shard_id).await?;
        assert_eq!(shard.manifest.len(), input.len());
        for (i, frag) in input.iter().enumerate() {
            assert_eq!(shard.manifest.get_fragments()[i].frag_path, frag.frag_path);
        }
        if let Some((_, c)) = self.collect_worker_nodes().await.first() {
            assert_eq!(c.compaction_len().await, 0);
        }
        assert_eq!(shard.progress.compaction_state.len(), 0);
        Ok(())
    }

    pub async fn must_compact_all(
        &self,
        index_key: &IndexKey,
        shard_id: u64,
        new_prefix: &str,
    ) -> Result<()> {
        let shard = self
            .get_shard_from_index(index_key, shard_id)
            .await
            .unwrap();
        let input = shard.read().await.manifest.delta_slice().to_vec();
        let new_fragment_range = shard.read().await.range.debug_derive_frag_range();
        let output = vec![make_frag_meta_by_range(
            &mock_frag_path(format!("{}{}", new_prefix, shard_id)),
            &new_fragment_range,
        )];
        let state = self
            .do_compact_scheduled(index_key, shard_id, input, true)
            .await?;
        self.do_compact_finished(shard_id, state.unwrap(), output, true)
            .await?;
        Ok(())
    }
}