// Copyright (C) Parity Technologies (UK) Ltd. and Dijital Kurdistan Tech Institute
// This file is part of Pezkuwi.

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

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

//! A utility for tracking groups and their members within a session.

use pezkuwi_primitives::{effective_minimum_backing_votes, GroupIndex, IndexedVec, ValidatorIndex};

use std::collections::HashMap;

/// Validator groups within a session, plus some helpful indexing for
/// looking up groups by validator indices or authority discovery ID.
#[derive(Debug, Clone)]
pub struct Groups {
	groups: IndexedVec<GroupIndex, Vec<ValidatorIndex>>,
	by_validator_index: HashMap<ValidatorIndex, GroupIndex>,
	backing_threshold: u32,
}

impl Groups {
	/// Create a new [`Groups`] tracker with the groups and discovery keys
	/// from the session.
	pub fn new(
		groups: IndexedVec<GroupIndex, Vec<ValidatorIndex>>,
		backing_threshold: u32,
	) -> Self {
		let mut by_validator_index = HashMap::new();

		for (i, group) in groups.iter().enumerate() {
			let index = GroupIndex(i as _);
			for v in group {
				by_validator_index.insert(*v, index);
			}
		}

		Groups { groups, by_validator_index, backing_threshold }
	}

	/// Access all the underlying groups.
	pub fn all(&self) -> &IndexedVec<GroupIndex, Vec<ValidatorIndex>> {
		&self.groups
	}

	/// Get the underlying group validators by group index.
	pub fn get(&self, group_index: GroupIndex) -> Option<&[ValidatorIndex]> {
		self.groups.get(group_index).map(|x| &x[..])
	}

	/// Get the backing group size and backing threshold.
	pub fn get_size_and_backing_threshold(
		&self,
		group_index: GroupIndex,
	) -> Option<(usize, usize)> {
		self.get(group_index)
			.map(|g| (g.len(), effective_minimum_backing_votes(g.len(), self.backing_threshold)))
	}

	/// Get the group index for a validator by index.
	pub fn by_validator_index(&self, validator_index: ValidatorIndex) -> Option<GroupIndex> {
		self.by_validator_index.get(&validator_index).map(|x| *x)
	}
}