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

use itertools::Itertools;
use log::{debug, info};
use tici_system::readable::ReadableDuration;
use tokio::sync::oneshot;

use crate::meta::reader::cache_scheduler::ReaderCacheScheduler;

/// Cache allocation functionality
#[derive(Clone)]
pub struct CacheAllocator {
    pub cache_scheduler: Arc<dyn ReaderCacheScheduler>,
    pub s3_bucket: String,
    pub allocate_parallelism: usize,
    pub max_reader_attempts: usize,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AllocateReaderReason {
    Split,
    ReadRequest,
    EagerSchedule,
    ReallocateOnReaderFailOver,
    Merge,
    Manual,
}

#[derive(Clone)]
pub struct ShardCacheAllocationElem {
    pub(crate) reader_addr: String,
    pub(crate) tiflash_server_addr: String,
    pub(crate) observed_seq: Option<u64>,
    pub(crate) ready_at: Instant,
}

impl ShardCacheAllocationElem {
    pub fn new(reader_addr: String, tiflash_server_addr: String, ready_at: Instant) -> Self {
        Self::with_observed_seq(reader_addr, tiflash_server_addr, None, ready_at)
    }

    pub fn with_observed_seq(
        reader_addr: String,
        tiflash_server_addr: String,
        observed_seq: Option<u64>,
        ready_at: Instant,
    ) -> Self {
        ShardCacheAllocationElem {
            reader_addr,
            tiflash_server_addr,
            observed_seq,
            ready_at,
        }
    }
}

impl std::fmt::Debug for ShardCacheAllocationElem {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!(
            "addr={}/tiflash_addr={}/seq={:?}/elapsed={}",
            self.reader_addr,
            self.tiflash_server_addr,
            self.observed_seq,
            self.ready_at.elapsed().as_millis()
        ))
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ManualRemoveResult {
    Success,
    ReaderExpired,
}

#[derive(Clone)]
pub struct RemovingReaderEntry {
    pub(crate) elem: ShardCacheAllocationElem,
    callbacks: Arc<Mutex<Vec<oneshot::Sender<ManualRemoveResult>>>>,
}

impl RemovingReaderEntry {
    pub fn new(elem: ShardCacheAllocationElem) -> Self {
        Self {
            elem,
            callbacks: Arc::new(Mutex::new(vec![])),
        }
    }

    pub fn add_callback(&self, callback: oneshot::Sender<ManualRemoveResult>) {
        if let Ok(mut callbacks) = self.callbacks.lock() {
            callbacks.push(callback);
        }
    }

    pub fn take_callbacks(&self) -> Vec<oneshot::Sender<ManualRemoveResult>> {
        if let Ok(mut callbacks) = self.callbacks.lock() {
            std::mem::take(&mut *callbacks)
        } else {
            vec![]
        }
    }
}

impl std::fmt::Debug for RemovingReaderEntry {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let callback_count = self
            .callbacks
            .lock()
            .map(|callbacks| callbacks.len())
            .unwrap_or_default();
        f.write_fmt(format_args!("{:?}/callbacks={}", self.elem, callback_count))
    }
}

/// Cache allocation details for a shard
#[derive(Clone)]
pub struct ShardCacheAllocation {
    ready_readers: HashMap<String, ShardCacheAllocationElem>,
    removing_readers: HashMap<String, RemovingReaderEntry>,
}

impl std::fmt::Debug for ShardCacheAllocation {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!(
            "ready=({}), removing=({})",
            self.ready_readers
                .iter()
                .map(|x| format!("{:?}", x.1))
                .join("&&"),
            self.removing_readers
                .iter()
                .map(|x| format!("{:?}", x.1))
                .join("&&")
        ))
    }
}

#[derive(PartialEq, Eq, PartialOrd, Ord)]
pub enum ShardCacheAllocationState {
    Timeout,
    Again(Duration),
    NotReady,
}

impl ShardCacheAllocationState {
    pub fn need_remove(&self) -> bool {
        match self {
            Self::Timeout | Self::NotReady => true,
            Self::Again(_) => false,
        }
    }
}

impl ShardCacheAllocation {
    pub fn not_allocated() -> Self {
        ShardCacheAllocation {
            ready_readers: HashMap::new(),
            removing_readers: HashMap::new(),
        }
    }

    pub fn new(
        reader_addrs: Vec<String>,
        tiflash_server_addrs: Vec<String>,
        ready_at: Instant,
    ) -> Self {
        Self {
            ready_readers: reader_addrs
                .into_iter()
                .zip(tiflash_server_addrs)
                .map(|(reader_addr, tiflash_server_addr)| {
                    ShardCacheAllocationElem::new(reader_addr, tiflash_server_addr, ready_at)
                })
                .map(|x| (x.reader_addr.clone(), x))
                .collect(),
            removing_readers: HashMap::new(),
        }
    }

    pub fn has_ready_readers(&self) -> bool {
        !self.ready_readers.is_empty()
    }

    pub fn get_reader_addrs(&self) -> impl Iterator<Item = &String> {
        self.ready_readers.keys()
    }

    pub fn cloned_reader_addrs(&self) -> Vec<String> {
        let mut elems = self.ready_readers.values().collect::<Vec<_>>();
        elems.sort_by(|a, b| {
            a.ready_at
                .cmp(&b.ready_at)
                .then_with(|| a.reader_addr.cmp(&b.reader_addr))
        });
        elems.into_iter().map(|x| x.reader_addr.clone()).collect()
    }

    pub fn cloned_tiflash_server_addrs(&self) -> Vec<String> {
        let mut elems = self.ready_readers.values().collect::<Vec<_>>();
        elems.sort_by(|a, b| {
            a.ready_at
                .cmp(&b.ready_at)
                .then_with(|| a.reader_addr.cmp(&b.reader_addr))
        });
        elems
            .into_iter()
            .map(|x| x.tiflash_server_addr.clone())
            .collect()
    }

    pub fn has_reader_addr(&self, s: &String) -> bool {
        self.ready_readers.contains_key(s) || self.removing_readers.contains_key(s)
    }

    pub fn get_reader_observed_seq(&self, reader_addr: &str) -> Option<u64> {
        self.ready_readers
            .get(reader_addr)
            .and_then(|elem| elem.observed_seq)
            .or_else(|| {
                self.removing_readers
                    .get(reader_addr)
                    .and_then(|entry| entry.elem.observed_seq)
            })
    }

    pub fn ready_reader_count(&self) -> usize {
        self.ready_readers.len()
    }

    pub fn is_reader_removing(&self, reader_addr: &str) -> bool {
        self.removing_readers.contains_key(reader_addr)
    }

    pub fn ready_reader_elems(&self) -> Vec<ShardCacheAllocationElem> {
        self.ready_readers.values().cloned().collect()
    }

    pub fn removing_reader_elems(&self) -> Vec<ShardCacheAllocationElem> {
        self.removing_readers
            .values()
            .map(|entry| entry.elem.clone())
            .collect()
    }

    pub fn add_ready_reader_elem(&mut self, elem: ShardCacheAllocationElem) -> bool {
        self.removing_readers.remove(&elem.reader_addr);
        if !self.ready_readers.contains_key(&elem.reader_addr) {
            self.ready_readers.insert(elem.reader_addr.clone(), elem);
            true
        } else {
            false
        }
    }

    pub fn add_ready_reader(&mut self, reader_addr: String, tiflash_server_addr: String) -> bool {
        self.add_ready_reader_with_observed_seq(reader_addr, tiflash_server_addr, None)
    }

    pub fn add_ready_reader_with_observed_seq(
        &mut self,
        reader_addr: String,
        tiflash_server_addr: String,
        observed_seq: Option<u64>,
    ) -> bool {
        if let Some(existing) = self.ready_readers.get_mut(&reader_addr) {
            existing.tiflash_server_addr = tiflash_server_addr;
            existing.observed_seq = observed_seq;
            return false;
        }
        self.add_ready_reader_elem(ShardCacheAllocationElem::with_observed_seq(
            reader_addr,
            tiflash_server_addr,
            observed_seq,
            Instant::now(),
        ))
    }

    pub fn remove_ready_reader(&mut self, reader_addr: &str) -> bool {
        self.ready_readers.remove(reader_addr).is_some()
    }

    pub fn mark_ready_reader_removing(&mut self, reader_addr: &str) -> bool {
        if let Some(elem) = self.ready_readers.remove(reader_addr) {
            self.removing_readers
                .insert(reader_addr.to_string(), RemovingReaderEntry::new(elem));
            true
        } else {
            false
        }
    }

    pub fn add_removing_reader_callback(
        &mut self,
        reader_addr: &str,
        callback: oneshot::Sender<ManualRemoveResult>,
    ) -> bool {
        if let Some(entry) = self.removing_readers.get(reader_addr) {
            entry.add_callback(callback);
            true
        } else {
            false
        }
    }

    pub fn take_remove_callbacks(
        &mut self,
        reader_addr: &str,
    ) -> Vec<oneshot::Sender<ManualRemoveResult>> {
        self.removing_readers
            .get(reader_addr)
            .map(RemovingReaderEntry::take_callbacks)
            .unwrap_or_default()
    }

    pub fn finalize_remove_reader(&mut self, reader_addr: &str) -> bool {
        self.ready_readers.remove(reader_addr).is_some()
            || self.removing_readers.remove(reader_addr).is_some()
    }

    pub fn update_removing_reader_with_observed_seq(
        &mut self,
        reader_addr: String,
        tiflash_server_addr: String,
        observed_seq: Option<u64>,
    ) -> bool {
        if let Some(existing) = self.removing_readers.get_mut(&reader_addr) {
            existing.elem.tiflash_server_addr = tiflash_server_addr;
            existing.elem.observed_seq = observed_seq;
            return true;
        }
        false
    }

    pub fn maybe_remove_ready_reader(
        &mut self,
        reader_addr: &str,
        tag: String,
        timeout: ReadableDuration,
    ) -> ShardCacheAllocationState {
        if let Some(alloc) = self.ready_readers.get(reader_addr) {
            let duration = Instant::now().duration_since(alloc.ready_at);
            if duration > timeout.0 {
                info!(
                    "maybe_remove_ready_reader for lost ready reader after given timeout, reader_addr={}, alloc={:?}, duration={:.4}>timeout={:.4}, unified_shard_info={}",
                    reader_addr,
                    alloc,
                    duration.as_secs_f32(),
                    timeout.0.as_secs_f64(),
                    tag,
                );
                self.remove_ready_reader(reader_addr);
                return ShardCacheAllocationState::Timeout;
            }
            return ShardCacheAllocationState::Again(duration);
        }

        debug!(
            "maybe_remove_ready_reader, but it has not been ready, reader_addr={}, unified_shard_info={}, timeout={:?}",
            reader_addr, tag, timeout
        );
        ShardCacheAllocationState::NotReady
    }
}

impl CacheAllocator {
    #[cfg(any(test, feature = "testexport"))]
    pub fn new_for_test() -> Self {
        use crate::meta::reader::RoundRobinCacheScheduler;

        Self {
            cache_scheduler: Arc::new(RoundRobinCacheScheduler::new()),
            s3_bucket: "".to_string(),
            allocate_parallelism: 1,
            max_reader_attempts: 1,
        }
    }

    pub fn new(
        cache_scheduler: Arc<dyn ReaderCacheScheduler>,
        s3_bucket: String,
        allocate_parallelism: usize,
        max_reader_attempts: usize,
    ) -> Self {
        Self {
            cache_scheduler,
            s3_bucket,
            allocate_parallelism,
            max_reader_attempts,
        }
    }
}