use std::{sync::Arc, time::Duration};

use anyhow::Result;
use either::Either;
use log::info;
use storage_backend::storage::create_backend_storage_for_test;
use tici::{
    HandleRange, ShardInfo,
    worker::{
        mocks::{MockShardWriterManager, make_frag_meta_by_range, make_frag_meta_extended},
        server::DispatchShardWriter,
    },
};
use tici_common::{codec::previous_key_helper, index::IndexKey};
use tici_grpc::{
    AppendBaseFragMetaRequest, AppendDeltaFragMetaRequest, AppendFragMetaResponse, ErrorCode,
    FragMeta, ShardProgress,
};

use crate::test_utils::{
    AppendExtraParam, create_cdc_file_name, mock_frag_path, must_finally_ex,
    test_suite::E2ETestSuite,
};

#[derive(Clone)]
pub struct ExpectedAppendFragRes {
    pub do_manifest_check: bool,
    pub expected_seq: u64,
    pub expected_delta_len: usize,
    pub expected_base_len: usize,
    pub appended: Vec<FragMeta>,
    pub shard_must_existed: bool,
}

pub struct RealAppendFragRes {
    pub shard_id: u64,
    pub index_key: IndexKey,
    pub frags: Vec<FragMeta>,
    pub original_shard: ShardInfo,
    pub compaction_finish_before_append: u64,
}

#[derive(Debug)]
pub enum AppendDeltaFragError {
    RpcCall {
        shard_id: u64,
        source: anyhow::Error,
    },
    ResponseCheck {
        shard_id: u64,
        source: anyhow::Error,
    },
    AppendEventually {
        shard_id: u64,
        last_error: i32,
    },
    PostCheck {
        shard_id: u64,
        source: anyhow::Error,
    },
}

impl std::fmt::Display for AppendDeltaFragError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::RpcCall { shard_id, source } => {
                write!(f, "append rpc failed, shard_id={}, err={source}", shard_id)
            }
            Self::ResponseCheck { shard_id, source } => write!(
                f,
                "append response check failed, shard_id={}, err={source}",
                shard_id
            ),
            Self::AppendEventually {
                shard_id,
                last_error,
            } => write!(
                f,
                "can't append eventually, shard_id={}, last_error={}",
                shard_id, last_error
            ),
            Self::PostCheck { shard_id, source } => write!(
                f,
                "append post-check failed, shard_id={}, err={source}",
                shard_id
            ),
        }
    }
}

impl std::error::Error for AppendDeltaFragError {}

impl AppendDeltaFragError {
    pub fn shard_id(&self) -> u64 {
        match self {
            Self::RpcCall { shard_id, .. }
            | Self::ResponseCheck { shard_id, .. }
            | Self::AppendEventually { shard_id, .. }
            | Self::PostCheck { shard_id, .. } => *shard_id,
        }
    }
}

impl RealAppendFragRes {
    pub fn new(
        shard_id: u64,
        index_key: IndexKey,
        frags: Vec<FragMeta>,
        original_shard: ShardInfo,
        compaction_finish_before_append: u64,
    ) -> Self {
        RealAppendFragRes {
            shard_id,
            index_key,
            frags,
            original_shard,
            compaction_finish_before_append,
        }
    }
}

impl E2ETestSuite {
    pub async fn check_append_frag(
        &self,
        shard_id: u64,
        index_key: &IndexKey,
        last_cdc_file: Option<String>,
        expected: ExpectedAppendFragRes,
    ) -> Result<()> {
        let m_s = match self.get_shard_from_restore(shard_id).await {
            Ok(o) => o,
            Err(e) => {
                if expected.shard_must_existed {
                    return Err(e);
                }
                return Ok(());
            }
        };
        if let Some(cdc_file) = last_cdc_file {
            assert_eq!(m_s.progress.cdc_s3_last_file, cdc_file);
        }
        if expected.do_manifest_check {
            assert_eq!(m_s.seq, expected.expected_seq);
            let m_manifest = m_s.manifest.clone();
            assert_eq!(m_manifest.delta_slice().len(), expected.expected_delta_len);
            assert_eq!(m_manifest.base_slice().len(), expected.expected_base_len);
        }
        #[cfg(feature = "get_manifest_from_writer")]
        {
            let range = m_s.range.clone();
            let debug_get_shard_manifest_req = DebugGetShardManifestRequest {
                keyspace_id: index_key.keyspace_id,
                table_id: index_key.table_id,
                index_id: index_key.index_id,
                key_range: Some(KeyRange {
                    start_key: range.encoded_start(),
                    end_key: range.encoded_end(),
                }),
            };
            let res = self
                .get_meta_internal_client()
                .debug_get_shard_manifest(debug_get_shard_manifest_req)
                .await?;
            let res = res.into_inner();
            if res.status == ErrorCode::Success.to_i32() {
                assert_ne!(res.frag_metas.len(), 0);
            }
        }
        #[cfg(not(feature = "get_manifest_from_writer"))]
        {
            let get_shard_manifest_req = tici_grpc::GetShardManifestRequest {
                keyspace_id: index_key.keyspace_id,
                table_id: index_key.table_id,
                index_id: index_key.index_id,
                shard_id,
            };
            let check = must_finally_ex(async |_| {
                let res = self
                    .get_meta_internal_client()
                    .get_shard_manifest(get_shard_manifest_req.clone())
                    .await
                    .unwrap();
                if res.get_ref().status == ErrorCode::ShardNotFound as i32 {
                    // already splitted, skip check.
                    return Either::Left(());
                }
                let res = res.into_inner();
                if res.status != ErrorCode::Success.to_i32() {
                    return Either::Right(format!(
                        "manifest not correct: res={:?}, shard_id={}",
                        res, shard_id
                    ));
                }
                if res.frag_metas.is_empty() {
                    return Either::Right(format!(
                        "manifest not correct: res={:?}, shard_id={}",
                        res, shard_id
                    ));
                }
                for frag in expected.appended.iter() {
                    if !res.frag_metas.iter().any(|f| f.frag_path == frag.frag_path) {
                        return Either::Right(format!(
                            "manifest not correct: res={:?}, shard_id={}",
                            res, shard_id
                        ));
                    }
                }
                Either::Left(())
            })
            .await;
            match check {
                Either::Left(_) => {}
                Either::Right(info) => anyhow::bail!("{info}"),
            }
        }
        Ok(())
    }

    pub async fn do_append_delta_frag(
        &self,
        addr: &str,
        shard_id: u64,
        index_key: &IndexKey,
        last_cdc_file: String,
        frag_metas: Vec<FragMeta>,
    ) -> Result<(AppendDeltaFragMetaRequest, AppendFragMetaResponse)> {
        let wn = self.workers.get(addr);
        let wn = if let Some(w) = wn {
            w
        } else {
            return Err(anyhow::anyhow!(
                "do_append_delta_frag: not found wn={}",
                addr,
            ));
        };
        let r = match &wn.shard_write_manager {
            DispatchShardWriter::Mock(mock) => {
                mock.mock_send_append_delta_frag(
                    index_key,
                    shard_id,
                    last_cdc_file,
                    frag_metas.clone(),
                )
                .await?
            }
            DispatchShardWriter::Real(_) => {
                // Directly send a request to meta service
                let append_req = AppendDeltaFragMetaRequest {
                    keyspace_id: index_key.keyspace_id,
                    table_id: index_key.table_id,
                    index_id: index_key.index_id,
                    shard_id,
                    worker_node_addr: self
                        .get_worker_node_config(&addr.to_string())
                        .await
                        .addr
                        .clone(),
                    frag_metas: frag_metas.clone(),
                    progress: Some(ShardProgress { last_cdc_file }),
                };
                let res = self
                    .get_meta_internal_client()
                    .append_delta_frag_meta(append_req.clone())
                    .await?;
                assert_eq!(res.get_ref().status, 0);
                (append_req, res.into_inner())
            }
        };
        Ok(r)
    }

    pub async fn do_append_base_frag(
        &self,
        addr: &String,
        index_key: &IndexKey,
        frag_metas: Vec<FragMeta>,
    ) -> Result<()> {
        let wn = self.workers.get(addr);
        let wn = if let Some(w) = wn {
            w
        } else {
            return Err(anyhow::anyhow!(
                "do_append_base_frag: not found wn={}",
                addr,
            ));
        };
        match &wn.shard_write_manager {
            DispatchShardWriter::Mock(mock) => {
                mock.mock_send_append_base_frag(index_key, frag_metas.clone())
                    .await?;
            }
            DispatchShardWriter::Real(_) => {
                // Directly send a request to meta service
                let append_req = AppendBaseFragMetaRequest {
                    keyspace_id: index_key.keyspace_id,
                    table_id: index_key.table_id,
                    index_id: index_key.index_id,
                    frag_metas: frag_metas.clone(),
                };
                let res = self
                    .get_meta_internal_client()
                    .append_base_frag_meta(append_req)
                    .await?;
                assert_eq!(res.get_ref().status, 0);
            }
        }
        Ok(())
    }

    async fn mock_upload_frag_to_s3(&self, frag_prefix: &str, frag_metas: &[FragMeta]) {
        let config = &self.get_meta_handle().config;
        let storage = create_backend_storage_for_test(&config.s3).await.unwrap();
        // TODO(P2) This only upload empty files, which can't be read.
        for frag in frag_metas {
            let remote_path = format!("{}/{}/{}", config.s3.prefix, frag_prefix, frag.frag_path);
            storage
                .put_data(&config.s3.bucket, &remote_path, Vec::new())
                .await
                .unwrap();
            info!(
                "try append frag and upload {}/{}",
                config.s3.bucket, remote_path
            );
        }
    }

    pub fn check_append_response(shard_id: u64, status: i32) -> Result<bool> {
        if matches!(status, const { tici_grpc::ErrorCode::Success as i32 }) {
            Ok(true)
        } else if matches!(
            status,
            const { tici_grpc::ErrorCode::FragRangeMismatch as i32 }
                | const { tici_grpc::ErrorCode::ShardNotFound as i32 }
                | const { tici_grpc::ErrorCode::ShardNotScheduled as i32 }
        ) {
            Ok(false)
        } else {
            anyhow::bail!(
                "critical append error, status={}, shard_id={}",
                status,
                shard_id
            );
        }
    }

    async fn must_do_append_delta_frag<F, Fut>(
        &self,
        addr: &str,
        last_cdc_file: String,
        generate: F,
        upload_to_s3: bool,
    ) -> std::result::Result<RealAppendFragRes, AppendDeltaFragError>
    where
        F: FnOnce() -> Fut + Send + 'static + Clone,
        Fut: std::future::Future<Output = (u64, IndexKey, Vec<FragMeta>)> + Send + 'static,
    {
        let mut att = 0;
        loop {
            let gen_once = generate.clone();
            let (shard_id, index_key, frag_metas) = gen_once().await;
            let s_origin = match self.get_shard_from_index(&index_key, shard_id).await {
                Ok(s) => s.read().await.clone(),
                Err(_) => {
                    info!("can't find shard {} in index {:?}", shard_id, &index_key);
                    tokio::time::sleep(Duration::from_millis(100)).await;
                    continue;
                }
            };

            let compaction_finish_before_append = s_origin.stats.get_compaction_finish();
            if upload_to_s3 {
                self.mock_upload_frag_to_s3("frags", frag_metas.as_slice())
                    .await;
            } else {
                info!("try append frag {:?}", &frag_metas);
            }
            let mut last_error = 0;
            let done = match self
                .do_append_delta_frag(
                    addr,
                    shard_id,
                    &index_key,
                    last_cdc_file.clone(),
                    frag_metas.clone(),
                )
                .await
            {
                Err(e) => {
                    // Rare errors directly raised by the grpc call.
                    return Err(AppendDeltaFragError::RpcCall {
                        shard_id,
                        source: e,
                    });
                }
                Ok((_, resp)) => {
                    // We check different ErrorCodes.
                    let r = match Self::check_append_response(shard_id, resp.status) {
                        Ok(ok) => ok,
                        Err(e) => {
                            return Err(AppendDeltaFragError::ResponseCheck {
                                shard_id,
                                source: e,
                            });
                        }
                    };
                    if !r {
                        last_error = resp.status;
                    }
                    r
                }
            };
            att += 1;
            if att > 10 {
                return Err(AppendDeltaFragError::AppendEventually {
                    shard_id,
                    last_error,
                });
            }
            if done {
                return Ok(RealAppendFragRes::new(
                    shard_id,
                    index_key,
                    frag_metas,
                    s_origin,
                    compaction_finish_before_append,
                ));
            }
        }
    }

    pub async fn must_append_delta_by_shard(
        &self,
        shard_info: &ShardInfo,
        path: &str,
        cdc_index: u64,
        extra: AppendExtraParam,
    ) -> FragMeta {
        let shard_id = shard_info.shard_id;
        let index_key = shard_info.index_key.clone();
        let path = mock_frag_path(path);
        let new_range_clone = HandleRange::new(
            shard_info.range.encoded_start(),
            previous_key_helper(shard_info.range.encoded_end().as_slice(), 1),
        );
        let frag = if extra.count.is_some() {
            make_frag_meta_extended(
                path.as_str(),
                &new_range_clone,
                extra.size.unwrap(),
                extra.count.unwrap(),
                extra.segs_count.unwrap(),
            )
        } else {
            make_frag_meta_by_range(path.as_str(), &new_range_clone)
        };
        let frag_clone = frag.clone();
        self.must_append_delta_frag(
            shard_info.writer.as_str(),
            create_cdc_file_name(cdc_index).to_string(),
            async move || (shard_id, index_key, vec![frag_clone]),
            extra,
        )
        .await
        .unwrap();
        frag
    }

    pub async fn must_append_delta_by_mocker(
        &self,
        mock: Arc<MockShardWriterManager>,
        path: Either<String, String>,
        cdc_index: u64,
        extra: AppendExtraParam,
        default_size_batch: u64,
        default_rows_batch: u64,
    ) -> std::result::Result<RealAppendFragRes, AppendDeltaFragError> {
        let append_addr = self.get_one_worker_node_addr();
        self.must_append_delta_frag(
            append_addr.as_str(),
            create_cdc_file_name(cdc_index).to_string(),
            async move || {
                let (shard_meta, frag_metas) = mock
                    .mock_new_frags(path, default_size_batch, default_rows_batch, 4)
                    .await;
                (
                    shard_meta.shard_header.shard_id,
                    shard_meta.index_key.clone(),
                    frag_metas,
                )
            },
            extra,
        )
        .await
    }

    async fn must_append_delta_frag<F, Fut>(
        &self,
        addr: &str,
        last_cdc_file: String,
        generate: F,
        extra: AppendExtraParam,
    ) -> std::result::Result<RealAppendFragRes, AppendDeltaFragError>
    where
        F: FnOnce() -> Fut + Send + 'static + Clone,
        Fut: std::future::Future<Output = (u64, IndexKey, Vec<FragMeta>)> + Send + 'static,
    {
        let r = self
            .must_do_append_delta_frag(addr, last_cdc_file.clone(), generate, extra.upload_to_s3)
            .await;
        if extra.call_check_append_frag {
            let dedup = extra.dedup;
            match r {
                Ok(real_res) => {
                    let shard_id = real_res.shard_id;
                    let frags_len = real_res.frags.len();
                    let finish_after_append = match self
                        .get_shard_from_index(&real_res.index_key, shard_id)
                        .await
                    {
                        Ok(shard) => shard.read().await.stats.get_compaction_finish(),
                        Err(_) => {
                            // Shard could be split/merged; skip strict check.
                            return Ok(real_res);
                        }
                    };
                    if finish_after_append != real_res.compaction_finish_before_append {
                        // A compaction finished concurrently and may have changed shard
                        // topology/manifest/seq. Skip strict append checks.
                        return Ok(real_res);
                    }
                    let s_origin = real_res.original_shard.clone();
                    let appended_count = frags_len.saturating_sub(dedup.0);
                    let expected = ExpectedAppendFragRes {
                        do_manifest_check: extra.do_manifest_check,
                        expected_seq: s_origin.seq + if appended_count > 0 { 1 } else { 0 },
                        expected_delta_len: s_origin.manifest.delta_slice().len() + appended_count,
                        expected_base_len: s_origin.manifest.base_slice().len() - dedup.1,
                        appended: real_res.frags.clone(),
                        shard_must_existed: false,
                    };
                    if let Err(e) = self
                        .check_append_frag(
                            shard_id,
                            &real_res.index_key,
                            Some(last_cdc_file),
                            expected,
                        )
                        .await
                    {
                        return Err(AppendDeltaFragError::PostCheck {
                            shard_id,
                            source: e,
                        });
                    }
                    Ok(real_res)
                }
                Err(e) => Err(e),
            }
        } else {
            r
        }
    }

    pub async fn must_append_base_frag(
        &self,
        addr: &String,
        shard_id: u64,
        index_key: &IndexKey,
        frag_metas: Vec<FragMeta>,
    ) -> Result<()> {
        self.must_append_base_frag_impl(addr, shard_id, index_key, frag_metas, (0, 0))
            .await
    }

    pub async fn must_append_base_frag_impl(
        &self,
        addr: &String,
        shard_id: u64,
        index_key: &IndexKey,
        frag_metas: Vec<FragMeta>,
        dedup: (usize, usize),
    ) -> Result<()> {
        let s_origin = self.get_shard_from_restore(shard_id).await?;
        self.do_append_base_frag(addr, index_key, frag_metas.clone())
            .await?;
        let len = frag_metas.len();
        let appended_count = len.saturating_sub(dedup.1);
        let expected = ExpectedAppendFragRes {
            do_manifest_check: true,
            expected_seq: s_origin.seq + if appended_count > 0 { 1 } else { 0 },
            expected_delta_len: s_origin.manifest.delta_slice().len() - dedup.0,
            expected_base_len: s_origin.manifest.base_slice().len() + appended_count,
            appended: frag_metas,
            shard_must_existed: false,
        };
        self.check_append_frag(shard_id, index_key, None, expected)
            .await?;
        Ok(())
    }
}