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

//! This module is responsible for maintaining a consistent initialization order for all other
//! teyrchains modules. It's also responsible for finalization and session change notifications.
//!
//! This module can throw fatal errors if session-change notifications are received after
//! initialization.

use crate::{
	configuration::{self, HostConfiguration},
	disputes::{self, DisputesHandler as _, SlashingHandler as _},
	dmp, hrmp, inclusion, paras, scheduler, session_info, shared,
};
use alloc::vec::Vec;
use codec::{Decode, Encode};
use pezframe_support::{
	traits::{OneSessionHandler, Randomness},
	weights::Weight,
};
use pezframe_system::limits::BlockWeights;
use pezkuwi_primitives::{BlockNumber, ConsensusLog, SessionIndex, ValidatorId};
use scale_info::TypeInfo;

#[cfg(test)]
mod tests;

#[cfg(feature = "runtime-benchmarks")]
mod benchmarking;

pub use pezpallet::*;

/// Information about a session change that has just occurred.
#[derive(Clone)]
pub struct SessionChangeNotification<BlockNumber> {
	/// The new validators in the session.
	pub validators: Vec<ValidatorId>,
	/// The queued validators for the following session.
	pub queued: Vec<ValidatorId>,
	/// The configuration before handling the session change
	pub prev_config: HostConfiguration<BlockNumber>,
	/// The configuration after handling the session change.
	pub new_config: HostConfiguration<BlockNumber>,
	/// A secure random seed for the session, gathered from BABE.
	pub random_seed: [u8; 32],
	/// New session index.
	pub session_index: SessionIndex,
}

/// Inform something about a new session.
pub trait OnNewSession<N> {
	/// A new session was started.
	fn on_new_session(notification: &SessionChangeNotification<N>);
}

impl<N> OnNewSession<N> for () {
	fn on_new_session(_: &SessionChangeNotification<N>) {}
}

/// Number of validators (not only teyrchain) in a session.
pub type ValidatorSetCount = u32;

impl<BlockNumber: Default + From<u32>> Default for SessionChangeNotification<BlockNumber> {
	fn default() -> Self {
		Self {
			validators: Vec::new(),
			queued: Vec::new(),
			prev_config: HostConfiguration::default(),
			new_config: HostConfiguration::default(),
			random_seed: Default::default(),
			session_index: Default::default(),
		}
	}
}

#[derive(Encode, Decode, TypeInfo)]
pub(crate) struct BufferedSessionChange {
	pub validators: Vec<ValidatorId>,
	pub queued: Vec<ValidatorId>,
	pub session_index: SessionIndex,
}

pub trait WeightInfo {
	fn force_approve(d: u32) -> Weight;
}

impl WeightInfo for () {
	fn force_approve(_: u32) -> Weight {
		BlockWeights::default().max_block
	}
}

#[pezframe_support::pezpallet]
pub mod pezpallet {
	use super::*;
	use pezframe_support::pezpallet_prelude::*;
	use pezframe_system::pezpallet_prelude::*;

	#[pezpallet::pezpallet]
	#[pezpallet::without_storage_info]
	pub struct Pezpallet<T>(_);

	#[pezpallet::config]
	pub trait Config:
		pezframe_system::Config
		+ configuration::Config
		+ shared::Config
		+ paras::Config
		+ scheduler::Config
		+ inclusion::Config
		+ session_info::Config
		+ disputes::Config
		+ dmp::Config
		+ hrmp::Config
	{
		/// A randomness beacon.
		type Randomness: Randomness<Self::Hash, BlockNumberFor<Self>>;
		/// An origin which is allowed to force updates to teyrchains.
		type ForceOrigin: EnsureOrigin<<Self as pezframe_system::Config>::RuntimeOrigin>;
		/// Temporary hack to call `Coretime::on_new_session` on chains that support `Coretime` or
		/// to disable it on the ones that don't support it. Can be removed and replaced by a simple
		/// bound to `coretime::Config` once all chains support it.
		type CoretimeOnNewSession: OnNewSession<BlockNumberFor<Self>>;
		/// Weight information for extrinsics in this pezpallet.
		type WeightInfo: WeightInfo;
	}

	/// Whether the teyrchains modules have been initialized within this block.
	///
	/// Semantically a `bool`, but this guarantees it should never hit the trie,
	/// as this is cleared in `on_finalize` and Frame optimizes `None` values to be empty values.
	///
	/// As a `bool`, `set(false)` and `remove()` both lead to the next `get()` being false, but one
	/// of them writes to the trie and one does not. This confusion makes `Option<()>` more suitable
	/// for the semantics of this variable.
	#[pezpallet::storage]
	pub(super) type HasInitialized<T: Config> = StorageValue<_, ()>;

	/// Buffered session changes.
	///
	/// Typically this will be empty or one element long. Apart from that this item never hits
	/// the storage.
	///
	/// However this is a `Vec` regardless to handle various edge cases that may occur at runtime
	/// upgrade boundaries or if governance intervenes.
	#[pezpallet::storage]
	pub(crate) type BufferedSessionChanges<T: Config> =
		StorageValue<_, Vec<BufferedSessionChange>, ValueQuery>;

	#[pezpallet::hooks]
	impl<T: Config> Hooks<BlockNumberFor<T>> for Pezpallet<T> {
		fn on_initialize(now: BlockNumberFor<T>) -> Weight {
			// The other modules are initialized in this order:
			// - Configuration
			// - Paras
			// - Scheduler
			// - Inclusion
			// - `SessionInfo`
			// - Disputes
			// - DMP
			// - UMP
			// - HRMP
			let total_weight = configuration::Pezpallet::<T>::initializer_initialize(now)
				+ shared::Pezpallet::<T>::initializer_initialize(now)
				+ paras::Pezpallet::<T>::initializer_initialize(now)
				+ scheduler::Pezpallet::<T>::initializer_initialize(now)
				+ inclusion::Pezpallet::<T>::initializer_initialize(now)
				+ session_info::Pezpallet::<T>::initializer_initialize(now)
				+ T::DisputesHandler::initializer_initialize(now)
				+ T::SlashingHandler::initializer_initialize(now)
				+ dmp::Pezpallet::<T>::initializer_initialize(now)
				+ hrmp::Pezpallet::<T>::initializer_initialize(now);

			HasInitialized::<T>::set(Some(()));

			total_weight
		}

		fn on_finalize(now: BlockNumberFor<T>) {
			// reverse initialization order.
			hrmp::Pezpallet::<T>::initializer_finalize();
			dmp::Pezpallet::<T>::initializer_finalize();
			T::SlashingHandler::initializer_finalize();
			T::DisputesHandler::initializer_finalize();
			session_info::Pezpallet::<T>::initializer_finalize();
			inclusion::Pezpallet::<T>::initializer_finalize();
			scheduler::Pezpallet::<T>::initializer_finalize();
			paras::Pezpallet::<T>::initializer_finalize(now);
			shared::Pezpallet::<T>::initializer_finalize();
			configuration::Pezpallet::<T>::initializer_finalize();

			// Apply buffered session changes as the last thing. This way the runtime APIs and the
			// next block will observe the next session.
			//
			// Note that we only apply the last session as all others lasted less than a block
			// (weirdly).
			if let Some(BufferedSessionChange { session_index, validators, queued }) =
				BufferedSessionChanges::<T>::take().pop()
			{
				Self::apply_new_session(session_index, validators, queued);
			}

			HasInitialized::<T>::take();
		}
	}

	#[pezpallet::call]
	impl<T: Config> Pezpallet<T> {
		/// Issue a signal to the consensus engine to forcibly act as though all teyrchain
		/// blocks in all relay chain blocks up to and including the given number in the current
		/// chain are valid and should be finalized.
		#[pezpallet::call_index(0)]
		#[pezpallet::weight((
			<T as Config>::WeightInfo::force_approve(
				pezframe_system::Pezpallet::<T>::digest().logs.len() as u32,
			),
			DispatchClass::Operational,
		))]
		pub fn force_approve(origin: OriginFor<T>, up_to: BlockNumber) -> DispatchResult {
			T::ForceOrigin::ensure_origin(origin)?;

			pezframe_system::Pezpallet::<T>::deposit_log(ConsensusLog::ForceApprove(up_to).into());
			Ok(())
		}
	}
}

impl<T: Config> Pezpallet<T> {
	fn apply_new_session(
		session_index: SessionIndex,
		all_validators: Vec<ValidatorId>,
		queued: Vec<ValidatorId>,
	) {
		let random_seed = {
			let mut buf = [0u8; 32];
			// TODO: audit usage of randomness API
			// https://github.com/pezkuwichain/pezkuwi-sdk/issues/139
			let (random_hash, _) = T::Randomness::random(&b"paras"[..]);
			let len = core::cmp::min(32, random_hash.as_ref().len());
			buf[..len].copy_from_slice(&random_hash.as_ref()[..len]);
			buf
		};

		let configuration::SessionChangeOutcome { prev_config, new_config } =
			configuration::Pezpallet::<T>::initializer_on_new_session(&session_index);
		let new_config = new_config.unwrap_or_else(|| prev_config.clone());

		let validators = shared::Pezpallet::<T>::initializer_on_new_session(
			session_index,
			random_seed,
			&new_config,
			all_validators,
		);

		let notification = SessionChangeNotification {
			validators,
			queued,
			prev_config,
			new_config,
			random_seed,
			session_index,
		};

		let outgoing_paras = paras::Pezpallet::<T>::initializer_on_new_session(&notification);
		scheduler::Pezpallet::<T>::initializer_on_new_session(&notification);
		inclusion::Pezpallet::<T>::initializer_on_new_session(&notification, &outgoing_paras);
		session_info::Pezpallet::<T>::initializer_on_new_session(&notification);
		T::DisputesHandler::initializer_on_new_session(&notification);
		T::SlashingHandler::initializer_on_new_session(session_index);
		dmp::Pezpallet::<T>::initializer_on_new_session(&notification, &outgoing_paras);
		hrmp::Pezpallet::<T>::initializer_on_new_session(&notification, &outgoing_paras);
		T::CoretimeOnNewSession::on_new_session(&notification);
	}

	/// Should be called when a new session occurs. Buffers the session notification to be applied
	/// at the end of the block. If `queued` is `None`, the `validators` are considered queued.
	fn on_new_session<'a, I: 'a>(
		_changed: bool,
		session_index: SessionIndex,
		validators: I,
		queued: Option<I>,
	) where
		I: Iterator<Item = (&'a T::AccountId, ValidatorId)>,
	{
		let validators: Vec<_> = validators.map(|(_, v)| v).collect();
		let queued: Vec<_> = if let Some(queued) = queued {
			queued.map(|(_, v)| v).collect()
		} else {
			validators.clone()
		};

		if session_index == 0 {
			// Genesis session should be immediately enacted.
			Self::apply_new_session(0, validators, queued);
		} else {
			BufferedSessionChanges::<T>::mutate(|v| {
				v.push(BufferedSessionChange { validators, queued, session_index })
			});
		}
	}

	// Allow to trigger `on_new_session` in tests, this is needed as long as `pezpallet_session` is
	// not implemented in mock.
	#[cfg(any(test, feature = "runtime-benchmarks"))]
	pub(crate) fn test_trigger_on_new_session<'a, I: 'a>(
		changed: bool,
		session_index: SessionIndex,
		validators: I,
		queued: Option<I>,
	) where
		I: Iterator<Item = (&'a T::AccountId, ValidatorId)>,
	{
		Self::on_new_session(changed, session_index, validators, queued)
	}

	/// Return whether at the end of this block a new session will be initialized.
	pub(crate) fn upcoming_session_change() -> bool {
		!BufferedSessionChanges::<T>::get().is_empty()
	}
}

impl<T: Config> pezsp_runtime::BoundToRuntimeAppPublic for Pezpallet<T> {
	type Public = ValidatorId;
}

impl<T: pezpallet_session::Config + Config> OneSessionHandler<T::AccountId> for Pezpallet<T> {
	type Key = ValidatorId;

	fn on_genesis_session<'a, I: 'a>(validators: I)
	where
		I: Iterator<Item = (&'a T::AccountId, Self::Key)>,
	{
		Pezpallet::<T>::on_new_session(false, 0, validators, None);
	}

	fn on_new_session<'a, I: 'a>(changed: bool, validators: I, queued: I)
	where
		I: Iterator<Item = (&'a T::AccountId, Self::Key)>,
	{
		let session_index = pezpallet_session::Pezpallet::<T>::current_index();
		Pezpallet::<T>::on_new_session(changed, session_index, validators, Some(queued));
	}

	fn on_disabled(_i: u32) {}
}