#![allow(clippy::arithmetic_side_effects)]

use {
    serial_test::serial,
    solana_account::{Account, AccountSharedData},
    solana_bench_tps::{
        bench::{do_bench_tps, generate_and_fund_keypairs},
        cli::{Config, InstructionPaddingConfig},
        send_batch::generate_durable_nonce_accounts,
    },
    solana_commitment_config::CommitmentConfig,
    solana_connection_cache::connection_cache::NewConnectionConfig,
    solana_core::validator::ValidatorConfig,
    solana_faucet::faucet::run_local_faucet_for_tests,
    solana_fee_calculator::FeeRateGovernor,
    solana_keypair::Keypair,
    solana_local_cluster::{
        cluster::Cluster,
        local_cluster::{ClusterConfig, LocalCluster},
        validator_configs::make_identical_validator_configs,
    },
    solana_quic_client::{QuicConfig, QuicConnectionManager},
    solana_rent::Rent,
    solana_rpc::rpc::JsonRpcConfig,
    solana_rpc_client::rpc_client::RpcClient,
    solana_signer::Signer,
    solana_streamer::socket::SocketAddrSpace,
    solana_test_validator::TestValidatorGenesis,
    solana_tpu_client::tpu_client::{TpuClient, TpuClientConfig},
    std::{sync::Arc, time::Duration},
};

fn program_account(program_data: &[u8]) -> AccountSharedData {
    AccountSharedData::from(Account {
        lamports: Rent::default().minimum_balance(program_data.len()).min(1),
        data: program_data.to_vec(),
        owner: solana_sdk_ids::bpf_loader::id(),
        executable: true,
        rent_epoch: 0,
    })
}

fn test_bench_tps_local_cluster(config: Config) {
    let additional_accounts = vec![(
        spl_instruction_padding_interface::ID,
        program_account(include_bytes!("fixtures/spl_instruction_padding.so")),
    )];

    solana_logger::setup();

    let faucet_keypair = Keypair::new();
    let faucet_pubkey = faucet_keypair.pubkey();
    let faucet_addr = run_local_faucet_for_tests(
        faucet_keypair,
        None, /* per_time_cap */
        0,    /* port */
    );

    const NUM_NODES: usize = 1;
    let cluster = LocalCluster::new(
        &mut ClusterConfig {
            node_stakes: vec![999_990; NUM_NODES],
            mint_lamports: 200_000_000,
            validator_configs: make_identical_validator_configs(
                &ValidatorConfig {
                    rpc_config: JsonRpcConfig {
                        faucet_addr: Some(faucet_addr),
                        ..JsonRpcConfig::default_for_test()
                    },
                    ..ValidatorConfig::default_for_test()
                },
                NUM_NODES,
            ),
            additional_accounts,
            ..ClusterConfig::default()
        },
        SocketAddrSpace::Unspecified,
    );

    cluster.transfer(&cluster.funding_keypair, &faucet_pubkey, 100_000_000);

    let client = Arc::new(
        cluster
            .build_validator_tpu_quic_client(cluster.entry_point_info.pubkey())
            .unwrap_or_else(|err| {
                panic!("Could not create TpuClient with Quic Cache {err:?}");
            }),
    );

    let lamports_per_account = 100;

    let keypair_count = config.tx_count * config.keypair_multiplier;
    let keypairs = generate_and_fund_keypairs(
        client.clone(),
        &config.id,
        keypair_count,
        lamports_per_account,
        false,
        false,
    )
    .unwrap();

    let _total = do_bench_tps(client, config, keypairs, None);

    #[cfg(not(debug_assertions))]
    assert!(_total > 100);
}

fn test_bench_tps_test_validator(config: Config) {
    solana_logger::setup();

    let mint_keypair = Keypair::new();
    let mint_pubkey = mint_keypair.pubkey();
    let faucet_addr = run_local_faucet_for_tests(
        mint_keypair,
        None, /* per_time_cap */
        0,    /* port */
    );

    let test_validator = TestValidatorGenesis::default()
        .fee_rate_governor(FeeRateGovernor::new(0, 0))
        .rent(Rent {
            lamports_per_byte_year: 1,
            exemption_threshold: 1.0,
            ..Rent::default()
        })
        .faucet_addr(Some(faucet_addr))
        .add_program(
            "spl_instruction_padding",
            spl_instruction_padding_interface::ID,
        )
        .start_with_mint_address(mint_pubkey, SocketAddrSpace::Unspecified)
        .expect("validator start failed");

    let rpc_client = Arc::new(RpcClient::new_with_commitment(
        test_validator.rpc_url(),
        CommitmentConfig::processed(),
    ));
    let websocket_url = test_validator.rpc_pubsub_url();

    let client = Arc::new(
        TpuClient::new(
            "tpu_client_quic_bench_tps",
            rpc_client,
            &websocket_url,
            TpuClientConfig::default(),
            QuicConnectionManager::new_with_connection_config(QuicConfig::new().unwrap()),
        )
        .expect("Should build Quic Tpu Client."),
    );

    let lamports_per_account = 1000;

    let keypair_count = config.tx_count * config.keypair_multiplier;
    let keypairs = generate_and_fund_keypairs(
        client.clone(),
        &config.id,
        keypair_count,
        lamports_per_account,
        false,
        false,
    )
    .unwrap();
    let nonce_keypairs = if config.use_durable_nonce {
        Some(generate_durable_nonce_accounts(client.clone(), &keypairs))
    } else {
        None
    };

    let _total = do_bench_tps(client, config, keypairs, nonce_keypairs);

    #[cfg(not(debug_assertions))]
    assert!(_total > 100);
}

#[test]
#[serial]
fn test_bench_tps_local_cluster_solana() {
    test_bench_tps_local_cluster(Config {
        tx_count: 100,
        duration: Duration::from_secs(10),
        ..Config::default()
    });
}

#[test]
#[serial]
fn test_bench_tps_tpu_client() {
    test_bench_tps_test_validator(Config {
        tx_count: 100,
        duration: Duration::from_secs(10),
        ..Config::default()
    });
}

#[test]
#[serial]
fn test_bench_tps_tpu_client_nonce() {
    test_bench_tps_test_validator(Config {
        tx_count: 100,
        duration: Duration::from_secs(10),
        use_durable_nonce: true,
        ..Config::default()
    });
}

#[test]
#[serial]
fn test_bench_tps_local_cluster_with_padding() {
    test_bench_tps_local_cluster(Config {
        tx_count: 100,
        duration: Duration::from_secs(10),
        instruction_padding_config: Some(InstructionPaddingConfig {
            program_id: spl_instruction_padding_interface::ID,
            data_size: 0,
        }),
        ..Config::default()
    });
}

#[test]
#[serial]
fn test_bench_tps_tpu_client_with_padding() {
    test_bench_tps_test_validator(Config {
        tx_count: 100,
        duration: Duration::from_secs(10),
        instruction_padding_config: Some(InstructionPaddingConfig {
            program_id: spl_instruction_padding_interface::ID,
            data_size: 0,
        }),
        ..Config::default()
    });
}
