// This file is part of Substrate.

// Copyright (C) Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: Apache-2.0

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Multi-block Migration framework for pallet-contracts.
//!
//! This module allows us to define a migration as a sequence of [`MigrationStep`]s that can be
//! executed across multiple blocks.
//!
//! # Usage
//!
//! A migration step is defined under `src/migration/vX.rs`, where `X` is the version number.
//! For example, `vX.rs` defines a migration from version `X - 1` to version `X`.
//!
//! ## Example:
//!
//! To configure a migration to `v11` for a runtime using `v10` of pallet-contracts on the chain,
//! you would set the `Migrations` type as follows:
//!
//! ```
//! use pallet_contracts::migration::{v10, v11};
//! # pub enum Runtime {};
//! # struct Currency;
//! type Migrations = (v10::Migration<Runtime, Currency>, v11::Migration<Runtime>);
//! ```
//!
//! ## Notes:
//!
//! - Migrations should always be tested with `try-runtime` before being deployed.
//! - By testing with `try-runtime` against a live network, you ensure that all migration steps work
//!   and that you have included the required steps.
//!
//! ## Low Level / Implementation Details
//!
//! When a migration starts and [`OnRuntimeUpgrade::on_runtime_upgrade`] is called, instead of
//! performing the actual migration, we set a custom storage item [`MigrationInProgress`].
//! This storage item defines a [`Cursor`] for the current migration.
//!
//! If the [`MigrationInProgress`] storage item exists, it means a migration is in progress, and its
//! value holds a cursor for the current migration step. These migration steps are executed during
//! [`Hooks<BlockNumber>::on_idle`] or when the [`Pallet::migrate`] dispatchable is
//! called.
//!
//! While the migration is in progress, all dispatchables except `migrate`, are blocked, and returns
//! a `MigrationInProgress` error.

pub mod v09;
pub mod v10;
pub mod v11;
pub mod v12;
pub mod v13;
pub mod v14;
pub mod v15;
pub mod v16;
include!(concat!(env!("OUT_DIR"), "/migration_codegen.rs"));

use crate::{weights::WeightInfo, Config, Error, MigrationInProgress, Pallet, Weight, LOG_TARGET};
use codec::{Codec, Decode};
use core::marker::PhantomData;
use frame_support::{
	pallet_prelude::*,
	traits::{ConstU32, OnRuntimeUpgrade},
	weights::WeightMeter,
};
use sp_runtime::Saturating;

#[cfg(feature = "try-runtime")]
use alloc::vec::Vec;
#[cfg(feature = "try-runtime")]
use sp_runtime::TryRuntimeError;

const PROOF_ENCODE: &str = "Tuple::max_encoded_len() < Cursor::max_encoded_len()` is verified in `Self::integrity_test()`; qed";
const PROOF_DECODE: &str =
	"We encode to the same type in this trait only. No other code touches this item; qed";

fn invalid_version(version: StorageVersion) -> ! {
	panic!("Required migration {version:?} not supported by this runtime. This is a bug.");
}

/// The cursor used to encode the position (usually the last iterated key) of the current migration
/// step.
pub type Cursor = BoundedVec<u8, ConstU32<1024>>;

/// IsFinished describes whether a migration is finished or not.
pub enum IsFinished {
	Yes,
	No,
}

/// A trait that allows to migrate storage from one version to another.
///
/// The migration is done in steps. The migration is finished when
/// `step()` returns `IsFinished::Yes`.
pub trait MigrationStep: Codec + MaxEncodedLen + Default {
	/// Returns the version of the migration.
	const VERSION: u16;

	/// Returns the maximum weight that can be consumed in a single step.
	fn max_step_weight() -> Weight;

	/// Process one step of the migration.
	///
	/// Returns whether the migration is finished.
	fn step(&mut self, meter: &mut WeightMeter) -> IsFinished;

	/// Verify that the migration step fits into `Cursor`, and that `max_step_weight` is not greater
	/// than `max_block_weight`.
	fn integrity_test(max_block_weight: Weight) {
		if Self::max_step_weight().any_gt(max_block_weight) {
			panic!(
				"Invalid max_step_weight for Migration {}. Value should be lower than {}",
				Self::VERSION,
				max_block_weight
			);
		}

		let len = <Self as MaxEncodedLen>::max_encoded_len();
		let max = Cursor::bound();
		if len > max {
			panic!(
				"Migration {} has size {} which is bigger than the maximum of {}",
				Self::VERSION,
				len,
				max,
			);
		}
	}

	/// Execute some pre-checks prior to running the first step of this migration.
	#[cfg(feature = "try-runtime")]
	fn pre_upgrade_step() -> Result<Vec<u8>, TryRuntimeError> {
		Ok(Vec::new())
	}

	/// Execute some post-checks after running the last step of this migration.
	#[cfg(feature = "try-runtime")]
	fn post_upgrade_step(_state: Vec<u8>) -> Result<(), TryRuntimeError> {
		Ok(())
	}
}

/// A noop migration that can be used when there is no migration to be done for a given version.
#[doc(hidden)]
#[derive(frame_support::DefaultNoBound, Encode, Decode, MaxEncodedLen)]
pub struct NoopMigration<const N: u16>;

impl<const N: u16> MigrationStep for NoopMigration<N> {
	const VERSION: u16 = N;
	fn max_step_weight() -> Weight {
		Weight::zero()
	}
	fn step(&mut self, _meter: &mut WeightMeter) -> IsFinished {
		log::debug!(target: LOG_TARGET, "Noop migration for version {}", N);
		IsFinished::Yes
	}
}

mod private {
	use crate::migration::MigrationStep;
	pub trait Sealed {}
	#[impl_trait_for_tuples::impl_for_tuples(10)]
	#[tuple_types_custom_trait_bound(MigrationStep)]
	impl Sealed for Tuple {}
}

/// Defines a sequence of migrations.
///
/// The sequence must be defined by a tuple of migrations, each of which must implement the
/// `MigrationStep` trait. Migrations must be ordered by their versions with no gaps.
pub trait MigrateSequence: private::Sealed {
	/// Returns the range of versions that this migrations sequence can handle.
	/// Migrations must be ordered by their versions with no gaps.
	///
	/// The following code will fail to compile:
	///
	/// ```compile_fail
	///     # use pallet_contracts::{NoopMigration, MigrateSequence};
	/// 	let _ = <(NoopMigration<1>, NoopMigration<3>)>::VERSION_RANGE;
	/// ```
	/// The following code will compile:
	/// ```
	///     # use pallet_contracts::{NoopMigration, MigrateSequence};
	/// 	let _ = <(NoopMigration<1>, NoopMigration<2>)>::VERSION_RANGE;
	/// ```
	const VERSION_RANGE: (u16, u16);

	/// Returns the default cursor for the given version.
	fn new(version: StorageVersion) -> Cursor;

	#[cfg(feature = "try-runtime")]
	fn pre_upgrade_step(_version: StorageVersion) -> Result<Vec<u8>, TryRuntimeError> {
		Ok(Vec::new())
	}

	#[cfg(feature = "try-runtime")]
	fn post_upgrade_step(_version: StorageVersion, _state: Vec<u8>) -> Result<(), TryRuntimeError> {
		Ok(())
	}

	/// Execute the migration step until the available weight is consumed.
	fn steps(version: StorageVersion, cursor: &[u8], meter: &mut WeightMeter) -> StepResult;

	/// Verify that the migration step fits into `Cursor`, and that `max_step_weight` is not greater
	/// than `max_block_weight`.
	fn integrity_test(max_block_weight: Weight);

	/// Returns whether migrating from `in_storage` to `target` is supported.
	///
	/// A migration is supported if `VERSION_RANGE` is (in_storage + 1, target).
	fn is_upgrade_supported(in_storage: StorageVersion, target: StorageVersion) -> bool {
		let (low, high) = Self::VERSION_RANGE;
		target == high && in_storage + 1 == low
	}
}

/// Performs all necessary migrations based on `StorageVersion`.
///
/// If `TEST_ALL_STEPS == true` and `try-runtime` is enabled, this will run all the migrations
/// inside `on_runtime_upgrade`. This should be set to false in tests that want to ensure the step
/// by step migration works.
pub struct Migration<T: Config, const TEST_ALL_STEPS: bool = true>(PhantomData<T>);

#[cfg(feature = "try-runtime")]
impl<T: Config, const TEST_ALL_STEPS: bool> Migration<T, TEST_ALL_STEPS> {
	fn run_all_steps() -> Result<(), TryRuntimeError> {
		let mut meter = &mut WeightMeter::new();
		let name = <Pallet<T>>::name();
		loop {
			let in_progress_version = <Pallet<T>>::on_chain_storage_version() + 1;
			let state = T::Migrations::pre_upgrade_step(in_progress_version)?;
			let before = meter.consumed();
			let status = Self::migrate(&mut meter);
			log::info!(
				target: LOG_TARGET,
				"{name}: Migration step {:?} weight = {}",
				in_progress_version,
				meter.consumed() - before
			);
			T::Migrations::post_upgrade_step(in_progress_version, state)?;
			if matches!(status, MigrateResult::Completed) {
				break
			}
		}

		let name = <Pallet<T>>::name();
		log::info!(target: LOG_TARGET, "{name}: Migration steps weight = {}", meter.consumed());
		Ok(())
	}
}

impl<T: Config, const TEST_ALL_STEPS: bool> OnRuntimeUpgrade for Migration<T, TEST_ALL_STEPS> {
	fn on_runtime_upgrade() -> Weight {
		let name = <Pallet<T>>::name();
		let in_code_version = <Pallet<T>>::in_code_storage_version();
		let on_chain_version = <Pallet<T>>::on_chain_storage_version();

		if on_chain_version == in_code_version {
			log::warn!(
				target: LOG_TARGET,
				"{name}: No Migration performed storage_version = latest_version = {:?}",
				&on_chain_version
			);
			return T::WeightInfo::on_runtime_upgrade_noop()
		}

		// In case a migration is already in progress we create the next migration
		// (if any) right when the current one finishes.
		if Self::in_progress() {
			log::warn!(
				target: LOG_TARGET,
				"{name}: Migration already in progress {:?}",
				&on_chain_version
			);

			return T::WeightInfo::on_runtime_upgrade_in_progress()
		}

		log::info!(
			target: LOG_TARGET,
			"{name}: Upgrading storage from {on_chain_version:?} to {in_code_version:?}.",
		);

		let cursor = T::Migrations::new(on_chain_version + 1);
		MigrationInProgress::<T>::set(Some(cursor));

		#[cfg(feature = "try-runtime")]
		if TEST_ALL_STEPS {
			Self::run_all_steps().unwrap();
		}

		T::WeightInfo::on_runtime_upgrade()
	}

	#[cfg(feature = "try-runtime")]
	fn pre_upgrade() -> Result<Vec<u8>, TryRuntimeError> {
		// We can't really do much here as our migrations do not happen during the runtime upgrade.
		// Instead, we call the migrations `pre_upgrade` and `post_upgrade` hooks when we iterate
		// over our migrations.
		let on_chain_version = <Pallet<T>>::on_chain_storage_version();
		let in_code_version = <Pallet<T>>::in_code_storage_version();

		if on_chain_version == in_code_version {
			return Ok(Default::default())
		}

		log::debug!(
			target: LOG_TARGET,
			"Requested migration of {} from {:?}(on-chain storage version) to {:?}(in-code storage version)",
			<Pallet<T>>::name(), on_chain_version, in_code_version
		);

		ensure!(
			T::Migrations::is_upgrade_supported(on_chain_version, in_code_version),
			"Unsupported upgrade: VERSION_RANGE should be (on-chain storage version + 1, in-code storage version)"
		);

		Ok(Default::default())
	}

	#[cfg(feature = "try-runtime")]
	fn post_upgrade(_state: Vec<u8>) -> Result<(), TryRuntimeError> {
		if !TEST_ALL_STEPS {
			return Ok(())
		}

		log::info!(target: LOG_TARGET, "=== POST UPGRADE CHECKS ===");

		// Ensure that the hashing algorithm is correct for each storage map.
		if let Some(hash) = crate::CodeInfoOf::<T>::iter_keys().next() {
			crate::CodeInfoOf::<T>::get(hash).expect("CodeInfo exists for hash; qed");
		}
		if let Some(hash) = crate::PristineCode::<T>::iter_keys().next() {
			crate::PristineCode::<T>::get(hash).expect("PristineCode exists for hash; qed");
		}
		if let Some(account_id) = crate::ContractInfoOf::<T>::iter_keys().next() {
			crate::ContractInfoOf::<T>::get(account_id)
				.expect("ContractInfo exists for account_id; qed");
		}
		if let Some(nonce) = crate::DeletionQueue::<T>::iter_keys().next() {
			crate::DeletionQueue::<T>::get(nonce).expect("DeletionQueue exists for nonce; qed");
		}

		Ok(())
	}
}

/// The result of running the migration.
#[derive(Debug, PartialEq)]
pub enum MigrateResult {
	/// No migration was performed
	NoMigrationPerformed,
	/// No migration currently in progress
	NoMigrationInProgress,
	/// A migration is in progress
	InProgress { steps_done: u32 },
	/// All migrations are completed
	Completed,
}

/// The result of running a migration step.
#[derive(Debug, PartialEq)]
pub enum StepResult {
	InProgress { cursor: Cursor, steps_done: u32 },
	Completed { steps_done: u32 },
}

impl<T: Config, const TEST_ALL_STEPS: bool> Migration<T, TEST_ALL_STEPS> {
	/// Verify that each migration's step of the [`Config::Migrations`] sequence fits into
	/// `Cursor`.
	pub(crate) fn integrity_test() {
		let max_weight = <T as frame_system::Config>::BlockWeights::get().max_block;
		T::Migrations::integrity_test(max_weight)
	}

	/// Execute the multi-step migration.
	/// Returns whether or not a migration is in progress
	pub(crate) fn migrate(mut meter: &mut WeightMeter) -> MigrateResult {
		let name = <Pallet<T>>::name();

		if meter.try_consume(T::WeightInfo::migrate()).is_err() {
			return MigrateResult::NoMigrationPerformed
		}

		MigrationInProgress::<T>::mutate_exists(|progress| {
			let Some(cursor_before) = progress.as_mut() else {
				meter.consume(T::WeightInfo::migration_noop());
				return MigrateResult::NoMigrationInProgress
			};

			// if a migration is running it is always upgrading to the next version
			let storage_version = <Pallet<T>>::on_chain_storage_version();
			let in_progress_version = storage_version + 1;

			log::info!(
				target: LOG_TARGET,
				"{name}: Migrating from {:?} to {:?},",
				storage_version,
				in_progress_version,
			);

			let result =
				match T::Migrations::steps(in_progress_version, cursor_before.as_ref(), &mut meter)
				{
					StepResult::InProgress { cursor, steps_done } => {
						*progress = Some(cursor);
						MigrateResult::InProgress { steps_done }
					},
					StepResult::Completed { steps_done } => {
						in_progress_version.put::<Pallet<T>>();
						if <Pallet<T>>::in_code_storage_version() != in_progress_version {
							log::info!(
								target: LOG_TARGET,
								"{name}: Next migration is {:?},",
								in_progress_version + 1
							);
							*progress = Some(T::Migrations::new(in_progress_version + 1));
							MigrateResult::InProgress { steps_done }
						} else {
							log::info!(
								target: LOG_TARGET,
								"{name}: All migrations done. At version {:?},",
								in_progress_version
							);
							*progress = None;
							MigrateResult::Completed
						}
					},
				};

			result
		})
	}

	pub(crate) fn ensure_migrated() -> DispatchResult {
		if Self::in_progress() {
			Err(Error::<T>::MigrationInProgress.into())
		} else {
			Ok(())
		}
	}

	pub(crate) fn in_progress() -> bool {
		MigrationInProgress::<T>::exists()
	}
}

#[impl_trait_for_tuples::impl_for_tuples(10)]
#[tuple_types_custom_trait_bound(MigrationStep)]
impl MigrateSequence for Tuple {
	const VERSION_RANGE: (u16, u16) = {
		let mut versions: (u16, u16) = (0, 0);
		for_tuples!(
			#(
				match versions {
					(0, 0) => {
						versions = (Tuple::VERSION, Tuple::VERSION);
					},
					(min_version, last_version) if Tuple::VERSION == last_version + 1 => {
						versions = (min_version, Tuple::VERSION);
					},
					_ => panic!("Migrations must be ordered by their versions with no gaps.")
				}
			)*
		);
		versions
	};

	fn new(version: StorageVersion) -> Cursor {
		for_tuples!(
			#(
				if version == Tuple::VERSION {
					return Tuple::default().encode().try_into().expect(PROOF_ENCODE)
				}
			)*
		);
		invalid_version(version)
	}

	#[cfg(feature = "try-runtime")]
	/// Execute the pre-checks of the step associated with this version.
	fn pre_upgrade_step(version: StorageVersion) -> Result<Vec<u8>, TryRuntimeError> {
		for_tuples!(
			#(
				if version == Tuple::VERSION {
					return Tuple::pre_upgrade_step()
				}
			)*
		);
		invalid_version(version)
	}

	#[cfg(feature = "try-runtime")]
	/// Execute the post-checks of the step associated with this version.
	fn post_upgrade_step(version: StorageVersion, state: Vec<u8>) -> Result<(), TryRuntimeError> {
		for_tuples!(
			#(
				if version == Tuple::VERSION {
					return Tuple::post_upgrade_step(state)
				}
			)*
		);
		invalid_version(version)
	}

	fn steps(version: StorageVersion, mut cursor: &[u8], meter: &mut WeightMeter) -> StepResult {
		for_tuples!(
			#(
				if version == Tuple::VERSION {
					let mut migration = <Tuple as Decode>::decode(&mut cursor)
						.expect(PROOF_DECODE);
					let max_weight = Tuple::max_step_weight();
					let mut steps_done = 0;
					while meter.can_consume(max_weight) {
						steps_done.saturating_accrue(1);
						if matches!(migration.step(meter), IsFinished::Yes) {
							return StepResult::Completed{ steps_done }
						}
					}
					return StepResult::InProgress{cursor: migration.encode().try_into().expect(PROOF_ENCODE), steps_done }
				}
			)*
		);
		invalid_version(version)
	}

	fn integrity_test(max_block_weight: Weight) {
		for_tuples!(
			#(
				Tuple::integrity_test(max_block_weight);
			)*
		);
	}
}

#[cfg(test)]
mod test {
	use super::*;
	use crate::{
		migration::codegen::LATEST_MIGRATION_VERSION,
		tests::{ExtBuilder, Test},
	};

	#[derive(Default, Encode, Decode, MaxEncodedLen)]
	struct MockMigration<const N: u16> {
		// MockMigration<N> needs `N` steps to finish
		count: u16,
	}

	impl<const N: u16> MigrationStep for MockMigration<N> {
		const VERSION: u16 = N;
		fn max_step_weight() -> Weight {
			Weight::from_all(1)
		}
		fn step(&mut self, meter: &mut WeightMeter) -> IsFinished {
			assert!(self.count != N);
			self.count += 1;
			meter.consume(Weight::from_all(1));
			if self.count == N {
				IsFinished::Yes
			} else {
				IsFinished::No
			}
		}
	}

	#[test]
	fn test_storage_version_matches_last_migration_file() {
		assert_eq!(StorageVersion::new(LATEST_MIGRATION_VERSION), crate::pallet::STORAGE_VERSION);
	}

	#[test]
	fn version_range_works() {
		let range = <(MockMigration<1>, MockMigration<2>)>::VERSION_RANGE;
		assert_eq!(range, (1, 2));
	}

	#[test]
	fn is_upgrade_supported_works() {
		type Migrations = (MockMigration<9>, MockMigration<10>, MockMigration<11>);
		assert!(Migrations::is_upgrade_supported(StorageVersion::new(8), StorageVersion::new(11)));
		assert!(!Migrations::is_upgrade_supported(StorageVersion::new(9), StorageVersion::new(11)));
		assert!(!Migrations::is_upgrade_supported(StorageVersion::new(8), StorageVersion::new(12)));
	}

	#[test]
	fn steps_works() {
		type Migrations = (MockMigration<2>, MockMigration<3>);
		let version = StorageVersion::new(2);
		let mut cursor = Migrations::new(version);

		let mut meter = WeightMeter::with_limit(Weight::from_all(1));
		let result = Migrations::steps(version, &cursor, &mut meter);
		cursor = alloc::vec![1u8, 0].try_into().unwrap();
		assert_eq!(result, StepResult::InProgress { cursor: cursor.clone(), steps_done: 1 });
		assert_eq!(meter.consumed(), Weight::from_all(1));

		let mut meter = WeightMeter::with_limit(Weight::from_all(1));
		assert_eq!(
			Migrations::steps(version, &cursor, &mut meter),
			StepResult::Completed { steps_done: 1 }
		);
	}

	#[test]
	fn no_migration_in_progress_works() {
		type TestMigration = Migration<Test>;

		ExtBuilder::default().build().execute_with(|| {
			assert_eq!(StorageVersion::get::<Pallet<Test>>(), LATEST_MIGRATION_VERSION);
			assert_eq!(
				TestMigration::migrate(&mut WeightMeter::new()),
				MigrateResult::NoMigrationInProgress
			)
		});
	}

	#[test]
	fn migration_works() {
		type TestMigration = Migration<Test, false>;

		ExtBuilder::default()
			.set_storage_version(LATEST_MIGRATION_VERSION - 2)
			.build()
			.execute_with(|| {
				assert_eq!(StorageVersion::get::<Pallet<Test>>(), LATEST_MIGRATION_VERSION - 2);
				TestMigration::on_runtime_upgrade();
				for (version, status) in [
					(LATEST_MIGRATION_VERSION - 1, MigrateResult::InProgress { steps_done: 1 }),
					(LATEST_MIGRATION_VERSION, MigrateResult::Completed),
				] {
					assert_eq!(TestMigration::migrate(&mut WeightMeter::new()), status);
					assert_eq!(
						<Pallet<Test>>::on_chain_storage_version(),
						StorageVersion::new(version)
					);
				}

				assert_eq!(
					TestMigration::migrate(&mut WeightMeter::new()),
					MigrateResult::NoMigrationInProgress
				);
				assert_eq!(StorageVersion::get::<Pallet<Test>>(), LATEST_MIGRATION_VERSION);
			});
	}
}