use std::{collections::HashSet, sync::Arc, time::Instant};

use anyhow::{Result, bail};
use async_channel::{Receiver, Sender, bounded};
use log::{debug, error, info, warn};
use scopeguard::defer;
use tici_grpc::{ErrorCode, ReaderNodeHeartbeatRequest, ReaderNodeStatus, ShardManifestHeader};
use tici_shard::{
    metadata::metadata_store::MetadataStore,
    reader::{frag_data_cache::FragDataCache, shard_mirror_manager::ShardMirrorManager},
    utils::meta_client_pool::MetaClientPool,
};
use tokio::{task::JoinHandle, time::interval};

use crate::{
    config::Config,
    metrics::{
        begin_shard_refresh_metrics, decrement_heartbeat_queue_backlog, inc_shard_refresh_event,
        increment_heartbeat_queue_backlog, observe_shard_refresh_queue_duration,
        reset_heartbeat_queue_backlog,
    },
};

#[derive(Debug)]
struct RespondedShardInfo {
    shard_header: ShardManifestHeader,
    enqueued_at: Instant,
}

pub struct HeartbeatManager {
    heartbeat_task: JoinHandle<()>,
    response_worker_tasks: Vec<JoinHandle<()>>,
}

impl HeartbeatManager {
    pub async fn new(
        config: Config,
        frag_data_cache: Arc<FragDataCache>,
        meta_client_pool: Arc<dyn MetaClientPool>,
        metadata_store: Arc<MetadataStore>,
        shard_mirror_manager: Arc<ShardMirrorManager>,
    ) -> Result<Self> {
        reset_heartbeat_queue_backlog();
        let (resp_tx, resp_rx) =
            bounded::<RespondedShardInfo>(config.tici.reader_node.heartbeat_queue_size as usize);
        let heartbeat_task = Self::spawn_heartbeat_loop(
            config.clone(),
            meta_client_pool.clone(),
            resp_tx,
            metadata_store.clone(),
            shard_mirror_manager.clone(),
        )
        .await?;
        let response_worker_tasks = Self::spawn_response_workers(
            resp_rx,
            frag_data_cache,
            meta_client_pool.clone(),
            &config,
            metadata_store,
            shard_mirror_manager.clone(),
        )
        .await?;
        Ok(Self {
            heartbeat_task,
            response_worker_tasks,
        })
    }

    async fn spawn_heartbeat_loop(
        config: Config,
        meta_client_pool: Arc<dyn MetaClientPool>,
        resp_tx: Sender<RespondedShardInfo>,
        metadata_store: Arc<MetadataStore>,
        shard_mirror_manager: Arc<ShardMirrorManager>,
    ) -> Result<JoinHandle<()>> {
        info!(
            "Starting reader node heartbeat manager for config={:?}",
            config.tici,
        );
        let handle = tokio::spawn(async move {
            let mut interval = interval(config.tici.reader_node.heartbeat_interval.0);
            loop {
                if let Err(e) = Self::send_heartbeat(
                    &config,
                    &meta_client_pool,
                    &resp_tx,
                    &metadata_store,
                    shard_mirror_manager.clone(),
                )
                .await
                {
                    error!("Heartbeat err={:?}", e);
                    // Continue without trying to reconnect immediately
                }
                // Wait for the next interval after sending heartbeat
                interval.tick().await;
            }
        });
        Ok(handle)
    }

    async fn spawn_response_workers(
        resp_rx: Receiver<RespondedShardInfo>,
        frag_data_cache: Arc<FragDataCache>,
        meta_client_pool: Arc<dyn MetaClientPool>,
        config: &Config,
        metadata_store: Arc<MetadataStore>,
        shard_mirror_manager: Arc<ShardMirrorManager>,
    ) -> Result<Vec<JoinHandle<()>>> {
        let pending_shards = Arc::new(parking_lot::Mutex::new(HashSet::new()));
        let mut handles = Vec::new();
        for _i in 0..config.tici.reader_node.heartbeat_worker_count {
            let resp_rx = resp_rx.clone();
            let frag_data_cache = frag_data_cache.clone();
            let meta_client_pool = meta_client_pool.clone();
            let pending_shards = pending_shards.clone();
            let metadata_store = metadata_store.clone();
            let shard_mirror_manager_clone = shard_mirror_manager.clone();
            let handle = tokio::spawn(async move {
                while let Ok(shard_info) = resp_rx.recv().await {
                    decrement_heartbeat_queue_backlog();
                    let queue_wait = shard_info.enqueued_at.elapsed();
                    observe_shard_refresh_queue_duration(queue_wait);
                    {
                        let mut guard = pending_shards.lock();
                        if !guard.insert(shard_info.shard_header.shard_id) {
                            inc_shard_refresh_event("skip_inflight");
                            info!(
                                "Shard {} is already processing, skip. target_seq={}, queued_ms={}, pending_shards={:?}",
                                shard_info.shard_header.shard_id,
                                shard_info.shard_header.seq,
                                queue_wait.as_millis(),
                                guard
                            );
                            continue;
                        }
                    }
                    defer!({
                        let mut guard = pending_shards.lock();
                        guard.remove(&shard_info.shard_header.shard_id);
                    });
                    let refresh_metrics = begin_shard_refresh_metrics("refresh_shard");
                    let refresh_start = Instant::now();
                    let _ = shard_mirror_manager_clone
                        .try_refresh_shard(
                            &shard_info.shard_header,
                            frag_data_cache.clone(),
                            meta_client_pool.clone(),
                            metadata_store.clone(),
                        )
                        .await
                        .inspect_err(|e| {
                            refresh_metrics.mark_error();
                            error!(
                                "Refresh shard failed, shard_id={}, target_seq={}, queued_ms={}, refresh_ms={}, err={:?}",
                                shard_info.shard_header.shard_id,
                                shard_info.shard_header.seq,
                                queue_wait.as_millis(),
                                refresh_start.elapsed().as_millis(),
                                e
                            );
                        });
                }
            });
            handles.push(handle);
        }
        Ok(handles)
    }

    /// Stop the heartbeat process
    pub async fn stop(&self) {
        self.heartbeat_task.abort();
        for handle in self.response_worker_tasks.iter() {
            handle.abort();
        }
        reset_heartbeat_queue_backlog();
        info!("Stopping reader node heartbeat manager");
    }

    /// Send heartbeat to meta service
    async fn send_heartbeat(
        config: &Config,
        meta_client_pool: &Arc<dyn MetaClientPool>,
        resp_tx: &Sender<RespondedShardInfo>,
        metadata_store: &Arc<MetadataStore>,
        shard_mirror_manager: Arc<ShardMirrorManager>,
    ) -> Result<()> {
        let mut client = meta_client_pool.get_meta_client().await?;
        let advertise_addr = if config.tici.reader_node.advertise_addr.is_empty() {
            config.tici.reader_node.addr.clone()
        } else {
            config.tici.reader_node.advertise_addr.clone()
        };
        let mut metrics = crate::metrics::SystemMetrics::new();
        let cached_shards = shard_mirror_manager.get_cached_shards();
        let cached_shard_ids = cached_shards.iter().map(|shard| shard.shard_id).collect();
        let requests_in_cycle = crate::metrics::get_requests_in_cycle();
        let request = ReaderNodeHeartbeatRequest {
            addr: advertise_addr,
            status: Some(ReaderNodeStatus {
                cpu_usage: metrics.get_cpu_usage(),
                memory_usage: metrics.get_memory_usage(),
                requests_in_cycle,
                cache_hit_rate: metrics.get_cache_hit_rate(),
            }),
            tiflash_server_addr: config.flash.service_addr.clone(),
            cached_shard_ids,
            cached_shards,
        };
        debug!("Heartbeat request={:?}", request);
        let resp = client.reader_node_heartbeat(request).await?;
        if resp.status != ErrorCode::Success as i32 {
            bail!(
                "Heartbeat failed with status={}, message={}, meta_addr={}",
                resp.status,
                resp.error_message,
                client.get_addr()
            );
        }
        // Reset the requests counter after successful heartbeat
        crate::metrics::reset_requests_in_cycle();

        Self::handle_heartbeat_resp(
            resp.shards_to_remove.as_ref(),
            resp.shards_to_keep.as_ref(),
            resp_tx,
            metadata_store,
            shard_mirror_manager,
        )
        .await;
        Ok(())
    }

    async fn handle_heartbeat_resp(
        shards_to_remove: &[ShardManifestHeader],
        shards_to_keep: &[ShardManifestHeader],
        resp_tx: &Sender<RespondedShardInfo>,
        metadata_store: &Arc<MetadataStore>,
        shard_mirror_manager: Arc<ShardMirrorManager>,
    ) {
        let mut shard_ids_to_remove = vec![];
        for shard in shards_to_remove {
            let shard_id = shard.shard_id;
            let remove_res: Result<()> = async {
                let rx = metadata_store
                    .remove_reader_local_shard_metadata(shard_id)
                    .await?;
                rx.await??;
                Ok(())
            }
            .await;
            if let Err(err) = remove_res {
                warn!(
                    "Failed to remove local reader shard metadata atomically, keep reporting shard_id={}, err={:?}",
                    shard_id, err
                );
                continue;
            }
            shard_ids_to_remove.push(shard_id);
        }
        shard_mirror_manager.remove_shards(&shard_ids_to_remove);

        let mut fresh_shards = vec![];
        let mut stale_shards = vec![];
        for shard in shards_to_keep {
            let shard_mirror = shard_mirror_manager.get_shard(shard.shard_id);
            if let Some(shard_mirror) = shard_mirror {
                let curr_seq = shard_mirror.get_curr_seq();
                let new_seq = shard.seq;
                match curr_seq.cmp(&new_seq) {
                    std::cmp::Ordering::Greater | std::cmp::Ordering::Equal => {
                        // If concurrent refresh is happened, curr_seq can be greater than new_seq.
                        fresh_shards.push((shard.shard_id, curr_seq));
                    }
                    std::cmp::Ordering::Less => {
                        stale_shards.push((shard.shard_id, curr_seq, new_seq));
                        if let Err(e) = resp_tx.try_send(RespondedShardInfo {
                            shard_header: shard.clone(),
                            enqueued_at: Instant::now(),
                        }) {
                            inc_shard_refresh_event("enqueue_failed");
                            warn!(
                                "Failed to send shard info to resp_tx, shard={:?}, err={:?}",
                                shard, e
                            );
                        } else {
                            increment_heartbeat_queue_backlog();
                            inc_shard_refresh_event("enqueued");
                        }
                    }
                }
            }
        }
        info!(
            "Removed: {:?}, Stale: {:?}",
            shard_ids_to_remove, stale_shards
        );
        debug!("Fresh shards: {:?}", fresh_shards);
    }
}