// Copyright (C) Parity Technologies (UK) Ltd.
// This file is part of Pezcumulus.

// Pezcumulus is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Pezcumulus is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Pezcumulus.  If not, see <http://www.gnu.org/licenses/>.

use crate::Client;
use codec::Encode;
use pezcumulus_primitives_core::{PersistedValidationData, TeyrchainBlockData};
use pezcumulus_primitives_teyrchain_inherent::{TeyrchainInherentData, INHERENT_IDENTIFIER};
use pezcumulus_test_relay_sproof_builder::RelayStateSproofBuilder;
use pezcumulus_test_runtime::{Block, GetLastTimestamp, Hash, Header};
use pezkuwi_primitives::{BlockNumber as PBlockNumber, Hash as PHash};
use pezsc_block_builder::BlockBuilderBuilder;
use pezsp_api::{ProofRecorder, ProofRecorderIgnoredNodes, ProvideRuntimeApi};
use pezsp_consensus_aura::{AuraApi, Slot};
use pezsp_runtime::{traits::Header as HeaderT, Digest, DigestItem};

/// A struct containing a block builder and support data required to build test scenarios.
pub struct BlockBuilderAndSupportData<'a> {
	pub block_builder: pezsc_block_builder::BlockBuilder<'a, Block, Client>,
	pub persisted_validation_data: PersistedValidationData<PHash, PBlockNumber>,
}

/// An extension for the Pezcumulus test client to init a block builder.
pub trait InitBlockBuilder {
	/// Init a specific block builder that works for the test runtime.
	///
	/// This will automatically create and push the inherents for you to make the block
	/// valid for the test runtime.
	///
	/// You can use the relay chain state sproof builder to arrange required relay chain state or
	/// just use a default one. The relay chain slot in the storage proof
	/// will be adjusted to align with the teyrchain slot to pass validation.
	///
	/// Returns the block builder and validation data for further usage.
	fn init_block_builder(
		&self,
		validation_data: Option<PersistedValidationData<PHash, PBlockNumber>>,
		relay_sproof_builder: RelayStateSproofBuilder,
	) -> BlockBuilderAndSupportData<'_>;

	/// Init a specific block builder at a specific block that works for the test runtime.
	///
	/// Same as [`InitBlockBuilder::init_block_builder`] besides that it takes a
	/// [`type@Hash`] to say which should be the parent block of the block that is being build.
	fn init_block_builder_at(
		&self,
		at: Hash,
		validation_data: Option<PersistedValidationData<PHash, PBlockNumber>>,
		relay_sproof_builder: RelayStateSproofBuilder,
	) -> BlockBuilderAndSupportData<'_>;

	/// Init a specific block builder using the given pre-digests.
	///
	/// Same as [`InitBlockBuilder::init_block_builder`] besides that it takes vector of
	/// [`DigestItem`]'s that are passed as pre-digest to the block builder.
	fn init_block_builder_with_pre_digests(
		&self,
		validation_data: Option<PersistedValidationData<PHash, PBlockNumber>>,
		relay_sproof_builder: RelayStateSproofBuilder,
		pre_digests: Vec<DigestItem>,
	) -> BlockBuilderAndSupportData<'_>;
	/// Init a specific block builder at a specific block that works for the test runtime.
	///
	/// Same as [`InitBlockBuilder::init_block_builder_with_timestamp`] besides that it takes
	/// `ignored_nodes` that instruct the proof recorder to not record these nodes.
	fn init_block_builder_with_ignored_nodes(
		&self,
		at: Hash,
		validation_data: Option<PersistedValidationData<PHash, PBlockNumber>>,
		relay_sproof_builder: RelayStateSproofBuilder,
		timestamp: u64,
		ignored_nodes: ProofRecorderIgnoredNodes<Block>,
	) -> BlockBuilderAndSupportData<'_>;

	/// Init a specific block builder that works for the test runtime.
	///
	/// Same as [`InitBlockBuilder::init_block_builder`] besides that it takes a
	/// [`type@Hash`] to say which should be the parent block of the block that is being build and
	/// it will use the given `timestamp` as input for the timestamp inherent.
	fn init_block_builder_with_timestamp(
		&self,
		at: Hash,
		validation_data: Option<PersistedValidationData<PHash, PBlockNumber>>,
		relay_sproof_builder: RelayStateSproofBuilder,
		timestamp: u64,
	) -> BlockBuilderAndSupportData<'_>;
}

fn init_block_builder(
	client: &Client,
	at: Hash,
	validation_data: Option<PersistedValidationData<PHash, PBlockNumber>>,
	mut relay_sproof_builder: RelayStateSproofBuilder,
	timestamp: Option<u64>,
	extra_pre_digests: Option<Vec<DigestItem>>,
	ignored_nodes: Option<ProofRecorderIgnoredNodes<Block>>,
) -> BlockBuilderAndSupportData<'_> {
	let timestamp = timestamp.unwrap_or_else(|| {
		let last_timestamp =
			client.runtime_api().get_last_timestamp(at).expect("Get last timestamp");

		if last_timestamp == 0 {
			if relay_sproof_builder.current_slot != 0u64 {
				*relay_sproof_builder.current_slot * 6_000
			} else {
				std::time::SystemTime::now()
					.duration_since(std::time::SystemTime::UNIX_EPOCH)
					.expect("Time is always after UNIX_EPOCH; qed")
					.as_millis() as u64
			}
		} else {
			last_timestamp + client.runtime_api().slot_duration(at).unwrap().as_millis()
		}
	});

	let slot: Slot =
		(timestamp / client.runtime_api().slot_duration(at).unwrap().as_millis()).into();

	if relay_sproof_builder.current_slot == 0u64 {
		relay_sproof_builder.current_slot = (timestamp / 6_000).into();
	}

	let pre_digests = Digest {
		logs: extra_pre_digests
			.unwrap_or_default()
			.into_iter()
			.chain(std::iter::once(DigestItem::PreRuntime(
				pezsp_consensus_aura::AURA_ENGINE_ID,
				slot.encode(),
			)))
			.collect::<Vec<_>>(),
	};

	let mut block_builder = BlockBuilderBuilder::new(client)
		.on_parent_block(at)
		.fetch_parent_block_number(client)
		.unwrap()
		.with_proof_recorder(Some(ProofRecorder::<Block>::with_ignored_nodes(
			ignored_nodes.unwrap_or_default(),
		)))
		.with_inherent_digests(pre_digests)
		.build()
		.expect("Creates new block builder for test runtime");

	let mut inherent_data = pezsp_inherents::InherentData::new();

	inherent_data
		.put_data(pezsp_timestamp::INHERENT_IDENTIFIER, &timestamp)
		.expect("Put timestamp failed");

	let (relay_parent_storage_root, relay_chain_state) =
		relay_sproof_builder.into_state_root_and_proof();

	let mut validation_data = validation_data.unwrap_or_default();
	validation_data.relay_parent_storage_root = relay_parent_storage_root;

	inherent_data
		.put_data(
			INHERENT_IDENTIFIER,
			&TeyrchainInherentData {
				validation_data: validation_data.clone(),
				relay_chain_state,
				downward_messages: Default::default(),
				horizontal_messages: Default::default(),
				relay_parent_descendants: Default::default(),
				collator_peer_id: None,
			},
		)
		.expect("Put validation function params failed");

	let inherents = block_builder.create_inherents(inherent_data).expect("Creates inherents");

	inherents
		.into_iter()
		.for_each(|ext| block_builder.push(ext).expect("Pushes inherent"));

	BlockBuilderAndSupportData { block_builder, persisted_validation_data: validation_data }
}

impl InitBlockBuilder for Client {
	fn init_block_builder(
		&self,
		validation_data: Option<PersistedValidationData<PHash, PBlockNumber>>,
		relay_sproof_builder: RelayStateSproofBuilder,
	) -> BlockBuilderAndSupportData<'_> {
		let chain_info = self.chain_info();
		self.init_block_builder_at(chain_info.best_hash, validation_data, relay_sproof_builder)
	}

	fn init_block_builder_with_pre_digests(
		&self,
		validation_data: Option<PersistedValidationData<PHash, PBlockNumber>>,
		relay_sproof_builder: RelayStateSproofBuilder,
		pre_digests: Vec<DigestItem>,
	) -> BlockBuilderAndSupportData<'_> {
		let chain_info = self.chain_info();
		init_block_builder(
			self,
			chain_info.best_hash,
			validation_data,
			relay_sproof_builder,
			None,
			Some(pre_digests),
			None,
		)
	}

	fn init_block_builder_at(
		&self,
		at: Hash,
		validation_data: Option<PersistedValidationData<PHash, PBlockNumber>>,
		relay_sproof_builder: RelayStateSproofBuilder,
	) -> BlockBuilderAndSupportData<'_> {
		init_block_builder(self, at, validation_data, relay_sproof_builder, None, None, None)
	}

	fn init_block_builder_with_ignored_nodes(
		&self,
		at: Hash,
		validation_data: Option<PersistedValidationData<PHash, PBlockNumber>>,
		relay_sproof_builder: RelayStateSproofBuilder,
		timestamp: u64,
		ignored_nodes: ProofRecorderIgnoredNodes<Block>,
	) -> BlockBuilderAndSupportData<'_> {
		init_block_builder(
			self,
			at,
			validation_data,
			relay_sproof_builder,
			Some(timestamp),
			None,
			Some(ignored_nodes),
		)
	}

	fn init_block_builder_with_timestamp(
		&self,
		at: Hash,
		validation_data: Option<PersistedValidationData<PHash, PBlockNumber>>,
		relay_sproof_builder: RelayStateSproofBuilder,
		timestamp: u64,
	) -> BlockBuilderAndSupportData<'_> {
		init_block_builder(
			self,
			at,
			validation_data,
			relay_sproof_builder,
			Some(timestamp),
			None,
			None,
		)
	}
}

/// Extension trait for the [`BlockBuilder`](pezsc_block_builder::BlockBuilder) to build directly a
/// [`TeyrchainBlockData`].
pub trait BuildTeyrchainBlockData {
	/// Directly build the [`TeyrchainBlockData`] from the block that comes out of the block
	/// builder.
	fn build_teyrchain_block(self, parent_state_root: Hash) -> TeyrchainBlockData<Block>;
}

impl<'a> BuildTeyrchainBlockData for pezsc_block_builder::BlockBuilder<'a, Block, Client> {
	fn build_teyrchain_block(self, parent_state_root: Hash) -> TeyrchainBlockData<Block> {
		let built_block = self.build().expect("Builds the block");

		let storage_proof = built_block
			.proof
			.expect("We enabled proof recording before.")
			.into_compact_proof::<<Header as HeaderT>::Hashing>(parent_state_root)
			.expect("Creates the compact proof");

		TeyrchainBlockData::new(vec![built_block.block], storage_proof)
	}
}