#![warn(missing_docs)]
use codec::Encode;
use sp_api::{
	ApiExt, ApiRef, Core, ProvideRuntimeApi, StorageChanges, StorageProof, TransactionOutcome,
};
use sp_blockchain::{ApplyExtrinsicFailed, Error};
use sp_core::traits::CallContext;
use sp_runtime::{
	legacy,
	traits::{Block as BlockT, Hash, HashingFor, Header as HeaderT, NumberFor, One},
	Digest,
};
use sc_client_api::backend;
pub use sp_block_builder::BlockBuilder as BlockBuilderApi;
#[derive(Copy, Clone, PartialEq)]
pub enum RecordProof {
	Yes,
	No,
}
impl RecordProof {
	pub fn yes(&self) -> bool {
		matches!(self, Self::Yes)
	}
}
impl Default for RecordProof {
	fn default() -> Self {
		Self::No
	}
}
impl From<bool> for RecordProof {
	fn from(val: bool) -> Self {
		if val {
			Self::Yes
		} else {
			Self::No
		}
	}
}
pub struct BuiltBlock<Block: BlockT> {
	pub block: Block,
	pub storage_changes: StorageChanges<Block>,
	pub proof: Option<StorageProof>,
}
impl<Block: BlockT> BuiltBlock<Block> {
	pub fn into_inner(self) -> (Block, StorageChanges<Block>, Option<StorageProof>) {
		(self.block, self.storage_changes, self.proof)
	}
}
pub trait BlockBuilderProvider<B, Block, RA>
where
	Block: BlockT,
	B: backend::Backend<Block>,
	Self: Sized,
	RA: ProvideRuntimeApi<Block>,
{
	fn new_block_at<R: Into<RecordProof>>(
		&self,
		parent: Block::Hash,
		inherent_digests: Digest,
		record_proof: R,
	) -> sp_blockchain::Result<BlockBuilder<Block, RA, B>>;
	fn new_block(
		&self,
		inherent_digests: Digest,
	) -> sp_blockchain::Result<BlockBuilder<Block, RA, B>>;
}
pub struct BlockBuilder<'a, Block: BlockT, A: ProvideRuntimeApi<Block>, B> {
	extrinsics: Vec<Block::Extrinsic>,
	api: ApiRef<'a, A::Api>,
	version: u32,
	parent_hash: Block::Hash,
	backend: &'a B,
	estimated_header_size: usize,
}
impl<'a, Block, A, B> BlockBuilder<'a, Block, A, B>
where
	Block: BlockT,
	A: ProvideRuntimeApi<Block> + 'a,
	A::Api: BlockBuilderApi<Block> + ApiExt<Block>,
	B: backend::Backend<Block>,
{
	pub fn new(
		api: &'a A,
		parent_hash: Block::Hash,
		parent_number: NumberFor<Block>,
		record_proof: RecordProof,
		inherent_digests: Digest,
		backend: &'a B,
	) -> Result<Self, Error> {
		let header = <<Block as BlockT>::Header as HeaderT>::new(
			parent_number + One::one(),
			Default::default(),
			Default::default(),
			parent_hash,
			inherent_digests,
		);
		let estimated_header_size = header.encoded_size();
		let mut api = api.runtime_api();
		if record_proof.yes() {
			api.record_proof();
		}
		api.set_call_context(CallContext::Onchain);
		api.initialize_block(parent_hash, &header)?;
		let version = api
			.api_version::<dyn BlockBuilderApi<Block>>(parent_hash)?
			.ok_or_else(|| Error::VersionInvalid("BlockBuilderApi".to_string()))?;
		Ok(Self {
			parent_hash,
			extrinsics: Vec::new(),
			api,
			version,
			backend,
			estimated_header_size,
		})
	}
	pub fn push(&mut self, xt: <Block as BlockT>::Extrinsic) -> Result<(), Error> {
		let parent_hash = self.parent_hash;
		let extrinsics = &mut self.extrinsics;
		let version = self.version;
		self.api.execute_in_transaction(|api| {
			let res = if version < 6 {
				#[allow(deprecated)]
				api.apply_extrinsic_before_version_6(parent_hash, xt.clone())
					.map(legacy::byte_sized_error::convert_to_latest)
			} else {
				api.apply_extrinsic(parent_hash, xt.clone())
			};
			match res {
				Ok(Ok(_)) => {
					extrinsics.push(xt);
					TransactionOutcome::Commit(Ok(()))
				},
				Ok(Err(tx_validity)) => TransactionOutcome::Rollback(Err(
					ApplyExtrinsicFailed::Validity(tx_validity).into(),
				)),
				Err(e) => TransactionOutcome::Rollback(Err(Error::from(e))),
			}
		})
	}
	pub fn build(mut self) -> Result<BuiltBlock<Block>, Error> {
		let header = self.api.finalize_block(self.parent_hash)?;
		debug_assert_eq!(
			header.extrinsics_root().clone(),
			HashingFor::<Block>::ordered_trie_root(
				self.extrinsics.iter().map(Encode::encode).collect(),
				sp_runtime::StateVersion::V0,
			),
		);
		let proof = self.api.extract_proof();
		let state = self.backend.state_at(self.parent_hash)?;
		let storage_changes = self
			.api
			.into_storage_changes(&state, self.parent_hash)
			.map_err(sp_blockchain::Error::StorageChanges)?;
		Ok(BuiltBlock {
			block: <Block as BlockT>::new(header, self.extrinsics),
			storage_changes,
			proof,
		})
	}
	pub fn create_inherents(
		&mut self,
		inherent_data: sp_inherents::InherentData,
	) -> Result<Vec<Block::Extrinsic>, Error> {
		let parent_hash = self.parent_hash;
		self.api
			.execute_in_transaction(move |api| {
				TransactionOutcome::Rollback(api.inherent_extrinsics(parent_hash, inherent_data))
			})
			.map_err(|e| Error::Application(Box::new(e)))
	}
	pub fn estimate_block_size(&self, include_proof: bool) -> usize {
		let size = self.estimated_header_size + self.extrinsics.encoded_size();
		if include_proof {
			size + self.api.proof_recorder().map(|pr| pr.estimate_encoded_size()).unwrap_or(0)
		} else {
			size
		}
	}
}
#[cfg(test)]
mod tests {
	use super::*;
	use sp_blockchain::HeaderBackend;
	use sp_core::Blake2Hasher;
	use sp_state_machine::Backend;
	use substrate_test_runtime_client::{
		runtime::ExtrinsicBuilder, DefaultTestClientBuilderExt, TestClientBuilderExt,
	};
	#[test]
	fn block_building_storage_proof_does_not_include_runtime_by_default() {
		let builder = substrate_test_runtime_client::TestClientBuilder::new();
		let backend = builder.backend();
		let client = builder.build();
		let genesis_hash = client.info().best_hash;
		let block = BlockBuilder::new(
			&client,
			genesis_hash,
			client.info().best_number,
			RecordProof::Yes,
			Default::default(),
			&*backend,
		)
		.unwrap()
		.build()
		.unwrap();
		let proof = block.proof.expect("Proof is build on request");
		let genesis_state_root = client.header(genesis_hash).unwrap().unwrap().state_root;
		let backend =
			sp_state_machine::create_proof_check_backend::<Blake2Hasher>(genesis_state_root, proof)
				.unwrap();
		assert!(backend
			.storage(&sp_core::storage::well_known_keys::CODE)
			.unwrap_err()
			.contains("Database missing expected key"),);
	}
	#[test]
	fn failing_extrinsic_rolls_back_changes_in_storage_proof() {
		let builder = substrate_test_runtime_client::TestClientBuilder::new();
		let backend = builder.backend();
		let client = builder.build();
		let mut block_builder = BlockBuilder::new(
			&client,
			client.info().best_hash,
			client.info().best_number,
			RecordProof::Yes,
			Default::default(),
			&*backend,
		)
		.unwrap();
		block_builder.push(ExtrinsicBuilder::new_read_and_panic(8).build()).unwrap_err();
		let block = block_builder.build().unwrap();
		let proof_with_panic = block.proof.expect("Proof is build on request").encoded_size();
		let mut block_builder = BlockBuilder::new(
			&client,
			client.info().best_hash,
			client.info().best_number,
			RecordProof::Yes,
			Default::default(),
			&*backend,
		)
		.unwrap();
		block_builder.push(ExtrinsicBuilder::new_read(8).build()).unwrap();
		let block = block_builder.build().unwrap();
		let proof_without_panic = block.proof.expect("Proof is build on request").encoded_size();
		let block = BlockBuilder::new(
			&client,
			client.info().best_hash,
			client.info().best_number,
			RecordProof::Yes,
			Default::default(),
			&*backend,
		)
		.unwrap()
		.build()
		.unwrap();
		let proof_empty_block = block.proof.expect("Proof is build on request").encoded_size();
		assert!(proof_without_panic > proof_with_panic);
		assert!(proof_without_panic > proof_empty_block);
		assert_eq!(proof_empty_block, proof_with_panic);
	}
}