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

/// `pezpallet-assets` has been enhanced with asset reserves information so that AH foreign assets
/// can be registered as either teleportable or reserve-based.
/// Originally, all foreign assets were exclusively teleportable, whereas now, after creation the
/// asset `Owner` also needs to set the asset's trusted reserves and teleport relationships.
///
/// This migration adds the origin chain of each existing foreign asset as a trusted reserve for it.
/// It marks Sibling Teyrchain foreign assets as teleportable between the local chain and the
/// sibling teyrchain.
/// It marks external ecosystems (bridged) foreign assets as non-teleportable between the local
/// chain and the external ecosystem.
/// For new assets, the reserve location(s) and teleport status need to be explicitly configured by
/// the asset's `Owner`.
///
/// See <https://github.com/pezkuwichain/pezkuwi-sdk/issues/129> for more info.
pub mod foreign_assets_reserves {
	use crate::*;
	use codec::{Decode, Encode, MaxEncodedLen};
	use core::{fmt::Debug, marker::PhantomData};
	use pezframe_support::{
		migrations::{MigrationId, SteppedMigration, SteppedMigrationError},
		weights::WeightMeter,
	};
	use pezpallet_assets::WeightInfo;

	const MIGRATIONS_ID: &[u8; 23] = b"foreign-assets-reserves";

	#[cfg(feature = "try-runtime")]
	#[derive(Encode, Decode)]
	struct TryRuntimeState<T: pezpallet_assets::Config<I>, I: 'static> {
		assets: Vec<T::AssetId>,
	}

	/// Progressive states of a migration. The migration starts with the first variant and ends with
	/// the last.
	#[derive(Decode, Encode, MaxEncodedLen, Eq, PartialEq)]
	pub enum MigrationState<A> {
		Asset(A),
		Finished,
	}

	/// Trait for plugging in the type that chooses the correct reserves information per asset_id
	/// for this migration.
	pub trait ForeignAssetsReservesProvider {
		type ReserveData: Debug;
		fn reserves_for(asset_id: &Location) -> Vec<Self::ReserveData>;
		#[cfg(feature = "try-runtime")]
		fn check_reserves_for(asset_id: &Location, reserves: Vec<Self::ReserveData>) -> bool;
	}

	/// The resulting state of the step and the actual weight consumed.
	type StepResultOf<T, I> = MigrationState<<T as pezpallet_assets::Config<I>>::AssetId>;

	/// This migration adds assets reserves information for Foreign Assets registered on Asset Hub.
	/// The migration can be used by multiple runtimes that need to populate reserves information
	/// for existing assets. The actual reserves information is provided by an external type
	/// `ReservesProvider` implementing the `ForeignAssetsReservesProvider` trait.
	pub struct ForeignAssetsReservesMigration<T, I, ReservesProvider>(
		PhantomData<(T, I, ReservesProvider)>,
	);
	impl<T, I, ReservesProvider> SteppedMigration
		for ForeignAssetsReservesMigration<T, I, ReservesProvider>
	where
		ReservesProvider: ForeignAssetsReservesProvider,
		T: pezpallet_assets::Config<
			I,
			AssetId = Location,
			ReserveData = ReservesProvider::ReserveData,
		>,
		I: 'static,
	{
		type Cursor = StepResultOf<T, I>;
		type Identifier = MigrationId<23>;

		fn id() -> Self::Identifier {
			// this migration doesn't change pezpallet storage version, from and to are both `1`
			MigrationId { pezpallet_id: *MIGRATIONS_ID, version_from: 1, version_to: 1 }
		}

		fn step(
			mut cursor: Option<Self::Cursor>,
			meter: &mut WeightMeter,
		) -> Result<Option<Self::Cursor>, SteppedMigrationError> {
			// Check that we have enough weight for at least the next step. If we don't, then the
			// migration cannot be complete.
			let required =
				<T as pezpallet_assets::Config<I>>::WeightInfo::migration_v2_foreign_asset_set_reserve_weight();
			tracing::debug!(target: "runtime::ForeignAssetsReservesMigration", ?meter, ?required);
			if !meter.can_consume(required) {
				return Err(SteppedMigrationError::InsufficientWeight { required });
			}

			loop {
				if !meter.can_consume(required) {
					break;
				}

				let next = match &cursor {
					// At first, start migrating assets.
					None => Self::asset_step(None),
					// Migrate any remaining assets.
					Some(MigrationState::Asset(maybe_last_asset)) =>
						Self::asset_step(Some(maybe_last_asset)),
					// After the last asset, migration is finished.
					Some(MigrationState::Finished) => {
						tracing::info!(target: "runtime::ForeignAssetsReservesMigration", "migration finished");
						return Ok(None);
					},
				};

				cursor = Some(next);
				meter.consume(required);
			}

			Ok(cursor)
		}

		#[cfg(feature = "try-runtime")]
		fn pre_upgrade() -> Result<Vec<u8>, pezsp_runtime::TryRuntimeError> {
			let assets = pezpallet_assets::Asset::<T, I>::iter_keys().collect();
			tracing::info!(target: "runtime::ForeignAssetsReservesMigration::pre_upgrade", ?assets);
			let state = TryRuntimeState::<T, I> { assets };
			Ok(state.encode())
		}

		#[cfg(feature = "try-runtime")]
		fn post_upgrade(state: Vec<u8>) -> Result<(), pezsp_runtime::TryRuntimeError> {
			let prev_state = TryRuntimeState::<T, I>::decode(&mut &state[..])
				.expect("Failed to decode the previous storage state");
			for id in prev_state.assets {
				let reserves = pezpallet_assets::Pezpallet::<T, I>::get_reserves_data(id.clone());
				tracing::info!(target: "runtime::ForeignAssetsReservesMigration::post_upgrade", ?id, ?reserves, "verify asset");
				assert!(ReservesProvider::check_reserves_for(&id, reserves));
			}
			Ok(())
		}
	}

	impl<T, I, ReservesProvider> ForeignAssetsReservesMigration<T, I, ReservesProvider>
	where
		ReservesProvider: ForeignAssetsReservesProvider,
		T: pezpallet_assets::Config<
			I,
			AssetId = Location,
			ReserveData = ReservesProvider::ReserveData,
		>,
		I: 'static,
	{
		fn asset_step(maybe_last_key: Option<&T::AssetId>) -> StepResultOf<T, I> {
			tracing::debug!(target: "runtime::ForeignAssetsReservesMigration::asset_step", ?maybe_last_key);
			let mut iter = if let Some(last_key) = maybe_last_key {
				pezpallet_assets::Asset::<T, I>::iter_keys_from(
					pezpallet_assets::Asset::<T, I>::hashed_key_for(last_key),
				)
			} else {
				pezpallet_assets::Asset::<T, I>::iter_keys()
			};
			if let Some(asset_id) = iter.next() {
				let reserves = ReservesProvider::reserves_for(&asset_id);
				tracing::info!(
					target: "runtime::ForeignAssetsReservesMigration::asset_step",
					?asset_id, ?reserves, "updating reserves for"
				);
				if let Err(e) = pezpallet_assets::Pezpallet::<T, I>::unchecked_update_reserves(
					asset_id.clone(),
					reserves,
				) {
					tracing::error!(
						target: "runtime::ForeignAssetsReservesMigration::asset_step",
						?e, ?asset_id, "failed migrating reserves for asset"
					);
				}
				MigrationState::Asset(asset_id)
			} else {
				MigrationState::Finished
			}
		}
	}
}