// Copyright 2025 PingCAP, Inc.
//! Utilities for handling TiDB (not SST) table IDs.

use std::{collections::HashSet, ops::Deref};

use bytes::BufMut;
use tidb_query_datatype::codec::table::{
    RECORD_PREFIX_SEP, TABLE_PREFIX, TABLE_PREFIX_KEY_LEN, decode_table_id,
};
use tikv_util::codec::number::NumberEncoder;

use crate::table::{DataBound, InnerKey, OwnedInnerKey, fts::TrackedFtsIndexMap};

pub fn merge_columnar_table_ids(
    source_table_ids: &[i64],
    target_table_ids: &[i64],
    source_bound: DataBound<'_>,
    target_bound: DataBound<'_>,
) -> Vec<i64> {
    let (min_source_table_id, max_source_table_id) = get_table_id_from_data_bound(source_bound);
    let (min_target_table_id, max_target_table_id) = get_table_id_from_data_bound(target_bound);
    let source_table_ids: HashSet<_> = source_table_ids.iter().cloned().collect();
    let target_table_ids: HashSet<_> = target_table_ids.iter().cloned().collect();
    let common_table_ids: HashSet<_> = source_table_ids
        .intersection(&target_table_ids)
        .cloned()
        .collect();
    let mut only_in_source: HashSet<_> = source_table_ids
        .difference(&target_table_ids)
        .cloned()
        .collect();
    only_in_source
        .retain(|&table_id| table_id < min_target_table_id || table_id > max_target_table_id);
    let mut only_in_target: HashSet<_> = target_table_ids
        .difference(&source_table_ids)
        .cloned()
        .collect();
    only_in_target
        .retain(|&table_id| table_id < min_source_table_id || table_id > max_source_table_id);
    let mut merged_table_ids = common_table_ids
        .union(&only_in_source)
        .cloned()
        .collect::<HashSet<_>>()
        .union(&only_in_target)
        .cloned()
        .collect::<Vec<_>>();
    merged_table_ids.sort();
    merged_table_ids
}

pub fn has_consistent_tracked_fts_indexes(
    source_indexes: &TrackedFtsIndexMap,
    target_indexes: &TrackedFtsIndexMap,
    source_bound: DataBound<'_>,
    target_bound: DataBound<'_>,
) -> bool {
    let (min_source_table_id, max_source_table_id) = get_table_id_from_data_bound(source_bound);
    let (min_target_table_id, max_target_table_id) = get_table_id_from_data_bound(target_bound);
    let min_table_id = std::cmp::max(min_source_table_id, min_target_table_id);
    let max_table_id = std::cmp::min(max_source_table_id, max_target_table_id);
    source_indexes
        .iter()
        .filter(|(&table_id, _)| table_id >= min_table_id && table_id <= max_table_id)
        .all(|(table_id, source_indexes)| target_indexes.get(table_id) == Some(source_indexes))
        && target_indexes
            .iter()
            .filter(|(&table_id, _)| table_id >= min_table_id && table_id <= max_table_id)
            .all(|(table_id, target_indexes)| source_indexes.get(table_id) == Some(target_indexes))
}

/// Try to get the table id from the data bound.
///
/// If the data bound is not a table id, min table id is 0 and max table id is
/// i64::MAX. If the upper bound is starts with "m", max table id should be 0.
pub fn get_table_id_from_data_bound(data_bound: DataBound<'_>) -> (i64, i64) {
    let data_upper_bound = data_bound.upper_bound.deref();
    let start_table_id = decode_table_id(data_bound.lower_bound.deref()).unwrap_or(0);
    let end_table_id = if !data_upper_bound.is_empty() && data_upper_bound < TABLE_PREFIX {
        0
    } else {
        let mut table_id = decode_table_id(data_upper_bound).unwrap_or(i64::MAX);
        if data_upper_bound.len() >= TABLE_PREFIX_KEY_LEN
            && &data_upper_bound[TABLE_PREFIX_KEY_LEN..] <= RECORD_PREFIX_SEP
        {
            table_id -= 1;
        }
        table_id
    };

    (start_table_id, end_table_id)
}

/// Return true if the range of data bound overlaps with any of the table ids.
pub(crate) fn is_bound_overlap_with_table_ids(
    data_bound: DataBound<'_>,
    table_ids: &[i64],
) -> bool {
    let (min_table_id, max_table_id) = get_table_id_from_data_bound(data_bound);
    table_ids
        .iter()
        .any(|&id| id >= min_table_id && id <= max_table_id)
}

/// Return true when both keys are table keys and they belong to the same table.
///
/// Should check `ApiV2::is_belongs_to_same_keyspace()` first if the keys can
/// belong to different keyspace.
pub fn keys_belong_to_same_table(k0: InnerKey<'_>, k1: InnerKey<'_>) -> bool {
    k0.len() >= TABLE_PREFIX_KEY_LEN
        && k1.len() >= TABLE_PREFIX_KEY_LEN
        && k0[..TABLE_PREFIX_KEY_LEN] == k1[..TABLE_PREFIX_KEY_LEN]
}

pub fn encode_table_prefix_key(table_id: i64) -> OwnedInnerKey {
    let mut key = Vec::with_capacity(TABLE_PREFIX_KEY_LEN);
    key.put(TABLE_PREFIX);
    key.encode_i64(table_id).unwrap();
    OwnedInnerKey::new(key.into())
}

/// Return whether the key is on a table boundary.
#[inline]
pub fn is_table_boundary_key(key: InnerKey<'_>) -> bool {
    key.len() == TABLE_PREFIX_KEY_LEN && decode_table_id(key.as_ref()).is_ok()
}

/// Get the table id of the ingest files changeset belongs to.
///
/// All ssts in changeset should be belongs to the same table.
pub fn get_table_id_from_ingest_files(ingest_files: &kvenginepb::IngestFiles) -> Option<i64> {
    let mut unique_table_id = None;
    for tbl in ingest_files.get_l0_creates() {
        let smallest = tbl.get_smallest();
        if let Ok(table_id) = decode_table_id(smallest) {
            match unique_table_id {
                None => unique_table_id = Some(table_id),
                Some(unique_table_id) => {
                    debug_assert_eq!(table_id, unique_table_id, "table id should be same");
                }
            }
        }
    }
    for tbl in ingest_files.get_table_creates() {
        let smallest = tbl.get_smallest();
        if let Ok(table_id) = decode_table_id(smallest) {
            match unique_table_id {
                None => unique_table_id = Some(table_id),
                Some(unique_table_id) => {
                    debug_assert_eq!(table_id, unique_table_id, "table id should be same");
                }
            }
        }
    }
    unique_table_id
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use super::*;
    use crate::table::{InnerKey, fts::new_legacy_tracked_index_meta};

    #[test]
    fn test_keys_belong_to_same_table() {
        let cases: Vec<(&'static str, &'static str, bool)> = vec![
            (
                "7800002A7480000000000000015F728000000000000001", // table id before `5F72`
                "7800002A7480000000000000015F728000000000000002",
                true,
            ),
            (
                "7800002A7480000000000000015F728000000000000001",
                "7800002A7480000000000000025F728000000000000001",
                false,
            ),
            (
                "7480000000000000015F728000000000000001",
                "7480000000000000015F728000000000000001",
                true,
            ),
            (
                "7480000000000000015F728000000000000001",
                "7480000000000000025F728000000000000001",
                false,
            ),
            (
                "7800002A",
                "7800002A7480000000000000115F728000000000000001",
                false,
            ),
            (
                "7800002A7480000000000000015F728000000000000001",
                "7800002AFF",
                false,
            ),
        ];

        for (k0, k1, expected) in cases {
            let key0 = hex::decode(k0).unwrap();
            let key1 = hex::decode(k1).unwrap();
            assert_eq!(
                keys_belong_to_same_table(
                    InnerKey::from_outer_key(&key0),
                    InnerKey::from_outer_key(&key1)
                ),
                expected
            );
        }
    }

    #[test]
    fn test_has_consistent_tracked_fts_indexes() {
        fn to_index_map(entries: &[(i64, &[i64])]) -> TrackedFtsIndexMap {
            let mut map: TrackedFtsIndexMap = HashMap::new();
            for &(table_id, index_ids) in entries {
                let indexes = map.entry(table_id).or_default();
                for index_id in index_ids {
                    indexes.insert(
                        *index_id,
                        new_legacy_tracked_index_meta(table_id, *index_id),
                    );
                }
            }
            map
        }

        let t1 = hex::decode("748000000000000001").unwrap();
        let t3 = hex::decode("748000000000000003").unwrap();
        let t5 = hex::decode("748000000000000005").unwrap();
        let t7 = hex::decode("748000000000000007").unwrap();

        // Bound covers table_id in [1, 2].
        let bound_1_2 = DataBound::new(
            InnerKey::from_inner_buf(&t1),
            InnerKey::from_inner_buf(&t3),
            true,
        );
        // Bound covers table_id in [5, 6].
        let bound_5_6 = DataBound::new(
            InnerKey::from_inner_buf(&t5),
            InnerKey::from_inner_buf(&t7),
            true,
        );

        // No overlap => vacuously consistent.
        assert!(has_consistent_tracked_fts_indexes(
            &to_index_map(&[(1, &[1, 2])]),
            &to_index_map(&[(5, &[9])]),
            bound_1_2,
            bound_5_6,
        ));

        // Overlap, exact match => consistent (entries outside overlap are ignored).
        assert!(has_consistent_tracked_fts_indexes(
            &to_index_map(&[(1, &[1, 2]), (2, &[3])]),
            &to_index_map(&[(1, &[2, 1]), (2, &[3]), (100, &[7])]),
            bound_1_2,
            bound_1_2,
        ));

        // Overlap, mismatched index set => inconsistent.
        assert!(!has_consistent_tracked_fts_indexes(
            &to_index_map(&[(1, &[1, 2]), (2, &[3])]),
            &to_index_map(&[(1, &[1, 2]), (2, &[4])]),
            bound_1_2,
            bound_1_2,
        ));

        // Overlap, missing table entry on either side => inconsistent.
        assert!(!has_consistent_tracked_fts_indexes(
            &to_index_map(&[(1, &[1, 2])]),
            &to_index_map(&[(2, &[3])]),
            bound_1_2,
            bound_1_2,
        ));
        assert!(!has_consistent_tracked_fts_indexes(
            &to_index_map(&[(2, &[3])]),
            &to_index_map(&[(1, &[1, 2])]),
            bound_1_2,
            bound_1_2,
        ));
    }

    // Adjacent regions may carry incompatible metadata after an interrupted
    // upgrade or restore. They must not be treated as merge-compatible.
    #[test]
    fn test_has_consistent_tracked_fts_indexes_rejects_mixed_field_layout() {
        use kvenginepb::fts::{FtsFieldLayout, FullTextIndexFieldDef};

        use crate::table::fts::new_tracked_index_meta;

        let t1 = hex::decode("748000000000000001").unwrap();
        let t3 = hex::decode("748000000000000003").unwrap();
        let bound = DataBound::new(
            InnerKey::from_inner_buf(&t1),
            InnerKey::from_inner_buf(&t3),
            true,
        );

        let legacy = new_legacy_tracked_index_meta(1, 7);
        // Same index id, only the layout flag differs (single-column STANDARD).
        let mut legacy_flag_only = legacy.clone();
        legacy_flag_only.set_field_layout(FtsFieldLayout::FtsFieldLayoutIndexFieldOrdinalString);

        let mut source: TrackedFtsIndexMap = HashMap::new();
        source.entry(1).or_default().insert(7, legacy.clone());
        let mut target_flag_only: TrackedFtsIndexMap = HashMap::new();
        target_flag_only
            .entry(1)
            .or_default()
            .insert(7, legacy_flag_only);
        assert!(
            !has_consistent_tracked_fts_indexes(&source, &target_flag_only, bound, bound),
            "same index_id with LegacyBody vs ordinal-string must block merge"
        );

        // Multi-column: legacy tracking truncates to the first field; ordinal keeps
        // all.
        let mut field_a = FullTextIndexFieldDef::new();
        field_a.set_col_id(101);
        let mut field_b = FullTextIndexFieldDef::new();
        field_b.set_col_id(202);
        let ordinal_multi = new_tracked_index_meta(
            1,
            7,
            vec![field_a, field_b],
            FtsFieldLayout::FtsFieldLayoutIndexFieldOrdinalString,
        );
        let mut target_multi: TrackedFtsIndexMap = HashMap::new();
        target_multi.entry(1).or_default().insert(7, ordinal_multi);
        assert!(
            !has_consistent_tracked_fts_indexes(&source, &target_multi, bound, bound),
            "legacy single-field vs ordinal multi-field meta must block merge"
        );

        let mut aligned: TrackedFtsIndexMap = HashMap::new();
        aligned.entry(1).or_default().insert(7, legacy);
        assert!(
            has_consistent_tracked_fts_indexes(&aligned, &aligned, bound, bound),
            "identical layout on both sides remains merge-compatible"
        );
    }

    #[test]
    fn test_is_table_boundary_key() {
        let cases: Vec<(&'static str, bool)> = vec![
            ("748000000000000001", true),
            ("7480000000000000015F728000000000000001", false),
            ("758000000000000001", false),
        ];

        for (k, expected) in cases {
            let key = hex::decode(k).unwrap();
            assert_eq!(
                is_table_boundary_key(InnerKey::from_inner_buf(&key)),
                expected
            );
        }
    }

    #[test]
    fn test_is_bound_overlap_with_table_ids() {
        let cases: Vec<(&'static str, &'static str, &[i64], bool)> = vec![
            (
                "748000000000000001",
                "7480000000000000015F728000000000000001",
                &[1, 2, 3],
                true,
            ),
            (
                "7480000000000000025F728000000000000001",
                "748000000000000003",
                &[4, 5, 6],
                false,
            ),
            (
                "748000000000000001",
                "7480000000000000045F728000000000000001",
                &[3],
                true,
            ),
            (
                "748000000000000004",
                "7480000000000000095F728000000000000001",
                &[1, 2, 3],
                false,
            ),
        ];

        for (lower_bound, upper_bound, table_ids, expected) in cases {
            let lower_bound = hex::decode(lower_bound).unwrap();
            let upper_bound = hex::decode(upper_bound).unwrap();
            assert_eq!(
                is_bound_overlap_with_table_ids(
                    DataBound::new(
                        InnerKey::from_inner_buf(&lower_bound),
                        InnerKey::from_inner_buf(&upper_bound),
                        true
                    ),
                    table_ids
                ),
                expected
            );
        }
    }
}