// Copyright 2015, 2016 Ethcore (UK) Ltd.
// This file is part of Parity.

// Parity 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.

// Parity 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 Parity.  If not, see <http://www.gnu.org/licenses/>.

//! Transaction Stats

use std::collections::{HashSet, HashMap};
use util::{H256, H512};
use util::hash::H256FastMap;

type NodeId = H512;
type BlockNumber = u64;

#[derive(Debug, Default, PartialEq)]
pub struct Stats {
	first_seen: BlockNumber,
	propagated_to: HashMap<NodeId, usize>,
}

#[derive(Debug, Default)]
pub struct TransactionsStats {
	pending_transactions: H256FastMap<Stats>,
}

impl TransactionsStats {
	/// Increases number of propagations to given `enodeid`.
	pub fn propagated(&mut self, hash: H256, enode_id: Option<NodeId>) {
		let enode_id = enode_id.unwrap_or_default();
		let mut stats = self.pending_transactions.entry(hash).or_insert_with(|| Stats::default());
		let mut count = stats.propagated_to.entry(enode_id).or_insert(0);
		*count = count.saturating_add(1);
	}

	/// Returns propagation stats for given hash or `None` if hash is not known.
	pub fn stats(&self, hash: &H256) -> Option<&Stats> {
		self.pending_transactions.get(hash)
	}

	/// Retains only transactions present in given `HashSet`.
	pub fn retain(&mut self, hashes: &HashSet<H256>) {
		let to_remove = self.pending_transactions.keys()
			.filter(|hash| !hashes.contains(hash))
			.cloned()
			.collect::<Vec<_>>();

		for hash in to_remove {
			self.pending_transactions.remove(&hash);
		}
	}
}

#[cfg(test)]
mod tests {

	use std::collections::{HashMap, HashSet};
	use super::{Stats, TransactionsStats};

	#[test]
	fn should_keep_track_of_propagations() {
		// given
		let mut stats = TransactionsStats::default();
		let hash = 5.into();
		let enodeid1 = 2.into();
		let enodeid2 = 5.into();

		// when
		stats.propagated(hash, Some(enodeid1));
		stats.propagated(hash, Some(enodeid1));
		stats.propagated(hash, Some(enodeid2));

		// then
		let stats = stats.stats(&hash);
		assert_eq!(stats, Some(&Stats {
			first_seen: 0,
			propagated_to: hash_map![
				enodeid1 => 2,
				enodeid2 => 1
			]
		}));
	}

	#[test]
	fn should_remove_hash_from_tracking() {
		// given
		let mut stats = TransactionsStats::default();
		let hash = 5.into();
		let enodeid1 = 5.into();
		stats.propagated(hash, Some(enodeid1));

		// when
		stats.retain(&HashSet::new());

		// then
		let stats = stats.stats(&hash);
		assert_eq!(stats, None);
	}
}