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

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

// Pezkuwi 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 Pezkuwi.  If not, see <http://www.gnu.org/licenses/>.

//! Pezkuwi test service only.

#![warn(missing_docs)]

pub mod chain_spec;

pub use chain_spec::*;
use futures::{future::Future, stream::StreamExt};
use pezkuwi_node_subsystem::messages::{CollationGenerationMessage, CollatorProtocolMessage};
use pezkuwi_overseer::Handle;
use pezkuwi_pez_node_primitives::{CollationGenerationConfig, CollatorFn};
use pezkuwi_primitives::{Balance, CollatorPair, HeadData, Id as ParaId, ValidationCode};
use pezkuwi_runtime_common::BlockHashCount;
use pezkuwi_runtime_teyrchains::paras::{ParaGenesisArgs, ParaKind};
use pezkuwi_service::{Error, FullClient, IsTeyrchainNode, NewFull, OverseerGen, PrometheusConfig};
use pezkuwi_test_runtime::{
	ParasCall, ParasSudoWrapperCall, Runtime, SignedPayload, SudoCall, TxExtension,
	UncheckedExtrinsic, VERSION,
};

use bizinikiwi_test_client::{
	BlockchainEventsExt, RpcHandlersExt, RpcTransactionError, RpcTransactionOutput,
};
use pezsc_chain_spec::ChainSpec;
use pezsc_client_api::BlockchainEvents;
use pezsc_network::{
	config::{NetworkConfiguration, TransportConfig},
	multiaddr,
	service::traits::NetworkService,
	NetworkStateInfo,
};
use pezsc_service::{
	config::{
		DatabaseSource, KeystoreConfig, MultiaddrWithPeerId, RpcBatchRequestConfig,
		WasmExecutionMethod, WasmtimeInstantiationStrategy,
	},
	BasePath, BlocksPruning, Configuration, Role, RpcHandlers, TaskManager,
};
use pezsp_arithmetic::traits::SaturatedConversion;
use pezsp_blockchain::HeaderBackend;
use pezsp_keyring::Sr25519Keyring;
use pezsp_runtime::{codec::Encode, generic, traits::IdentifyAccount, MultiSigner};
use pezsp_state_machine::BasicExternalities;
use std::{
	collections::HashSet,
	net::{Ipv4Addr, SocketAddr},
	path::PathBuf,
	sync::Arc,
};

/// The client type being used by the test service.
pub type Client = FullClient;

pub use pezkuwi_service::{FullBackend, GetLastTimestamp};
use pezsc_service::config::{ExecutorConfiguration, RpcConfiguration};

/// Create a new full node.
#[pezsc_tracing::logging::prefix_logs_with(custom_log_prefix.unwrap_or(config.network.node_name.as_str()))]
pub fn new_full<OverseerGenerator: OverseerGen>(
	config: Configuration,
	is_teyrchain_node: IsTeyrchainNode,
	workers_path: Option<PathBuf>,
	overseer_gen: OverseerGenerator,
	custom_log_prefix: Option<&'static str>,
) -> Result<NewFull, Error> {
	let workers_path = Some(workers_path.unwrap_or_else(get_relative_workers_path_for_test));

	let params = pezkuwi_service::NewFullParams {
		is_teyrchain_node,
		enable_beefy: true,
		force_authoring_backoff: false,
		telemetry_worker_handle: None,
		node_version: None,
		secure_validator_mode: false,
		workers_path,
		workers_names: None,
		overseer_gen,
		overseer_message_channel_capacity_override: None,
		malus_finality_delay: None,
		hwbench: None,
		execute_workers_max_num: None,
		prepare_workers_hard_max_num: None,
		prepare_workers_soft_max_num: None,
		keep_finalized_for: None,
		invulnerable_ah_collators: HashSet::new(),
		collator_protocol_hold_off: None,
	};

	match config.network.network_backend {
		pezsc_network::config::NetworkBackendType::Libp2p => {
			pezkuwi_service::new_full::<_, pezsc_network::NetworkWorker<_, _>>(config, params)
		},
		pezsc_network::config::NetworkBackendType::Litep2p => {
			pezkuwi_service::new_full::<_, pezsc_network::Litep2pNetworkBackend>(config, params)
		},
	}
}

fn get_relative_workers_path_for_test() -> PathBuf {
	// If no explicit worker path is passed in, we need to specify it ourselves as test binaries
	// are in the "deps/" directory, one level below where the worker binaries are generated.
	let mut exe_path = std::env::current_exe()
		.expect("for test purposes it's reasonable to expect that this will not fail");
	let _ = exe_path.pop();
	let _ = exe_path.pop();
	exe_path
}

/// Returns a prometheus config usable for testing.
pub fn test_prometheus_config(port: u16) -> PrometheusConfig {
	PrometheusConfig::new_with_default_registry(
		SocketAddr::new(Ipv4Addr::LOCALHOST.into(), port),
		"test-chain".to_string(),
	)
}

/// Create a Pezkuwi `Configuration`.
///
/// By default a TCP socket will be used, therefore you need to provide boot
/// nodes if you want the future node to be connected to other nodes.
///
/// The `storage_update_func` function will be executed in an externalities provided environment
/// and can be used to make adjustments to the runtime genesis storage.
pub fn node_config(
	storage_update_func: impl Fn(),
	tokio_handle: tokio::runtime::Handle,
	key: Sr25519Keyring,
	boot_nodes: Vec<MultiaddrWithPeerId>,
	is_validator: bool,
) -> Configuration {
	let base_path = BasePath::new_temp_dir().expect("could not create temporary directory");
	let root = base_path.path().join(key.to_string());
	let role = if is_validator { Role::Authority } else { Role::Full };
	let key_seed = key.to_seed();
	let mut spec = pezkuwi_local_testnet_config();
	let mut storage = spec.as_storage_builder().build_storage().expect("could not build storage");

	BasicExternalities::execute_with_storage(&mut storage, storage_update_func);
	spec.set_storage(storage);

	let mut network_config = NetworkConfiguration::new(
		key_seed.to_string(),
		"network/test/0.1",
		Default::default(),
		None,
	);

	network_config.boot_nodes = boot_nodes;

	network_config.allow_non_globals_in_dht = true;

	// Libp2p needs to know the local address on which it should listen for incoming connections,
	// while Litep2p will use `/ip4/127.0.0.1/tcp/0` by default.
	let addr: multiaddr::Multiaddr = "/ip4/127.0.0.1/tcp/0".parse().expect("valid address; qed");
	network_config.listen_addresses.push(addr.clone());
	network_config.public_addresses.push(addr);
	network_config.transport =
		TransportConfig::Normal { enable_mdns: false, allow_private_ip: true };

	Configuration {
		impl_name: "pezkuwi-test-node".to_string(),
		impl_version: "0.1".to_string(),
		role,
		tokio_handle,
		transaction_pool: Default::default(),
		network: network_config,
		keystore: KeystoreConfig::InMemory,
		database: DatabaseSource::RocksDb { path: root.join("db"), cache_size: 128 },
		trie_cache_maximum_size: Some(64 * 1024 * 1024),
		warm_up_trie_cache: None,
		state_pruning: Default::default(),
		blocks_pruning: BlocksPruning::KeepFinalized,
		chain_spec: Box::new(spec),
		executor: ExecutorConfiguration {
			wasm_method: WasmExecutionMethod::Compiled {
				instantiation_strategy: WasmtimeInstantiationStrategy::PoolingCopyOnWrite,
			},
			..ExecutorConfiguration::default()
		},
		wasm_runtime_overrides: Default::default(),
		rpc: RpcConfiguration {
			addr: Default::default(),
			max_request_size: Default::default(),
			max_response_size: Default::default(),
			max_connections: Default::default(),
			cors: None,
			methods: Default::default(),
			id_provider: None,
			max_subs_per_conn: Default::default(),
			port: 9944,
			message_buffer_capacity: Default::default(),
			batch_config: RpcBatchRequestConfig::Unlimited,
			rate_limit: None,
			rate_limit_whitelisted_ips: Default::default(),
			rate_limit_trust_proxy_headers: Default::default(),
			request_logger_limit: 1024,
		},
		prometheus_config: None,
		telemetry_endpoints: None,
		offchain_worker: Default::default(),
		force_authoring: false,
		disable_grandpa: false,
		dev_key_seed: Some(key_seed),
		tracing_targets: None,
		tracing_receiver: Default::default(),
		announce_block: true,
		data_path: root,
		base_path,
	}
}

/// Get the listen multiaddr from the network service.
///
/// The address is used to connect to the node.
pub async fn get_listen_address(network: Arc<dyn NetworkService>) -> pezsc_network::Multiaddr {
	loop {
		// Litep2p provides instantly the listen address of the TCP protocol and
		// ditched the `/0` port used by the `node_config` function.
		//
		// Libp2p backend needs to be polled in a separate tokio task a few times
		// before the listen address is available. The address is made available
		// through the `SwarmEvent::NewListenAddr` event.
		let listen_addresses = network.listen_addresses();

		// The network backend must produce a valid TCP port.
		match listen_addresses.into_iter().find(|addr| {
			addr.iter().any(|protocol| match protocol {
				multiaddr::Protocol::Tcp(port) => port > 0,
				_ => false,
			})
		}) {
			Some(multiaddr) => return multiaddr,
			None => {
				tokio::time::sleep(std::time::Duration::from_millis(500)).await;
				continue;
			},
		}
	}
}

/// Run a test validator node that uses the test runtime and specified `config`.
pub async fn run_validator_node(
	config: Configuration,
	worker_program_path: Option<PathBuf>,
) -> PezkuwiTestNode {
	let NewFull { task_manager, client, network, rpc_handlers, overseer_handle, .. } = new_full(
		config,
		IsTeyrchainNode::No,
		worker_program_path,
		pezkuwi_service::ValidatorOverseerGen,
		None,
	)
	.expect("could not create Pezkuwi test service");

	let overseer_handle = overseer_handle.expect("test node must have an overseer handle");
	let peer_id = network.local_peer_id();
	let multiaddr = get_listen_address(network).await;

	let addr = MultiaddrWithPeerId { multiaddr, peer_id };

	PezkuwiTestNode { task_manager, client, overseer_handle, addr, rpc_handlers }
}

/// Run a test collator node that uses the test runtime.
///
/// The node will be using an in-memory socket, therefore you need to provide boot nodes if you
/// want it to be connected to other nodes.
///
/// The `storage_update_func` function will be executed in an externalities provided environment
/// and can be used to make adjustments to the runtime genesis storage.
///
/// # Note
///
/// The collator functionality still needs to be registered at the node! This can be done using
/// [`PezkuwiTestNode::register_collator`].
pub async fn run_collator_node(
	tokio_handle: tokio::runtime::Handle,
	key: Sr25519Keyring,
	storage_update_func: impl Fn(),
	boot_nodes: Vec<MultiaddrWithPeerId>,
	collator_pair: CollatorPair,
) -> PezkuwiTestNode {
	let config = node_config(storage_update_func, tokio_handle, key, boot_nodes, false);
	let NewFull { task_manager, client, network, rpc_handlers, overseer_handle, .. } = new_full(
		config,
		IsTeyrchainNode::Collator(collator_pair),
		None,
		pezkuwi_service::CollatorOverseerGen,
		None,
	)
	.expect("could not create Pezkuwi test service");

	let overseer_handle = overseer_handle.expect("test node must have an overseer handle");
	let peer_id = network.local_peer_id();

	let multiaddr = get_listen_address(network).await;
	let addr = MultiaddrWithPeerId { multiaddr, peer_id };

	PezkuwiTestNode { task_manager, client, overseer_handle, addr, rpc_handlers }
}

/// A Pezkuwi test node instance used for testing.
pub struct PezkuwiTestNode {
	/// `TaskManager`'s instance.
	pub task_manager: TaskManager,
	/// Client's instance.
	pub client: Arc<Client>,
	/// A handle to Overseer.
	pub overseer_handle: Handle,
	/// The `MultiaddrWithPeerId` to this node. This is useful if you want to pass it as "boot
	/// node" to other nodes.
	pub addr: MultiaddrWithPeerId,
	/// `RPCHandlers` to make RPC queries.
	pub rpc_handlers: RpcHandlers,
}

impl PezkuwiTestNode {
	/// Send a sudo call to this node.
	async fn send_sudo(
		&self,
		call: impl Into<pezkuwi_test_runtime::RuntimeCall>,
		caller: Sr25519Keyring,
		nonce: u32,
	) -> Result<(), RpcTransactionError> {
		let sudo = SudoCall::sudo { call: Box::new(call.into()) };

		let extrinsic = construct_extrinsic(&self.client, sudo, caller, nonce);
		self.rpc_handlers.send_transaction(extrinsic.into()).await.map(drop)
	}

	/// Send an extrinsic to this node.
	pub async fn send_extrinsic(
		&self,
		function: impl Into<pezkuwi_test_runtime::RuntimeCall>,
		caller: Sr25519Keyring,
	) -> Result<RpcTransactionOutput, RpcTransactionError> {
		let extrinsic = construct_extrinsic(&self.client, function, caller, 0);

		self.rpc_handlers.send_transaction(extrinsic.into()).await
	}

	/// Register a teyrchain at this relay chain.
	pub async fn register_teyrchain(
		&self,
		id: ParaId,
		validation_code: impl Into<ValidationCode>,
		genesis_head: impl Into<HeadData>,
	) -> Result<(), RpcTransactionError> {
		let validation_code: ValidationCode = validation_code.into();
		let call = ParasSudoWrapperCall::sudo_schedule_para_initialize {
			id,
			genesis: ParaGenesisArgs {
				genesis_head: genesis_head.into(),
				validation_code: validation_code.clone(),
				para_kind: ParaKind::Teyrchain,
			},
		};

		self.send_sudo(call, Sr25519Keyring::Alice, 0).await?;

		// Bypass pvf-checking.
		let call = ParasCall::add_trusted_validation_code { validation_code };
		self.send_sudo(call, Sr25519Keyring::Alice, 1).await
	}

	/// Wait for `count` blocks to be imported in the node and then exit. This function will not
	/// return if no blocks are ever created, thus you should restrict the maximum amount of time of
	/// the test execution.
	pub fn wait_for_blocks(&self, count: usize) -> impl Future<Output = ()> {
		self.client.wait_for_blocks(count)
	}

	/// Wait for `count` blocks to be finalized and then exit. Similarly with `wait_for_blocks` this
	/// function will not return if no block are ever finalized.
	pub async fn wait_for_finalized_blocks(&self, count: usize) {
		let mut import_notification_stream = self.client.finality_notification_stream();
		let mut blocks = HashSet::new();

		while let Some(notification) = import_notification_stream.next().await {
			blocks.insert(notification.hash);
			if blocks.len() == count {
				break;
			}
		}
	}

	/// Register the collator functionality in the overseer of this node.
	pub async fn register_collator(
		&mut self,
		collator_key: CollatorPair,
		para_id: ParaId,
		collator: CollatorFn,
	) {
		let config =
			CollationGenerationConfig { key: collator_key, collator: Some(collator), para_id };

		self.overseer_handle
			.send_msg(CollationGenerationMessage::Initialize(config), "Collator")
			.await;

		self.overseer_handle
			.send_msg(CollatorProtocolMessage::CollateOn(para_id), "Collator")
			.await;
	}
}

/// Construct an extrinsic that can be applied to the test runtime.
pub fn construct_extrinsic(
	client: &Client,
	function: impl Into<pezkuwi_test_runtime::RuntimeCall>,
	caller: Sr25519Keyring,
	nonce: u32,
) -> UncheckedExtrinsic {
	let function = function.into();
	let current_block_hash = client.info().best_hash;
	let current_block = client.info().best_number.saturated_into();
	let genesis_block = client.hash(0).unwrap().unwrap();
	let period =
		BlockHashCount::get().checked_next_power_of_two().map(|c| c / 2).unwrap_or(2) as u64;
	let tip = 0;
	let tx_ext: TxExtension = (
		pezframe_system::AuthorizeCall::<Runtime>::new(),
		pezframe_system::CheckNonZeroSender::<Runtime>::new(),
		pezframe_system::CheckSpecVersion::<Runtime>::new(),
		pezframe_system::CheckTxVersion::<Runtime>::new(),
		pezframe_system::CheckGenesis::<Runtime>::new(),
		pezframe_system::CheckEra::<Runtime>::from(generic::Era::mortal(period, current_block)),
		pezframe_system::CheckNonce::<Runtime>::from(nonce),
		pezframe_system::CheckWeight::<Runtime>::new(),
		pezpallet_transaction_payment::ChargeTransactionPayment::<Runtime>::from(tip),
		pezframe_system::WeightReclaim::<Runtime>::new(),
	)
		.into();
	let raw_payload = SignedPayload::from_raw(
		function.clone(),
		tx_ext.clone(),
		(
			(),
			(),
			VERSION.spec_version,
			VERSION.transaction_version,
			genesis_block,
			current_block_hash,
			(),
			(),
			(),
			(),
		),
	);
	let signature = raw_payload.using_encoded(|e| caller.sign(e));
	UncheckedExtrinsic::new_signed(
		function.clone(),
		pezkuwi_test_runtime::Address::Id(caller.public().into()),
		pezkuwi_primitives::Signature::Sr25519(signature),
		tx_ext.clone(),
	)
}

/// Construct a transfer extrinsic.
pub fn construct_transfer_extrinsic(
	client: &Client,
	origin: pezsp_keyring::Sr25519Keyring,
	dest: pezsp_keyring::Sr25519Keyring,
	value: Balance,
) -> UncheckedExtrinsic {
	let function = pezkuwi_test_runtime::RuntimeCall::Balances(
		pezpallet_balances::Call::transfer_allow_death {
			dest: MultiSigner::from(dest.public()).into_account().into(),
			value,
		},
	);

	construct_extrinsic(client, function, origin, 0)
}