use std::{
    collections::BTreeMap,
    sync::{Arc, atomic::AtomicU64 as StdAtomicU64},
    time::Instant,
};

use portable_atomic::{AtomicU16, AtomicU64};
use tici_common::{
    codec::{table_end_key, table_start_key},
    index::IndexKey,
};
use tici_shard::test_utils::index_meta::TEST_FT_INDEX_INFO_STD_STR;
use tokio::sync::RwLock;

use crate::{
    HandleRange, RangeBound, ShardManifest,
    meta::{
        reader::{
            ShardCacheAllocation, ShardCacheManager, UnifiedShardInfo, UnifiedShardInfoInner,
        },
        writer::ShardEventIndexInfo,
    },
};

pub fn create_test_cache_allocation(reader_addr: String) -> ShardCacheAllocation {
    ShardCacheAllocation::new(
        vec![reader_addr.to_string()],
        vec!["127.0.0.1:3930".to_string()],
        Instant::now(),
    )
}

// Test helper functions to reduce redundancy
pub fn create_test_shard_cache_info(shard_id: u64, reader_addr: &str) -> UnifiedShardInfo {
    UnifiedShardInfo {
        shard_id,
        index_key: IndexKey::new_nullspace(1, 1),
        range: HandleRange::new(
            format!("start_{}", shard_id).into_bytes(),
            format!("end_{}", shard_id).into_bytes(),
        ),
        seq: Arc::new(AtomicU64::new(0)),
        epoch: 1,
        inner: Arc::new(RwLock::new(UnifiedShardInfoInner::new(
            create_test_cache_allocation(reader_addr.to_string()),
            1,
            ShardManifest::new_for_first(),
        ))),
        in_warming: Arc::new(AtomicU16::new(0)),
    }
}

/// Creates a test TableInfo with default values
pub fn create_test_table_info(index_key: &IndexKey) -> tici_grpc::TableInfo {
    tici_grpc::TableInfo {
        keyspace_id: index_key.keyspace_id,
        table_id: index_key.table_id,
        table_name: format!("test_table_{}", index_key.table_id),
        database_name: "test_db".to_string(),
        version: 1,
        columns: vec![],
        is_clustered: false,
        is_partition: false,
        logical_table_id: 0,
    }
}

/// Creates a test IndexInfo with default values
pub fn create_test_index_info(index_key: &IndexKey) -> tici_grpc::IndexInfo {
    tici_grpc::IndexInfo {
        keyspace_id: index_key.keyspace_id,
        table_id: index_key.table_id,
        index_id: index_key.index_id,
        index_name: format!("test_index_{}_{}", index_key.table_id, index_key.index_id),
        index_type: tici_grpc::IndexType::Fulltext as i32,
        columns: vec![],
        is_unique: false,
        parser_info: Some(tici_grpc::DEFAULT_PARSER_INFO.clone()),
        other_params: Default::default(),
        fulltext_index: TEST_FT_INDEX_INFO_STD_STR.to_string(),
        hybrid_index: String::new(),
        created_ts: 0,
    }
}

/// Creates a test ShardInfo with given parameters
pub fn create_test_shard_info(
    index_key: IndexKey,
    shard_id: u64,
    start_key: &[u8],
    end_key: &[u8],
    seq: u64,
) -> crate::ShardInfo {
    crate::ShardInfo {
        index_key,
        shard_id,
        range: crate::HandleRange::new(start_key.to_vec(), end_key.to_vec()),
        epoch: 1,
        seq,
        manifest: crate::ShardManifest::new_for_first(),
        writer: "test_writer".to_string(),
        progress: crate::ShardBuildProgress::default(),
        frag_upload_location: tici_grpc::S3Location {
            bucket: "test_bucket".to_string(),
            prefix: "test_prefix".to_string(),
        },
        stats: crate::ShardStats::default(),
        append_empty_persist_count: Arc::new(StdAtomicU64::new(0)),
    }
}

pub fn create_test_shard_info2(
    index_key: &IndexKey,
    shard_id: u64,
    start_key: &[u8],
    end_key: &[u8],
    seq: u64,
) -> crate::ShardInfo {
    let start_key: Vec<u8> = [table_start_key(index_key), start_key.to_vec()].concat();
    let end_key: Vec<u8> = if end_key.is_empty() {
        table_end_key(index_key)
    } else {
        [table_start_key(index_key), end_key.to_vec()].concat()
    };
    create_test_shard_info(
        index_key.clone(),
        shard_id,
        start_key.as_slice(),
        end_key.as_slice(),
        seq,
    )
}

/// Creates a test UnifiedShardInfo with given parameters
pub fn create_test_unified_shard_info(
    index_key: &IndexKey,
    shard_id: u64,
    start_key: &[u8],
    end_key: &[u8],
    seq: u64,
) -> UnifiedShardInfo {
    UnifiedShardInfo {
        shard_id,
        index_key: index_key.clone(),
        epoch: 1,
        range: HandleRange::new(start_key.to_vec(), end_key.to_vec()),
        seq: Arc::new(AtomicU64::new(seq)),
        inner: Arc::new(RwLock::new(UnifiedShardInfoInner::default())),
        in_warming: Arc::new(AtomicU16::new(0)),
    }
}

/// Creates a test ShardEventIndexInfo
pub fn create_test_shard_event_index_info(index_key: &IndexKey) -> ShardEventIndexInfo {
    ShardEventIndexInfo {
        table_info: create_test_table_info(index_key),
        index_info: create_test_index_info(index_key),
    }
}

pub async fn setup_test_metadata(
    manager: &ShardCacheManager,
    index_key: &IndexKey,
    shards: Vec<(u64, &[u8], &[u8], u64)>, // (shard_id, start_key, end_key, seq)
) {
    for (shard_id, start_key, end_key, seq) in shards {
        let _ = manager
            .handle_shard_add(
                ShardEventIndexInfo::new(
                    create_test_table_info(index_key),
                    create_test_index_info(index_key),
                ),
                &create_test_shard_info2(index_key, shard_id, start_key, end_key, seq),
            )
            .await;
    }
}

pub fn create_test_cache_allocations(
    v: Vec<(RangeBound, String)>,
) -> BTreeMap<RangeBound, ShardCacheAllocation> {
    let mut cache_infos = BTreeMap::new();
    for (start_key, addr) in v {
        cache_infos.insert(start_key, create_test_cache_allocation(addr));
    }
    cache_infos
}