use std::{collections::HashMap, sync::Arc};

use anyhow::Result;
use log::{error, info, warn};
use tici_common::{codec::parse_table_id_from_start_key, index::IndexKey, keyspace::NULLSPACEID};
use tici_grpc::{
    ErrorCode, GetShardLocalCacheRequest, GetShardLocalCacheResponse, ShardLocalCacheInfo,
};
use tokio::{sync::oneshot, time::timeout};

use crate::meta::reader::{
    AllocateReaderReason, ManualRemoveResult, ReaderPool, ShardCacheManager, UnifiedShardInfo,
};

impl ShardCacheManager {
    pub async fn on_msg_manual_reschedule_by_shard(
        manager: &Arc<Self>,
        reader_pool: Arc<ReaderPool>,
        index_key: &IndexKey,
        start_key: Vec<u8>,
        end_key: Vec<u8>,
    ) -> Result<(ErrorCode, String)> {
        let key_ranges = vec![tici_grpc::KeyRange {
            start_key: start_key.clone(),
            end_key: end_key.clone(),
        }];
        match ShardCacheManager::normalize_range_and_get_cache(
            manager,
            index_key,
            key_ranges.as_slice(),
            reader_pool,
            true,
            0,
            AllocateReaderReason::Manual,
        )
        .await
        {
            Ok(res) => match res {
                either::Left(vec) => {
                    info!(
                        "reader_manual_reschedule, succ: {:?}, start_key={:?}, end_key={:?}",
                        index_key, start_key, end_key
                    );
                    Ok((ErrorCode::Success, format!("scheduled={}", vec.len())))
                }
                either::Either::Right(ec) => {
                    info!(
                        "reader_manual_reschedule, succ: {:?}, start_key={:?}, end_key={:?}",
                        index_key, start_key, end_key
                    );
                    Ok((ec, format!("scheduled_error={:?}", ec)))
                }
            },
            Err(e) => {
                info!(
                    "reader_manual_reschedule, fail: {:?}, start_key={:?}, end_key={:?}",
                    index_key, start_key, end_key
                );
                Ok((ErrorCode::Success, format!("error={:?}", e)))
            }
        }
    }

    pub async fn on_msg_manual_add_replica_by_shard(
        manager: &Arc<Self>,
        reader_pool: Arc<ReaderPool>,
        index_key: &IndexKey,
        start_key: Vec<u8>,
        end_key: Vec<u8>,
        designated_reader_addr: Option<String>,
    ) -> Result<(ErrorCode, String)> {
        enum AddReplicaMode {
            Designated(String),
            Auto,
        }

        let (table_info, index_info, shards) = {
            let metadata_guard = manager.index_metadata.read().await;
            let Some(index_meta) = metadata_guard.get(index_key) else {
                info!(
                    "reader_manual_add_replica: index not found, index_key={:?}",
                    index_key
                );
                return Ok((
                    ErrorCode::IndexNotFound,
                    format!("index not found, index_key={:?}", index_key),
                ));
            };

            let key_ranges = [tici_grpc::KeyRange {
                start_key: start_key.clone(),
                end_key: end_key.clone(),
            }];
            let overlapping_shards = match index_meta
                .find_overlapping_shards(&key_ranges, AllocateReaderReason::Manual, None, false)
                .await
            {
                Ok((overlapping, ..)) => overlapping,
                Err(e) => {
                    info!(
                        "reader_manual_add_replica: failed to query overlapping shards, index_key={:?}, start_key={:?}, end_key={:?}, err={:?}",
                        index_key, start_key, end_key, e
                    );
                    return Ok((ErrorCode::UnknownError, format!("error={:?}", e)));
                }
            };
            if overlapping_shards.is_empty() {
                info!(
                    "reader_manual_add_replica: shard not found, index_key={:?}, start_key={:?}, end_key={:?}",
                    index_key, start_key, end_key
                );
                return Ok((
                    ErrorCode::ShardNotFound,
                    format!(
                        "no shard included, index_key={:?}, start_key={:?}, end_key={:?}",
                        index_key, start_key, end_key
                    ),
                ));
            }

            let shards: HashMap<u64, UnifiedShardInfo> = overlapping_shards
                .into_iter()
                .map(|shard| (shard.shard_id, shard))
                .collect();

            (
                index_meta.table_info.clone(),
                index_meta.index_info.clone(),
                shards,
            )
        };

        let mode = if let Some(addr) = designated_reader_addr {
            if !reader_pool.is_reader_active(&addr).await {
                info!(
                    "reader_manual_add_replica: designated reader not found or inactive, index_key={:?}, addr={}",
                    index_key, addr
                );
                return Ok((
                    ErrorCode::WorkerNotFound,
                    format!("designated reader not found or inactive, addr={}", addr),
                ));
            }
            AddReplicaMode::Designated(addr)
        } else {
            let active_readers = reader_pool.get_readers_for_schedule().await;
            if active_readers.is_empty() {
                info!(
                    "reader_manual_add_replica: failed for no readers, index_key={:?}, start_key={:?}, end_key={:?}",
                    index_key, start_key, end_key
                );
                return Ok((ErrorCode::WorkerNotFound, "no readers".to_string()));
            }
            AddReplicaMode::Auto
        };

        let mut already_present = 0usize;
        let mut to_wait = Vec::with_capacity(shards.len());
        for (_, shard) in shards {
            let before = match &mode {
                AddReplicaMode::Designated(addr) => {
                    if shard.has_reader_addr(addr).await && !shard.is_reader_removing(addr).await {
                        already_present += 1;
                        continue;
                    }
                    None
                }
                AddReplicaMode::Auto => Some(shard.get_reader_addrs().await.len()),
            };

            let handle_res = match &mode {
                AddReplicaMode::Designated(addr) => reader_pool.background_add_replica_designated(
                    &shard,
                    reader_pool.clone(),
                    table_info.clone(),
                    index_info.clone(),
                    AllocateReaderReason::Manual,
                    addr.clone(),
                ),
                AddReplicaMode::Auto => reader_pool.background_add_replica(
                    &shard,
                    reader_pool.clone(),
                    table_info.clone(),
                    index_info.clone(),
                    AllocateReaderReason::Manual,
                ),
            };

            let handle = match handle_res {
                Ok(handle) => handle,
                Err(e) => {
                    match &mode {
                        AddReplicaMode::Designated(addr) => warn!(
                            "reader_manual_add_replica: failed to start designated warmup, shard={}, addr={}, err={:?}",
                            shard.to_long_ident(),
                            addr,
                            e
                        ),
                        AddReplicaMode::Auto => warn!(
                            "reader_manual_add_replica: failed to start background warmup, shard={}, err={:?}",
                            shard.to_long_ident(),
                            e
                        ),
                    }
                    return Ok((ErrorCode::UnknownError, format!("error={:?}", e)));
                }
            };
            to_wait.push((shard, before, handle));
        }

        let scheduled = to_wait.len();
        let mut added = 0usize;
        let mut success = already_present;
        let mut increased_desired_count = 0usize;
        let mut failed = vec![];
        for (shard, before, handle) in to_wait {
            if let Err(e) = handle.await {
                failed.push((shard.shard_id, format!("join_error={}", e)));
                continue;
            }

            match &mode {
                AddReplicaMode::Designated(addr) => {
                    if shard.has_reader_addr(addr).await {
                        shard.inc_desired_replica_count().await;
                        added += 1;
                        success += 1;
                        increased_desired_count += 1;
                    } else {
                        failed.push((
                            shard.shard_id,
                            format!("designated reader not present after warmup, addr={}", addr),
                        ));
                    }
                }
                AddReplicaMode::Auto => {
                    let Some(before) = before else {
                        failed.push((
                            shard.shard_id,
                            "missing before reader count in auto mode".to_string(),
                        ));
                        continue;
                    };
                    let after = shard.get_reader_addrs().await.len();
                    if after > before {
                        shard.inc_desired_replica_count().await;
                        added += 1;
                        success += 1;
                        increased_desired_count += 1;
                    } else {
                        failed.push((
                            shard.shard_id,
                            format!(
                                "no new replica, before_readers={}, after_readers={}",
                                before, after
                            ),
                        ));
                    }
                }
            }
        }

        if failed.is_empty() {
            match &mode {
                AddReplicaMode::Designated(addr) => {
                    let shards_len = scheduled + already_present;
                    let msg = format!(
                        "addr={}, added={}, already_present={}, desired_count_updated={}, shards={}/{}",
                        addr, added, already_present, increased_desired_count, success, shards_len
                    );
                    info!(
                        "reader_manual_add_replica: succ: {:?}, start_key={:?}, end_key={:?}, {}",
                        index_key, start_key, end_key, msg
                    );
                    Ok((ErrorCode::Success, msg))
                }
                AddReplicaMode::Auto => {
                    info!(
                        "reader_manual_add_replica: succ: {:?}, start_key={:?}, end_key={:?}, shards={}/{}, desired_count_updated={}",
                        index_key, start_key, end_key, success, scheduled, increased_desired_count
                    );
                    Ok((
                        ErrorCode::Success,
                        format!(
                            "added={}, desired_count_updated={}",
                            added, increased_desired_count
                        ),
                    ))
                }
            }
        } else {
            match &mode {
                AddReplicaMode::Designated(addr) => {
                    let shards_len = scheduled + already_present;
                    let msg = format!(
                        "addr={}, added={}, already_present={}, desired_count_updated={}, shards={}/{}",
                        addr, added, already_present, increased_desired_count, success, shards_len
                    );
                    info!(
                        "reader_manual_add_replica: partial: {:?}, start_key={:?}, end_key={:?}, {} failed={:?}",
                        index_key, start_key, end_key, msg, failed
                    );
                    Ok((ErrorCode::TryAgain, format!("{msg} failed={:?}", failed)))
                }
                AddReplicaMode::Auto => {
                    info!(
                        "reader_manual_add_replica: partial: {:?}, start_key={:?}, end_key={:?}, succ={}/{}, desired_count_updated={} failed={:?}",
                        index_key,
                        start_key,
                        end_key,
                        success,
                        scheduled,
                        increased_desired_count,
                        failed
                    );
                    Ok((
                        ErrorCode::TryAgain,
                        format!(
                            "added={}/{}, desired_count_updated={} failed={:?}",
                            success, scheduled, increased_desired_count, failed
                        ),
                    ))
                }
            }
        }
    }

    /// Manually remove a shard replica from a specific reader. The replica is
    /// first marked as removing and will be finalized after the reader stops
    /// reporting it in heartbeat.
    pub async fn on_msg_manual_remove_replica_by_shard(
        manager: &Arc<Self>,
        _reader_pool: Arc<ReaderPool>,
        shard_id: u64,
        reader_addr: String,
    ) -> Result<(ErrorCode, String)> {
        let Some(shard) = manager.get_shard_by_id(&shard_id).await else {
            info!(
                "reader_manual_remove_replica: shard not found, shard_id={}",
                shard_id
            );
            return Ok((
                ErrorCode::ShardNotFound,
                format!("shard not found, shard_id={}", shard_id),
            ));
        };

        let (tx, rx) = oneshot::channel();
        let started = if shard.is_reader_removing(&reader_addr).await {
            shard.add_removing_reader_callback(&reader_addr, tx).await
        } else {
            let removed = shard.mark_ready_reader_removing(&reader_addr).await;
            if removed {
                shard.dec_desired_replica_count().await;
                let _ = shard.add_removing_reader_callback(&reader_addr, tx).await;
            }
            removed
        };

        if !started {
            info!(
                "reader_manual_remove_replica: addr={}, shard_id={}, removed=false, desired_replica_count={}",
                reader_addr,
                shard_id,
                shard.desired_replica_count().await
            );
            return Ok((
                ErrorCode::Success,
                format!(
                    "addr={}, shard_id={}, removed=false, desired_replica_count={}",
                    reader_addr,
                    shard_id,
                    shard.desired_replica_count().await
                ),
            ));
        }

        let wait_timeout = manager.config.cache_scheduling.manual_remove_wait_timeout.0;
        let wait_result = timeout(wait_timeout, rx).await;
        let (status, msg) = match wait_result {
            Ok(Ok(ManualRemoveResult::Success)) => (
                ErrorCode::Success,
                format!(
                    "addr={}, shard_id={}, removed=true, desired_replica_count={}",
                    reader_addr,
                    shard_id,
                    shard.desired_replica_count().await
                ),
            ),
            Ok(Ok(ManualRemoveResult::ReaderExpired)) => (
                ErrorCode::TryAgain,
                format!(
                    "addr={}, shard_id={}, remove_failed=reader_expired",
                    reader_addr, shard_id
                ),
            ),
            Ok(Err(_)) => (
                ErrorCode::TryAgain,
                format!(
                    "addr={}, shard_id={}, remove_failed=callback_cancelled",
                    reader_addr, shard_id
                ),
            ),
            Err(_) => (
                ErrorCode::TryAgain,
                format!(
                    "addr={}, shard_id={}, remove_failed=timeout",
                    reader_addr, shard_id
                ),
            ),
        };
        info!("reader_manual_remove_replica: {}", msg);
        Ok((status, msg))
    }

    /// Get shard local cache information for the given request
    pub async fn on_msg_get_shard_local_cache(
        manager: &Arc<Self>,
        request: &GetShardLocalCacheRequest,
        reader_pool: Arc<ReaderPool>,
    ) -> Result<GetShardLocalCacheResponse> {
        let mut result_infos: Vec<ShardLocalCacheInfo> = Vec::new();

        if request.key_ranges.is_empty() {
            error!(
                "get_shard_local_cache failed, key_ranges is empty, keyspace_id={:?}, table_id={:?}, index_id={:?}",
                request.keyspace_id, request.table_id, request.index_id
            );

            return Ok(GetShardLocalCacheResponse {
                status: ErrorCode::InvalidArguments.into(),
                shard_local_cache_infos: vec![],
            });
        }

        // 'key_ranges' may contain different table ids when using partition table
        for range in &request.key_ranges {
            let table_id = if range.start_key.is_empty() {
                #[cfg(any(test, feature = "testexport"))]
                {
                    request.table_id
                }
                #[cfg(not(any(test, feature = "testexport")))]
                {
                    error!(
                        "start_key is empty, keyspace_id={:?}, table_id={:?}, index_id={:?}, range={:?}",
                        request.keyspace_id, request.table_id, request.index_id, range
                    );
                    return Ok(GetShardLocalCacheResponse {
                        status: ErrorCode::InvalidArguments.into(),
                        shard_local_cache_infos: vec![],
                    });
                }
            } else {
                parse_table_id_from_start_key(&range.start_key).ok_or_else(|| {
                    anyhow::anyhow!(
                        "failed to get table_id from start_key={:?}",
                        range.start_key
                    )
                })?
            };

            let index_key = IndexKey::new(request.keyspace_id, table_id, request.index_id);
            let result_info = match manager
                .normalize_range_and_get_cache(
                    &index_key,
                    &[range.clone()],
                    reader_pool.clone(),
                    false,
                    request.limit,
                    AllocateReaderReason::ReadRequest,
                )
                .await
            {
                Ok(ok) => match ok {
                    either::Either::Left(vec) => vec,
                    either::Either::Right(ec) => {
                        return Ok(GetShardLocalCacheResponse {
                            status: ec.into(),
                            shard_local_cache_infos: vec![],
                        });
                    }
                },
                Err(_) => {
                    return Ok(GetShardLocalCacheResponse {
                        status: ErrorCode::UnknownError.into(),
                        shard_local_cache_infos: vec![],
                    });
                }
            };

            result_infos.extend(result_info);
        }

        // TiDB takes only non-prefixed key ranges.
        if request.keyspace_id != NULLSPACEID {
            for shard in result_infos.iter_mut() {
                if let Some(header) = shard.shard.as_mut() {
                    header.start_key.drain(..4);
                    header.end_key.drain(..4);
                } else {
                    warn!("meet error header, {:?}", shard.shard);
                }
            }
        }
        Ok(GetShardLocalCacheResponse {
            status: ErrorCode::Success.into(),
            shard_local_cache_infos: result_infos,
        })
    }

    /// Get shard local cache information for the given request
    pub async fn normalize_range_and_get_cache(
        &self,
        index_key: &IndexKey,
        key_ranges: &[tici_grpc::KeyRange],
        reader_pool: Arc<ReaderPool>,
        tidb_x_already_encoded: bool,
        limit: i32,
        reason: AllocateReaderReason,
    ) -> Result<either::Either<Vec<ShardLocalCacheInfo>, tici_grpc::ErrorCode>> {
        let normal_ranges = self.normalize_keyrange(index_key, tidb_x_already_encoded, key_ranges);

        info!(
            "get_shard_local_cache: ranges={:?}, index_key={:?}, normal_ranges={:?}",
            key_ranges, index_key, normal_ranges
        );

        let result_infos = match self
            .get_or_allocate_cache(
                index_key,
                normal_ranges,
                reader_pool,
                limit as usize,
                reason.clone(),
                None,
            )
            .await
        {
            Ok(infos) => infos,
            Err(e) => {
                info!("get_or_allocate_cache error: {:?}", e);
                return Err(e);
            }
        };

        info!(
            "get_shard_local_cache: result, ranges={:?}, index_key={:?}, result_infos={:?}",
            key_ranges, index_key, result_infos
        );
        Ok(result_infos)
    }
}