use crate::{BlockNumberFor, Config, Message, MessagesV2, Pallet, SchemaId, LOG_TARGET};
use frame_support::{pallet_prelude::*, storage_alias, traits::OnRuntimeUpgrade, weights::Weight};
use log;
use sp_runtime::Saturating;

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

/// old structures and storages
pub mod old {
	use super::*;
	use common_primitives::msa::MessageSourceId;
	use sp_std::fmt::Debug;

	/// old message structure that was stored
	#[derive(Default, Encode, Decode, PartialEq, Debug, TypeInfo, Eq, MaxEncodedLen)]
	#[scale_info(skip_type_params(MaxDataSize))]
	#[codec(mel_bound(MaxDataSize: MaxEncodedLen))]
	pub struct OldMessage<MaxDataSize>
	where
		MaxDataSize: Get<u32> + Debug,
	{
		///  Data structured by the associated schema's model.
		pub payload: BoundedVec<u8, MaxDataSize>,
		/// Message source account id of the Provider. This may be the same id as contained in `msa_id`,
		/// indicating that the original source MSA is acting as its own provider. An id differing from that
		/// of `msa_id` indicates that `provider_msa_id` was delegated by `msa_id` to send this message on
		/// its behalf.
		pub provider_msa_id: MessageSourceId,
		///  Message source account id (the original source).
		pub msa_id: Option<MessageSourceId>,
		///  Stores index of message in block to keep total order.
		pub index: u16,
	}

	/// old permanent storage for messages mapped by block number and schema id.
	#[storage_alias]
	pub(crate) type Messages<T: Config> = StorageDoubleMap<
		Pallet<T>,
		Twox64Concat,
		BlockNumberFor<T>,
		Twox64Concat,
		SchemaId,
		BoundedVec<
			OldMessage<<T as crate::pallet::Config>::MessagesMaxPayloadSizeBytes>,
			ConstU32<200>,
		>,
		ValueQuery,
	>;
}
/// migration to v2 implementation
pub struct MigrateToV2<T>(PhantomData<T>);

impl<T: Config> OnRuntimeUpgrade for MigrateToV2<T> {
	fn on_runtime_upgrade() -> Weight {
		migrate_to_v2::<T>()
	}

	#[cfg(feature = "try-runtime")]
	fn pre_upgrade() -> Result<Vec<u8>, TryRuntimeError> {
		use frame_support::storage::generator::StorageDoubleMap;
		log::info!(target: LOG_TARGET, "Running pre_upgrade...");
		let onchain_version = Pallet::<T>::on_chain_storage_version();
		if onchain_version < 2 {
			let pallet_prefix = old::Messages::<T>::pallet_prefix();
			let storage_prefix = old::Messages::<T>::storage_prefix();
			assert_eq!(&b"Messages"[..], pallet_prefix);
			assert_eq!(&b"Messages"[..], storage_prefix);

			let mut count = 0u32;
			for (_, _, messages) in old::Messages::<T>::iter() {
				count += messages.len() as u32;
			}
			log::info!(target: LOG_TARGET, "Finish pre_upgrade for {:?}", count);
			return Ok(count.encode())
		}
		Ok(Vec::new())
	}

	#[cfg(feature = "try-runtime")]
	fn post_upgrade(state: Vec<u8>) -> Result<(), TryRuntimeError> {
		log::info!(target: LOG_TARGET, "Running post_upgrade...");
		let onchain_version = Pallet::<T>::on_chain_storage_version();
		if onchain_version < 2 {
			let old_count: u32 = Decode::decode(&mut state.as_slice()).expect(
				"the state parameter should be something that was generated by pre_upgrade",
			);

			let count = old::Messages::<T>::iter().count();
			let moved_count = MessagesV2::<T>::iter().count();

			log::info!(target: LOG_TARGET, "Finish post_upgrade for {:?}", moved_count);
			let onchain_version = Pallet::<T>::on_chain_storage_version();

			assert_eq!(count, 0usize);
			assert_eq!(moved_count, old_count as usize);
			assert_eq!(onchain_version, crate::pallet::STORAGE_VERSION);
		}
		Ok(())
	}
}
/// migrating to v2
pub fn migrate_to_v2<T: Config>() -> Weight {
	log::info!(target: LOG_TARGET, "Running storage migration...");
	let onchain_version = Pallet::<T>::on_chain_storage_version();
	let current_version = Pallet::<T>::in_code_storage_version();
	log::info!(target: LOG_TARGET, "onchain_version= {:?}, current_version={:?}", onchain_version, current_version);

	if onchain_version < 2 {
		let mut reads = 1u64;
		let mut writes = 0u64;
		let mut bytes = 0u64;
		for (block_number, schema_id, messages) in old::Messages::<T>::drain() {
			bytes = bytes.saturating_add(messages.encode().len() as u64);

			for message in &messages {
				let new_msg = Message {
					provider_msa_id: message.provider_msa_id,
					msa_id: message.msa_id,
					payload: message.payload.clone(),
				};
				bytes = bytes.saturating_add(new_msg.encode().len() as u64);
				MessagesV2::<T>::insert((block_number, schema_id, message.index), new_msg);
			}

			reads.saturating_inc();
			writes = writes.saturating_add(messages.len() as u64 + 1);
		}

		// Set storage version to `2`.
		StorageVersion::new(2).put::<Pallet<T>>();
		writes.saturating_inc();

		log::info!(target: LOG_TARGET, "Storage migrated to version 2  read={:?}, write={:?}, bytes={:?}", reads, writes, bytes);
		let weights = T::DbWeight::get().reads_writes(reads, writes).add_proof_size(bytes);
		log::info!(target: LOG_TARGET, "Migration Calculated weights={:?}",weights);
		weights
	} else {
		log::info!(
			target: LOG_TARGET,
			"Migration did not execute. This probably should be removed onchain:{:?}, current:{:?}",
			onchain_version,
			current_version
		);
		T::DbWeight::get().reads(1)
	}
}