use std::{
    collections::HashMap,
    sync::{
        Arc,
        atomic::{AtomicU32, Ordering},
    },
    time::Duration,
};

use anyhow::Result;
use async_trait::async_trait;
use log::info;
use portable_atomic::AtomicUsize;
use tici_common::index::IndexKey;
pub use tici_grpc::test_utils::*;
use tici_grpc::{
    GetDebugCacheInfoRequest, GrpcMetaServiceInternalClient, GrpcReaderService,
    GrpcReaderServiceServer, ReaderNodeHeartbeatRequest, ShardManifestHeader, WarmupShardRequest,
    WarmupShardResponse,
};
use tokio::sync::{Mutex, RwLock};
use tonic::{
    Request, Response, Status,
    transport::{Channel, Server},
};

/// Waits for and validates heartbeat response
pub async fn send_and_validate_heartbeat(
    client: &mut tici_grpc::GrpcMetaServiceInternalClient<tonic::transport::Channel>,
    request: ReaderNodeHeartbeatRequest,
    expected_status: i32,
) -> Result<tici_grpc::ReaderNodeHeartbeatResponse> {
    info!("reader send heartbeat, req={:?}", request);
    let response = client.reader_node_heartbeat(request.clone()).await.unwrap();
    let response_inner = response.into_inner();

    assert_eq!(
        response_inner.status, expected_status,
        "Heartbeat should return status {} for reader {}, but got {}",
        expected_status, request.addr, response_inner.status
    );

    Ok(response_inner)
}

#[derive(Debug, Clone)]
pub struct MockReader {
    warmup_count: Arc<AtomicU32>,
    pub addr: String,
    pub tiflash_addr: String,
    pub cached_shard_seqs: Arc<Mutex<HashMap<u64, u64>>>,
    warmup_count_by_index: Arc<Mutex<HashMap<IndexKey, u32>>>,
    pub meta_client: GrpcMetaServiceInternalClient<Channel>,
    heartbeat_stoppers: Arc<RwLock<Option<tokio::sync::oneshot::Sender<()>>>>,
    heartbeat_handles: Arc<RwLock<Option<tokio::task::JoinHandle<()>>>>,
}

impl MockReader {
    pub fn new(
        addr: String,
        tiflash_addr: String,
        meta_client: GrpcMetaServiceInternalClient<Channel>,
    ) -> Self {
        Self {
            warmup_count: Arc::new(AtomicU32::new(0)),
            addr,
            tiflash_addr,
            cached_shard_seqs: Arc::new(Mutex::new(HashMap::new())),
            warmup_count_by_index: Arc::new(Mutex::new(HashMap::new())),
            meta_client,
            heartbeat_stoppers: Arc::new(RwLock::new(None)),
            heartbeat_handles: Arc::new(RwLock::new(None)),
        }
    }

    pub fn start_server(&self) {
        let reader_addr = self.addr.parse().unwrap();
        let cloned = self.clone();
        tokio::spawn(async move {
            let _ = Server::builder()
                .add_service(GrpcReaderServiceServer::new(cloned))
                .serve(reader_addr)
                .await;
        });
    }

    pub async fn start_heartbeat(&mut self, heartbeat_interval: Duration) {
        let heartbeat_req =
            create_test_heartbeat_request(self.addr.clone(), self.tiflash_addr.clone(), vec![]);

        send_and_validate_heartbeat(&mut self.meta_client, heartbeat_req.clone(), 0)
            .await
            .unwrap();

        // Start heartbeat task
        let (stop_tx, mut stop_rx) = tokio::sync::oneshot::channel::<()>();
        let _ = self.heartbeat_stoppers.write().await.replace(stop_tx);

        let heartbeat_req_template = heartbeat_req.clone();
        let reader_service = self.clone();

        let mut internal_client_clone = self.meta_client.clone();
        let handle = tokio::spawn(async move {
            let mut interval = tokio::time::interval(heartbeat_interval);
            loop {
                let errcount = Arc::new(AtomicUsize::new(0));
                tokio::select! {
                    _ = interval.tick() => {
                        errcount.store(0, Ordering::Relaxed);
                        log::debug!("ticker");
                        let cached_shards_before = reader_service.get_cached_shards().await;
                        let cached_shard_ids_before = cached_shards_before
                            .iter()
                            .map(|shard| shard.shard_id)
                            .collect::<Vec<_>>();
                        let heartbeat_req = ReaderNodeHeartbeatRequest {
                            cached_shard_ids: cached_shard_ids_before.clone(),
                            cached_shards: cached_shards_before,
                            ..heartbeat_req_template.clone()
                        };

                        match internal_client_clone.reader_node_heartbeat(heartbeat_req).await {
                            Ok(response) => {
                                let response_inner = response.into_inner();
                                log::debug!("ticker finish, response={:?}", response_inner);

                                // Process cache synchronization based on heartbeat response
                                reader_service.process_heartbeat_response(&response_inner, &heartbeat_req_template.addr).await;

                                let cached_shard_ids_after = reader_service.get_cached_shard_ids().await;
                                if cached_shard_ids_after.len() != cached_shard_ids_before.len() {
                                    info!(
                                        "reader {} cache updated, {} -> {} shards, {:?} ==> {:?}",
                                        heartbeat_req_template.addr,
                                        cached_shard_ids_before.len(),
                                        cached_shard_ids_after.len(),
                                        cached_shard_ids_before,
                                        cached_shard_ids_after
                                    );
                                }
                            }
                            Err(e) => {
                                errcount.fetch_add(1, Ordering::Relaxed);
                                log::debug!("ticker err, e={:?}", e);
                                if errcount.load(Ordering::Relaxed) > 10 {
                                    break;
                                }
                            },
                        }
                    }
                    _ = &mut stop_rx => {
                        info!("stop heartbeat");
                        break
                    },
                }
            }
        });
        let _ = self.heartbeat_handles.write().await.replace(handle);
    }

    pub async fn shutdown(&mut self) {
        if let Some(s) = self.heartbeat_stoppers.write().await.take() {
            s.send(()).unwrap();
        }
        if let Some(s) = self.heartbeat_handles.write().await.take() {
            s.abort();
            let _ = s.await;
        }
    }

    pub async fn get_warmup_count_for_index(&self, index_key: &IndexKey) -> u32 {
        self.warmup_count_by_index
            .lock()
            .await
            .get(index_key)
            .copied()
            .unwrap_or(0)
    }

    pub async fn get_cached_shard_ids(&self) -> Vec<u64> {
        self.cached_shard_seqs
            .lock()
            .await
            .keys()
            .copied()
            .collect()
    }

    pub async fn get_cached_shards(&self) -> Vec<tici_grpc::ReaderCachedShard> {
        self.cached_shard_seqs
            .lock()
            .await
            .iter()
            .map(|(&shard_id, &seq)| tici_grpc::ReaderCachedShard { shard_id, seq })
            .collect()
    }

    pub async fn clear_all_shards(&self) {
        self.cached_shard_seqs.lock().await.clear();
    }

    pub async fn has_shard(&self, shard_id: u64) -> bool {
        self.cached_shard_seqs.lock().await.contains_key(&shard_id)
    }

    pub async fn remove_shard(&self, shard_id: u64) -> bool {
        self.cached_shard_seqs
            .lock()
            .await
            .remove(&shard_id)
            .is_some()
    }

    pub async fn get_shard_count(&self) -> usize {
        self.cached_shard_seqs.lock().await.len()
    }

    pub async fn process_heartbeat_response(
        &self,
        response: &tici_grpc::ReaderNodeHeartbeatResponse,
        reader_addr: &str,
    ) -> (Vec<ShardManifestHeader>, Vec<ShardManifestHeader>) {
        let shards_to_keep = &response.shards_to_keep;
        let shards_to_remove = &response.shards_to_remove;

        let shards_to_keep_ids: Vec<u64> = shards_to_keep.iter().map(|h| h.shard_id).collect();
        let shards_to_remove_ids: Vec<u64> = shards_to_remove.iter().map(|h| h.shard_id).collect();

        // Update reader's cached shards based on meta's response
        {
            let mut cached_shard_seqs = self.cached_shard_seqs.lock().await;

            // Remove shards that meta says should be removed
            for shard_header in shards_to_remove {
                cached_shard_seqs.remove(&shard_header.shard_id);
            }

            // Ensure shards that meta says should be kept are present
            for shard_header in shards_to_keep {
                let entry = cached_shard_seqs.entry(shard_header.shard_id).or_insert(0);
                *entry = (*entry).max(shard_header.seq);
            }
        }

        if !shards_to_remove.is_empty() || !shards_to_keep.is_empty() {
            info!(
                "reader {} processed heartbeat keep=[{}]({:?}), remove=[{}]{:?}",
                reader_addr,
                shards_to_keep.len(),
                shards_to_keep_ids,
                shards_to_remove.len(),
                shards_to_remove_ids
            );
        }

        (shards_to_keep.clone(), shards_to_remove.clone())
    }
}

#[async_trait]
impl GrpcReaderService for MockReader {
    async fn warmup_shard(
        &self,
        request: Request<WarmupShardRequest>,
    ) -> Result<Response<WarmupShardResponse>, Status> {
        let req = request.into_inner();

        // Validate required fields
        let shard = req
            .shard
            .as_ref()
            .ok_or_else(|| Status::invalid_argument("Shard information is required"))?;

        let table_info = req
            .table_info
            .as_ref()
            .ok_or_else(|| Status::invalid_argument("Table information is required"))?;

        let index_info = req
            .index_info
            .as_ref()
            .ok_or_else(|| Status::invalid_argument("Index information is required"))?;

        if index_info.table_id != table_info.table_id {
            return Err(Status::invalid_argument(format!(
                "Table ID mismatch: index_info.table_id={}, table_info.table_id={}",
                index_info.table_id, table_info.table_id
            )));
        }

        let shard_id = shard.shard_id;
        let index_key = IndexKey::new(
            table_info.keyspace_id,
            table_info.table_id,
            index_info.index_id,
        );

        let warmup_count = self.warmup_count.fetch_add(1, Ordering::Relaxed) + 1;

        // Update per-index warmup count
        {
            let mut count_map = self.warmup_count_by_index.lock().await;
            *count_map.entry(index_key.clone()).or_insert(0) += 1;
        }

        // Add shard to cached set
        self.cached_shard_seqs
            .lock()
            .await
            .insert(shard_id, shard.seq);

        info!(
            "Reader {} warmup shard {} for index_key={:?}, warmup_count={}, index_count={}",
            self.addr,
            shard_id,
            index_key,
            warmup_count,
            self.get_warmup_count_for_index(&index_key).await
        );

        Ok(Response::new(WarmupShardResponse {
            status: 0,
            error_message: String::new(),
        }))
    }

    async fn get_debug_cache_info(
        &self,
        _request: Request<GetDebugCacheInfoRequest>,
    ) -> Result<Response<tici_grpc::GetDebugCacheInfoResponse>, Status> {
        let cached_ids = self.get_cached_shard_ids().await;

        Ok(Response::new(tici_grpc::GetDebugCacheInfoResponse {
            status: 0,
            error_message: String::new(),
            cached_shard_count: cached_ids.len() as i32,
            cached_shard_ids: cached_ids,
        }))
    }
}