pub mod cache_scheduler;
pub mod reader_meta;
pub mod warmup;

// TODO: Add monitoring and statistics functionality later
// TODO: Add advanced load balancing strategies (weighted, round-robin, etc.)
use std::{
    collections::{HashMap, HashSet},
    sync::Arc,
};

use anyhow::{Result, anyhow};
pub use cache_scheduler::*;
use log::{debug, error, info, warn};
pub use reader_meta::*;
use tici_common::config::GrpcConfig;
use tici_grpc::{ReaderNodeHeartbeatRequest, ReaderNodeHeartbeatResponse};
use tici_system::readable::ReadableDuration;
use tokio::{sync::RwLock, task::JoinHandle, time::interval};

use crate::{
    config::{MetaServiceConfig, ReaderPoolConfig},
    meta::{
        metrics::SHARD_CACHE_READER_COUNT_GAUGE,
        reader::{
            AllocateReaderReason, UnifiedShardInfo,
            shard_cache::ShardCacheManager,
            warmup::{WarmupMgr, WarmupTask},
        },
    },
};

/// Pool of reader nodes with automatic cleanup and load balancing
/// TODO(P2) Make this struct accessible from ShardCacheManager.
pub struct ReaderPool {
    readers: RwLock<HashMap<String, Arc<ReaderNodeMeta>>>,
    blocked_read_addrs: RwLock<HashSet<String>>,
    config: ReaderPoolConfig,
    grpc_client_config: GrpcConfig,
    // TODO(P2) Make this a message sender to ShardCacheManager
    pub shard_cache_manager: Option<Arc<ShardCacheManager>>,
    warmup_mgr: Arc<WarmupMgr>,
}

impl ReaderPool {
    /// Create a new ReaderPool with configuration and shard cache manager
    pub fn new(meta_config: &MetaServiceConfig) -> Self {
        Self::new_with_warmup_mgr(meta_config, WarmupMgr::new())
    }

    pub fn new_with_warmup_mgr(meta_config: &MetaServiceConfig, warmup_mgr: WarmupMgr) -> Self {
        let config = meta_config.reader_pool.clone();
        Self {
            readers: RwLock::new(HashMap::new()),
            blocked_read_addrs: RwLock::new(HashSet::new()),
            config,
            grpc_client_config: meta_config.server.grpc_client.clone(),
            shard_cache_manager: None,
            warmup_mgr: Arc::new(warmup_mgr),
        }
    }

    pub fn init(&mut self, shard_cache_manager: Arc<ShardCacheManager>) {
        let _ = self.shard_cache_manager.insert(shard_cache_manager);
    }

    pub fn warmup_mgr(&self) -> Arc<WarmupMgr> {
        self.warmup_mgr.clone()
    }

    pub fn background_warmup(
        &self,
        unified_shard: &UnifiedShardInfo,
        reader_pool: Arc<ReaderPool>,
        table_info: tici_grpc::TableInfo,
        index_info: tici_grpc::IndexInfo,
        reason: AllocateReaderReason,
    ) -> Result<JoinHandle<()>> {
        let in_warming = unified_shard.in_warming.clone();
        let header = unified_shard.get_header();
        let reader_pool_clone = reader_pool.clone();
        let unified_shard = unified_shard.clone();
        let shard_cache_manager = self
            .shard_cache_manager
            .clone()
            .ok_or(anyhow!("shard cache manager not inited"))?;
        let warmup_task = WarmupTask::new(
            unified_shard.clone(),
            index_info.clone(),
            table_info.clone(),
            reason.clone(),
        );
        let warmup_mgr = self.warmup_mgr.clone();
        Ok(tokio::spawn(async move {
            warmup_mgr
                .do_warmup(
                    reader_pool_clone.clone(),
                    warmup_task,
                    header,
                    in_warming,
                    shard_cache_manager,
                )
                .await;
        }))
    }

    pub fn background_add_replica(
        &self,
        unified_shard: &UnifiedShardInfo,
        reader_pool: Arc<ReaderPool>,
        table_info: tici_grpc::TableInfo,
        index_info: tici_grpc::IndexInfo,
        reason: AllocateReaderReason,
    ) -> Result<JoinHandle<()>> {
        let in_warming = unified_shard.in_warming.clone();
        let header = unified_shard.get_header();
        let reader_pool_clone = reader_pool.clone();
        let unified_shard = unified_shard.clone();
        let shard_cache_manager = self
            .shard_cache_manager
            .clone()
            .ok_or(anyhow!("shard cache manager not inited"))?;
        let warmup_task = WarmupTask::new(
            unified_shard.clone(),
            index_info.clone(),
            table_info.clone(),
            reason.clone(),
        );
        let warmup_mgr = self.warmup_mgr.clone();
        Ok(tokio::spawn(async move {
            warmup_mgr
                .add_new_peer(
                    reader_pool_clone.clone(),
                    warmup_task,
                    header,
                    in_warming,
                    shard_cache_manager,
                )
                .await;
        }))
    }

    pub fn background_add_replica_designated(
        &self,
        unified_shard: &UnifiedShardInfo,
        reader_pool: Arc<ReaderPool>,
        table_info: tici_grpc::TableInfo,
        index_info: tici_grpc::IndexInfo,
        reason: AllocateReaderReason,
        designated_reader_addr: String,
    ) -> Result<JoinHandle<()>> {
        let in_warming = unified_shard.in_warming.clone();
        let header = unified_shard.get_header();
        let reader_pool_clone = reader_pool.clone();
        let unified_shard = unified_shard.clone();
        let shard_cache_manager = self
            .shard_cache_manager
            .clone()
            .ok_or(anyhow!("shard cache manager not inited"))?;
        let warmup_task = WarmupTask::new(
            unified_shard.clone(),
            index_info.clone(),
            table_info.clone(),
            reason.clone(),
        );
        let warmup_mgr = self.warmup_mgr.clone();
        Ok(tokio::spawn(async move {
            warmup_mgr
                .add_new_peer_designated(
                    reader_pool_clone.clone(),
                    warmup_task,
                    header,
                    in_warming,
                    shard_cache_manager,
                    designated_reader_addr,
                )
                .await;
        }))
    }

    /// Start the reader pool and background cleanup task to remove expired
    /// readers
    pub async fn start(self: Arc<Self>) {
        let pool = self.clone();
        tokio::spawn(async move {
            let mut interval = interval(pool.config.cleanup_interval_duration.0);
            loop {
                interval.tick().await;
                if let Err(e) = pool.cleanup_expired_readers().await {
                    warn!("failed to cleanup expired readers, err={}", e);
                }
                // Periodically refresh active reader count metric
                let active = pool.get_active_readers().await.len() as i64;
                SHARD_CACHE_READER_COUNT_GAUGE.set(active);
            }
        });
    }

    pub async fn on_warmup_finished(&self, set: HashSet<u64>) {
        for (_, read_node) in self.readers.read().await.iter() {
            for s in set.iter() {
                read_node.remove_shard_cache(s).await
            }
        }
    }

    /// Register or update a reader node
    pub async fn register_or_update_reader(
        &self,
        req: &ReaderNodeHeartbeatRequest,
    ) -> Result<Arc<ReaderNodeMeta>> {
        let status = req
            .status
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("Missing reader node status"))?
            .clone();
        let mut readers = self.readers.write().await;
        match readers.get(&req.addr) {
            Some(reader) => {
                reader.on_heartbeat_update_status(status).await;
                Ok(reader.clone())
            }
            None => {
                if let Some(manager) = self.shard_cache_manager.as_ref() {
                    let reader = Arc::new(ReaderNodeMeta::new(
                        req.addr.clone(),
                        req.tiflash_server_addr.clone(),
                        status,
                        self.grpc_client_config.clone(),
                        manager.make_cache_allocator(),
                    ));
                    readers.insert(req.addr.clone(), reader.clone());
                    info!("reader node registered, addr={}", req.addr);
                    Ok(reader)
                } else {
                    Err(anyhow!("shard cache manager is not set"))
                }
            }
        }
    }

    /// Get all reader nodes (including expired ones)
    pub async fn get_all_readers(&self) -> Vec<Arc<ReaderNodeMeta>> {
        self.readers.read().await.values().cloned().collect()
    }

    pub async fn get_readers_len(&self) -> usize {
        self.readers.read().await.values().len()
    }

    pub async fn get_shard_assigned(&self, shard_id: &u64) -> Vec<Arc<ReaderNodeMeta>> {
        let mut res = vec![];
        for (_, node) in self.readers.read().await.iter() {
            if node.has_shard_cache(*shard_id).await {
                res.push(node.clone());
            }
        }
        res
    }

    pub async fn get_reader_schedule_load(&self, reader_addr: &str) -> Option<usize> {
        let reader = self.readers.read().await.get(reader_addr).cloned();
        match reader {
            Some(reader) => Some(reader.shard_cache_schedule_load().await),
            None => None,
        }
    }

    pub async fn get_schedulable_reader_schedule_load(&self, reader_addr: &str) -> Option<usize> {
        if self.is_reader_blocked(reader_addr).await {
            return None;
        }

        let reader = self.readers.read().await.get(reader_addr).cloned();
        let reader = match reader {
            Some(reader) => reader,
            None => return None,
        };
        if reader.is_expired(self.config.ttl_duration).await {
            return None;
        }
        Some(reader.shard_cache_schedule_load().await)
    }

    pub async fn block_readers(&self, addrs: &[String]) {
        let mut g = self.blocked_read_addrs.write().await;
        for addr in addrs {
            g.insert(addr.clone());
        }
    }

    pub async fn unblock_readers(&self, addrs: &[String]) {
        let mut g = self.blocked_read_addrs.write().await;
        for addr in addrs {
            g.remove(addr);
        }
    }

    pub async fn list_blocked_readers(&self) -> Vec<String> {
        let mut out: Vec<String> = self
            .blocked_read_addrs
            .read()
            .await
            .iter()
            .cloned()
            .collect();
        out.sort();
        out
    }

    pub async fn is_reader_blocked(&self, addr: &str) -> bool {
        self.blocked_read_addrs.read().await.contains(addr)
    }

    pub async fn remove_shards(&self, shard_ids: &[u64]) {
        for (_, node) in self.readers.read().await.iter() {
            node.remove_shard_caches(shard_ids).await;
        }
    }

    /// Remove shard allocations from a specific reader only.
    pub async fn remove_shards_from_reader(&self, addr: &str, shard_ids: &[u64]) {
        if let Some(reader) = self.readers.read().await.get(addr) {
            reader.remove_shard_caches(shard_ids).await;
        }
    }

    /// Get only active (non-expired) reader nodes
    pub async fn get_active_readers(&self) -> Vec<Arc<ReaderNodeMeta>> {
        let readers = self.readers.read().await;
        let blocked = self.blocked_read_addrs.read().await;
        let mut active_readers = Vec::new();

        for reader in readers.values() {
            if blocked.contains(reader.addr.as_str()) {
                continue;
            }
            if !reader.is_expired(self.config.ttl_duration).await {
                active_readers.push(reader.clone());
            }
        }

        active_readers
    }

    pub async fn get_readers_for_schedule(&self) -> Vec<Arc<ReaderNodeMeta>> {
        let blocked = self.blocked_read_addrs.read().await;
        let mut ans = vec![];
        for reader in self.get_active_readers().await {
            if blocked.contains(reader.addr.as_str()) {
                continue;
            }
            ans.push(reader);
        }
        ans
    }

    /// Remove expired readers and clean up associated cache allocations
    pub async fn cleanup_expired_readers(&self) -> Result<usize> {
        let expired_addrs = self.find_expired_readers().await;
        let removed_count = expired_addrs.len();

        if removed_count == 0 {
            return Ok(0);
        }

        // Remove from reader pool
        {
            let mut readers = self.readers.write().await;
            for addr in &expired_addrs {
                readers.remove(addr);
            }
        }

        // Clean up cache allocations
        let mut total_cache_removals = 0;
        for addr in &expired_addrs {
            if let Some(shard_cache_manager) = self.shard_cache_manager.as_ref() {
                let cache_removals = shard_cache_manager.on_reader_node_down(addr).await;
                total_cache_removals += cache_removals;
            } else {
                error!("shard_cache_manager not set");
            }
        }

        info!(
            "Removed {} expired reader nodes and {} cache allocations",
            removed_count, total_cache_removals
        );

        // Update active reader gauge after cleanup
        let active = self.get_readers_for_schedule().await.len() as i64;
        SHARD_CACHE_READER_COUNT_GAUGE.set(active);

        Ok(removed_count)
    }

    /// Find all expired reader addresses
    async fn find_expired_readers(&self) -> Vec<String> {
        let readers = self.readers.read().await;
        let mut expired_addrs = Vec::new();

        for (addr, reader) in readers.iter() {
            if reader.is_expired(self.config.ttl_duration).await {
                expired_addrs.push(addr.clone());
            }
        }

        expired_addrs
    }

    pub async fn on_msg_reader_heartbeat(
        &self,
        req: ReaderNodeHeartbeatRequest,
        dead_timeout: Option<ReadableDuration>,
    ) -> Result<ReaderNodeHeartbeatResponse> {
        let read_node = self.register_or_update_reader(&req).await?;

        let shard_cache_manager = self
            .shard_cache_manager
            .as_ref()
            .ok_or(anyhow!("shard cache not bootstrap"))?;

        let dead_timeout =
            dead_timeout.unwrap_or(shard_cache_manager.config.cache_scheduling.dead_timeout);

        let reported_shards_by_id: HashMap<u64, Option<u64>> = if req.cached_shards.is_empty() {
            req.cached_shard_ids
                .iter()
                .copied()
                .map(|shard_id| (shard_id, None))
                .collect()
        } else {
            req.cached_shards
                .iter()
                .map(|shard| (shard.shard_id, Some(shard.seq)))
                .collect()
        };

        let (shards_to_keep, shards_to_remove) = read_node
            .handle_hearbeart_for_shard_cache(
                &reported_shards_by_id,
                dead_timeout,
                self,
                shard_cache_manager.as_ref(),
            )
            .await;

        // Update active reader gauge after processing heartbeat
        let active = self.get_readers_for_schedule().await.len() as i64;
        SHARD_CACHE_READER_COUNT_GAUGE.set(active);

        if !shards_to_keep.is_empty() || !shards_to_remove.is_empty() {
            debug!(
                "on_msg_reader_heartbeat: shards_to_keep={:?}, shards_to_remove={:?}",
                shards_to_keep, shards_to_remove
            );
        }

        Ok(ReaderNodeHeartbeatResponse {
            status: 0,
            error_message: String::new(),
            shards_to_keep,
            shards_to_remove,
        })
    }

    /// Get reader pool configuration
    pub fn get_config(&self) -> &ReaderPoolConfig {
        &self.config
    }

    /// Check if a specific reader is active (not expired)
    pub async fn is_reader_active(&self, reader_addr: &str) -> bool {
        if self.is_reader_blocked(reader_addr).await {
            return false;
        }
        let readers = self.readers.read().await;
        match readers.get(reader_addr) {
            Some(reader) => !reader.is_expired(self.config.ttl_duration).await,
            None => false,
        }
    }

    /// Get reader statistics for monitoring
    pub async fn get_stats(&self) -> ReaderPoolStats {
        let readers = self.readers.read().await;
        let total_readers = readers.len();
        let mut active_readers = 0;
        let mut total_cpu_usage = 0.0;
        let mut total_memory_usage = 0;

        for reader in readers.values() {
            if !reader.is_expired(self.config.ttl_duration).await {
                active_readers += 1;
                let status = reader.get_status().await;
                total_cpu_usage += status.cpu_usage;
                total_memory_usage += status.memory_usage;
            }
        }

        ReaderPoolStats {
            total_readers,
            active_readers,
            expired_readers: total_readers - active_readers,
            avg_cpu_usage: if active_readers > 0 {
                total_cpu_usage / active_readers as f64
            } else {
                0.0
            },
            avg_memory_usage: if active_readers > 0 {
                total_memory_usage / active_readers as u64
            } else {
                0
            },
        }
    }
}

/// Statistics for the reader pool
#[derive(Debug, Clone)]
pub struct ReaderPoolStats {
    pub total_readers: usize,
    pub active_readers: usize,
    pub expired_readers: usize,
    pub avg_cpu_usage: f64,
    pub avg_memory_usage: u64,
}

#[cfg(test)]
mod tests {
    use tici_grpc::{ReaderNodeStatus, test_utils::create_test_heartbeat_request};
    use tokio::time::{Duration, sleep};

    use super::*;
    use crate::meta::reader::shard_cache::ShardCacheManager;

    #[tokio::test]
    async fn test_reader_pool_basic_operations() {
        let meta_config = MetaServiceConfig {
            reader_pool: ReaderPoolConfig {
                ttl_duration: ReadableDuration::secs(3),
                cleanup_interval_duration: ReadableDuration::secs(1),
                ..Default::default()
            },
            ..Default::default()
        };
        let shard_cache_manager = ShardCacheManager::new(&meta_config, false);
        let mut pool = ReaderPool::new(&meta_config);
        pool.init(shard_cache_manager);

        let status = ReaderNodeStatus {
            cpu_usage: 50.0,
            memory_usage: 1024,
            requests_in_cycle: 10,
            cache_hit_rate: 85.0,
        };
        let mut req = create_test_heartbeat_request(
            "127.0.0.1:8080".to_string(),
            "127.0.0.1:3930".to_string(),
            vec![],
        );
        req.status = Some(status);

        let reader = pool.register_or_update_reader(&req).await.unwrap();
        assert_eq!(reader.addr, "127.0.0.1:8080");

        let all_readers = pool.get_all_readers().await;
        assert_eq!(all_readers.len(), 1);
        assert_eq!(all_readers[0].addr, "127.0.0.1:8080");

        // Test that readers are active
        let active_readers = pool.get_readers_for_schedule().await;
        assert_eq!(active_readers.len(), 1);
        assert_eq!(active_readers[0].addr, "127.0.0.1:8080");

        // Test stats
        let stats = pool.get_stats().await;
        assert_eq!(stats.total_readers, 1);
        assert_eq!(stats.active_readers, 1);
        assert_eq!(stats.expired_readers, 0);
    }

    #[tokio::test]
    async fn test_reader_pool_ttl_expiration() {
        let meta_config = MetaServiceConfig {
            reader_pool: ReaderPoolConfig {
                ttl_duration: ReadableDuration::secs(1),
                cleanup_interval_duration: ReadableDuration::secs(1),
                ..Default::default()
            },
            ..Default::default()
        };
        let shard_cache_manager = ShardCacheManager::new(&meta_config, false);
        let mut pool = ReaderPool::new(&meta_config);
        pool.init(shard_cache_manager);
        pool.register_or_update_reader(&create_test_heartbeat_request(
            "127.0.0.1:8080".to_string(),
            "127.0.0.1:3930".to_string(),
            vec![],
        ))
        .await
        .unwrap();
        assert_eq!(pool.get_all_readers().await.len(), 1);

        // Wait for TTL to expire
        sleep(Duration::from_secs(2)).await;

        let removed_count = pool.cleanup_expired_readers().await.unwrap();
        assert_eq!(removed_count, 1);
        assert_eq!(pool.get_all_readers().await.len(), 0);
    }
}