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

use anyhow::Result;
use either::Either;
use log::{debug, info, warn};
use portable_atomic::{AtomicBool, AtomicU16};
use tici_grpc::{ShardManifestHeader, WarmupShardRequest};
use tokio::sync::{Mutex, RwLock};

use crate::meta::{
    metrics::{
        SHARD_CACHE_EVENT_COUNTER, SHARD_CACHE_WARM_UP_QUEUE_HISTOGRAM,
        SHARD_CACHE_WARM_UP_REQ_LATENCY_HISTOGRAM,
    },
    reader::{
        AllocateReaderReason, CacheAllocator, ReaderPool, ShardCacheAllocationElem,
        ShardCacheManager, UnifiedShardInfo, reader_pool::ReaderNodeMeta,
    },
};

pub struct WarmupTask {
    unified_shard_info: UnifiedShardInfo,
    index_info: tici_grpc::IndexInfo,
    table_info: tici_grpc::TableInfo,
    failed_nodes: HashSet<String>,
    reason: AllocateReaderReason,
}

impl std::fmt::Debug for WarmupTask {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("WarmupTask")
            .field_with("shard", |f| {
                f.write_fmt(format_args!("{}", &self.unified_shard_info.to_long_ident()))
            })
            .field("failed_nodes", &self.failed_nodes)
            .field("reason", &self.reason)
            .finish()
    }
}

impl WarmupTask {
    pub fn new(
        unified_shard_info: UnifiedShardInfo,
        index_info: tici_grpc::IndexInfo,
        table_info: tici_grpc::TableInfo,
        reason: AllocateReaderReason,
    ) -> Self {
        Self {
            unified_shard_info,
            index_info,
            table_info,
            failed_nodes: HashSet::new(),
            reason,
        }
    }
}

#[derive(Clone, Debug)]
pub enum WarmupTestAction {
    Succeed,
    SucceedAfter(Duration),
    SucceedCond(Arc<RwLock<Option<tokio::sync::oneshot::Receiver<()>>>>),
    Fail(String),
}

#[derive(Clone)]
/// Test-only hook to script warmup outcomes without relying on global
/// failpoints.
pub struct WarmupTestController {
    per_shard_actions: Arc<Mutex<HashMap<u64, VecDeque<WarmupTestAction>>>>,
    fallback_action: Arc<Mutex<Option<WarmupTestAction>>>,
    warm_derived_while_parent_warming: Arc<AtomicBool>,
}

impl Default for WarmupTestController {
    fn default() -> Self {
        Self {
            per_shard_actions: Arc::new(Mutex::new(HashMap::new())),
            fallback_action: Arc::new(Mutex::new(None)),
            warm_derived_while_parent_warming: Arc::new(AtomicBool::new(true)),
        }
    }
}

impl WarmupTestController {
    pub fn new() -> Self {
        Self::default()
    }

    pub async fn set_fallback_action(&self, action: Option<WarmupTestAction>) {
        *self.fallback_action.lock().await = action;
    }

    pub async fn set_shard_actions(
        &self,
        shard_id: u64,
        actions: impl IntoIterator<Item = WarmupTestAction>,
    ) {
        let mut guard = self.per_shard_actions.lock().await;
        guard.insert(shard_id, actions.into_iter().collect());
    }

    pub async fn clear(&self) {
        self.per_shard_actions.lock().await.clear();
        *self.fallback_action.lock().await = None;
        self.warm_derived_while_parent_warming
            .store(true, std::sync::atomic::Ordering::Relaxed);
    }

    pub async fn next_action(&self, shard_id: u64) -> Option<WarmupTestAction> {
        if let Some(action) = self
            .per_shard_actions
            .lock()
            .await
            .get_mut(&shard_id)
            .and_then(|queue| queue.pop_front())
        {
            return Some(action);
        }

        self.fallback_action.lock().await.clone()
    }

    #[cfg(test)]
    pub fn set_warm_derived_while_parent_warming(&self, enabled: bool) {
        self.warm_derived_while_parent_warming
            .store(enabled, std::sync::atomic::Ordering::Relaxed);
    }

    fn warm_derived_while_parent_warming(&self) -> bool {
        self.warm_derived_while_parent_warming
            .load(std::sync::atomic::Ordering::Relaxed)
    }
}

#[derive(Clone, Default)]
pub struct WarmupMgr {
    test_controller: Arc<RwLock<Option<Arc<WarmupTestController>>>>,
}

impl std::fmt::Debug for WarmupMgr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("WarmupMgr").finish()
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
enum WarmupOp {
    First,
    Normal,
}

impl WarmupMgr {
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        Self {
            test_controller: Arc::new(RwLock::new(None)),
        }
    }

    pub fn with_test_controller(test_controller: Arc<WarmupTestController>) -> Self {
        Self {
            test_controller: Arc::new(RwLock::new(Some(test_controller))),
        }
    }

    pub async fn set_test_controller(&self, test_controller: Option<Arc<WarmupTestController>>) {
        *self.test_controller.write().await = test_controller;
    }

    pub async fn test_controller(&self) -> Option<Arc<WarmupTestController>> {
        self.test_controller.read().await.clone()
    }

    pub async fn should_warm_derived_when_parent_warming(&self) -> bool {
        if let Some(controller) = self.test_controller().await {
            controller.warm_derived_while_parent_warming()
        } else {
            true
        }
    }

    pub async fn do_warmup(
        self: Arc<Self>,
        reader_pool_clone: Arc<ReaderPool>,
        warmup_task: WarmupTask,
        header: ShardManifestHeader,
        in_warming: Arc<AtomicU16>,
        shard_cache_manager: Arc<ShardCacheManager>,
    ) {
        if warmup_task.unified_shard_info.has_ready_readers().await {
            info!(
                "warmup skipped: shard already has ready readers, header={:?}",
                header
            );
            return;
        }
        self.warmup_impl(
            reader_pool_clone,
            warmup_task,
            header,
            in_warming,
            shard_cache_manager,
            HashSet::new(),
            None,
            WarmupOp::First,
        )
        .await;
    }

    pub async fn add_new_peer(
        self: Arc<Self>,
        reader_pool_clone: Arc<ReaderPool>,
        warmup_task: WarmupTask,
        header: ShardManifestHeader,
        in_warming: Arc<AtomicU16>,
        shard_cache_manager: Arc<ShardCacheManager>,
    ) {
        let existing = warmup_task.unified_shard_info.get_reader_addrs().await;
        let excluded_readers = existing.into_iter().collect::<HashSet<_>>();
        self.warmup_impl(
            reader_pool_clone,
            warmup_task,
            header,
            in_warming,
            shard_cache_manager,
            excluded_readers,
            None,
            WarmupOp::Normal,
        )
        .await;
    }

    pub async fn add_new_peer_designated(
        self: Arc<Self>,
        reader_pool_clone: Arc<ReaderPool>,
        warmup_task: WarmupTask,
        header: ShardManifestHeader,
        in_warming: Arc<AtomicU16>,
        shard_cache_manager: Arc<ShardCacheManager>,
        designated_reader_addr: String,
    ) {
        let existing = warmup_task.unified_shard_info.get_reader_addrs().await;
        let excluded_readers = existing.into_iter().collect::<HashSet<_>>();
        self.warmup_impl(
            reader_pool_clone,
            warmup_task,
            header,
            in_warming,
            shard_cache_manager,
            excluded_readers,
            Some(designated_reader_addr),
            WarmupOp::Normal,
        )
        .await;
    }

    async fn warmup_impl(
        &self,
        reader_pool_clone: Arc<ReaderPool>,
        warmup_task: WarmupTask,
        header: ShardManifestHeader,
        in_warming: Arc<AtomicU16>,
        shard_cache_manager: Arc<ShardCacheManager>,
        excluded_readers: HashSet<String>,
        designated_reader_addr: Option<String>,
        op: WarmupOp,
    ) {
        if in_warming
            .compare_exchange(
                0,
                1,
                std::sync::atomic::Ordering::SeqCst,
                std::sync::atomic::Ordering::SeqCst,
            )
            .is_err()
        {
            if let Some(designated_reader) = designated_reader_addr.as_deref() {
                debug!(
                    "ongoing warmup exists, op={:?}, header={:?}, designated_reader={}",
                    op, header, designated_reader
                );
            } else {
                debug!("ongoing warmup exists, op={:?}, header={:?}", op, header);
            }
            return;
        }
        scopeguard::defer!({
            in_warming.store(0, std::sync::atomic::Ordering::Relaxed);
        });

        if let Some(designated_reader) = designated_reader_addr.as_deref()
            && excluded_readers.contains(designated_reader)
        {
            debug!(
                "designated reader already has a ready replica, op={:?}, header={:?}, designated_reader={}",
                op, header, designated_reader
            );
            return;
        }

        let mut warmup_task = warmup_task;
        let cache_allocator = shard_cache_manager.make_cache_allocator();

        let designated_reader_node = if let Some(designated_reader) =
            designated_reader_addr.as_deref()
        {
            let Some(reader_node) = reader_pool_clone
                .get_readers_for_schedule()
                .await
                .into_iter()
                .find(|r| r.addr == designated_reader)
            else {
                SHARD_CACHE_EVENT_COUNTER
                    .with_label_values(&["warm_failed_no_reader"])
                    .inc();
                info!(
                    "{:?} failed on designated reader not found, shard={:?}, designated_reader={}, excluded_readers=[{}]",
                    op,
                    header,
                    designated_reader,
                    excluded_readers.len(),
                );
                return;
            };

            debug!(
                "start warmup on designated reader, op={:?}, header={:?}, designated_reader={}",
                op, header, reader_node.addr
            );
            Some(reader_node)
        } else {
            debug!("start warmup, op={:?}, header={:?}", op, header);
            None
        };

        let max_attempts = if designated_reader_node.is_some() {
            1
        } else {
            cache_allocator.max_reader_attempts
        };

        for attempt in 0..max_attempts {
            let chosen = match designated_reader_node.as_ref() {
                Some(reader_node) => Some(reader_node.clone()),
                None => {
                    WarmupMgr::allocate_reader(
                        &warmup_task,
                        &cache_allocator,
                        reader_pool_clone.as_ref(),
                        &excluded_readers,
                    )
                    .await
                }
            };

            if let Some(r) = chosen {
                // This means the shard is allocate to reader_node, but it is not ready yet.
                // When it is ready, we have then `unified_shard.has_reader_addr()` should be
                // true.
                if designated_reader_node.is_some() {
                    r.add_shard_cache(warmup_task.unified_shard_info.clone())
                        .await;
                } else {
                    r.add_reserved_shard_cache(warmup_task.unified_shard_info.clone())
                        .await;
                }
                let unified_shard = warmup_task.unified_shard_info.clone();
                let reason = warmup_task.reason.clone();
                match self
                    .call_warmup(warmup_task, r.clone(), &cache_allocator)
                    .await
                {
                    Either::Left(allocation) => {
                        info!("warmup succeed, op={:?}, shard={:?}", op, header);
                        let _ = unified_shard.add_ready_reader_elem(allocation).await;
                        match shard_cache_manager
                            .on_warmup_finished(&unified_shard, reason)
                            .await
                        {
                            Err(e) => {
                                // Warmup succeed, but failed to notify the index, this is not
                                // commonm maybe because the index is dropped.
                                if op == WarmupOp::First {
                                    SHARD_CACHE_EVENT_COUNTER
                                        .with_label_values(&["warm_f_err_update"])
                                        .inc();
                                } else {
                                    SHARD_CACHE_EVENT_COUNTER
                                        .with_label_values(&["warm_n_err_update"])
                                        .inc();
                                }
                                warn!(
                                    "`{:?}` update index failed, attempt={}, shard={}, e={:?}",
                                    op,
                                    attempt,
                                    unified_shard.to_long_ident(),
                                    e
                                );
                            }
                            Ok(ok) => {
                                if op == WarmupOp::First {
                                    SHARD_CACHE_EVENT_COUNTER
                                        .with_label_values(&["warm_f_suc"])
                                        .inc();
                                } else {
                                    SHARD_CACHE_EVENT_COUNTER
                                        .with_label_values(&["warm_n_suc"])
                                        .inc();
                                }
                                reader_pool_clone.on_warmup_finished(ok).await;
                                info!(
                                    "`{:?}` update index finished, attempt={}, shard={}",
                                    op,
                                    attempt,
                                    unified_shard.to_long_ident()
                                );
                            }
                        }
                        return;
                    }
                    Either::Right(new_warmup_task) => {
                        r.remove_shard_cache(&new_warmup_task.unified_shard_info.shard_id)
                            .await;
                        if designated_reader_node.is_some() {
                            info!(
                                "{:?} failed on designated reader, shard={:?}, reader={}",
                                op, header, &r.addr
                            );
                            return;
                        } else {
                            info!(
                                "{:?} failed on reader, shard={:?}, reader={}",
                                op, header, &r.addr
                            );
                            warmup_task = new_warmup_task;
                        }
                    }
                }
            } else {
                // There is no available readers, because there is no alive reader, or all nodes
                // are blacklisted.
                let readers = reader_pool_clone.get_all_readers().await;
                SHARD_CACHE_EVENT_COUNTER
                    .with_label_values(&["warm_failed_no_reader"])
                    .inc();
                // TODO(P1) WE can do nothing in this case, so we can tell or let
                // find_overlapping_shards to fail immediately instead of keep retrying.
                info!(
                    "`{:?}` failed on no reader, shard={:?}, task={:?}, excluded_readers=[{}] all_readers=[{}]{}",
                    op,
                    header,
                    warmup_task,
                    excluded_readers.len(),
                    readers.len(),
                    readers
                        .iter()
                        .map(|x| x.addr.to_string())
                        .collect::<Vec<_>>()
                        .join("|")
                );
            };
        }
        SHARD_CACHE_EVENT_COUNTER
            .with_label_values(&["warm_failed_retry_limit"])
            .inc();
    }

    /// Send warmup request to a reader node
    async fn send_warmup_request(
        &self,
        s3_bucket: String,
        shard_header: &ShardManifestHeader,
        reader_node: Arc<ReaderNodeMeta>,
        index_info: &tici_grpc::IndexInfo,
        table_info: &tici_grpc::TableInfo,
    ) -> Result<()> {
        #[cfg(any(test, feature = "testexport"))]
        {
            if let Some(controller) = self.test_controller.read().await.clone() {
                if let Some(action) = controller.next_action(shard_header.shard_id).await {
                    match action {
                        WarmupTestAction::Succeed => return Ok(()),
                        WarmupTestAction::SucceedAfter(delay) => {
                            tokio::time::sleep(delay).await;
                            return Ok(());
                        }
                        WarmupTestAction::SucceedCond(cond) => {
                            if let Some(k) = cond.write().await.take() {
                                log::info!("mock: await warmup");
                                let _ = k.await;
                                log::info!("mock: warmup finish");
                            }
                            return Ok(());
                        }
                        WarmupTestAction::Fail(msg) => {
                            return Err(anyhow::anyhow!(msg));
                        }
                    }
                }
            }
        }
        let start = Instant::now();
        let warmup_request = WarmupShardRequest {
            shard: Some(shard_header.clone()),
            index_info: Some(index_info.clone()),
            table_info: Some(table_info.clone()),
            frag_upload_bucket: s3_bucket.clone(),
            fragments: vec![], // Not used in current implementation
        };
        let mut client = reader_node.get_client().await?;
        let result = match client.warmup_shard(warmup_request.clone()).await {
            Ok(response) => Ok(response.into_inner()),
            Err(e) => Err(anyhow::anyhow!("warmup_shard gRPC err={}", e)),
        };
        let elapsed = start.elapsed().as_secs_f64();
        SHARD_CACHE_WARM_UP_REQ_LATENCY_HISTOGRAM.observe(elapsed);

        match result {
            Ok(response) if response.status == 0 => {
                info!(
                    "warmup successful, shard={}, addr={}",
                    shard_header.to_long_ident(),
                    reader_node.addr
                );
                Ok(())
            }
            Ok(response) => {
                let error_msg = response.error_message.clone();
                Err(anyhow::anyhow!(
                    "warmup received error response, shard_header={:?}, addr={}, err={}",
                    shard_header,
                    reader_node.addr,
                    error_msg
                ))
            }
            Err(e) => Err(anyhow::anyhow!(
                "failed to send warmup to reader, shard_header={:?}, addr={}, err={}",
                shard_header,
                reader_node.addr,
                e
            )),
        }
    }

    async fn allocate_reader(
        warmup_task: &WarmupTask,
        cache_allocator: &CacheAllocator,
        reader_pool: &ReaderPool,
        excluded_readers: &HashSet<String>,
    ) -> Option<Arc<ReaderNodeMeta>> {
        let active_readers = reader_pool.get_readers_for_schedule().await;
        let candidates = active_readers
            .into_iter()
            // If the shard has failed to be warmed up at these nodes.
            .filter(|reader_node| !warmup_task.failed_nodes.contains(&reader_node.addr))
            // If the shard is allocated on these readers.
            .filter(|reader_node| !excluded_readers.contains(&reader_node.addr))
            .collect::<Vec<_>>();
        let reader_node = cache_allocator
            .cache_scheduler
            .select_reader(&candidates)
            .await?;

        info!(
            "allocated reader, addr={}, shard_id={}, reason={:?}, excluded_readers={}, failed_on={:?}",
            reader_node.addr,
            warmup_task.unified_shard_info.shard_id,
            warmup_task.reason,
            excluded_readers.len(),
            warmup_task.failed_nodes
        );
        Some(reader_node)
    }

    async fn call_warmup(
        &self,
        warmup_task: WarmupTask,
        reader_node: Arc<ReaderNodeMeta>,
        cache_allocator: &CacheAllocator,
    ) -> Either<ShardCacheAllocationElem, WarmupTask> {
        info!("acquire sem begin, warmup={:?}", warmup_task);
        let start = Instant::now();
        let _permit = match reader_node.warmup_sem.acquire().await {
            Ok(ok) => ok,
            Err(e) => {
                log::error!("acuqire sem error, warmup={:?}, e={:?}", warmup_task, e);
                let mut warmup_task = warmup_task;
                warmup_task.failed_nodes.insert(reader_node.addr.clone());
                return Either::Right(warmup_task);
            }
        };
        if start.elapsed() > std::time::Duration::from_millis(100) {
            info!(
                "acquire sem delayed, warmup={:?}, time={:.4}",
                warmup_task,
                start.elapsed().as_secs_f64()
            );
            SHARD_CACHE_WARM_UP_QUEUE_HISTOGRAM.observe(start.elapsed().as_secs_f64());
        }
        match self
            .send_warmup_request(
                cache_allocator.s3_bucket.clone(),
                &warmup_task.unified_shard_info.get_header(),
                reader_node.clone(),
                &warmup_task.index_info,
                &warmup_task.table_info,
            )
            .await
        {
            Ok(()) => {
                let allocation = ShardCacheAllocationElem::with_observed_seq(
                    reader_node.addr.clone(),
                    reader_node.tiflash_server_addr.clone(),
                    Some(warmup_task.unified_shard_info.get_header().seq),
                    Instant::now(),
                );
                info!(
                    "finished warmup, prepare allocate, shard={}, addr={}, reason={:?}, allocation={:?}",
                    warmup_task.unified_shard_info.to_long_ident(),
                    reader_node.addr,
                    warmup_task.reason,
                    allocation,
                );
                Either::Left(allocation)
            }
            Err(e) => {
                // NOTE We must use round robin to ensure this behavior.
                warn!(
                    "failed to warmup shard on reader, shard={}, addr={}, reason={:?}, err={}",
                    warmup_task.unified_shard_info.to_long_ident(),
                    reader_node.addr,
                    warmup_task.reason,
                    e
                );
                let mut warmup_task = warmup_task;
                warmup_task.failed_nodes.insert(reader_node.addr.clone());
                Either::Right(warmup_task)
            }
        }
    }
}