// This file is part of Bizinikiwi.

// Copyright (C) Parity Technologies (UK) Ltd. and Dijital Kurdistan Tech Institute
// 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.

use super::*;
use crate::types::RegionRecord;
use codec::{Decode, Encode};
use core::marker::PhantomData;
use pezframe_support::traits::{Get, UncheckedOnRuntimeUpgrade};
use pezsp_runtime::Saturating;

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

mod v1 {
	use super::*;

	/// V0 region record.
	#[derive(Encode, Decode)]
	struct RegionRecordV0<AccountId, Balance> {
		/// The end of the Region.
		pub end: Timeslice,
		/// The owner of the Region.
		pub owner: AccountId,
		/// The amount paid to Pezkuwi for this Region, or `None` if renewal is not allowed.
		pub paid: Option<Balance>,
	}

	pub struct MigrateToV1Impl<T>(PhantomData<T>);

	impl<T: Config> UncheckedOnRuntimeUpgrade for MigrateToV1Impl<T> {
		fn on_runtime_upgrade() -> pezframe_support::weights::Weight {
			let mut count: u64 = 0;

			<Regions<T>>::translate::<RegionRecordV0<T::AccountId, BalanceOf<T>>, _>(|_, v0| {
				count.saturating_inc();
				Some(RegionRecord { end: v0.end, owner: Some(v0.owner), paid: v0.paid })
			});

			log::info!(
				target: LOG_TARGET,
				"Storage migration v1 for pezpallet-broker finished.",
			);

			// calculate and return migration weights
			T::DbWeight::get().reads_writes(count as u64 + 1, count as u64 + 1)
		}

		#[cfg(feature = "try-runtime")]
		fn pre_upgrade() -> Result<Vec<u8>, pezsp_runtime::TryRuntimeError> {
			Ok((Regions::<T>::iter_keys().count() as u32).encode())
		}

		#[cfg(feature = "try-runtime")]
		fn post_upgrade(state: Vec<u8>) -> Result<(), pezsp_runtime::TryRuntimeError> {
			let old_count = u32::decode(&mut &state[..]).expect("Known good");
			let new_count = Regions::<T>::iter_values().count() as u32;

			ensure!(old_count == new_count, "Regions count should not change");
			Ok(())
		}
	}
}

mod v2 {
	use super::*;
	use pezframe_support::{
		pezpallet_prelude::{OptionQuery, Twox64Concat},
		storage_alias,
	};

	#[storage_alias]
	pub type AllowedRenewals<T: Config> = StorageMap<
		Pezpallet<T>,
		Twox64Concat,
		PotentialRenewalId,
		PotentialRenewalRecordOf<T>,
		OptionQuery,
	>;

	pub struct MigrateToV2Impl<T>(PhantomData<T>);

	impl<T: Config> UncheckedOnRuntimeUpgrade for MigrateToV2Impl<T> {
		fn on_runtime_upgrade() -> pezframe_support::weights::Weight {
			let mut count = 0;
			for (renewal_id, renewal) in AllowedRenewals::<T>::drain() {
				PotentialRenewals::<T>::insert(renewal_id, renewal);
				count += 1;
			}

			log::info!(
				target: LOG_TARGET,
				"Storage migration v2 for pezpallet-broker finished.",
			);

			// calculate and return migration weights
			T::DbWeight::get().reads_writes(count as u64 + 1, count as u64 + 1)
		}

		#[cfg(feature = "try-runtime")]
		fn pre_upgrade() -> Result<Vec<u8>, pezsp_runtime::TryRuntimeError> {
			Ok((AllowedRenewals::<T>::iter_keys().count() as u32).encode())
		}

		#[cfg(feature = "try-runtime")]
		fn post_upgrade(state: Vec<u8>) -> Result<(), pezsp_runtime::TryRuntimeError> {
			let old_count = u32::decode(&mut &state[..]).expect("Known good");
			let new_count = PotentialRenewals::<T>::iter_values().count() as u32;

			ensure!(old_count == new_count, "Renewal count should not change");
			Ok(())
		}
	}
}

mod v3 {
	use super::*;
	use codec::MaxEncodedLen;
	use pezframe_support::{
		pezpallet_prelude::{OptionQuery, RuntimeDebug, TypeInfo},
		storage_alias,
	};
	use pezframe_system::Pezpallet as System;
	use pezsp_arithmetic::Perbill;

	pub struct MigrateToV3Impl<T>(PhantomData<T>);

	impl<T: Config> UncheckedOnRuntimeUpgrade for MigrateToV3Impl<T> {
		fn on_runtime_upgrade() -> pezframe_support::weights::Weight {
			let acc = Pezpallet::<T>::account_id();
			System::<T>::inc_providers(&acc);
			// calculate and return migration weights
			T::DbWeight::get().writes(1)
		}

		#[cfg(feature = "try-runtime")]
		fn pre_upgrade() -> Result<Vec<u8>, pezsp_runtime::TryRuntimeError> {
			Ok(System::<T>::providers(&Pezpallet::<T>::account_id()).encode())
		}

		#[cfg(feature = "try-runtime")]
		fn post_upgrade(state: Vec<u8>) -> Result<(), pezsp_runtime::TryRuntimeError> {
			let old_providers = u32::decode(&mut &state[..]).expect("Known good");
			let new_providers = System::<T>::providers(&Pezpallet::<T>::account_id()) as u32;

			ensure!(new_providers == old_providers + 1, "Providers count should increase by one");
			Ok(())
		}
	}

	#[storage_alias]
	pub type Configuration<T: Config> = StorageValue<Pezpallet<T>, ConfigRecordOf<T>, OptionQuery>;
	pub type ConfigRecordOf<T> =
		ConfigRecord<pezframe_system::pezpallet_prelude::BlockNumberFor<T>, RelayBlockNumberOf<T>>;

	// types added here for v4 migration
	#[derive(Encode, Decode, Clone, PartialEq, Eq, RuntimeDebug, TypeInfo, MaxEncodedLen)]
	pub struct ConfigRecord<BlockNumber, RelayBlockNumber> {
		/// The number of Relay-chain blocks in advance which scheduling should be fixed and the
		/// `Coretime::assign` API used to inform the Relay-chain.
		pub advance_notice: RelayBlockNumber,
		/// The length in blocks of the Interlude Period for forthcoming sales.
		pub interlude_length: BlockNumber,
		/// The length in blocks of the Leadin Period for forthcoming sales.
		pub leadin_length: BlockNumber,
		/// The length in timeslices of Regions which are up for sale in forthcoming sales.
		pub region_length: Timeslice,
		/// The proportion of cores available for sale which should be sold in order for the price
		/// to remain the same in the next sale.
		pub ideal_bulk_proportion: Perbill,
		/// An artificial limit to the number of cores which are allowed to be sold. If `Some` then
		/// no more cores will be sold than this.
		pub limit_cores_offered: Option<CoreIndex>,
		/// The amount by which the renewal price increases each sale period.
		pub renewal_bump: Perbill,
		/// The duration by which rewards for contributions to the InstaPool must be collected.
		pub contribution_timeout: Timeslice,
	}

	#[storage_alias]
	pub type SaleInfo<T: Config> = StorageValue<Pezpallet<T>, SaleInfoRecordOf<T>, OptionQuery>;
	pub type SaleInfoRecordOf<T> =
		SaleInfoRecord<BalanceOf<T>, pezframe_system::pezpallet_prelude::BlockNumberFor<T>>;

	/// The status of a Bulk Coretime Sale.
	#[derive(Encode, Decode, Clone, PartialEq, Eq, RuntimeDebug, TypeInfo, MaxEncodedLen)]
	pub struct SaleInfoRecord<Balance, BlockNumber> {
		/// The relay block number at which the sale will/did start.
		pub sale_start: BlockNumber,
		/// The length in relay chain blocks of the Leadin Period (where the price is decreasing).
		pub leadin_length: BlockNumber,
		/// The price of Bulk Coretime after the Leadin Period.
		pub price: Balance,
		/// The first timeslice of the Regions which are being sold in this sale.
		pub region_begin: Timeslice,
		/// The timeslice on which the Regions which are being sold in the sale terminate. (i.e.
		/// One after the last timeslice which the Regions control.)
		pub region_end: Timeslice,
		/// The number of cores we want to sell, ideally. Selling this amount would result in no
		/// change to the price for the next sale.
		pub ideal_cores_sold: CoreIndex,
		/// Number of cores which are/have been offered for sale.
		pub cores_offered: CoreIndex,
		/// The index of the first core which is for sale. Core of Regions which are sold have
		/// incrementing indices from this.
		pub first_core: CoreIndex,
		/// The latest price at which Bulk Coretime was purchased until surpassing the ideal number
		/// of cores were sold.
		pub sellout_price: Option<Balance>,
		/// Number of cores which have been sold; never more than cores_offered.
		pub cores_sold: CoreIndex,
	}
}

pub mod v4 {
	use super::*;

	type BlockNumberFor<T> = pezframe_system::pezpallet_prelude::BlockNumberFor<T>;

	pub trait BlockToRelayHeightConversion<T: Config> {
		/// Converts absolute value of teyrchain block number to relay chain block number
		fn convert_block_number_to_relay_height(
			block_number: BlockNumberFor<T>,
		) -> RelayBlockNumberOf<T>;

		/// Converts teyrchain block length into equivalent relay chain block length
		fn convert_block_length_to_relay_length(
			block_number: BlockNumberFor<T>,
		) -> RelayBlockNumberOf<T>;
	}

	pub struct MigrateToV4Impl<T, BlockConversion>(PhantomData<T>, PhantomData<BlockConversion>);
	impl<T: Config, BlockConversion: BlockToRelayHeightConversion<T>> UncheckedOnRuntimeUpgrade
		for MigrateToV4Impl<T, BlockConversion>
	{
		#[cfg(feature = "try-runtime")]
		fn pre_upgrade() -> Result<Vec<u8>, pezsp_runtime::TryRuntimeError> {
			let (interlude_length, configuration_leadin_length) =
				if let Some(config_record) = v3::Configuration::<T>::get() {
					(config_record.interlude_length, config_record.leadin_length)
				} else {
					((0 as u32).into(), (0 as u32).into())
				};

			let updated_interlude_length: RelayBlockNumberOf<T> =
				BlockConversion::convert_block_length_to_relay_length(interlude_length);
			let updated_leadin_length: RelayBlockNumberOf<T> =
				BlockConversion::convert_block_length_to_relay_length(configuration_leadin_length);
			log::info!(target: LOG_TARGET, "Configuration Pre-Migration: Interlude Length {:?}->{:?} Leadin Length {:?}->{:?}", interlude_length, updated_interlude_length, configuration_leadin_length, updated_leadin_length);

			let (sale_start, sale_info_leadin_length) =
				if let Some(sale_info_record) = v3::SaleInfo::<T>::get() {
					(sale_info_record.sale_start, sale_info_record.leadin_length)
				} else {
					((0 as u32).into(), (0 as u32).into())
				};

			let updated_sale_start: RelayBlockNumberOf<T> =
				BlockConversion::convert_block_number_to_relay_height(sale_start);
			let updated_sale_info_leadin_length: RelayBlockNumberOf<T> =
				BlockConversion::convert_block_length_to_relay_length(sale_info_leadin_length);
			log::info!(target: LOG_TARGET, "SaleInfo Pre-Migration: Sale Start {:?}->{:?} Interlude Length {:?}->{:?}", sale_start, updated_sale_start, sale_info_leadin_length, updated_sale_info_leadin_length);

			Ok((interlude_length, configuration_leadin_length, sale_start, sale_info_leadin_length)
				.encode())
		}

		fn on_runtime_upgrade() -> pezframe_support::weights::Weight {
			let mut weight = T::DbWeight::get().reads(1);

			if let Some(config_record) = v3::Configuration::<T>::take() {
				log::info!(target: LOG_TARGET, "migrating Configuration record");

				let updated_interlude_length: RelayBlockNumberOf<T> =
					BlockConversion::convert_block_length_to_relay_length(
						config_record.interlude_length,
					);
				let updated_leadin_length: RelayBlockNumberOf<T> =
					BlockConversion::convert_block_length_to_relay_length(
						config_record.leadin_length,
					);

				let updated_config_record = ConfigRecord {
					interlude_length: updated_interlude_length,
					leadin_length: updated_leadin_length,
					advance_notice: config_record.advance_notice,
					region_length: config_record.region_length,
					ideal_bulk_proportion: config_record.ideal_bulk_proportion,
					limit_cores_offered: config_record.limit_cores_offered,
					renewal_bump: config_record.renewal_bump,
					contribution_timeout: config_record.contribution_timeout,
				};
				Configuration::<T>::put(updated_config_record);
			}
			weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));

			if let Some(sale_info) = v3::SaleInfo::<T>::take() {
				log::info!(target: LOG_TARGET, "migrating SaleInfo record");

				let updated_sale_start: RelayBlockNumberOf<T> =
					BlockConversion::convert_block_number_to_relay_height(sale_info.sale_start);
				let updated_leadin_length: RelayBlockNumberOf<T> =
					BlockConversion::convert_block_length_to_relay_length(sale_info.leadin_length);

				let updated_sale_info = SaleInfoRecord {
					sale_start: updated_sale_start,
					leadin_length: updated_leadin_length,
					end_price: sale_info.price,
					region_begin: sale_info.region_begin,
					region_end: sale_info.region_end,
					ideal_cores_sold: sale_info.ideal_cores_sold,
					cores_offered: sale_info.cores_offered,
					first_core: sale_info.first_core,
					sellout_price: sale_info.sellout_price,
					cores_sold: sale_info.cores_sold,
				};
				SaleInfo::<T>::put(updated_sale_info);
			}

			weight.saturating_add(T::DbWeight::get().reads_writes(1, 2))
		}

		#[cfg(feature = "try-runtime")]
		fn post_upgrade(state: Vec<u8>) -> Result<(), pezsp_runtime::TryRuntimeError> {
			let (
				old_interlude_length,
				old_configuration_leadin_length,
				old_sale_start,
				old_sale_info_leadin_length,
			): (BlockNumberFor<T>, BlockNumberFor<T>, BlockNumberFor<T>, BlockNumberFor<T>) =
				Decode::decode(&mut &state[..]).expect("pre_upgrade provides a valid state; qed");

			if let Some(config_record) = Configuration::<T>::get() {
				ensure!(
					Self::verify_updated_block_length(
						old_configuration_leadin_length,
						config_record.leadin_length
					),
					"must migrate configuration leadin_length"
				);

				ensure!(
					Self::verify_updated_block_length(
						old_interlude_length,
						config_record.interlude_length
					),
					"must migrate configuration interlude_length"
				);
			}

			if let Some(sale_info) = SaleInfo::<T>::get() {
				ensure!(
					Self::verify_updated_block_time(old_sale_start, sale_info.sale_start),
					"must migrate sale info sale_start"
				);

				ensure!(
					Self::verify_updated_block_length(
						old_sale_info_leadin_length,
						sale_info.leadin_length
					),
					"must migrate sale info leadin_length"
				);
			}

			Ok(())
		}
	}

	#[cfg(feature = "try-runtime")]
	impl<T: Config, BlockConversion: BlockToRelayHeightConversion<T>>
		MigrateToV4Impl<T, BlockConversion>
	{
		fn verify_updated_block_time(
			old_value: BlockNumberFor<T>,
			new_value: RelayBlockNumberOf<T>,
		) -> bool {
			BlockConversion::convert_block_number_to_relay_height(old_value) == new_value
		}

		fn verify_updated_block_length(
			old_value: BlockNumberFor<T>,
			new_value: RelayBlockNumberOf<T>,
		) -> bool {
			BlockConversion::convert_block_length_to_relay_length(old_value) == new_value
		}
	}
}

/// Migrate the pezpallet storage from `0` to `1`.
pub type MigrateV0ToV1<T> = pezframe_support::migrations::VersionedMigration<
	0,
	1,
	v1::MigrateToV1Impl<T>,
	Pezpallet<T>,
	<T as pezframe_system::Config>::DbWeight,
>;

pub type MigrateV1ToV2<T> = pezframe_support::migrations::VersionedMigration<
	1,
	2,
	v2::MigrateToV2Impl<T>,
	Pezpallet<T>,
	<T as pezframe_system::Config>::DbWeight,
>;

pub type MigrateV2ToV3<T> = pezframe_support::migrations::VersionedMigration<
	2,
	3,
	v3::MigrateToV3Impl<T>,
	Pezpallet<T>,
	<T as pezframe_system::Config>::DbWeight,
>;

pub type MigrateV3ToV4<T, BlockConversion> = pezframe_support::migrations::VersionedMigration<
	3,
	4,
	v4::MigrateToV4Impl<T, BlockConversion>,
	Pezpallet<T>,
	<T as pezframe_system::Config>::DbWeight,
>;