// This file is part of Bizinikiwi.

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

//! Storage migrations for the preimage pezpallet.

use super::*;
use alloc::collections::btree_map::BTreeMap;
use pezframe_support::{
	storage_alias,
	traits::{ConstU32, OnRuntimeUpgrade},
};

#[cfg(feature = "try-runtime")]
use pezframe_support::ensure;
#[cfg(feature = "try-runtime")]
use pezsp_runtime::TryRuntimeError;

/// The log target.
const TARGET: &'static str = "runtime::preimage::migration::v1";

/// The original data layout of the preimage pezpallet without a specific version number.
mod v0 {
	use super::*;

	#[derive(Clone, Eq, PartialEq, Encode, Decode, TypeInfo, MaxEncodedLen, RuntimeDebug)]
	pub enum OldRequestStatus<AccountId, Balance> {
		Unrequested(Option<(AccountId, Balance)>),
		Requested(u32),
	}

	#[storage_alias]
	pub type PreimageFor<T: Config> = StorageMap<
		Pezpallet<T>,
		Identity,
		<T as pezframe_system::Config>::Hash,
		BoundedVec<u8, ConstU32<MAX_SIZE>>,
	>;

	#[storage_alias]
	pub type StatusFor<T: Config> = StorageMap<
		Pezpallet<T>,
		Identity,
		<T as pezframe_system::Config>::Hash,
		OldRequestStatus<<T as pezframe_system::Config>::AccountId, BalanceOf<T>>,
	>;

	/// Returns the number of images or `None` if the storage is corrupted.
	#[cfg(feature = "try-runtime")]
	pub fn image_count<T: Config>() -> Option<u32> {
		let images = v0::PreimageFor::<T>::iter_values().count() as u32;
		let status = v0::StatusFor::<T>::iter_values().count() as u32;

		if images == status {
			Some(images)
		} else {
			None
		}
	}
}

pub mod v1 {
	use super::*;

	/// Migration for moving preimage from V0 to V1 storage.
	///
	/// Note: This needs to be run with the same hashing algorithm as before
	/// since it is not re-hashing the preimages.
	pub struct Migration<T>(core::marker::PhantomData<T>);

	impl<T: Config> OnRuntimeUpgrade for Migration<T> {
		#[cfg(feature = "try-runtime")]
		fn pre_upgrade() -> Result<Vec<u8>, TryRuntimeError> {
			ensure!(StorageVersion::get::<Pezpallet<T>>() == 0, "can only upgrade from version 0");

			let images = v0::image_count::<T>().expect("v0 storage corrupted");
			log::info!(target: TARGET, "Migrating {} images", &images);
			Ok((images as u32).encode())
		}

		fn on_runtime_upgrade() -> Weight {
			let mut weight = T::DbWeight::get().reads(1);
			if StorageVersion::get::<Pezpallet<T>>() != 0 {
				log::warn!(
					target: TARGET,
					"skipping MovePreimagesIntoBuckets: executed on wrong storage version.\
				Expected version 0"
				);
				return weight;
			}

			let status = v0::StatusFor::<T>::drain().collect::<Vec<_>>();
			weight.saturating_accrue(T::DbWeight::get().reads(status.len() as u64));

			let preimages = v0::PreimageFor::<T>::drain().collect::<BTreeMap<_, _>>();
			weight.saturating_accrue(T::DbWeight::get().reads(preimages.len() as u64));

			for (hash, status) in status.into_iter() {
				let preimage = if let Some(preimage) = preimages.get(&hash) {
					preimage
				} else {
					log::error!(target: TARGET, "preimage not found for hash {:?}", &hash);
					continue;
				};
				let len = preimage.len() as u32;
				if len > MAX_SIZE {
					log::error!(
						target: TARGET,
						"preimage too large for hash {:?}, len: {}",
						&hash,
						len
					);
					continue;
				}

				let status = match status {
					v0::OldRequestStatus::Unrequested(deposit) => match deposit {
						Some(deposit) => OldRequestStatus::Unrequested { deposit, len },
						// `None` depositor becomes system-requested.
						None =>
							OldRequestStatus::Requested { deposit: None, count: 1, len: Some(len) },
					},
					v0::OldRequestStatus::Requested(0) => {
						log::error!(target: TARGET, "preimage has counter of zero: {:?}", hash);
						continue;
					},
					v0::OldRequestStatus::Requested(count) =>
						OldRequestStatus::Requested { deposit: None, count, len: Some(len) },
				};
				log::trace!(target: TARGET, "Moving preimage {:?} with len {}", hash, len);

				#[allow(deprecated)]
				crate::StatusFor::<T>::insert(hash, status);
				crate::PreimageFor::<T>::insert(&(hash, len), preimage);

				weight.saturating_accrue(T::DbWeight::get().writes(2));
			}
			StorageVersion::new(1).put::<Pezpallet<T>>();

			weight.saturating_add(T::DbWeight::get().writes(1))
		}

		#[cfg(feature = "try-runtime")]
		fn post_upgrade(state: Vec<u8>) -> DispatchResult {
			let old_images: u32 =
				Decode::decode(&mut &state[..]).expect("pre_upgrade provides a valid state; qed");
			let new_images = image_count::<T>().expect("V1 storage corrupted");

			if new_images != old_images {
				log::error!(
					target: TARGET,
					"migrated {} images, expected {}",
					new_images,
					old_images
				);
			}
			ensure!(StorageVersion::get::<Pezpallet<T>>() == 1, "must upgrade");
			Ok(())
		}
	}

	/// Returns the number of images or `None` if the storage is corrupted.
	#[cfg(feature = "try-runtime")]
	pub fn image_count<T: Config>() -> Option<u32> {
		// Use iter_values() to ensure that the values are decodable.
		let images = crate::PreimageFor::<T>::iter_values().count() as u32;
		#[allow(deprecated)]
		let status = crate::StatusFor::<T>::iter_values().count() as u32;

		if images == status {
			Some(images)
		} else {
			None
		}
	}
}

#[cfg(test)]
#[cfg(feature = "try-runtime")]
mod test {
	#![allow(deprecated)]
	use super::*;
	use crate::mock::{Test as T, *};

	use pezsp_runtime::bounded_vec;

	#[test]
	fn migration_works() {
		new_test_ext().execute_with(|| {
			assert_eq!(StorageVersion::get::<Pezpallet<T>>(), 0);
			// Insert some preimages into the v0 storage:

			// Case 1: Unrequested without deposit
			let (p, h) = preimage::<T>(128);
			v0::PreimageFor::<T>::insert(h, p);
			v0::StatusFor::<T>::insert(h, v0::OldRequestStatus::Unrequested(None));
			// Case 2: Unrequested with deposit
			let (p, h) = preimage::<T>(1024);
			v0::PreimageFor::<T>::insert(h, p);
			v0::StatusFor::<T>::insert(h, v0::OldRequestStatus::Unrequested(Some((1, 1))));
			// Case 3: Requested by 0 (invalid)
			let (p, h) = preimage::<T>(8192);
			v0::PreimageFor::<T>::insert(h, p);
			v0::StatusFor::<T>::insert(h, v0::OldRequestStatus::Requested(0));
			// Case 4: Requested by 10
			let (p, h) = preimage::<T>(65536);
			v0::PreimageFor::<T>::insert(h, p);
			v0::StatusFor::<T>::insert(h, v0::OldRequestStatus::Requested(10));

			assert_eq!(v0::image_count::<T>(), Some(4));
			assert_eq!(v1::image_count::<T>(), None, "V1 storage should be corrupted");

			let state = v1::Migration::<T>::pre_upgrade().unwrap();
			let _w = v1::Migration::<T>::on_runtime_upgrade();
			v1::Migration::<T>::post_upgrade(state).unwrap();

			// V0 and V1 share the same prefix, so `iter_values` still counts the same.
			assert_eq!(v0::image_count::<T>(), Some(3));
			assert_eq!(v1::image_count::<T>(), Some(3)); // One gets skipped therefore 3.
			assert_eq!(StorageVersion::get::<Pezpallet<T>>(), 1);

			// Case 1: Unrequested without deposit becomes system-requested
			let (p, h) = preimage::<T>(128);
			assert_eq!(crate::PreimageFor::<T>::get(&(h, 128)), Some(p));
			assert_eq!(
				crate::StatusFor::<T>::get(h),
				Some(OldRequestStatus::Requested { deposit: None, count: 1, len: Some(128) })
			);
			// Case 2: Unrequested with deposit becomes unrequested
			let (p, h) = preimage::<T>(1024);
			assert_eq!(crate::PreimageFor::<T>::get(&(h, 1024)), Some(p));
			assert_eq!(
				crate::StatusFor::<T>::get(h),
				Some(OldRequestStatus::Unrequested { deposit: (1, 1), len: 1024 })
			);
			// Case 3: Requested by 0 should be skipped
			let (_, h) = preimage::<T>(8192);
			assert_eq!(crate::PreimageFor::<T>::get(&(h, 8192)), None);
			assert_eq!(crate::StatusFor::<T>::get(h), None);
			// Case 4: Requested by 10 becomes requested by 10
			let (p, h) = preimage::<T>(65536);
			assert_eq!(crate::PreimageFor::<T>::get(&(h, 65536)), Some(p));
			assert_eq!(
				crate::StatusFor::<T>::get(h),
				Some(OldRequestStatus::Requested { deposit: None, count: 10, len: Some(65536) })
			);
		});
	}

	/// Returns a preimage with a given size and its hash.
	fn preimage<T: Config>(
		len: usize,
	) -> (BoundedVec<u8, ConstU32<MAX_SIZE>>, <T as pezframe_system::Config>::Hash) {
		let p = bounded_vec![1; len];
		let h = <T as pezframe_system::Config>::Hashing::hash_of(&p);
		(p, h)
	}
}