use std::{
    collections::{HashMap, HashSet},
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
    time::{Duration, SystemTime},
};

use anyhow::Result;
use log::{debug, info};
use parking_lot::Mutex;
use tici_common::config::GrpcConfig;
use tici_grpc::{GrpcReaderServiceClient, ReaderNodeStatus, ShardManifestHeader};
use tici_system::readable::ReadableDuration;
use tokio::sync::{OnceCell, RwLock, Semaphore};

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

/// Metadata for a reader node including connection and status information
#[derive(Debug)]
pub struct ReaderNodeMeta {
    pub addr: String,
    pub tiflash_server_addr: String,
    latest_status: Mutex<ReaderNodeStatus>,
    last_heartbeat: Mutex<SystemTime>,
    client: OnceCell<GrpcReaderServiceClient<tonic::transport::Channel>>,
    grpc_client_config: GrpcConfig,
    allocated_shards: RwLock<HashMap<u64, UnifiedShardInfo>>,
    // Scheduling currently uses a count-based load model. This counter tracks
    // both shards already assigned to the reader and pending reservations made
    // during reader selection, so `shard_cache_schedule_load()` intentionally
    // returns the same value as `shard_cache_count()`.
    scheduled_shard_cache_count: AtomicUsize,
    pub warmup_sem: Arc<Semaphore>,
}

impl PartialEq for ReaderNodeMeta {
    fn eq(&self, other: &Self) -> bool {
        self.addr == other.addr
    }
}

unsafe impl Send for ReaderNodeMeta {}
unsafe impl Sync for ReaderNodeMeta {}

impl ReaderNodeMeta {
    pub fn new(
        addr: String,
        tiflash_server_addr: String,
        status: ReaderNodeStatus,
        grpc_client_config: GrpcConfig,
        cache_allocator: CacheAllocator,
    ) -> Self {
        Self {
            addr,
            tiflash_server_addr,
            latest_status: Mutex::new(status),
            last_heartbeat: Mutex::new(SystemTime::now()),
            client: OnceCell::new(),
            grpc_client_config,
            allocated_shards: RwLock::new(HashMap::new()),
            scheduled_shard_cache_count: AtomicUsize::new(0),
            warmup_sem: Arc::new(Semaphore::new(cache_allocator.allocate_parallelism)),
        }
    }

    /// Update reader status and heartbeat timestamp
    pub async fn on_heartbeat_update_status(&self, status: ReaderNodeStatus) {
        *self.latest_status.lock() = status;
        *self.last_heartbeat.lock() = SystemTime::now();
    }

    /// Check if this reader has expired based on TTL
    pub async fn is_expired(&self, ttl_duration: ReadableDuration) -> bool {
        let last_heartbeat = *self.last_heartbeat.lock();
        let now = SystemTime::now();

        match now.duration_since(last_heartbeat) {
            Ok(duration) => duration > ttl_duration.0,
            Err(_) => true,
        }
    }

    /// Get the current status for load balancing decisions
    pub async fn get_status(&self) -> ReaderNodeStatus {
        self.latest_status.lock().clone()
    }

    async fn create_client(&self) -> Result<GrpcReaderServiceClient<tonic::transport::Channel>> {
        const MAX_ATTEMPTS: u32 = 3;
        const RETRY_DELAY: Duration = Duration::from_secs(1);

        let connect_addr = self.format_connection_address();

        for attempt in 1..=MAX_ATTEMPTS {
            match GrpcReaderServiceClient::connect(connect_addr.clone()).await {
                Ok(client) => return Ok(client),
                Err(_) if attempt < MAX_ATTEMPTS => {
                    tokio::time::sleep(RETRY_DELAY).await;
                }
                Err(e) => {
                    return Err(anyhow::anyhow!(
                        "Failed to connect to reader server at {} after {} attempts={}",
                        connect_addr,
                        MAX_ATTEMPTS,
                        e
                    ));
                }
            }
        }

        unreachable!()
    }

    fn format_connection_address(&self) -> String {
        if self.addr.starts_with("http://") {
            self.addr.clone()
        } else {
            format!("http://{}", self.addr)
        }
    }

    pub async fn get_client(&self) -> Result<GrpcReaderServiceClient<tonic::transport::Channel>> {
        Ok(self
            .client
            .get_or_try_init(|| self.create_client())
            .await?
            .clone()
            .max_decoding_message_size(self.grpc_client_config.max_decoding_message_size.0 as usize)
            .max_encoding_message_size(
                self.grpc_client_config.max_encoding_message_size.0 as usize,
            ))
    }

    pub async fn add_shard_cache(&self, shard: UnifiedShardInfo) {
        let old = self
            .allocated_shards
            .write()
            .await
            .insert(shard.shard_id, shard);
        if old.is_none() {
            self.scheduled_shard_cache_count
                .fetch_add(1, Ordering::Relaxed);
        }
    }

    pub async fn add_reserved_shard_cache(&self, shard: UnifiedShardInfo) {
        let old = self
            .allocated_shards
            .write()
            .await
            .insert(shard.shard_id, shard);
        if old.is_some() {
            self.release_pending_shard_cache_allocation();
        }
    }

    pub async fn has_shard_cache(&self, shard_id: u64) -> bool {
        self.allocated_shards.read().await.contains_key(&shard_id)
    }

    pub async fn shard_cache_count(&self) -> usize {
        self.scheduled_shard_cache_count.load(Ordering::Relaxed)
    }

    pub async fn shard_cache_schedule_load(&self) -> usize {
        // The current scheduler only reasons about "how many shard slots are
        // already on this reader", including pending warmup reservations.
        self.shard_cache_count().await
    }

    pub fn reserve_pending_shard_cache_allocation(&self) {
        self.scheduled_shard_cache_count
            .fetch_add(1, Ordering::Relaxed);
    }

    pub fn release_pending_shard_cache_allocation(&self) {
        let _ = self.scheduled_shard_cache_count.fetch_update(
            Ordering::Relaxed,
            Ordering::Relaxed,
            |count| count.checked_sub(1),
        );
    }

    pub async fn allocated_shards_snapshot(&self) -> HashMap<u64, UnifiedShardInfo> {
        self.allocated_shards.read().await.clone()
    }

    pub async fn remove_shard_cache(&self, shard_id: &u64) {
        let old = self.allocated_shards.write().await.remove(shard_id);
        if old.is_some() {
            self.release_pending_shard_cache_allocation();
        }
    }

    pub async fn remove_shard_caches(&self, shard_ids: &[u64]) {
        let mut allocated_shards = self.allocated_shards.write().await;
        for shard_id in shard_ids {
            if allocated_shards.remove(shard_id).is_some() {
                self.release_pending_shard_cache_allocation();
            }
        }
    }
}

impl ReaderNodeMeta {
    async fn reconcile_redundant_ready_readers_batch(
        &self,
        affected_shards: &[(u64, UnifiedShardInfo)],
        reader_pool: &ReaderPool,
    ) -> HashSet<u64> {
        let mut planned_load = HashMap::<String, isize>::new();
        let mut should_keep_shards = HashSet::new();

        for (shard_id, unified_shard_info) in affected_shards {
            if self
                .reconcile_redundant_ready_readers_for_shard(
                    unified_shard_info,
                    reader_pool,
                    &mut planned_load,
                )
                .await
            {
                should_keep_shards.insert(*shard_id);
            }
        }

        should_keep_shards
    }

    async fn reconcile_redundant_ready_readers_for_shard(
        &self,
        unified_shard_info: &UnifiedShardInfo,
        reader_pool: &ReaderPool,
        planned_load: &mut HashMap<String, isize>,
    ) -> bool {
        if unified_shard_info.is_reader_removing(&self.addr).await {
            return false;
        }

        let desired_replica_count = unified_shard_info.desired_replica_count().await as usize;
        let allocation = unified_shard_info.cloned_cache_allocation().await;
        let ready_readers = allocation.ready_reader_elems();
        if ready_readers.len() <= desired_replica_count {
            return true;
        }

        let mut ranked_readers = Vec::with_capacity(ready_readers.len());
        for elem in &ready_readers {
            let Some(observed_seq) = elem.observed_seq else {
                debug!(
                    "skip redundant cleanup for shard={} because reader {} has unknown observed seq",
                    unified_shard_info.to_long_ident(),
                    elem.reader_addr
                );
                return true;
            };
            let load = if let Some(load) = planned_load.get(&elem.reader_addr) {
                *load
            } else {
                let Some(load) = reader_pool
                    .get_schedulable_reader_schedule_load(&elem.reader_addr)
                    .await
                else {
                    debug!(
                        "skip reader {} from redundant cleanup ranking for shard={} because it is not schedulable",
                        elem.reader_addr,
                        unified_shard_info.to_long_ident(),
                    );
                    continue;
                };
                let load = load as isize;
                planned_load.insert(elem.reader_addr.clone(), load);
                load
            };
            ranked_readers.push((elem.clone(), observed_seq, load));
        }

        if ranked_readers.len() < desired_replica_count {
            debug!(
                "skip redundant cleanup for shard={} because only {} schedulable readers are available for desired_replica_count={}",
                unified_shard_info.to_long_ident(),
                ranked_readers.len(),
                desired_replica_count
            );
            return true;
        }

        ranked_readers.sort_by(
            |(lhs_elem, lhs_seq, lhs_load), (rhs_elem, rhs_seq, rhs_load)| {
                lhs_load
                    .cmp(rhs_load)
                    .then_with(|| rhs_seq.cmp(lhs_seq))
                    .then_with(|| lhs_elem.ready_at.cmp(&rhs_elem.ready_at))
                    .then_with(|| lhs_elem.reader_addr.cmp(&rhs_elem.reader_addr))
            },
        );

        let target_readers = ranked_readers
            .iter()
            .take(desired_replica_count)
            .map(|(elem, seq, ..)| (elem.reader_addr.clone(), *seq))
            .collect::<Vec<_>>();
        let max_seq = ranked_readers
            .iter()
            .map(|(_, seq, _)| *seq)
            .max()
            .unwrap_or_default();

        let mut keep_reader_addrs = target_readers
            .iter()
            .map(|(reader_addr, _)| reader_addr.clone())
            .collect::<HashSet<_>>();
        if target_readers.iter().any(|(_, seq)| *seq < max_seq) {
            for (elem, seq, _) in &ranked_readers {
                if *seq == max_seq {
                    keep_reader_addrs.insert(elem.reader_addr.clone());
                }
            }
        }

        for elem in ready_readers {
            if !keep_reader_addrs.contains(&elem.reader_addr)
                && unified_shard_info
                    .mark_ready_reader_removing(&elem.reader_addr)
                    .await
            {
                if let Some(load) = planned_load.get_mut(&elem.reader_addr) {
                    *load = load.saturating_sub(1);
                }
            }
        }

        keep_reader_addrs.contains(&self.addr)
    }

    pub async fn handle_hearbeart_for_shard_cache(
        &self,
        reported_shards_by_id: &HashMap<u64, Option<u64>>,
        dead_timeout: ReadableDuration,
        reader_pool: &ReaderPool,
        shard_cache_manager: &ShardCacheManager,
    ) -> (
        Vec<tici_grpc::ShardManifestHeader>,
        Vec<tici_grpc::ShardManifestHeader>,
    ) {
        let mut all_shards_to_keep = HashMap::<u64, ShardManifestHeader>::new();
        let mut all_shards_to_remove = HashMap::<u64, ShardManifestHeader>::new();
        let mut not_allocated_ids = reported_shards_by_id
            .keys()
            .copied()
            .collect::<HashSet<u64>>();
        let allocated_shards = self.allocated_shards_snapshot().await;
        let mut affected_shards = HashMap::<u64, UnifiedShardInfo>::new();

        let mut to_be_removed = vec![];
        for (shard_id, unified_shard_info) in &allocated_shards {
            affected_shards.insert(*shard_id, unified_shard_info.clone());
            if let Some(reported_seq) = reported_shards_by_id.get(shard_id) {
                not_allocated_ids.remove(shard_id);
                if unified_shard_info.is_reader_removing(&self.addr).await {
                    unified_shard_info
                        .update_removing_reader_with_observed_seq(
                            self.addr.clone(),
                            self.tiflash_server_addr.clone(),
                            *reported_seq,
                        )
                        .await;
                    all_shards_to_keep.remove(shard_id);
                    all_shards_to_remove.insert(*shard_id, unified_shard_info.get_header());
                } else {
                    let _ = unified_shard_info
                        .add_ready_reader_with_observed_seq(
                            self.addr.clone(),
                            self.tiflash_server_addr.clone(),
                            *reported_seq,
                        )
                        .await;
                    if !all_shards_to_remove.contains_key(shard_id) {
                        all_shards_to_keep.insert(*shard_id, unified_shard_info.get_header());
                    }
                }
            } else {
                if unified_shard_info.is_reader_removing(&self.addr).await {
                    info!(
                        "finalize removing reader replica, shard={}, addr={}",
                        unified_shard_info.to_long_ident(),
                        self.addr,
                    );
                    let callbacks = unified_shard_info.take_remove_callbacks(&self.addr).await;
                    unified_shard_info.finalize_remove_reader(&self.addr).await;
                    for callback in callbacks {
                        let _ = callback.send(ManualRemoveResult::Success);
                    }
                    to_be_removed.push(unified_shard_info.shard_id);
                    all_shards_to_keep.remove(shard_id);
                    all_shards_to_remove.remove(shard_id);
                    continue;
                }
                match unified_shard_info
                    .maybe_remove_ready_reader(&self.addr, dead_timeout)
                    .await
                {
                    ShardCacheAllocationState::Again(d) => {
                        debug!(
                            "shard cache temporarily missing, reader_addr={}, shard={}, timeout={:?}, elapsed={:.4?}",
                            self.addr,
                            unified_shard_info.to_long_ident(),
                            dead_timeout,
                            d.as_millis_f64()
                        );
                        if !all_shards_to_remove.contains_key(shard_id) {
                            all_shards_to_keep.insert(*shard_id, unified_shard_info.get_header());
                        }
                    }
                    ShardCacheAllocationState::NotReady => {
                        if !all_shards_to_remove.contains_key(shard_id) {
                            all_shards_to_keep.insert(*shard_id, unified_shard_info.get_header());
                        }
                    }
                    ShardCacheAllocationState::Timeout => {
                        info!(
                            "remove reader for timeout, unified_shard_info={}, readers={:?}, addr={}, reported_shards={:?}",
                            unified_shard_info.to_long_ident(),
                            unified_shard_info.get_reader_addrs().await,
                            &self.addr,
                            reported_shards_by_id
                        );
                        to_be_removed.push(unified_shard_info.shard_id);
                        all_shards_to_keep.remove(shard_id);
                    }
                }
            }
        }
        for shard_id in to_be_removed {
            self.remove_shard_cache(&shard_id).await;
        }

        let (need_to_remedy, nonexist) = shard_cache_manager
            .remedy_shard_caches(not_allocated_ids)
            .await;

        for shard_id in nonexist {
            debug!(
                "shard nonexist in shard cache, shard_id={}, addr={}, tiflash_addr={}",
                shard_id, self.addr, self.tiflash_server_addr,
            );
            all_shards_to_keep.remove(&shard_id);
            all_shards_to_remove.insert(
                shard_id,
                ShardManifestHeader {
                    shard_id,
                    ..Default::default()
                },
            );
        }

        for shard in need_to_remedy {
            affected_shards.insert(shard.shard_id, shard.clone());
            debug!(
                "reconcile reader for remedy, shard={}, addr={}, tiflash_addr={}",
                shard.to_long_ident(),
                self.addr,
                self.tiflash_server_addr,
            );
            let observed_seq = reported_shards_by_id
                .get(&shard.shard_id)
                .copied()
                .flatten();
            if shard.is_reader_removing(&self.addr).await {
                shard
                    .update_removing_reader_with_observed_seq(
                        self.addr.clone(),
                        self.tiflash_server_addr.clone(),
                        observed_seq,
                    )
                    .await;
                all_shards_to_keep.remove(&shard.shard_id);
                all_shards_to_remove.insert(shard.shard_id, shard.get_header());
            } else {
                shard
                    .add_ready_reader_with_observed_seq(
                        self.addr.clone(),
                        self.tiflash_server_addr.clone(),
                        observed_seq,
                    )
                    .await;
                if !all_shards_to_remove.contains_key(&shard.shard_id) {
                    all_shards_to_keep.insert(shard.shard_id, shard.get_header());
                }
            }
            self.add_shard_cache(shard).await;
        }

        let mut affected_shards = affected_shards.into_iter().collect::<Vec<_>>();
        affected_shards.sort_by_key(|(shard_id, _)| *shard_id);

        let should_keep_shards = self
            .reconcile_redundant_ready_readers_batch(&affected_shards, reader_pool)
            .await;
        for (shard_id, unified_shard_info) in affected_shards {
            if unified_shard_info.is_reader_removing(&self.addr).await
                || !should_keep_shards.contains(&shard_id)
            {
                all_shards_to_keep.remove(&shard_id);
                all_shards_to_remove.insert(shard_id, unified_shard_info.get_header());
            }
        }

        (
            all_shards_to_keep.into_values().collect(),
            all_shards_to_remove.into_values().collect(),
        )
    }
}