// This file is part of Substrate.

// Copyright (C) 2018-2020 Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

use std::{
	collections::HashMap,
	sync::{Arc, atomic::{AtomicIsize, Ordering as AtomicOrdering}},
};
use parking_lot::{RwLock, RwLockWriteGuard, RwLockReadGuard};

/// Something that can report it's size.
pub trait Size {
	fn size(&self) -> usize;
}

/// Map with size tracking.
///
/// Size reported might be slightly off and only approximately true.
#[derive(Debug, parity_util_mem::MallocSizeOf)]
pub struct TrackedMap<K, V> {
	index: Arc<RwLock<HashMap<K, V>>>,
	bytes: AtomicIsize,
	length: AtomicIsize,
}

impl<K, V> Default for TrackedMap<K, V> {
	fn default() -> Self {
		Self {
			index: Arc::new(HashMap::default().into()),
			bytes: 0.into(),
			length: 0.into(),
		}
	}
}

impl<K, V> TrackedMap<K, V> {
	/// Current tracked length of the content.
	pub fn len(&self) -> usize {
		std::cmp::max(self.length.load(AtomicOrdering::Relaxed), 0) as usize
	}

	/// Current sum of content length.
	pub fn bytes(&self) -> usize {
		std::cmp::max(self.bytes.load(AtomicOrdering::Relaxed), 0) as usize
	}

	/// Read-only clone of the interior.
	pub fn clone(&self) -> ReadOnlyTrackedMap<K, V> {
		ReadOnlyTrackedMap(self.index.clone())
	}

	/// Lock map for read.
	pub fn read<'a>(&'a self) -> TrackedMapReadAccess<'a, K, V> {
		TrackedMapReadAccess {
			inner_guard: self.index.read(),
		}
	}

	/// Lock map for write.
	pub fn write<'a>(&'a self) -> TrackedMapWriteAccess<'a, K, V> {
		TrackedMapWriteAccess {
			inner_guard: self.index.write(),
			bytes: &self.bytes,
			length: &self.length,
		}
	}
}

/// Read-only access to map.
///
/// The only thing can be done is .read().
pub struct ReadOnlyTrackedMap<K, V>(Arc<RwLock<HashMap<K, V>>>);

impl<K, V> ReadOnlyTrackedMap<K, V>
where
	K: Eq + std::hash::Hash
{
	/// Lock map for read.
	pub fn read<'a>(&'a self) -> TrackedMapReadAccess<'a, K, V> {
		TrackedMapReadAccess {
			inner_guard: self.0.read(),
		}
	}
}

pub struct TrackedMapReadAccess<'a, K, V> {
	inner_guard: RwLockReadGuard<'a, HashMap<K, V>>,
}

impl<'a, K, V> TrackedMapReadAccess<'a, K, V>
where
	K: Eq + std::hash::Hash
{
	/// Returns true if map contains key.
	pub fn contains_key(&self, key: &K) -> bool {
		self.inner_guard.contains_key(key)
	}

	/// Returns reference to the contained value by key, if exists.
	pub fn get(&self, key: &K) -> Option<&V> {
		self.inner_guard.get(key)
	}

	/// Returns iterator over all values.
	pub fn values(&self) -> std::collections::hash_map::Values<K, V> {
		self.inner_guard.values()
	}
}

pub struct TrackedMapWriteAccess<'a, K, V> {
	bytes: &'a AtomicIsize,
	length: &'a AtomicIsize,
	inner_guard: RwLockWriteGuard<'a, HashMap<K, V>>,
}

impl<'a, K, V> TrackedMapWriteAccess<'a, K, V>
where
	K: Eq + std::hash::Hash, V: Size
{
	/// Insert value and return previous (if any).
	pub fn insert(&mut self, key: K, val: V) -> Option<V> {
		let new_bytes = val.size();
		self.bytes.fetch_add(new_bytes as isize, AtomicOrdering::Relaxed);
		self.length.fetch_add(1, AtomicOrdering::Relaxed);
		self.inner_guard.insert(key, val).and_then(|old_val| {
			self.bytes.fetch_sub(old_val.size() as isize, AtomicOrdering::Relaxed);
			self.length.fetch_sub(1, AtomicOrdering::Relaxed);
			Some(old_val)
		})
	}

	/// Remove value by key.
	pub fn remove(&mut self, key: &K) -> Option<V> {
		let val = self.inner_guard.remove(key);
		if let Some(size) = val.as_ref().map(Size::size) {
			self.bytes.fetch_sub(size as isize, AtomicOrdering::Relaxed);
			self.length.fetch_sub(1, AtomicOrdering::Relaxed);
		}
		val
	}

	/// Returns mutable reference to the contained value by key, if exists.
	pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
		self.inner_guard.get_mut(key)
	}
}

#[cfg(test)]
mod tests {

	use super::*;

	impl Size for i32 {
		fn size(&self) -> usize { *self as usize / 10 }
	}

	#[test]
	fn basic() {
		let map = TrackedMap::default();
		map.write().insert(5, 10);
		map.write().insert(6, 20);

		assert_eq!(map.bytes(), 3);
		assert_eq!(map.len(), 2);

		map.write().insert(6, 30);

		assert_eq!(map.bytes(), 4);
		assert_eq!(map.len(), 2);

		map.write().remove(&6);
		assert_eq!(map.bytes(), 1);
		assert_eq!(map.len(), 1);
	}
}