// Copyright (C) Parity Technologies (UK) Ltd.
// This file is part of Polkadot.

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

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

use crate::{
	approval::{ApprovalsOptions, BlockTestData, CandidateTestData},
	configuration::TestAuthorities,
};
use codec::{Decode, Encode};
use itertools::Itertools;
use polkadot_node_network_protocol::v3 as protocol_v3;
use polkadot_primitives::{CandidateIndex, Hash, ValidatorIndex};
use sc_network_types::PeerId;
use std::collections::{HashMap, HashSet};

#[derive(Debug, Clone, Encode, Decode, PartialEq, Eq)]
pub struct TestMessageInfo {
	/// The actual message
	pub msg: protocol_v3::ApprovalDistributionMessage,
	/// The list of peers that would sends this message in a real topology.
	/// It includes both the peers that would send the message because of the topology
	/// or because of randomly choosing so.
	pub sent_by: Vec<ValidatorIndex>,
	/// The tranche at which this message should be sent.
	pub tranche: u32,
	/// The block hash this message refers to.
	pub block_hash: Hash,
}

impl std::hash::Hash for TestMessageInfo {
	fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
		match &self.msg {
			protocol_v3::ApprovalDistributionMessage::Assignments(assignments) => {
				for (assignment, candidates) in assignments {
					(assignment.block_hash, assignment.validator).hash(state);
					candidates.hash(state);
				}
			},
			protocol_v3::ApprovalDistributionMessage::Approvals(approvals) => {
				for approval in approvals {
					(approval.block_hash, approval.validator).hash(state);
					approval.candidate_indices.hash(state);
				}
			},
		};
	}
}

#[derive(Debug, Clone, Encode, Decode, PartialEq, Eq)]
/// A list of messages that depend of each-other, approvals cover one of the assignments and
/// vice-versa.
pub struct MessagesBundle {
	pub assignments: Vec<TestMessageInfo>,
	pub approvals: Vec<TestMessageInfo>,
}

impl MessagesBundle {
	/// The tranche when this bundle can be sent correctly, so no assignments or approvals will be
	/// from the future.
	pub fn tranche_to_send(&self) -> u32 {
		self.assignments
			.iter()
			.chain(self.approvals.iter())
			.max_by(|a, b| a.tranche.cmp(&b.tranche))
			.unwrap()
			.tranche
	}

	/// The min tranche in the bundle.
	pub fn min_tranche(&self) -> u32 {
		self.assignments
			.iter()
			.chain(self.approvals.iter())
			.min_by(|a, b| a.tranche.cmp(&b.tranche))
			.unwrap()
			.tranche
	}

	/// Tells if the bundle is needed for sending.
	/// We either send it because we need more assignments and approvals to approve the candidates
	/// or because we configured the test to send messages until a given tranche.
	pub fn should_send(
		&self,
		candidates_test_data: &HashMap<(Hash, CandidateIndex), CandidateTestData>,
		options: &ApprovalsOptions,
	) -> bool {
		self.needed_for_approval(candidates_test_data) ||
			(!options.stop_when_approved &&
				self.min_tranche() <= options.last_considered_tranche)
	}

	/// Tells if the bundle is needed because we need more messages to approve the candidates.
	pub fn needed_for_approval(
		&self,
		candidates_test_data: &HashMap<(Hash, CandidateIndex), CandidateTestData>,
	) -> bool {
		self.assignments
			.iter()
			.any(|message| message.needed_for_approval(candidates_test_data))
	}

	/// Mark the assignments in the bundle as sent.
	pub fn record_sent_assignment(
		&self,
		candidates_test_data: &mut HashMap<(Hash, CandidateIndex), CandidateTestData>,
	) {
		self.assignments
			.iter()
			.for_each(|assignment| assignment.record_sent_assignment(candidates_test_data));
	}
}

impl TestMessageInfo {
	/// Tells if the message is an approval.
	fn is_approval(&self) -> bool {
		match self.msg {
			protocol_v3::ApprovalDistributionMessage::Assignments(_) => false,
			protocol_v3::ApprovalDistributionMessage::Approvals(_) => true,
		}
	}

	/// Records an approval.
	/// We use this to check after all messages have been processed that we didn't loose any
	/// message.
	pub fn record_vote(&self, state: &BlockTestData) {
		if self.is_approval() {
			match &self.msg {
				protocol_v3::ApprovalDistributionMessage::Assignments(_) => todo!(),
				protocol_v3::ApprovalDistributionMessage::Approvals(approvals) =>
					for approval in approvals {
						for candidate_index in approval.candidate_indices.iter_ones() {
							state
								.votes
								.get(approval.validator.0 as usize)
								.unwrap()
								.get(candidate_index)
								.unwrap()
								.store(true, std::sync::atomic::Ordering::SeqCst);
						}
					},
			}
		}
	}

	/// Mark the assignments in the message as sent.
	pub fn record_sent_assignment(
		&self,
		candidates_test_data: &mut HashMap<(Hash, CandidateIndex), CandidateTestData>,
	) {
		match &self.msg {
			protocol_v3::ApprovalDistributionMessage::Assignments(assignments) => {
				for (assignment, candidate_indices) in assignments {
					for candidate_index in candidate_indices.iter_ones() {
						let candidate_test_data = candidates_test_data
							.get_mut(&(assignment.block_hash, candidate_index as CandidateIndex))
							.unwrap();
						candidate_test_data.mark_sent_assignment(self.tranche)
					}
				}
			},
			protocol_v3::ApprovalDistributionMessage::Approvals(_approvals) => todo!(),
		}
	}

	/// Returns a list of candidates indices in this message
	pub fn candidate_indices(&self) -> HashSet<usize> {
		let mut unique_candidate_indices = HashSet::new();
		match &self.msg {
			protocol_v3::ApprovalDistributionMessage::Assignments(assignments) =>
				for (_assignment, candidate_indices) in assignments {
					for candidate_index in candidate_indices.iter_ones() {
						unique_candidate_indices.insert(candidate_index);
					}
				},
			protocol_v3::ApprovalDistributionMessage::Approvals(approvals) =>
				for approval in approvals {
					for candidate_index in approval.candidate_indices.iter_ones() {
						unique_candidate_indices.insert(candidate_index);
					}
				},
		}
		unique_candidate_indices
	}

	/// Marks this message as no-shows if the number of configured no-shows is above the registered
	/// no-shows.
	/// Returns true if the message is a no-show.
	pub fn no_show_if_required(
		&self,
		assignments: &[TestMessageInfo],
		candidates_test_data: &mut HashMap<(Hash, CandidateIndex), CandidateTestData>,
	) -> bool {
		let mut should_no_show = false;
		if self.is_approval() {
			let covered_candidates = assignments
				.iter()
				.map(|assignment| (assignment, assignment.candidate_indices()))
				.collect_vec();

			match &self.msg {
				protocol_v3::ApprovalDistributionMessage::Assignments(_) => todo!(),
				protocol_v3::ApprovalDistributionMessage::Approvals(approvals) => {
					assert_eq!(approvals.len(), 1);

					for approval in approvals {
						should_no_show = should_no_show ||
							approval.candidate_indices.iter_ones().all(|candidate_index| {
								let candidate_test_data = candidates_test_data
									.get_mut(&(
										approval.block_hash,
										candidate_index as CandidateIndex,
									))
									.unwrap();
								let assignment = covered_candidates
									.iter()
									.find(|(_assignment, candidates)| {
										candidates.contains(&candidate_index)
									})
									.unwrap();
								candidate_test_data.should_no_show(assignment.0.tranche)
							});

						if should_no_show {
							for candidate_index in approval.candidate_indices.iter_ones() {
								let candidate_test_data = candidates_test_data
									.get_mut(&(
										approval.block_hash,
										candidate_index as CandidateIndex,
									))
									.unwrap();
								let assignment = covered_candidates
									.iter()
									.find(|(_assignment, candidates)| {
										candidates.contains(&candidate_index)
									})
									.unwrap();
								candidate_test_data.record_no_show(assignment.0.tranche)
							}
						}
					}
				},
			}
		}
		should_no_show
	}

	/// Tells if a message is needed for approval
	pub fn needed_for_approval(
		&self,
		candidates_test_data: &HashMap<(Hash, CandidateIndex), CandidateTestData>,
	) -> bool {
		match &self.msg {
			protocol_v3::ApprovalDistributionMessage::Assignments(assignments) =>
				assignments.iter().any(|(assignment, candidate_indices)| {
					candidate_indices.iter_ones().any(|candidate_index| {
						candidates_test_data
							.get(&(assignment.block_hash, candidate_index as CandidateIndex))
							.map(|data| data.should_send_tranche(self.tranche))
							.unwrap_or_default()
					})
				}),
			protocol_v3::ApprovalDistributionMessage::Approvals(approvals) =>
				approvals.iter().any(|approval| {
					approval.candidate_indices.iter_ones().any(|candidate_index| {
						candidates_test_data
							.get(&(approval.block_hash, candidate_index as CandidateIndex))
							.map(|data| data.should_send_tranche(self.tranche))
							.unwrap_or_default()
					})
				}),
		}
	}

	/// Splits a message into multiple messages based on what peers should send this message.
	/// It build a HashMap of messages that should be sent by each peer.
	pub fn split_by_peer_id(
		self,
		authorities: &TestAuthorities,
	) -> HashMap<(ValidatorIndex, PeerId), Vec<TestMessageInfo>> {
		let mut result: HashMap<(ValidatorIndex, PeerId), Vec<TestMessageInfo>> = HashMap::new();

		for validator_index in &self.sent_by {
			let peer = authorities.peer_ids.get(validator_index.0 as usize).unwrap();
			result.entry((*validator_index, *peer)).or_default().push(TestMessageInfo {
				msg: self.msg.clone(),
				sent_by: Default::default(),
				tranche: self.tranche,
				block_hash: self.block_hash,
			});
		}
		result
	}
}