use std::{
    collections::HashSet,
    sync::{
        Arc,
        atomic::{AtomicBool, AtomicU64, Ordering},
    },
};

use anyhow::Result;
use dashmap::{DashMap, DashSet};
use either::Either;
use log::info;
use rand::{SeedableRng, seq::IteratorRandom};
use tici_common::{
    codec::{previous_key_helper, table_end_key, table_start_key},
    config::ServerConfig,
    index::IndexKey,
};
use tici_grpc::{
    AddShardRequest, AppendBaseFragMetaRequest, AppendDeltaFragMetaRequest, AppendFragMetaResponse,
    CompactFragRequest, ErrorCode, FinishCompactFragRequest, FragMeta, FragProperty, ShardProgress,
    WorkerNodeShardStatus, WorkerNodeStatus,
};
use tici_shard::{ShardWriterMeta, test_utils::test_mock_frag_storage_refs_total};
use tici_system::cpu_time::ProcessStat;
use tokio::sync::RwLock;

use crate::{
    HandleRange, RangeBound,
    meta::{MetaServiceCore, writer::FRAG_FILE_PREFIX},
    worker::WorkerNodeServiceCore,
};

pub fn midpoint_bytes(a: &[u8], b: &[u8]) -> Vec<u8> {
    let mut a = a.to_vec();
    let mut b = b.to_vec();
    if b.len() < a.len() {
        b.resize(a.len(), 0);
    }
    if a.len() < b.len() {
        a.resize(b.len(), 0);
    }
    let mut result = a.clone();
    for i in 0..result.len() {
        if result[i] != b[i] {
            let mid = ((result[i] as u16 + b[i] as u16) / 2) as u8;
            if mid == result[i] {
                if i + 1 < result.len() {
                    result[i + 1] = ((result[i + 1] as u16 + 256) / 2) as u8;
                    if result[i + 1] == a[i + 1] {
                        // [1,1,255], [1,2,0]
                        break;
                    } else {
                        // [1,1,1], [1,2,0]
                        return result;
                    }
                } else {
                    // [1,1], [1,2]
                    break;
                }
            } else {
                result[i] = mid;
                return result;
            }
        }
    }

    result.push(128);
    result
}

pub fn make_frag_meta(index_key: &IndexKey, frag_path: &str) -> FragMeta {
    let range = HandleRange::new(
        table_start_key(index_key).to_vec(),
        previous_key_helper(table_end_key(index_key).as_slice(), 1),
    );
    make_frag_meta_extended(frag_path, &range, 1, 1, 1)
}

pub fn make_frag_meta_by_range(frag_path: &str, range: &HandleRange) -> FragMeta {
    make_frag_meta_extended(frag_path, range, 1, 1, 1)
}

pub fn make_frag_meta_extended(
    frag_path: &str,
    range: &HandleRange,
    size: u64,
    count: u64,
    segs_count: usize,
) -> FragMeta {
    if !range.is_valid() {
        panic!("invalid range {:?}", range);
    }
    let mid_handle = midpoint_bytes(&range.encoded_start(), &range.encoded_end());
    if range
        .sortable_start_ref()
        .effective_eq(&RangeBound::Key(mid_handle.clone()))
        || range
            .sortable_end_ref()
            .effective_eq(&RangeBound::Key(mid_handle.clone()))
    {
        panic!(
            "make_frag_meta_extended compute error: [{:?} {:?}) => {:?}",
            range.encoded_start(),
            range.encoded_end(),
            mid_handle
        );
    }
    let split_handles = vec![mid_handle];
    FragMeta::new(
        999,
        frag_path.to_string(),
        FragProperty {
            size,
            count,
            min_handle: range.encoded_start(),
            split_handles,
            max_handle: range.encoded_end(),
            delete_count: 0,
        },
        test_mock_frag_storage_refs_total(size),
        segs_count as u32,
    )
    .expect("test frag path must be a valid f_{tso} path")
}

#[derive(Clone)]
pub struct MockedFragment {
    pub frag_path: String,
    // A set of keys which shards the same prefix.
    // We use this as an estimate because we can't actually store all keys.
    pub keys: Vec<(String, Vec<u8>)>,
}

impl MockedFragment {
    pub fn from_multiple(frag_path: String, frags: &[MockedFragment]) -> Self {
        let mut s = vec![];
        for f in frags.iter() {
            for key in f.keys.iter() {
                s.push(key.clone());
            }
        }
        Self { frag_path, keys: s }
    }

    pub fn get_min_key(&self) -> (String, Vec<u8>) {
        self.keys
            .iter()
            .min_by_key(|x| x.1.clone())
            .unwrap()
            .clone()
    }

    pub fn get_max_key(&self) -> (String, Vec<u8>) {
        self.keys
            .iter()
            .max_by_key(|x| x.1.clone())
            .unwrap()
            .clone()
    }
}

pub struct MockShardWriterManager {
    meta: RwLock<Arc<MetaServiceCore>>,
    pub config: ServerConfig,
    pub skip_heartbeat: AtomicBool,
    pub skipped_heartbeat_shards: DashSet<u64>,
    // If set to true: the compactor sleep for some time, and send a response
    // Otherwise: the test shall manually send the response
    pub mock_compact_process: AtomicBool,
    pub check_to_add_shard: AtomicBool,
    pub to_add_shards: DashMap<u64, AtomicU64>,
    pub shards: DashMap<u64, ShardWriterMeta>,
    counter: AtomicU64,
    pub compaction_counter: AtomicU64,
    pub respond_compact_delay: AtomicU64,
    // How much data left after the compaction
    pub after_compaction_ratio: portable_atomic::AtomicF64,
    pub frags: DashMap<String, MockedFragment>,
    pub expected_frags: DashMap<String, MockedFragment>,
    rng: tokio::sync::RwLock<rand::rngs::StdRng>,
}

impl MockShardWriterManager {
    pub fn new(meta: Arc<MetaServiceCore>, config: ServerConfig) -> Self {
        Self {
            meta: RwLock::new(meta),
            config,
            skip_heartbeat: AtomicBool::new(false),
            skipped_heartbeat_shards: DashSet::new(),
            mock_compact_process: AtomicBool::new(false),
            check_to_add_shard: AtomicBool::new(false),
            to_add_shards: DashMap::new(),
            shards: DashMap::new(),
            counter: AtomicU64::new(0),
            compaction_counter: AtomicU64::new(0),
            respond_compact_delay: AtomicU64::new(500),
            after_compaction_ratio: portable_atomic::AtomicF64::new(0.8),
            frags: DashMap::new(),
            expected_frags: DashMap::new(),
            rng: tokio::sync::RwLock::new(Self::gen_rng()),
        }
    }

    pub fn gen_rng() -> rand::rngs::StdRng {
        use rand::rngs::StdRng;
        let seed = [
            1, 0, 0, 0, 23, 0, 0, 0, 200, 1, 0, 0, 210, 30, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0,
        ];
        StdRng::from_seed(seed)
    }

    pub async fn reset_meta(&self, meta: Arc<MetaServiceCore>) {
        *self.meta.write().await = meta;
    }

    pub async fn add_shard(&self, shard: &ShardWriterMeta) {
        info!("mocker_worker: add_shard, shard={:?}", shard);
        self.shards
            .insert(shard.shard_header.shard_id, shard.clone());
    }

    pub fn shard_statuses(&self) -> Vec<WorkerNodeShardStatus> {
        self.shards
            .iter()
            .map(|entry| {
                let shard_header = &entry.value().shard_header;
                WorkerNodeShardStatus {
                    start_key: shard_header.start_key.clone(),
                    shard_id: shard_header.shard_id,
                }
            })
            .collect()
    }

    pub async fn remove_shard(&self, shard_id: &u64) {
        info!("mocker_worker: remove_shard, shard_id={}", shard_id);
        self.shards.remove(shard_id);
    }

    pub async fn get_total_rows_by_frags(&self) -> u64 {
        let mut keys = HashSet::new();
        for f in self.frags.iter() {
            for key in f.value().keys.iter() {
                keys.insert(key.clone());
            }
        }
        keys.len() as u64
    }

    pub async fn get_total_expected_rows_by_frags(&self) -> u64 {
        let mut keys = HashSet::new();
        for f in self.expected_frags.iter() {
            for key in f.value().keys.iter() {
                keys.insert(key.clone());
            }
        }
        keys.len() as u64
    }

    pub async fn compact_fragments(
        &self,
        req: CompactFragRequest,
    ) -> Result<FinishCompactFragRequest> {
        tokio::time::sleep(std::time::Duration::from_millis(
            self.respond_compact_delay.load(Ordering::Relaxed),
        ))
        .await;
        let shard = req.shard.as_ref().unwrap().clone();
        let range = HandleRange::new(shard.start_key.clone(), shard.end_key.clone());
        let ratio = self.after_compaction_ratio.load(Ordering::Relaxed);
        let frag_path = format!(
            "c_{}_/test/{}/t_{}/i_{}/f_9999999999999999999",
            self.compaction_counter.fetch_add(1, Ordering::Relaxed),
            req.table_id,
            req.index_id,
            FRAG_FILE_PREFIX
        );
        let mut count: u64 = 0;
        let mut total_count: u64 = 0;
        let mut mocked_frags = vec![];
        for f in req.input_frags.iter() {
            let mock_frag = match self.frags.get(&f.frag_path) {
                Some(x) => x.clone(),
                None => {
                    panic!(
                        "frag {} is not registered, total={:?}",
                        f.frag_path,
                        self.frags
                            .iter()
                            .map(|x| x.value().frag_path.clone())
                            .collect::<Vec<_>>()
                            .join(";")
                    );
                }
            };
            // We don't use result of `.bytes()` and `.rows()` in `input_frags`,
            // because compaction will remove all out-of-bound data.
            // We have to look in to the fragment, to see how many keys are to be removed.
            for key in mock_frag.keys.iter() {
                if range.contains_key(&crate::RangeBound::Key(key.1.clone())) {
                    count += 1;
                }
            }
            total_count += mock_frag.keys.len() as u64;
            mocked_frags.push(mock_frag);
        }
        let new_start = mocked_frags
            .iter()
            .map(|x| x.get_min_key().1.clone())
            .min()
            .unwrap();
        let new_end = mocked_frags
            .iter()
            .map(|x| x.get_max_key().1.clone())
            .max()
            .unwrap();
        let new_range = HandleRange::new(new_start, new_end);
        let output_frags = if new_range
            .sortable_start_ref()
            .effective_eq(new_range.sortable_end_ref())
        {
            let trim_ratio: f64 = (count as f64) / (total_count as f64);
            let size =
                req.input_frags.iter().map(|x| x.bytes()).sum::<u64>() as f64 * ratio * trim_ratio;
            let count =
                req.input_frags.iter().map(|x| x.rows()).sum::<u64>() as f64 * ratio * trim_ratio;
            let out_frag =
                make_frag_meta_extended(&frag_path, &new_range, size as u64, count as u64, 1);
            self.frags.insert(
                out_frag.frag_path.clone(),
                MockedFragment::from_multiple(out_frag.frag_path.clone(), mocked_frags.as_slice()),
            );
            info!("mock: finished compaction, shard={:?}", &req.shard);
            vec![out_frag]
        } else {
            info!(
                "mock: finished compaction with no data, shard={:?}",
                &req.shard
            );
            Vec::new()
        };
        Ok(FinishCompactFragRequest {
            addr: self.config.get_advertise_addr().to_string(),
            shard: req.shard.clone(),
            keyspace_id: req.keyspace_id,
            table_id: req.table_id,
            index_id: req.index_id,
            seq: req.seq,
            timestamp: req.timestamp,
            leader_term: req.leader_term.clone(),
            output_frags,
            input_frags: req.input_frags.clone(),
        })
    }
}

impl MockShardWriterManager {
    // Maintain all rows which are successfully added.
    async fn add_frag_info(&self, frag: &FragMeta) {
        self.frags.insert(
            frag.frag_path.clone(),
            MockedFragment {
                frag_path: frag.frag_path.clone(),
                keys: vec![(
                    frag.frag_path.clone(),
                    frag.property
                        .as_ref()
                        .unwrap()
                        .split_handles
                        .first()
                        .unwrap()
                        .clone(),
                )],
            },
        );
        self.expected_frags.insert(
            frag.frag_path.clone(),
            MockedFragment {
                frag_path: frag.frag_path.clone(),
                keys: vec![(
                    frag.frag_path.clone(),
                    frag.property
                        .as_ref()
                        .unwrap()
                        .split_handles
                        .first()
                        .unwrap()
                        .clone(),
                )],
            },
        );
    }

    pub async fn mock_send_append_delta_frag(
        &self,
        index_key: &IndexKey,
        shard_id: u64,
        last_cdc_file: String,
        frag_metas: Vec<FragMeta>,
    ) -> Result<(AppendDeltaFragMetaRequest, AppendFragMetaResponse)> {
        let req = AppendDeltaFragMetaRequest {
            keyspace_id: index_key.keyspace_id,
            table_id: index_key.table_id,
            index_id: index_key.index_id,
            shard_id,
            worker_node_addr: self.config.get_advertise_addr().to_string(),
            frag_metas,
            progress: Some(ShardProgress { last_cdc_file }),
        };
        let resp = self.meta.read().await.append_delta_frag_meta(&req).await?;
        if resp.status == ErrorCode::Success as i32 {
            for frag in req.frag_metas.iter() {
                self.add_frag_info(frag).await;
            }
        }
        Ok((req, resp))
    }

    pub async fn mock_send_append_base_frag(
        &self,
        index_key: &IndexKey,
        frag_metas: Vec<FragMeta>,
    ) -> Result<()> {
        let req = AppendBaseFragMetaRequest {
            keyspace_id: index_key.keyspace_id,
            table_id: index_key.table_id,
            index_id: index_key.index_id,
            frag_metas,
        };
        self.meta.read().await.append_base_frag_meta(&req).await?;
        for frag in req.frag_metas.iter() {
            self.add_frag_info(frag).await;
        }
        Ok(())
    }

    pub async fn mock_send_finish_compact_frags(
        &mut self,
        req: FinishCompactFragRequest,
    ) -> Result<()> {
        self.meta
            .read()
            .await
            .finish_compact_fragments(&req)
            .await?;
        Ok(())
    }

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

    pub fn mock_skip_shard_heartbeat(&self, shard_id: u64) {
        self.skipped_heartbeat_shards.insert(shard_id);
    }

    pub fn mock_report_shard_heartbeat(&self, shard_id: u64) {
        self.skipped_heartbeat_shards.remove(&shard_id);
    }

    pub fn mock_should_skip_shard_heartbeat(&self, shard_id: &u64) -> bool {
        self.skipped_heartbeat_shards.contains(shard_id)
    }

    pub fn mock_collect_worker_node_status(
        &self,
        proc_stats: &mut ProcessStat,
    ) -> WorkerNodeStatus {
        WorkerNodeServiceCore::collect_worker_node_status(proc_stats)
    }

    pub async fn maybe_check_to_add_shards(&self, shards: &[AddShardRequest]) -> bool {
        if self
            .check_to_add_shard
            .load(std::sync::atomic::Ordering::Relaxed)
        {
            for shard in shards.iter() {
                match self
                    .to_add_shards
                    .entry(shard.shard.clone().unwrap().shard_id)
                {
                    dashmap::mapref::entry::Entry::Occupied(c) => {
                        if c.get().fetch_add(1, Ordering::Relaxed) > 2 {
                            return true;
                        }
                    }
                    dashmap::mapref::entry::Entry::Vacant(c) => {
                        c.insert(AtomicU64::new(0));
                    }
                };
            }
        }
        false
    }

    pub async fn mock_new_frags(
        &self,
        path: Either<String, String>,
        size: u64,
        count: u64,
        segs_count: usize,
    ) -> (ShardWriterMeta, Vec<FragMeta>) {
        let shard = {
            let shards = self.shards.iter();
            let item = shards.choose(&mut *self.rng.write().await).unwrap();
            item.value().clone()
        };
        let range = HandleRange::new(
            shard.shard_header.start_key.clone(),
            previous_key_helper(shard.shard_header.end_key.as_slice(), 1),
        );
        if !range.is_valid() {
            panic!(
                "invalid range {:?}, shard={:?}..{:?}",
                range, shard.shard_header.start_key, shard.shard_header.end_key
            );
        }
        let frag_path = match path {
            Either::Right(x) => x,
            Either::Left(prefix) => {
                format!(
                    "{}_{}/f_1",
                    prefix,
                    self.counter.fetch_add(1, Ordering::Relaxed)
                )
            }
        };
        // By test, found seg_count from append is 4.
        let frag_metas = vec![make_frag_meta_extended(
            &frag_path, &range, size, count, segs_count,
        )];
        (shard, frag_metas)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_midpoint_bytes() {
        assert_eq!(midpoint_bytes(&[1, 1, 1], &[1, 1, 2]), &[1, 1, 1, 128]);
        assert_eq!(midpoint_bytes(&[1, 1, 1], &[1, 2]), &[1, 1, 128]);
        assert_eq!(midpoint_bytes(&[1, 1], &[1, 2]), &[1, 1, 128]);
        assert_eq!(
            midpoint_bytes(
                &[116, 128, 0, 0, 0, 0, 0, 0, 1],
                &[116, 128, 0, 0, 0, 0, 0, 0, 1, 128]
            ),
            &[116, 128, 0, 0, 0, 0, 0, 0, 1, 64]
        );
        assert_eq!(
            midpoint_bytes(
                hex::decode(b"748000000000000001ad6e").unwrap().as_slice(),
                hex::decode(b"748000000000000001ad6f").unwrap().as_slice()
            ),
            [116, 128, 0, 0, 0, 0, 0, 0, 1, 173, 110, 128]
        );
        assert_eq!(
            midpoint_bytes(
                &[116, 128, 0, 0, 0, 0, 0, 0, 1, 255],
                &[116, 128, 0, 0, 0, 0, 0, 0, 2],
            ),
            [116, 128, 0, 0, 0, 0, 0, 0, 1, 255, 128]
        );
        assert_eq!(
            midpoint_bytes(
                &[116, 128, 0, 0, 0, 0, 0, 0, 1, 127, 31, 0],
                &[116, 128, 0, 0, 0, 0, 0, 0, 1, 127, 31, 0, 0]
            ),
            &[116, 128, 0, 0, 0, 0, 0, 0, 1, 127, 31, 0, 0, 128],
        );
    }
}