use std::{fmt::Debug, sync::Arc};

use portable_atomic::{AtomicU16, AtomicU64};
use tici_common::index::IndexKey;
use tici_grpc::{FragMeta, ShardManifestHeader};
use tici_system::readable::ReadableDuration;
use tokio::sync::{RwLock, oneshot};

use crate::{
    HandleRange, RangeBound, ShardInfo, ShardInfoAppendUpdate, ShardManifest,
    meta::reader::{
        ManualRemoveResult, ShardCacheAllocation, ShardCacheAllocationElem,
        ShardCacheAllocationState,
    },
};

#[derive(Clone, Debug)]
pub struct UnifiedShardInfoInner {
    // Cache allocation information (optional - Some if allocated to readers)
    pub(crate) cache_allocation: ShardCacheAllocation,
    pub(crate) desired_replica_count: u8,
    pub manifest: ShardManifest,
}

impl UnifiedShardInfoInner {
    pub fn new(
        cache_allocation: ShardCacheAllocation,
        desired_replica_count: u8,
        manifest: ShardManifest,
    ) -> Self {
        UnifiedShardInfoInner {
            cache_allocation,
            desired_replica_count,
            manifest,
        }
    }
}

impl Default for UnifiedShardInfoInner {
    fn default() -> Self {
        UnifiedShardInfoInner {
            cache_allocation: ShardCacheAllocation::not_allocated(),
            desired_replica_count: 1,
            manifest: ShardManifest::new_for_first(),
        }
    }
}

/// Unified shard information that includes both metadata and cache allocation
// TODO(P2) Split the mutable and immutable part
#[derive(Clone)]
pub struct UnifiedShardInfo {
    // Basic shard metadata
    pub shard_id: u64,
    pub index_key: IndexKey,
    pub epoch: u64,
    pub range: HandleRange,
    pub seq: Arc<AtomicU64>,

    pub in_warming: Arc<AtomicU16>,

    // TODO(P2) This field can be made exclusive.
    #[cfg(any(test, feature = "testexport"))]
    pub inner: Arc<RwLock<UnifiedShardInfoInner>>,
    #[cfg(not(any(test, feature = "testexport")))]
    pub(super) inner: Arc<RwLock<UnifiedShardInfoInner>>,
}

impl std::fmt::Debug for UnifiedShardInfo {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("UShardInfo")
            .field_with("shard", |f| {
                f.write_fmt(format_args!(
                    "shard_id={}/{:?}/epoch={}/seq={}/range={:?}/in_warming={}",
                    self.shard_id,
                    self.index_key,
                    self.epoch,
                    self.seq.load(std::sync::atomic::Ordering::SeqCst),
                    self.range,
                    self.in_warming.load(std::sync::atomic::Ordering::SeqCst)
                ))
            })
            .field("inner", &self.inner)
            .finish()
    }
}

impl UnifiedShardInfo {
    pub fn to_long_ident(&self) -> String {
        format!(
            "({:?}/shard_id={}/start_key={:?}/seq={})",
            self.index_key,
            self.shard_id,
            self.range.sortable_start_ref(),
            self.seq.load(std::sync::atomic::Ordering::SeqCst),
        )
    }

    pub fn is_warming(&self) -> bool {
        self.in_warming.load(std::sync::atomic::Ordering::SeqCst) == 1
    }

    pub async fn apply_append(&self, append: &ShardInfoAppendUpdate) {
        // Check `ShardInfo::apply_append`.
        if append.is_base {
            self.inner.write().await.manifest.append_base_fragments(
                append.frags.clone(),
                append.shard_id,
                append.new_seq,
            );
            self.seq
                .store(append.new_seq, std::sync::atomic::Ordering::SeqCst);
        } else {
            self.inner.write().await.manifest.append_delta_fragments(
                append.frags.clone(),
                append.shard_id,
                append.new_seq,
            );
            self.seq
                .store(append.new_seq, std::sync::atomic::Ordering::SeqCst);
        }
    }

    pub async fn apply_update(&self, new_shard: &ShardInfo) {
        self.seq
            .store(new_shard.seq, std::sync::atomic::Ordering::SeqCst);
        self.inner.write().await.manifest = new_shard.manifest.clone();
    }

    pub async fn cloned_fragments(&self) -> Vec<FragMeta> {
        self.inner.read().await.manifest.cloned_fragments()
    }
}

impl UnifiedShardInfo {
    pub fn from_shard_info(shard_info: &ShardInfo) -> Self {
        Self {
            shard_id: shard_info.shard_id,
            index_key: shard_info.index_key.clone(),
            epoch: shard_info.epoch,
            seq: Arc::new(AtomicU64::new(shard_info.seq)),
            range: shard_info.range.clone(),
            inner: Arc::new(RwLock::new(UnifiedShardInfoInner {
                manifest: shard_info.manifest.clone(),
                ..Default::default()
            })),
            in_warming: Arc::new(AtomicU16::new(0)),
        }
    }

    pub fn start_key(&self) -> RangeBound {
        self.range.sortable_start()
    }

    pub fn end_key(&self) -> RangeBound {
        self.range.sortable_end()
    }

    pub fn start_key_ref(&self) -> &RangeBound {
        self.range.sortable_start_ref()
    }

    pub fn end_key_ref(&self) -> &RangeBound {
        self.range.sortable_end_ref()
    }

    pub fn get_header(&self) -> ShardManifestHeader {
        ShardManifestHeader {
            shard_id: self.shard_id,
            start_key: self.range.encoded_start(),
            end_key: self.range.encoded_end(),
            epoch: self.epoch,
            seq: self.seq.load(std::sync::atomic::Ordering::SeqCst),
        }
    }

    pub fn overlaps_with_key_range(&self, key_range: &tici_grpc::KeyRange) -> bool {
        self.range
            .range_overlaps_range(&key_range.start_key, &key_range.end_key)
    }

    pub async fn has_ready_readers(&self) -> bool {
        self.inner.read().await.cache_allocation.has_ready_readers()
    }

    pub async fn clear_all_readers(&self) {
        self.inner.write().await.cache_allocation = ShardCacheAllocation::not_allocated();
    }

    pub async fn remove_ready_reader(&self, expired_reader_addr: &str) -> bool {
        self.inner
            .write()
            .await
            .cache_allocation
            .remove_ready_reader(expired_reader_addr)
    }

    pub async fn finalize_remove_reader(&self, reader_addr: &str) -> bool {
        self.inner
            .write()
            .await
            .cache_allocation
            .finalize_remove_reader(reader_addr)
    }

    pub async fn mark_ready_reader_removing(&self, reader_addr: &str) -> bool {
        self.inner
            .write()
            .await
            .cache_allocation
            .mark_ready_reader_removing(reader_addr)
    }

    pub async fn add_removing_reader_callback(
        &self,
        reader_addr: &str,
        callback: oneshot::Sender<ManualRemoveResult>,
    ) -> bool {
        self.inner
            .write()
            .await
            .cache_allocation
            .add_removing_reader_callback(reader_addr, callback)
    }

    pub async fn take_remove_callbacks(
        &self,
        reader_addr: &str,
    ) -> Vec<oneshot::Sender<ManualRemoveResult>> {
        self.inner
            .write()
            .await
            .cache_allocation
            .take_remove_callbacks(reader_addr)
    }

    pub async fn get_reader_addrs(&self) -> Vec<String> {
        self.inner
            .read()
            .await
            .cache_allocation
            .cloned_reader_addrs()
    }

    pub async fn get_removing_reader_addrs(&self) -> Vec<String> {
        let mut addrs = self
            .inner
            .read()
            .await
            .cache_allocation
            .removing_reader_elems()
            .into_iter()
            .map(|elem| elem.reader_addr)
            .collect::<Vec<_>>();
        addrs.sort();
        addrs
    }

    pub async fn has_reader_addr(&self, s: &String) -> bool {
        self.inner.read().await.cache_allocation.has_reader_addr(s)
    }

    pub async fn is_reader_removing(&self, s: &str) -> bool {
        self.inner
            .read()
            .await
            .cache_allocation
            .is_reader_removing(s)
    }

    pub async fn get_reader_observed_seq(&self, reader_addr: &str) -> Option<u64> {
        self.inner
            .read()
            .await
            .cache_allocation
            .get_reader_observed_seq(reader_addr)
    }

    pub async fn maybe_remove_ready_reader(
        &self,
        reader_addr: &str,
        timeout: ReadableDuration,
    ) -> ShardCacheAllocationState {
        self.inner
            .write()
            .await
            .cache_allocation
            .maybe_remove_ready_reader(reader_addr, self.to_long_ident(), timeout)
    }

    pub async fn add_ready_reader(&self, reader_addr: String, tiflash_server_addr: String) -> bool {
        self.inner
            .write()
            .await
            .cache_allocation
            .add_ready_reader(reader_addr, tiflash_server_addr)
    }

    pub async fn add_ready_reader_with_observed_seq(
        &self,
        reader_addr: String,
        tiflash_server_addr: String,
        observed_seq: Option<u64>,
    ) -> bool {
        self.inner
            .write()
            .await
            .cache_allocation
            .add_ready_reader_with_observed_seq(reader_addr, tiflash_server_addr, observed_seq)
    }

    pub async fn add_ready_reader_elem(&self, allocation: ShardCacheAllocationElem) -> bool {
        self.inner
            .write()
            .await
            .cache_allocation
            .add_ready_reader_elem(allocation)
    }

    pub async fn update_removing_reader_with_observed_seq(
        &self,
        reader_addr: String,
        tiflash_server_addr: String,
        observed_seq: Option<u64>,
    ) -> bool {
        self.inner
            .write()
            .await
            .cache_allocation
            .update_removing_reader_with_observed_seq(
                reader_addr,
                tiflash_server_addr,
                observed_seq,
            )
    }

    pub async fn cloned_cache_allocation(&self) -> ShardCacheAllocation {
        self.inner.read().await.cache_allocation.clone()
    }

    pub async fn cloned_tiflash_server_addrs(&self) -> Vec<String> {
        self.inner
            .read()
            .await
            .cache_allocation
            .cloned_tiflash_server_addrs()
    }

    pub async fn desired_replica_count(&self) -> u8 {
        self.inner.read().await.desired_replica_count
    }

    pub async fn set_desired_replica_count(&self, desired_replica_count: u8) {
        self.inner.write().await.desired_replica_count = desired_replica_count;
    }

    pub async fn inc_desired_replica_count(&self) -> u8 {
        let mut inner = self.inner.write().await;
        inner.desired_replica_count = inner.desired_replica_count.saturating_add(1);
        inner.desired_replica_count
    }

    pub async fn dec_desired_replica_count(&self) -> u8 {
        let mut inner = self.inner.write().await;
        inner.desired_replica_count = inner.desired_replica_count.saturating_sub(1);
        inner.desired_replica_count
    }
}