use std::sync::{
    Arc,
    atomic::{AtomicUsize, Ordering},
};

use async_trait::async_trait;

use crate::meta::reader::reader_pool::ReaderNodeMeta;

/// Trait for reader cache scheduling strategies
/// NOTE Such a scheduler must ensure that it can return different reader per
/// call by their best effort
#[async_trait]
pub trait ReaderCacheScheduler: Send + Sync + std::fmt::Debug {
    /// Select the best reader from the given list of active readers and reserve
    /// one schedule slot on it.
    async fn select_reader(&self, readers: &[Arc<ReaderNodeMeta>]) -> Option<Arc<ReaderNodeMeta>>;

    /// Get cache scheduler name for logging
    fn name(&self) -> &'static str;
}

/// Round-robin cache scheduler
#[derive(Debug)]
pub struct RoundRobinCacheScheduler {
    counter: AtomicUsize,
}

impl Default for RoundRobinCacheScheduler {
    fn default() -> Self {
        Self::new()
    }
}

impl RoundRobinCacheScheduler {
    pub fn new() -> Self {
        Self {
            counter: AtomicUsize::new(0),
        }
    }
}

#[async_trait]
impl ReaderCacheScheduler for RoundRobinCacheScheduler {
    async fn select_reader(&self, readers: &[Arc<ReaderNodeMeta>]) -> Option<Arc<ReaderNodeMeta>> {
        if readers.is_empty() {
            return None;
        }

        let index = self.counter.fetch_add(1, Ordering::Relaxed) % readers.len();
        let reader = readers[index].clone();
        reader.reserve_pending_shard_cache_allocation();
        Some(reader)
    }

    fn name(&self) -> &'static str {
        "round-robin"
    }
}

/// Balanced-count cache scheduler.
///
/// This scheduler selects the reader with the fewest assigned shards,
/// including ready, warming, and just-selected shards that have not been
/// inserted into the reader allocation map yet.
#[derive(Debug)]
pub struct BalancedCountCacheScheduler {
    counter: AtomicUsize,
}

impl Default for BalancedCountCacheScheduler {
    fn default() -> Self {
        Self::new()
    }
}

impl BalancedCountCacheScheduler {
    pub fn new() -> Self {
        Self {
            counter: AtomicUsize::new(0),
        }
    }
}

#[async_trait]
impl ReaderCacheScheduler for BalancedCountCacheScheduler {
    async fn select_reader(&self, readers: &[Arc<ReaderNodeMeta>]) -> Option<Arc<ReaderNodeMeta>> {
        if readers.is_empty() {
            return None;
        }

        let mut min_count = usize::MAX;
        let mut min_readers: Vec<Arc<ReaderNodeMeta>> = Vec::new();

        for reader in readers {
            let assigned = reader.shard_cache_schedule_load().await;
            match assigned.cmp(&min_count) {
                std::cmp::Ordering::Less => {
                    min_count = assigned;
                    min_readers.clear();
                    min_readers.push(reader.clone());
                }
                std::cmp::Ordering::Equal => {
                    min_readers.push(reader.clone());
                }
                std::cmp::Ordering::Greater => {}
            }
        }

        if min_readers.is_empty() {
            return None;
        }

        let idx = self.counter.fetch_add(1, Ordering::Relaxed) % min_readers.len();
        let reader = min_readers[idx].clone();
        reader.reserve_pending_shard_cache_allocation();
        Some(reader)
    }

    fn name(&self) -> &'static str {
        "balanced-count"
    }
}

/// Factory for creating cache schedulers based on configuration
pub struct CacheSchedulerFactory;

impl CacheSchedulerFactory {
    pub fn create_cache_scheduler(strategy: &str) -> Arc<dyn ReaderCacheScheduler> {
        match strategy.to_lowercase().as_str() {
            "round-robin" | "round_robin" | "roundrobin" => {
                Arc::new(RoundRobinCacheScheduler::new())
            }
            "balanced-count" | "balanced_count" | "balancedcount" => {
                Arc::new(BalancedCountCacheScheduler::new())
            }
            _ => Arc::new(BalancedCountCacheScheduler::new()),
        }
    }
}

#[cfg(test)]
mod tests {
    use tici_common::config::GrpcConfig;
    use tici_grpc::ReaderNodeStatus;

    use super::*;
    use crate::meta::reader::{CacheAllocator, ReaderNodeMeta, ShardCacheManager};

    fn create_test_readers(count: usize) -> Vec<Arc<ReaderNodeMeta>> {
        (0..count)
            .map(|i| {
                Arc::new(ReaderNodeMeta::new(
                    format!("reader_{}", i),
                    format!("tiflash_{}", i),
                    ReaderNodeStatus::default(),
                    GrpcConfig::default(),
                    CacheAllocator::new_for_test(),
                ))
            })
            .collect()
    }

    fn create_test_unified_shard(shard_id: u64) -> crate::meta::reader::UnifiedShardInfo {
        let mut info = crate::ShardInfo::new_for_test(crate::HandleRange::new(vec![], vec![1]));
        info.shard_id = shard_id;
        crate::meta::reader::UnifiedShardInfo::from_shard_info(&info)
    }

    #[tokio::test]
    async fn test_round_robin_cache_scheduler() {
        let scheduler = RoundRobinCacheScheduler::new();
        let readers = create_test_readers(3);

        // Test round-robin behavior
        let selections = vec![
            scheduler
                .select_reader(&readers)
                .await
                .unwrap()
                .addr
                .clone(),
            scheduler
                .select_reader(&readers)
                .await
                .unwrap()
                .addr
                .clone(),
            scheduler
                .select_reader(&readers)
                .await
                .unwrap()
                .addr
                .clone(),
            scheduler
                .select_reader(&readers)
                .await
                .unwrap()
                .addr
                .clone(),
        ];

        // Should cycle through readers in order
        assert_eq!(
            selections,
            vec!["reader_0", "reader_1", "reader_2", "reader_0"]
        );
    }

    #[tokio::test]
    async fn test_empty_readers() {
        let round_robin = RoundRobinCacheScheduler::new();
        let empty_readers = vec![];

        assert!(round_robin.select_reader(&empty_readers).await.is_none());
    }

    #[test]
    fn test_cache_scheduler_factory() {
        let rr_scheduler = CacheSchedulerFactory::create_cache_scheduler("round-robin");
        assert_eq!(rr_scheduler.name(), "round-robin");

        let bc_scheduler = CacheSchedulerFactory::create_cache_scheduler("balanced-count");
        assert_eq!(bc_scheduler.name(), "balanced-count");

        let default_scheduler = CacheSchedulerFactory::create_cache_scheduler("unknown");
        assert_eq!(default_scheduler.name(), "balanced-count");
    }

    #[tokio::test]
    async fn test_balanced_count_cache_scheduler_selects_least_assigned() {
        let scheduler = BalancedCountCacheScheduler::new();
        let readers = create_test_readers(2);

        readers[0]
            .add_shard_cache(create_test_unified_shard(1))
            .await;
        readers[0]
            .add_shard_cache(create_test_unified_shard(2))
            .await;

        let selected = scheduler.select_reader(&readers).await.unwrap();
        assert_eq!(selected.addr, "reader_1");
    }

    #[tokio::test]
    async fn test_cache_scheduler_integration() {
        // Test round-robin cache scheduler
        let mut config = crate::config::make_default_test_meta_config(false);
        config.cache_scheduling.strategy = "round-robin".to_string();
        let manager = ShardCacheManager::new(&config, false);

        // Create mock readers
        let readers = vec![
            Arc::new(ReaderNodeMeta::new(
                "127.0.0.1:8080".to_string(),
                "127.0.0.1:3930".to_string(),
                tici_grpc::ReaderNodeStatus {
                    cpu_usage: 50.0,
                    memory_usage: 1024,
                    requests_in_cycle: 10,
                    cache_hit_rate: 85.0,
                },
                GrpcConfig::default(),
                CacheAllocator::new_for_test(),
            )),
            Arc::new(ReaderNodeMeta::new(
                "127.0.0.1:8081".to_string(),
                "127.0.0.1:3931".to_string(),
                tici_grpc::ReaderNodeStatus {
                    cpu_usage: 45.0,
                    memory_usage: 1536,
                    requests_in_cycle: 12,
                    cache_hit_rate: 88.0,
                },
                GrpcConfig::default(),
                CacheAllocator::new_for_test(),
            )),
        ];

        // Test round-robin selection
        let mut selections = Vec::new();
        for _ in 0..4 {
            let selected = manager
                .cache_scheduler
                .select_reader(readers.as_slice())
                .await
                .unwrap();
            selections.push(selected.addr.clone());
        }

        // Should cycle through readers consistently
        assert_eq!(selections.len(), 4);
        assert_eq!(selections[0], selections[2]); // First and third should be same reader
        assert_eq!(selections[1], selections[3]); // Second and fourth should be same reader
        assert_ne!(selections[0], selections[1]); // First two should be different
    }
}