// This file is part of Substrate.

// Copyright (C) 2020-2021 Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0

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

// This program 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 this program. If not, see <https://www.gnu.org/licenses/>.

mod common;
mod construct;
#[macro_use]
mod core;
mod generator;
mod import;
mod simple_trie;
mod state_sizes;
mod tempdb;
mod trie;
mod txpool;

use structopt::StructOpt;

use node_testing::bench::{BlockType, DatabaseType as BenchDataBaseType, KeyTypes, Profile};

use crate::{
	common::SizeType,
	construct::ConstructionBenchmarkDescription,
	core::{run_benchmark, Mode as BenchmarkMode},
	import::ImportBenchmarkDescription,
	tempdb::DatabaseType,
	trie::{DatabaseSize, TrieReadBenchmarkDescription, TrieWriteBenchmarkDescription},
	txpool::PoolBenchmarkDescription,
};

#[derive(Debug, StructOpt)]
#[structopt(name = "node-bench", about = "Node integration benchmarks")]
struct Opt {
	/// Show list of all available benchmarks.
	///
	/// Will output ("name", "path"). Benchmarks can then be filtered by path.
	#[structopt(short, long)]
	list: bool,

	/// Machine readable json output.
	///
	/// This also suppresses all regular output (except to stderr)
	#[structopt(short, long)]
	json: bool,

	/// Filter benchmarks.
	///
	/// Run with `--list` for the hint of what to filter.
	filter: Option<String>,

	/// Number of transactions for block import with `custom` size.
	#[structopt(long)]
	transactions: Option<usize>,

	/// Mode
	///
	/// "regular" for regular benchmark
	///
	/// "profile" mode adds pauses between measurable runs,
	/// so that actual interval can be selected in the profiler of choice.
	#[structopt(short, long, default_value = "regular")]
	mode: BenchmarkMode,
}

fn main() {
	let opt = Opt::from_args();

	if !opt.json {
		sp_tracing::try_init_simple();
	}

	let mut import_benchmarks = Vec::new();

	for profile in [Profile::Wasm, Profile::Native].iter() {
		for size in [
			SizeType::Empty,
			SizeType::Small,
			SizeType::Medium,
			SizeType::Large,
			SizeType::Full,
			SizeType::Custom(opt.transactions.unwrap_or(0)),
		]
		.iter()
		{
			for block_type in [
				BlockType::RandomTransfersKeepAlive,
				BlockType::RandomTransfersReaping,
				BlockType::Noop,
			]
			.iter()
			{
				for database_type in
					[BenchDataBaseType::RocksDb, BenchDataBaseType::ParityDb].iter()
				{
					import_benchmarks.push((
						profile,
						size.clone(),
						block_type.clone(),
						database_type,
					));
				}
			}
		}
	}

	let benchmarks = matrix!(
		(profile, size, block_type, database_type) in import_benchmarks.into_iter() =>
			ImportBenchmarkDescription {
				profile: *profile,
				key_types: KeyTypes::Sr25519,
				size: size,
				block_type: block_type,
				database_type: *database_type,
			},
		(size, db_type) in
			[
				DatabaseSize::Empty, DatabaseSize::Smallest, DatabaseSize::Small,
				DatabaseSize::Medium, DatabaseSize::Large, DatabaseSize::Huge,
			]
			.iter().flat_map(|size|
			[
				DatabaseType::RocksDb, DatabaseType::ParityDb
			]
			.iter().map(move |db_type| (size, db_type)))
			=> TrieReadBenchmarkDescription { database_size: *size, database_type: *db_type },
		(size, db_type) in
			[
				DatabaseSize::Empty, DatabaseSize::Smallest, DatabaseSize::Small,
				DatabaseSize::Medium, DatabaseSize::Large, DatabaseSize::Huge,
			]
			.iter().flat_map(|size|
			[
				DatabaseType::RocksDb, DatabaseType::ParityDb
			]
			.iter().map(move |db_type| (size, db_type)))
			=> TrieWriteBenchmarkDescription { database_size: *size, database_type: *db_type },
		ConstructionBenchmarkDescription {
			profile: Profile::Wasm,
			key_types: KeyTypes::Sr25519,
			block_type: BlockType::RandomTransfersKeepAlive,
			size: SizeType::Medium,
			database_type: BenchDataBaseType::RocksDb,
		},
		ConstructionBenchmarkDescription {
			profile: Profile::Wasm,
			key_types: KeyTypes::Sr25519,
			block_type: BlockType::RandomTransfersKeepAlive,
			size: SizeType::Large,
			database_type: BenchDataBaseType::RocksDb,
		},
		PoolBenchmarkDescription { database_type: BenchDataBaseType::RocksDb },
	);

	if opt.list {
		println!("Available benchmarks:");
		if let Some(filter) = opt.filter.as_ref() {
			println!("\t(filtered by \"{}\")", filter);
		}
		for benchmark in benchmarks.iter() {
			if opt.filter.as_ref().map(|f| benchmark.path().has(f)).unwrap_or(true) {
				println!("{}: {}", benchmark.name(), benchmark.path().full())
			}
		}
		return
	}

	let mut results = Vec::new();
	for benchmark in benchmarks {
		if opt.filter.as_ref().map(|f| benchmark.path().has(f)).unwrap_or(true) {
			log::info!("Starting {}", benchmark.name());
			let result = run_benchmark(benchmark, opt.mode);
			log::info!("{}", result);

			results.push(result);
		}
	}

	if results.is_empty() {
		eprintln!("No benchmark was found for query");
		std::process::exit(1);
	}

	if opt.json {
		let json_result: String =
			serde_json::to_string(&results).expect("Failed to construct json");
		println!("{}", json_result);
	}
}