// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, RwLock};

use store_api::region_engine::RegionRole;
use store_api::storage::RegionId;

use crate::DatanodeId;

/// Tracks the operating(i.e., creating, opening, dropping) regions.
#[derive(Debug, Clone)]
pub struct OperatingRegionGuard {
    datanode_id: DatanodeId,
    region_id: RegionId,
    inner: Arc<RwLock<HashMap<(DatanodeId, RegionId), RegionRole>>>,
}

impl Drop for OperatingRegionGuard {
    fn drop(&mut self) {
        let mut inner = self.inner.write().unwrap();
        inner.remove(&(self.datanode_id, self.region_id));
    }
}

impl OperatingRegionGuard {
    /// Returns opening region info.
    pub fn info(&self) -> (DatanodeId, RegionId) {
        (self.datanode_id, self.region_id)
    }
}

pub type MemoryRegionKeeperRef = Arc<MemoryRegionKeeper>;

/// Tracks regions in memory.
///
/// It used in following cases:
/// - Tracks the opening regions before the corresponding metadata is created.
/// - Tracks the deleting regions after the corresponding metadata is deleted.
#[derive(Debug, Clone, Default)]
pub struct MemoryRegionKeeper {
    inner: Arc<RwLock<HashMap<(DatanodeId, RegionId), RegionRole>>>,
}

impl MemoryRegionKeeper {
    pub fn new() -> Self {
        Default::default()
    }

    /// Returns [OperatingRegionGuard] if Region(`region_id`) on Peer(`datanode_id`) does not exist.
    pub fn register_with_role(
        &self,
        datanode_id: DatanodeId,
        region_id: RegionId,
        role: RegionRole,
    ) -> Option<OperatingRegionGuard> {
        let mut inner = self.inner.write().unwrap();

        match inner.entry((datanode_id, region_id)) {
            Entry::Occupied(_) => None,
            Entry::Vacant(vacant_entry) => {
                vacant_entry.insert(role);
                Some(OperatingRegionGuard {
                    datanode_id,
                    region_id,
                    inner: self.inner.clone(),
                })
            }
        }
    }

    /// Returns true if the keeper contains a (`datanoe_id`, `region_id`) tuple.
    pub fn contains(&self, datanode_id: DatanodeId, region_id: RegionId) -> bool {
        let inner = self.inner.read().unwrap();
        inner.contains_key(&(datanode_id, region_id))
    }

    /// Extracts all operating regions with roles from `region_ids`.
    pub fn extract_operating_region_roles(
        &self,
        datanode_id: DatanodeId,
        region_ids: &HashSet<RegionId>,
    ) -> HashMap<RegionId, RegionRole> {
        let inner = self.inner.read().unwrap();
        region_ids
            .iter()
            .filter_map(|region_id| {
                inner
                    .get(&(datanode_id, *region_id))
                    .map(|role| (*region_id, *role))
            })
            .collect()
    }

    /// Returns number of element in tracking set.
    pub fn len(&self) -> usize {
        let inner = self.inner.read().unwrap();
        inner.len()
    }

    /// Returns true if it's empty.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    #[cfg(test)]
    pub fn clear(&self) {
        let mut inner = self.inner.write().unwrap();
        inner.clear();
    }
}

#[cfg(test)]
mod tests {
    use std::collections::{HashMap, HashSet};

    use store_api::region_engine::RegionRole;
    use store_api::storage::RegionId;

    use crate::region_keeper::MemoryRegionKeeper;

    #[test]
    fn test_opening_region_keeper() {
        let keeper = MemoryRegionKeeper::new();

        let guard = keeper
            .register_with_role(1, RegionId::from_u64(1), RegionRole::Leader)
            .unwrap();
        assert!(
            keeper
                .register_with_role(1, RegionId::from_u64(1), RegionRole::Leader)
                .is_none()
        );
        let guard2 = keeper
            .register_with_role(1, RegionId::from_u64(2), RegionRole::Follower)
            .unwrap();

        let regions = HashSet::from([
            RegionId::from_u64(1),
            RegionId::from_u64(2),
            RegionId::from_u64(3),
        ]);
        let output = keeper.extract_operating_region_roles(1, &regions);
        assert_eq!(output.len(), 2);

        assert!(output.contains_key(&RegionId::from_u64(1)));
        assert!(output.contains_key(&RegionId::from_u64(2)));
        assert_eq!(keeper.len(), 2);

        let regions = HashSet::from([
            RegionId::from_u64(1),
            RegionId::from_u64(2),
            RegionId::from_u64(3),
        ]);
        let output = keeper.extract_operating_region_roles(1, &regions);
        assert_eq!(
            output,
            HashMap::from([
                (RegionId::from_u64(1), RegionRole::Leader),
                (RegionId::from_u64(2), RegionRole::Follower),
            ])
        );
        assert_eq!(keeper.len(), 2);

        drop(guard);
        assert_eq!(keeper.len(), 1);
        assert!(keeper.contains(1, RegionId::from_u64(2)));

        drop(guard2);
        assert!(keeper.is_empty());
    }
}
