use std::{
    fmt::Debug,
    sync::{
        Arc,
        atomic::{AtomicU64, Ordering},
    },
    time::{Duration, SystemTime},
};

use anyhow::anyhow;
use derivative::Derivative;
use either::Either;
use heapsize::HeapSizeOf;
use log::{debug, error, info, warn};
use tici_common::{cdc_file_name::older_cdc_file, config::S3Config, index::IndexKey};
use tici_grpc::{AppendDeltaFragMetaRequest, ErrorCode, FragMeta, S3Location, ShardManifestHeader};

use crate::{
    HandleRange, RangeBound, ShardManifest, SplittableFrag,
    config::{ShardConfig, UniCompactionConfig},
    meta::{
        MetaStore,
        metrics::SHARD_EVENT_COUNTER_VEC,
        writer::{
            CompactShardReason, CompactionState, CompactionStateKey, FRAG_FILE_PREFIX, FTSIndex,
            ShardBuildProgress, ShardSplitReason, ShardStats,
        },
    },
};

#[derive(Clone, Derivative)]
#[derivative(PartialEq)]
pub struct ShardInfo {
    pub index_key: IndexKey,
    // NOTE `shard_id` is not guaranteed to be increment, although it is marked as AUTO_IMCREMENT
    // in the storage table. We have to use `epoch` to get the partial ordering of two shards
    // if they overlapped with each other.
    pub shard_id: u64,
    pub range: HandleRange,
    pub epoch: u64,
    pub seq: u64,
    pub manifest: ShardManifest,
    // Used by `add_shard` and `get_index_progress`.
    pub writer: String,
    pub progress: ShardBuildProgress,
    pub frag_upload_location: S3Location,
    #[derivative(PartialEq = "ignore")]
    pub stats: ShardStats,
    // Runtime-only counter for empty appends since the last persisted append progress.
    #[derivative(PartialEq = "ignore")]
    pub append_empty_persist_count: Arc<AtomicU64>,
}

impl heapsize::HeapSizeOf for ShardInfo {
    fn heap_size_of_children(&self) -> usize {
        self.manifest.heap_size_of_children()
            + self.writer.capacity()
            + self.progress.heap_size_of_children()
            + self.frag_upload_location.heap_size_of_children()
    }
}

impl Debug for ShardInfo {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ShardInfo")
            .field("index", &self.index_key)
            .field_with("shard", |f| {
                f.write_fmt(format_args!(
                    "shard_id={}/range={:?}/epoch={}/seq={}",
                    &self.shard_id, self.range, self.epoch, self.seq
                ))
            })
            .field("manifest", &self.manifest.len())
            .field_with("writer", |f| f.write_fmt(format_args!("{}", &self.writer)))
            .field("progress", &self.progress)
            .field("frag_upload_location", &self.frag_upload_location)
            .field("stats", &self.stats)
            .field_with("total_(bytes|rows)", |f| {
                f.write_fmt(format_args!(
                    "{}/{}",
                    self.get_total_estimated_bytes(),
                    self.get_total_estimated_rows()
                ))
            })
            .finish()
    }
}

impl ShardInfo {
    fn new_empty_append_persist_count() -> Arc<AtomicU64> {
        Arc::new(AtomicU64::new(0))
    }

    fn reset_runtime_state_for_new_shard(&mut self) {
        self.append_empty_persist_count = Self::new_empty_append_persist_count();
    }

    pub fn new_for_test(range: HandleRange) -> Self {
        let default_s3 = S3Config::default();
        Self {
            index_key: IndexKey::new_nullspace(0, 0),
            shard_id: 0,
            range,
            epoch: 0,
            seq: 0,
            manifest: ShardManifest::new_for_first(),
            writer: String::new(),
            progress: ShardBuildProgress::default(),
            frag_upload_location: S3Location {
                bucket: default_s3.bucket,
                prefix: format!("{}/{}", default_s3.prefix, FRAG_FILE_PREFIX),
            },
            stats: ShardStats::default(),
            append_empty_persist_count: Self::new_empty_append_persist_count(),
        }
    }

    pub fn to_long_ident(&self) -> String {
        format!(
            "({:?}, shard_id={}, start_key={:?})",
            self.index_key,
            self.shard_id,
            self.range.sortable_start_ref()
        )
    }

    pub fn keyspace_id(&self) -> u32 {
        self.index_key.keyspace_id
    }

    pub fn table_id(&self) -> i64 {
        self.index_key.table_id
    }

    pub fn index_id(&self) -> i64 {
        self.index_key.index_id
    }

    pub fn get_upload_location(&self) -> S3Location {
        self.frag_upload_location.clone()
    }

    pub fn find_overlapping_compaction(
        &self,
        input_fragments: &[SplittableFrag],
    ) -> Option<&CompactionState> {
        for prev_compaction in self.progress.compaction_state.values() {
            if prev_compaction.overlap_with(input_fragments) {
                debug!(
                    "Overlap detected: shard_id={}, input={:?}, conflicting={:?}",
                    self.shard_id,
                    input_fragments
                        .iter()
                        .map(|f| &f.frag_path)
                        .collect::<Vec<_>>(),
                    prev_compaction
                        .input_fragments
                        .iter()
                        .map(|f| &f.frag_path)
                        .collect::<Vec<_>>()
                );
                return Some(prev_compaction);
            }
        }
        None
    }

    // Keep backward compatibility helper
    pub fn is_overlapped_with_running_compaction(
        &self,
        input_fragments: &[SplittableFrag],
    ) -> bool {
        self.find_overlapping_compaction(input_fragments).is_some()
    }
}

impl ShardInfo {
    pub fn current_ts() -> Duration {
        SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap_or_default()
    }

    pub fn get_header(&self) -> ShardManifestHeader {
        ShardManifestHeader {
            shard_id: self.shard_id,
            start_key: self.range.encoded_start(),
            end_key: self.range.encoded_end(),
            epoch: self.epoch,
            seq: self.seq,
        }
    }

    pub fn get_total_estimated_bytes(&self) -> u64 {
        self.manifest.get_total_estimated_bytes()
    }

    pub fn get_total_estimated_rows(&self) -> u64 {
        self.manifest.get_total_estimated_rows()
    }

    pub fn get_total_datafile_count(&self) -> usize {
        self.manifest.get_total_datafile_count()
    }

    pub fn get_total_fragment_count(&self) -> usize {
        self.manifest.get_fragments().len()
    }

    pub fn should_persist_empty_append(&self, update: &ShardInfoUpdate, freq: u64) -> bool {
        let Some(appended_frags) = update.appended_frags.as_ref() else {
            return true;
        };
        if !appended_frags.is_empty() || update.progress.is_none() || update.seq.is_some() {
            return true;
        }
        if freq <= 1 {
            return true;
        }

        let count = self
            .append_empty_persist_count
            .fetch_add(1, Ordering::Relaxed)
            + 1;
        if count >= freq {
            self.reset_empty_append_persist_count();
            true
        } else {
            false
        }
    }

    pub fn reset_empty_append_persist_count(&self) {
        self.append_empty_persist_count.store(0, Ordering::Relaxed);
    }

    pub fn should_split(&self, config: &ShardConfig, reason: ShardSplitReason) -> bool {
        match reason {
            ShardSplitReason::NewAppend => {
                self.get_total_estimated_bytes() >= config.get_split_size()
            }
            ShardSplitReason::WritePressure => {
                self.get_total_estimated_bytes() as f64
                    >= config.get_split_size() as f64 * config.split.split_by_append_factor
            }
            ShardSplitReason::Manual => true,
        }
    }

    pub fn can_compaction(&self, config: &ShardConfig) -> bool {
        // Can't have more than `compaction_concurrency_shard` concurrent compaction
        // tasks.
        self.progress.compaction_state.len()
            < config.compactor.compaction_concurrency_shard as usize
    }
}

impl ShardInfo {
    fn is_frag_in_existing_task(&self, frag_path: &str) -> bool {
        self.progress.compaction_state.values().any(|state| {
            state
                .input_fragments
                .iter()
                .any(|f| f.frag_path == frag_path)
        })
    }

    fn is_better_fallback_candidate(
        candidate_left: usize,
        candidate_max: u64,
        candidate_sum: u64,
        candidate_count: usize,
        best_left: usize,
        best_max: u64,
        best_sum: u64,
        best_count: usize,
    ) -> bool {
        // Compare largest_ratio = max / sum without floats.
        // Treat sum==0 as ratio==1.0 so that it never wins over sum>0.
        let (candidate_max, candidate_sum) = if candidate_sum == 0 {
            (1u64, 1u64)
        } else {
            (candidate_max, candidate_sum)
        };
        let (best_max, best_sum) = if best_sum == 0 {
            (1u64, 1u64)
        } else {
            (best_max, best_sum)
        };

        let left_is_better_ratio = (candidate_max as u128) * (best_sum as u128)
            < (best_max as u128) * (candidate_sum as u128);
        let left_is_equal_ratio = (candidate_max as u128) * (best_sum as u128)
            == (best_max as u128) * (candidate_sum as u128);

        left_is_better_ratio
            || (left_is_equal_ratio && candidate_count > best_count)
            || (left_is_equal_ratio && candidate_count == best_count && candidate_left < best_left)
    }

    pub fn merge_overlap_compaction_inputs(&self, policy: &str) -> Option<Vec<SplittableFrag>> {
        let delta = self.manifest.delta_slice();
        if delta.len() <= 1 {
            return None;
        }
        match policy {
            "uni" => Some(delta.to_vec()),
            _ => None,
        }
    }

    /// Pick a compaction window from the delta manifest tail based on the
    /// universal compaction policy described in
    /// `prompts/20260226-new-compaction.md`.
    ///
    /// - Fix the right boundary at the newest delta fragment.
    /// - Enumerate windows by extending to the left and select the widest
    ///   window with `largest_ratio <= max_largest_ratio`.
    /// - While extending, stop when encountering a fragment that belongs to an
    ///   existing task.
    /// - Fallback: if no window satisfies the ratio constraint, pick the window
    ///   with minimal `largest_ratio` to avoid stalling.
    /// - Only windows with at least `min_compaction_frags` are considered as
    ///   candidates.
    pub fn pick_uni_compaction(&self, config: &UniCompactionConfig) -> Vec<SplittableFrag> {
        let delta = self.manifest.delta_slice();
        let Some(right) = delta.len().checked_sub(1) else {
            return vec![];
        };
        self.pick_uni_compaction_with_right(config, right, true)
    }

    /// Pick a compaction window with the right boundary fixed to a specific
    /// fragment path. The boundary must be in the delta layer.
    pub fn pick_uni_compaction_at(
        &self,
        config: &UniCompactionConfig,
        frag_path: &str,
        allow_fallback: bool,
    ) -> Vec<SplittableFrag> {
        use crate::common::shard::FragmentLocation;

        let frag_path = frag_path.to_string();
        let Some(location) = self.manifest.find_pos_by_frag_path(&frag_path) else {
            return vec![];
        };
        let right = match location {
            FragmentLocation::Delta(pos) => pos,
            FragmentLocation::Base(_) => return vec![],
        };
        self.pick_uni_compaction_with_right(config, right, allow_fallback)
    }

    fn pick_uni_compaction_with_right(
        &self,
        config: &UniCompactionConfig,
        right: usize,
        allow_fallback: bool,
    ) -> Vec<SplittableFrag> {
        if config.min_compaction_frags == 0 {
            return vec![];
        }
        if !(0.0 < config.max_largest_ratio && config.max_largest_ratio <= 1.0) {
            return vec![];
        }

        let delta = self.manifest.delta_slice();
        if right >= delta.len() {
            return vec![];
        }

        let mut sum_bytes: u64 = 0;
        let mut max_bytes: u64 = 0;
        let mut count: usize = 0;

        let mut best_passing_left: Option<usize> = None;
        let mut best_fallback_left: Option<(usize, u64, u64, usize)> = None;

        for left in (0..=right).rev() {
            if self.is_frag_in_existing_task(delta[left].frag_path.as_str()) {
                break;
            }
            let size = delta[left].weighted_bytes();
            sum_bytes = sum_bytes.saturating_add(size);
            max_bytes = max_bytes.max(size);
            count += 1;

            let ratio = if sum_bytes == 0 {
                1.0
            } else {
                max_bytes as f64 / sum_bytes as f64
            };

            if count >= config.min_compaction_frags {
                if ratio <= config.max_largest_ratio {
                    best_passing_left = Some(left);
                }

                if allow_fallback {
                    best_fallback_left = match best_fallback_left {
                        None => Some((left, max_bytes, sum_bytes, count)),
                        Some((best_left, best_max, best_sum, best_count)) => {
                            if Self::is_better_fallback_candidate(
                                left, max_bytes, sum_bytes, count, best_left, best_max, best_sum,
                                best_count,
                            ) {
                                Some((left, max_bytes, sum_bytes, count))
                            } else {
                                Some((best_left, best_max, best_sum, best_count))
                            }
                        }
                    };
                }
            }
        }

        let selected_left = if let Some(left) = best_passing_left {
            Some(left)
        } else if allow_fallback {
            best_fallback_left.map(|(left, ..)| left)
        } else {
            None
        };

        let Some(left) = selected_left else {
            return vec![];
        };

        let picked = delta[left..=right].to_vec();
        if picked.len() < config.min_compaction_frags {
            vec![]
        } else {
            picked
        }
    }
}

impl ShardInfo {
    pub fn new_from_merge(
        writer: String,
        input: &[&ShardInfo],
    ) -> anyhow::Result<Either<Self, (SplittableFrag, SplittableFrag)>> {
        #[allow(suspicious_double_ref_op)]
        let mut new_shard = if let Some(first) = input.first() {
            first.clone().clone()
        } else {
            return Err(anyhow!("empty shards"));
        };
        new_shard.epoch = input
            .iter()
            .map(|x| x.epoch)
            .max()
            .ok_or(anyhow!("can not find max epoch, {:?}", input))?
            + 1;
        new_shard.seq = 0;
        new_shard.range = HandleRange::new_from_sorted(
            input
                .first()
                .ok_or(anyhow!("can not find first, {:?}", input))?
                .range
                .sortable_start(),
            input
                .last()
                .ok_or(anyhow!("can not find last, {:?}", input))?
                .range
                .sortable_end(),
        );
        new_shard.writer = writer.clone();
        for shard in input.iter().skip(1) {
            new_shard.progress.merge(&shard.progress);
        }
        // Must be called from the left side shard to the right side shard to keep
        // order.
        for shard in input.iter().skip(1) {
            // We can merge by stack only if fragments do not overlap between shards.
            // Consider two shards:
            // (Older fragments/indexes come first)
            // (del: [a] means to delete all `a`s before this fragment)
            // - s1: [f0, f1], range: [a, c], s2: [f0, f1], range: [d, e]
            // - f0: index: [a,b,c,d,e], del: [a]
            // - f1: index: [a], del: [a]
            // Then s1 compacted f0 while s2 compacted f1:
            // - s1: [f0_1, f1], s2: [f0, f1_1]
            // - f0_1: index: [a,b,c], del: [a]
            // - f1_1: index: [], del: []
            // Then s1 and s2 merged:
            // - s1.f0_1: index: [a,b,c], del: [a]
            // - s1.f1: index: [a], del: [a]
            // - s2.f0: index: [a,b,c,d,e], del: [a]
            // The `del: [a]` in s2.f0 will kill all `a`s in s1.f0_1 and s1.f1.
            if let Some((left, right)) = new_shard.manifest.merge(&shard.manifest) {
                SHARD_EVENT_COUNTER_VEC
                    .with_label_values(&["merge_failed_overlap"])
                    .inc();
                // TODO(P1) Too many logs
                info!(
                    "merge_shards, fragments overlapped: {:?} <====> {:?}, new_shard={}, old_shard={:?}, old_shards={:?}",
                    left,
                    right,
                    new_shard.to_long_ident(),
                    shard.shard_id,
                    input
                        .iter()
                        .map(|f| format!("{}", f.shard_id))
                        .collect::<Vec<_>>()
                        .join("|"),
                );
                return Ok(Either::Right((left, right)));
            }
        }
        new_shard.reset_runtime_state_for_new_shard();
        Ok(Either::Left(new_shard))
    }

    pub fn new_from_split(
        original: &ShardInfo,
        new_ranges: &[HandleRange],
        writers: Vec<String>,
    ) -> Vec<ShardInfo> {
        let mut shard_infos = vec![];
        let new_epoch: u64 = original.epoch + new_ranges.len() as u64;
        let mut writers = writers.into_iter();
        for range in new_ranges.iter() {
            let mut new_shard_info = original.clone();
            new_shard_info
                .manifest
                .trim_fragments(range, new_ranges.len() as u64);
            new_shard_info.epoch = new_epoch;
            new_shard_info.range = range.clone();
            new_shard_info.seq = 0;
            new_shard_info.stats = ShardStats::split(&original.stats, new_ranges.len() as u64);
            new_shard_info.writer = match writers.next() {
                Some(x) => x,
                None => {
                    warn!(
                        "can't assign new worker for shard, use original, {:?}",
                        new_shard_info
                    );
                    original.writer.to_string()
                }
            };
            new_shard_info.reset_runtime_state_for_new_shard();
            shard_infos.push(new_shard_info);
        }
        shard_infos
    }
}

impl ShardInfo {
    pub fn apply_update(&mut self, update: ShardInfoUpdate) {
        if let Some(s) = update.seq {
            self.seq = s;
        }
        if let Some(s) = update.progress {
            self.progress = s;
        }
        if let Some(s) = update.manifest {
            self.manifest = s;
        }
        if let Some(s) = update.writer {
            self.writer = s;
        }
    }

    // Only used by shard events to replay
    pub fn apply_append(&mut self, append: ShardInfoAppendUpdate) {
        if append.is_base {
            self.manifest
                .append_base_fragments(append.frags, append.shard_id, append.new_seq);
            self.seq = append.new_seq;
            self.progress.cdc_s3_last_file = append.cdc_s3_last_file.clone();
            self.progress.mark_cdc_last_file_persisted();
        } else {
            self.manifest
                .append_delta_fragments(append.frags, append.shard_id, append.new_seq);
            self.seq = append.new_seq;
            self.progress.cdc_s3_last_file = append.cdc_s3_last_file.clone();
            self.progress.mark_cdc_last_file_persisted();
        }
    }

    pub fn process_resche(&self, new_writer: String) -> anyhow::Result<ShardInfoUpdate> {
        Ok(ShardInfoUpdate {
            shard_id: self.shard_id,
            progress: None,
            manifest: None,
            appended_frags: None,
            seq: None,
            compaction: None,
            writer: Some(new_writer),
            size_added: 0,
        })
    }

    pub fn process_append_base(&self, frags: &[FragMeta]) -> ShardInfoUpdate {
        let mut new_manifest = self.manifest.clone();
        let new_seq = self.seq + 1;
        let frags = new_manifest.make_deduplicated(
            frags.to_vec(),
            "append_base_fragments",
            self.shard_id,
            new_seq,
        );
        let mut update = ShardInfoUpdate::from(self);
        if frags.is_empty() {
            info!(
                "append_base_frag_meta: skip empty append shard={}, seq={}",
                self.to_long_ident(),
                self.seq,
            );
        } else {
            new_manifest.append_base_fragments(frags.clone(), self.shard_id, new_seq);
            SHARD_EVENT_COUNTER_VEC
                .with_label_values(&["append_base"])
                .inc();
            update.manifest = Some(new_manifest.clone());
            update.seq = Some(new_seq);
            info!(
                "append_base_frag_meta: start append frag shard={}, new_seq={}, len(frags)={}",
                self.to_long_ident(),
                new_seq,
                frags.len(),
            );
        }
        update.appended_frags = Some(frags);
        update
    }

    fn process_switch_base_to_delta(&self) -> Option<ShardInfoUpdate> {
        let mut new_manifest = self.manifest.clone();
        if !new_manifest.switch_base_to_delta() {
            return None;
        }

        let mut update = ShardInfoUpdate::from(self);
        update.manifest = Some(new_manifest);
        update.seq = Some(self.seq + 1);
        Some(update)
    }

    pub async fn switch_base_to_delta(
        &mut self,
        meta_store: &Arc<dyn MetaStore>,
    ) -> anyhow::Result<bool> {
        let Some(mut update) = self.process_switch_base_to_delta() else {
            return Ok(false);
        };
        meta_store
            .update_shard("switch_append_base", self, update.clone())
            .await?;
        update.mark_progress_cdc_persisted();
        self.apply_update(update);
        Ok(true)
    }

    pub fn process_append_delta(
        &self,
        request: &AppendDeltaFragMetaRequest,
    ) -> anyhow::Result<MaybeShardInfoUpdate> {
        if request.worker_node_addr != self.writer {
            warn!(
                "append_frag_meta: reject delta append for writer mismatch, shard={}, writer={}, request_writer={}",
                self.to_long_ident(),
                self.writer,
                request.worker_node_addr,
            );
            return Ok(MaybeShardInfoUpdate::ErrorCode(
                ErrorCode::ShardNotScheduled,
            ));
        }
        let new_frags = request.frag_metas.clone();
        let mut new_manifest = self.manifest.clone();
        let new_frags = new_manifest.make_deduplicated(
            new_frags,
            "process_append_delta",
            request.shard_id,
            self.seq,
        );
        let mut size_added = 0;
        let mut ranges = vec![];
        for f in new_frags.iter() {
            let prop = f.validate().map_err(|e| {
                anyhow!(
                    "append_frag_meta: reject invalid fragments, fragment={:?}, shard={:?}, err={:?}",
                    f,
                    self,
                    e,
                )
            })?;

            size_added += prop.size;
            // According to the logic of fragment generation, the `r` here means
            // the fragment with keys in [start_key, end_key]. So should use
            // `range_contains_minmax`.
            let r = HandleRange::new(prop.min_handle.clone(), prop.max_handle.clone());
            if !self.range.range_contains_minmax(&r) {
                info!(
                    "append_frag_meta: reject fragments for bound, fragment={:?}, expected={:?}, epoch={}, shard={}",
                    f,
                    self.range,
                    self.epoch,
                    self.to_long_ident(),
                );
                return Ok(MaybeShardInfoUpdate::ErrorCode(
                    ErrorCode::FragRangeMismatch,
                ));
            }
            ranges.push(r);
        }
        let range_str = ranges
            .iter()
            .map(|x| format!("{:?}", x))
            .collect::<Vec<_>>()
            .join("&");

        let update_progress = request.progress.clone();
        let append_frag_size: u64 = new_frags.iter().map(|x| x.bytes()).sum();
        let append_frag_count: u64 = new_frags.iter().map(|x| x.rows()).sum();
        let append_frag_segs: usize = new_frags.iter().map(|x| x.seg_count()).sum();
        let append_frag_storage_refs: usize = new_frags.iter().map(|x| x.storage_refs.len()).sum();
        let append_frag_mem: usize =
            new_frags.heap_size_of_children() + std::mem::size_of_val(&new_frags);

        let progress = if let Some(update_progress) = update_progress {
            let mut progress = self.progress.clone();
            #[allow(clippy::collapsible_if)]
            if update_progress.last_cdc_file != progress.cdc_s3_last_file {
                if older_cdc_file(
                    update_progress.last_cdc_file.as_str(),
                    progress.cdc_s3_last_file.as_str(),
                ) == update_progress.last_cdc_file.as_str()
                {
                    error!(
                        "suspect rewind of cdc log file, current={}, new={}, shard={}",
                        progress.cdc_s3_last_file.as_str(),
                        update_progress.last_cdc_file.as_str(),
                        self.to_long_ident()
                    );
                    return Ok(MaybeShardInfoUpdate::ErrorCode(ErrorCode::ProgressUnwinded));
                }
            }
            progress.cdc_s3_last_file = update_progress.last_cdc_file;
            Some(progress)
        } else {
            None
        };

        if new_frags.is_empty() {
            SHARD_EVENT_COUNTER_VEC
                .with_label_values(&["append_empty"])
                .inc();
            info!(
                "append_frag_meta: skip empty append and keep seq, shard={}, seq={}, range={}, manifest_len={}, update_progress={:?}, location={:?}",
                self.to_long_ident(),
                self.seq,
                range_str,
                self.manifest.len(),
                request.progress,
                self.frag_upload_location,
            );
            return Ok(MaybeShardInfoUpdate::Update(Box::new(ShardInfoUpdate {
                shard_id: self.shard_id,
                progress,
                manifest: None,
                appended_frags: Some(new_frags),
                seq: None,
                compaction: None,
                writer: None,
                size_added,
            })));
        }

        let new_seq = self.seq + 1;
        new_manifest.append_delta_fragments(new_frags.clone(), self.shard_id, new_seq);
        SHARD_EVENT_COUNTER_VEC
            .with_label_values(&["append_frag"])
            .inc();
        info!(
            "append_frag_meta: start append frag noempty, shard={}, new_seq={}, range={}, append=(len={}, size={}, rows={}, segs={}, storage_refs={}, frag_mem={}), new_manifest_len={}, update_progress={:?}, location={:?}, appended={}, writer={}",
            self.to_long_ident(),
            new_seq,
            range_str,
            new_frags.len(),
            append_frag_size,
            append_frag_count,
            append_frag_segs,
            append_frag_storage_refs,
            append_frag_mem,
            new_manifest.len(),
            request.progress,
            self.frag_upload_location,
            new_frags
                .iter()
                .map(|x| x.frag_path.as_ref())
                .collect::<Vec<_>>()
                .join(","),
            self.writer,
        );

        Ok(MaybeShardInfoUpdate::Update(Box::new(ShardInfoUpdate {
            shard_id: self.shard_id,
            progress,
            manifest: Some(new_manifest),
            appended_frags: Some(new_frags),
            seq: Some(new_seq),
            compaction: None,
            writer: None,
            size_added,
        })))
    }

    pub fn process_add_compaction(
        &mut self,
        index: Arc<FTSIndex>,
        input_fragments: Vec<SplittableFrag>,
        min_split_row_count: u64,
        max_split_point_count: u64,
        reason: CompactShardReason,
        remaining_iterations: usize,
    ) -> anyhow::Result<ShardInfoUpdate> {
        let input_start = if let Some(x) = input_fragments.first() {
            x
        } else {
            return Err(anyhow!(
                "schedule_compaction: empty input fragment {:?}, shard={:?}",
                input_fragments,
                self
            ));
        };
        let location =
            if let Some(loc) = self.manifest.find_pos_by_frag_path(&input_start.frag_path) {
                loc
            } else {
                return Err(anyhow!(
                    "schedule_compaction: can't find corresponding fragment {:?}, shard={:?}",
                    input_fragments,
                    self
                ));
            };

        // Extract position and determine level from location
        use crate::common::shard::FragmentLocation;
        let level = match location {
            FragmentLocation::Base(pos) => {
                let end = pos + input_fragments.len();
                if end <= self.manifest.base_slice().len() {
                    Some(2u8) // All in base layer
                } else {
                    None // Spans across boundary
                }
            }
            FragmentLocation::Delta(_) => None,
        };

        let compaction_state = CompactionState::new(
            self,
            index,
            input_fragments,
            level,
            min_split_row_count,
            max_split_point_count,
            reason,
            remaining_iterations,
        );
        let compaction_key = compaction_state.get_compaction_state_key();
        let mut new_progress = self.progress.clone();
        new_progress
            .compaction_state
            .insert(compaction_key.clone(), compaction_state.clone());

        Ok(ShardInfoUpdate {
            shard_id: self.shard_id,
            progress: Some(new_progress),
            manifest: None,
            appended_frags: None,
            seq: None,
            compaction: Some((compaction_key, compaction_state)),
            writer: None,
            size_added: 0,
        })
    }

    pub fn process_apply_compaction(
        &self,
        compaction_key: &CompactionStateKey,
        compaction_state: &CompactionState,
        output_frags: &[FragMeta],
    ) -> anyhow::Result<ShardInfoUpdate> {
        let mut new_progress = self.progress.clone();
        if new_progress
            .compaction_state
            .remove(compaction_key)
            .is_none()
        {
            log::warn!(
                "process_apply_compaction: can't remove compaction key in shard, key:{:?}",
                compaction_key
            );
        }
        let mut new_manifest = self.manifest.clone();
        let new_seq = self.seq + 1;
        let output_frags = output_frags
            .iter()
            .map(|f| SplittableFrag::new(f.clone(), self.shard_id, new_seq))
            .collect::<Vec<_>>();
        new_manifest.splice(compaction_state.input_fragments.as_ref(), output_frags)?;
        Ok(ShardInfoUpdate {
            shard_id: self.shard_id,
            progress: Some(new_progress),
            manifest: Some(new_manifest),
            appended_frags: None,
            seq: Some(new_seq),
            compaction: None,
            writer: None,
            size_added: 0,
        })
    }
}

#[derive(Debug, Clone)]
pub struct ShardInfoUpdate {
    pub shard_id: u64,
    pub progress: Option<ShardBuildProgress>,
    pub manifest: Option<ShardManifest>,
    // As extra informations,
    pub appended_frags: Option<Vec<FragMeta>>,
    pub seq: Option<u64>,
    pub compaction: Option<(CompactionStateKey, CompactionState)>,
    pub writer: Option<String>,
    pub size_added: u64,
}

impl From<&ShardInfo> for ShardInfoUpdate {
    fn from(value: &ShardInfo) -> Self {
        Self {
            shard_id: value.shard_id,
            progress: None,
            manifest: None,
            appended_frags: None,
            seq: None,
            compaction: None,
            writer: None,
            size_added: 0,
        }
    }
}

impl ShardInfoUpdate {
    pub fn mark_progress_cdc_persisted(&mut self) {
        if let Some(progress) = self.progress.as_mut() {
            progress.mark_cdc_last_file_persisted();
        }
    }
}

pub enum MaybeShardInfoUpdate {
    Update(Box<ShardInfoUpdate>),
    ErrorCode(ErrorCode),
}

// Specification for append update, in order to reduce memory usage
#[derive(Clone, Debug)]
pub struct ShardInfoAppendUpdate {
    pub shard_id: u64,
    pub start_key: RangeBound,
    pub frags: Vec<FragMeta>,
    pub new_seq: u64,
    pub is_base: bool,
    pub cdc_s3_last_file: String,
}

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

    use tici_grpc::{AppendDeltaFragMetaRequest, ErrorCode, FragProperty, ShardProgress};
    use tici_shard::test_utils::test_mock_frag_storage_refs_total;

    use super::*;
    use crate::common::test_utils::*;

    fn test_frag_path(path: &str) -> String {
        format!("{path}/f_1")
    }

    fn make_frag_with_path(path: &str, size: u64) -> FragMeta {
        FragMeta::new(
            0,
            test_frag_path(path),
            FragProperty {
                min_handle: vec![0],
                split_handles: vec![vec![0]],
                max_handle: vec![0],
                size,
                count: 1,
                delete_count: 0,
            },
            test_mock_frag_storage_refs_total(size),
            1,
        )
        .unwrap()
    }

    #[test]
    fn test_process_append_delta_rejects_writer_mismatch() {
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"B".to_vec()));
        shard.shard_id = 1;
        shard.writer = "writer-a".to_string();

        let request = AppendDeltaFragMetaRequest {
            keyspace_id: 0,
            table_id: 0,
            index_id: 0,
            shard_id: shard.shard_id,
            worker_node_addr: "writer-b".to_string(),
            frag_metas: vec![],
            progress: None,
        };

        let result = shard.process_append_delta(&request).unwrap();
        match result {
            MaybeShardInfoUpdate::ErrorCode(code) => {
                assert_eq!(code, ErrorCode::ShardNotScheduled);
            }
            _ => panic!("expected shard not scheduled"),
        }
    }

    #[test]
    fn test_process_append_delta_rejects_empty_storage_refs() {
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"B".to_vec()));
        shard.shard_id = 1;
        shard.writer = "writer-a".to_string();

        let request = AppendDeltaFragMetaRequest {
            keyspace_id: 0,
            table_id: 0,
            index_id: 0,
            shard_id: shard.shard_id,
            worker_node_addr: shard.writer.clone(),
            frag_metas: vec![FragMeta {
                frag_path: "f_1".to_string(),
                property: Some(FragProperty {
                    min_handle: b"A".to_vec(),
                    split_handles: vec![b"A".to_vec()],
                    max_handle: b"B".to_vec(),
                    ..Default::default()
                }),
                storage_refs: Vec::new(),
                ..Default::default()
            }],
            progress: None,
        };

        assert!(shard.process_append_delta(&request).is_err());
    }

    #[test]
    fn test_splitted_size() {
        let mut manifest = ShardManifest::default();
        let size_factor = 320;
        let test_fn = |n: u64| {
            let n_rows: u64 = 10000 + 5000 * n;
            for i in 0..n_rows {
                let s = format!("A{:06}A", i);
                let m = format!("A{:06}M", i);
                let e = format!("A{:06}Z", i);
                let frags = vec![make_sfsc(
                    s.as_bytes(),
                    m.as_bytes(),
                    e.as_bytes(),
                    size_factor,
                    size_factor,
                )];
                manifest.append_delta_fragments(frags, 1, 16);
            }
            let mut first = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"B".to_vec()));
            first.manifest = manifest;
            let mut new_ranges = vec![
                HandleRange::new(format!("A{:06}", 0).into(), format!("A{:06}", 5000).into()),
                HandleRange::new(
                    format!("A{:06}", 5000).into(),
                    format!("A{:06}", 10000).into(),
                ),
            ];
            for i in 0..n {
                new_ranges.push(HandleRange::new(
                    format!("A{:06}", 10000 + i * 5000).into(),
                    format!("A{:06}", 15000 + i * 5000).into(),
                ));
            }
            let splitted = ShardInfo::new_from_split(
                &first,
                new_ranges.as_slice(),
                std::iter::repeat_n("W".to_string(), 2 + n as usize).collect::<Vec<_>>(),
            );
            let size: u64 = splitted.iter().map(|x| x.get_total_estimated_bytes()).sum();
            assert_eq!(size, size_factor * n_rows);
        };
        (test_fn.clone())(0);
        (test_fn.clone())(3);
    }

    #[test]
    fn test_split_resets_empty_append_persist_count() {
        let original = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"Z".to_vec()));
        original
            .append_empty_persist_count
            .store(5, Ordering::Relaxed);

        let new_ranges = vec![
            HandleRange::new(b"A".to_vec(), b"M".to_vec()),
            HandleRange::new(b"M".to_vec(), b"Z".to_vec()),
        ];
        let split = ShardInfo::new_from_split(
            &original,
            &new_ranges,
            vec!["worker-a".to_string(), "worker-b".to_string()],
        );

        assert_eq!(split.len(), 2);
        for shard in &split {
            assert_eq!(shard.append_empty_persist_count.load(Ordering::Relaxed), 0);
            assert!(!Arc::ptr_eq(
                &shard.append_empty_persist_count,
                &original.append_empty_persist_count
            ));
        }

        split[0]
            .append_empty_persist_count
            .store(9, Ordering::Relaxed);
        assert_eq!(
            original.append_empty_persist_count.load(Ordering::Relaxed),
            5
        );
        assert_eq!(
            split[1].append_empty_persist_count.load(Ordering::Relaxed),
            0
        );
    }

    #[test]
    fn test_merge_resets_empty_append_persist_count() -> anyhow::Result<()> {
        let left = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"M".to_vec()));
        let right = ShardInfo::new_for_test(HandleRange::new(b"M".to_vec(), b"Z".to_vec()));
        left.append_empty_persist_count.store(5, Ordering::Relaxed);
        right.append_empty_persist_count.store(6, Ordering::Relaxed);

        let merged = match ShardInfo::new_from_merge("worker-c".to_string(), &[&left, &right])? {
            Either::Left(shard) => shard,
            Either::Right((left, right)) => {
                panic!("unexpected merge overlap: {left:?}, {right:?}");
            }
        };

        assert_eq!(merged.append_empty_persist_count.load(Ordering::Relaxed), 0);
        assert!(!Arc::ptr_eq(
            &merged.append_empty_persist_count,
            &left.append_empty_persist_count
        ));
        assert!(!Arc::ptr_eq(
            &merged.append_empty_persist_count,
            &right.append_empty_persist_count
        ));

        merged
            .append_empty_persist_count
            .store(9, Ordering::Relaxed);
        assert_eq!(left.append_empty_persist_count.load(Ordering::Relaxed), 5);
        assert_eq!(right.append_empty_persist_count.load(Ordering::Relaxed), 6);

        Ok(())
    }

    #[test]
    fn test_pick_uni_compaction_can_dilute_to_pass_ratio() {
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"B".to_vec()));
        shard.manifest.append_delta_fragments(
            vec![
                make_frag_with_path("F0", 9000),
                make_frag_with_path("F1", 1),
                make_frag_with_path("F2", 9000),
            ],
            1,
            1,
        );

        let config = UniCompactionConfig {
            max_largest_ratio: 0.5,
            min_compaction_frags: 2,
            max_compaction_iterations: 2,
        };

        let picked = shard.pick_uni_compaction(&config);
        let picked_paths = picked
            .iter()
            .map(|f| f.frag_path.clone())
            .collect::<Vec<_>>();
        assert_eq!(
            picked_paths,
            vec![
                test_frag_path("F0"),
                test_frag_path("F1"),
                test_frag_path("F2")
            ]
        );
    }

    #[test]
    fn test_pick_uni_compaction_avoids_rewriting_huge_frag() {
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"B".to_vec()));
        shard.manifest.append_delta_fragments(
            vec![
                make_frag_with_path("BIG", 10_000),
                make_frag_with_path("S1", 10),
                make_frag_with_path("S2", 10),
                make_frag_with_path("S3", 10),
            ],
            1,
            1,
        );

        let config = UniCompactionConfig {
            max_largest_ratio: 0.9,
            min_compaction_frags: 2,
            max_compaction_iterations: 2,
        };

        let picked = shard.pick_uni_compaction(&config);
        let picked_paths = picked
            .iter()
            .map(|f| f.frag_path.clone())
            .collect::<Vec<_>>();
        assert_eq!(
            picked_paths,
            vec![
                test_frag_path("S1"),
                test_frag_path("S2"),
                test_frag_path("S3")
            ]
        );
    }

    #[test]
    fn test_pick_uni_compaction_fallback_to_min_ratio_when_no_window_passes() {
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"B".to_vec()));
        shard.manifest.append_delta_fragments(
            vec![
                make_frag_with_path("F0", 1),
                make_frag_with_path("F1", 9000),
            ],
            1,
            1,
        );

        let config = UniCompactionConfig {
            max_largest_ratio: 0.5,
            min_compaction_frags: 2,
            max_compaction_iterations: 2,
        };

        let picked = shard.pick_uni_compaction(&config);
        let picked_paths = picked
            .iter()
            .map(|f| f.frag_path.clone())
            .collect::<Vec<_>>();
        assert_eq!(
            picked_paths,
            vec![test_frag_path("F0"), test_frag_path("F1")]
        );
    }

    #[test]
    fn test_pick_uni_compaction_uses_min_length_candidates() {
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"B".to_vec()));
        shard.manifest.append_delta_fragments(
            vec![
                make_frag_with_path("F0", 100),
                make_frag_with_path("F1", 1),
                make_frag_with_path("F2", 1),
                make_frag_with_path("F3", 1),
                make_frag_with_path("F4", 1),
            ],
            1,
            1,
        );

        let config = UniCompactionConfig {
            max_largest_ratio: 0.5,
            min_compaction_frags: 5,
            max_compaction_iterations: 2,
        };

        let picked = shard.pick_uni_compaction(&config);
        let picked_paths = picked
            .iter()
            .map(|f| f.frag_path.clone())
            .collect::<Vec<_>>();
        assert_eq!(
            picked_paths,
            vec![
                test_frag_path("F0"),
                test_frag_path("F1"),
                test_frag_path("F2"),
                test_frag_path("F3"),
                test_frag_path("F4")
            ]
        );
    }

    #[test]
    fn test_pick_uni_compaction_at_respects_right_boundary() {
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"B".to_vec()));
        shard.manifest.append_delta_fragments(
            vec![
                make_frag_with_path("F0", 10),
                make_frag_with_path("F1", 1),
                make_frag_with_path("F2", 1),
                make_frag_with_path("F3", 1),
            ],
            1,
            1,
        );

        let config = UniCompactionConfig {
            max_largest_ratio: 0.5,
            min_compaction_frags: 2,
            max_compaction_iterations: 2,
        };

        let picked = shard.pick_uni_compaction_at(&config, &test_frag_path("F2"), false);
        let picked_paths = picked
            .iter()
            .map(|f| f.frag_path.clone())
            .collect::<Vec<_>>();
        assert_eq!(
            picked_paths,
            vec![test_frag_path("F1"), test_frag_path("F2")]
        );
    }

    #[test]
    fn test_pick_uni_compaction_at_disables_fallback() {
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"B".to_vec()));
        shard.manifest.append_delta_fragments(
            vec![
                make_frag_with_path("F0", 1),
                make_frag_with_path("F1", 9000),
            ],
            1,
            1,
        );

        let config = UniCompactionConfig {
            max_largest_ratio: 0.5,
            min_compaction_frags: 2,
            max_compaction_iterations: 2,
        };

        assert!(
            shard
                .pick_uni_compaction_at(&config, &test_frag_path("F1"), false)
                .is_empty()
        );
    }

    #[test]
    fn test_pick_uni_compaction_stops_at_existing_task_fragment() {
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"B".to_vec()));
        shard.manifest.append_delta_fragments(
            vec![
                make_frag_with_path("F0", 10),
                make_frag_with_path("F1", 10),
                make_frag_with_path("F2", 10),
                make_frag_with_path("F3", 10),
            ],
            1,
            1,
        );

        let conflict = shard.manifest.delta_slice()[1].clone();
        let state = CompactionState::new(
            &shard,
            Arc::new(FTSIndex::new_for_test()),
            vec![conflict],
            None,
            0,
            0,
            CompactShardReason::Manual,
            0,
        );
        shard
            .progress
            .compaction_state
            .insert(state.get_compaction_state_key(), state);

        let config = UniCompactionConfig {
            max_largest_ratio: 0.5,
            min_compaction_frags: 2,
            max_compaction_iterations: 2,
        };

        let picked = shard.pick_uni_compaction(&config);
        let picked_paths = picked
            .iter()
            .map(|f| f.frag_path.clone())
            .collect::<Vec<_>>();
        assert_eq!(
            picked_paths,
            vec![test_frag_path("F2"), test_frag_path("F3")]
        );
    }

    #[test]
    fn test_pick_uni_compaction_empty_when_rightmost_in_existing_task() {
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"B".to_vec()));
        shard.manifest.append_delta_fragments(
            vec![make_frag_with_path("F0", 10), make_frag_with_path("F1", 10)],
            1,
            1,
        );

        let conflict = shard.manifest.delta_slice()[1].clone();
        let state = CompactionState::new(
            &shard,
            Arc::new(FTSIndex::new_for_test()),
            vec![conflict],
            None,
            0,
            0,
            CompactShardReason::Manual,
            0,
        );
        shard
            .progress
            .compaction_state
            .insert(state.get_compaction_state_key(), state);

        let config = UniCompactionConfig {
            max_largest_ratio: 1.0,
            min_compaction_frags: 2,
            max_compaction_iterations: 2,
        };

        assert!(shard.pick_uni_compaction(&config).is_empty());
    }

    #[test]
    fn test_merge_overlap_compaction_inputs_requires_multiple_delta_frags() {
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"B".to_vec()));

        assert!(shard.merge_overlap_compaction_inputs("uni").is_none());

        shard
            .manifest
            .append_delta_fragments(vec![make_frag_with_path("F0", 1)], 1, 1);
        assert!(shard.merge_overlap_compaction_inputs("uni").is_none());

        shard
            .manifest
            .append_delta_fragments(vec![make_frag_with_path("F1", 1)], 1, 1);

        let inputs = match shard.merge_overlap_compaction_inputs("uni") {
            Some(inputs) => inputs,
            None => panic!("expected merge overlap compaction inputs"),
        };
        let paths = inputs
            .iter()
            .map(|frag| frag.frag_path.clone())
            .collect::<Vec<_>>();
        assert_eq!(paths, vec![test_frag_path("F0"), test_frag_path("F1")]);
    }

    #[test]
    fn test_process_append_delta_rejects_progress_unwind() {
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"B".to_vec()));
        shard.progress.cdc_s3_last_file =
            "tici_default_prefix/cdc/10313/461869219056451591/CDC00000000000000000002.json"
                .to_string();

        let request = AppendDeltaFragMetaRequest {
            progress: Some(ShardProgress {
                last_cdc_file:
                    "tici_default_prefix/cdc/10313/461869219056451591/CDC00000000000000000001.json"
                        .to_string(),
            }),
            ..Default::default()
        };

        match shard.process_append_delta(&request).unwrap() {
            MaybeShardInfoUpdate::ErrorCode(ErrorCode::ProgressUnwinded) => {}
            _ => panic!("expected ProgressUnwinded"),
        }
    }

    #[test]
    fn test_process_append_delta_empty_append_keeps_seq() {
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"B".to_vec()));
        shard.shard_id = 1;
        shard.writer = "writer-a".to_string();
        shard.seq = 7;
        shard.progress.cdc_s3_last_file =
            "tici_default_prefix/cdc/10313/461869219056451591/CDC00000000000000000001.json"
                .to_string();

        let request = AppendDeltaFragMetaRequest {
            shard_id: shard.shard_id,
            worker_node_addr: shard.writer.clone(),
            progress: Some(ShardProgress {
                last_cdc_file:
                    "tici_default_prefix/cdc/10313/461869219056451591/CDC00000000000000000002.json"
                        .to_string(),
            }),
            ..Default::default()
        };

        match shard.process_append_delta(&request).unwrap() {
            MaybeShardInfoUpdate::Update(update) => {
                assert_eq!(update.seq, None);
                assert_eq!(update.manifest, None);
                assert_eq!(
                    update.progress.unwrap().cdc_s3_last_file,
                    "tici_default_prefix/cdc/10313/461869219056451591/CDC00000000000000000002.json"
                );
                assert!(update.appended_frags.unwrap().is_empty());
            }
            _ => panic!("expected update"),
        }
    }

    #[test]
    fn test_process_append_base_deduplicated_append_keeps_seq() {
        let mut shard = ShardInfo::new_for_test(HandleRange::new(b"A".to_vec(), b"B".to_vec()));
        shard.shard_id = 1;
        shard.seq = 5;
        shard.manifest.append_base_fragments(
            vec![make_frag_with_path("F0", 1)],
            shard.shard_id,
            shard.seq,
        );

        let update = shard.process_append_base(&[make_frag_with_path("F0", 1)]);
        assert_eq!(update.seq, None);
        assert_eq!(update.manifest, None);
        assert!(update.appended_frags.unwrap().is_empty());
    }

    #[test]
    fn test_mark_progress_cdc_persisted_updates_progress() {
        let mut update = ShardInfoUpdate {
            shard_id: 1,
            progress: Some(ShardBuildProgress {
                cdc_s3_last_file:
                    "tici_default_prefix/cdc/10313/461869219056451591/CDC00000000000000000002.json"
                        .to_string(),
                persisted_cdc_s3_last_file:
                    "tici_default_prefix/cdc/10313/461869219056451591/CDC00000000000000000001.json"
                        .to_string(),
                ..Default::default()
            }),
            manifest: None,
            appended_frags: None,
            seq: None,
            compaction: None,
            writer: None,
            size_added: 0,
        };

        update.mark_progress_cdc_persisted();

        let progress = update.progress.unwrap();
        assert_eq!(
            progress.persisted_cdc_s3_last_file,
            progress.cdc_s3_last_file
        );
    }
}