use {
    serde_json::{json, Value},
    serial_test::serial,
    solana_client::{
        pubsub_client::PubsubClient,
        rpc_client::RpcClient,
        rpc_config::{
            RpcAccountInfoConfig, RpcBlockSubscribeConfig, RpcBlockSubscribeFilter,
            RpcProgramAccountsConfig,
        },
        rpc_response::{RpcBlockUpdate, SlotInfo},
    },
    solana_ledger::{blockstore::Blockstore, get_tmp_ledger_path},
    solana_rpc::{
        optimistically_confirmed_bank_tracker::OptimisticallyConfirmedBank,
        rpc::create_test_transactions_and_populate_blockstore,
        rpc_pubsub_service::{PubSubConfig, PubSubService},
        rpc_subscriptions::RpcSubscriptions,
    },
    solana_runtime::{
        bank::Bank,
        bank_forks::BankForks,
        commitment::{BlockCommitmentCache, CommitmentSlots},
        genesis_utils::{create_genesis_config, GenesisConfigInfo},
    },
    solana_sdk::{
        clock::Slot,
        commitment_config::{CommitmentConfig, CommitmentLevel},
        native_token::sol_to_lamports,
        pubkey::Pubkey,
        rpc_port,
        signature::{Keypair, Signer},
        system_program, system_transaction,
    },
    solana_streamer::socket::SocketAddrSpace,
    solana_test_validator::TestValidator,
    solana_transaction_status::{TransactionDetails, UiTransactionEncoding},
    std::{
        collections::HashSet,
        net::{IpAddr, SocketAddr},
        sync::{
            atomic::{AtomicBool, AtomicU64, Ordering},
            Arc, RwLock,
        },
        thread::sleep,
        time::{Duration, Instant},
    },
    systemstat::Ipv4Addr,
};

#[test]
fn test_rpc_client() {
    solana_logger::setup();

    let alice = Keypair::new();
    let test_validator =
        TestValidator::with_no_fees(alice.pubkey(), None, SocketAddrSpace::Unspecified);

    let bob_pubkey = solana_sdk::pubkey::new_rand();

    let client = RpcClient::new(test_validator.rpc_url());

    assert_eq!(
        client.get_version().unwrap().solana_core,
        solana_version::semver!()
    );

    assert!(client.get_account(&bob_pubkey).is_err());

    assert_eq!(client.get_balance(&bob_pubkey).unwrap(), 0);

    let original_alice_balance = client.get_balance(&alice.pubkey()).unwrap();

    let blockhash = client.get_latest_blockhash().unwrap();

    let tx = system_transaction::transfer(&alice, &bob_pubkey, sol_to_lamports(20.0), blockhash);
    let signature = client.send_transaction(&tx).unwrap();

    let mut confirmed_tx = false;

    let now = Instant::now();
    while now.elapsed().as_secs() <= 20 {
        let response = client
            .confirm_transaction_with_commitment(&signature, CommitmentConfig::default())
            .unwrap();

        if response.value {
            confirmed_tx = true;
            break;
        }

        sleep(Duration::from_millis(500));
    }

    assert!(confirmed_tx);

    assert_eq!(
        client.get_balance(&bob_pubkey).unwrap(),
        sol_to_lamports(20.0)
    );
    assert_eq!(
        client.get_balance(&alice.pubkey()).unwrap(),
        original_alice_balance - sol_to_lamports(20.0)
    );
}

#[test]
#[serial]
fn test_account_subscription() {
    let pubsub_addr = SocketAddr::new(
        IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)),
        rpc_port::DEFAULT_RPC_PUBSUB_PORT,
    );
    let exit = Arc::new(AtomicBool::new(false));

    let GenesisConfigInfo {
        genesis_config,
        mint_keypair: alice,
        ..
    } = create_genesis_config(10_000);
    let bank = Bank::new_for_tests(&genesis_config);
    let blockhash = bank.last_blockhash();
    let bank_forks = Arc::new(RwLock::new(BankForks::new(bank)));
    let bank0 = bank_forks.read().unwrap().get(0).unwrap().clone();
    let bank1 = Bank::new_from_parent(&bank0, &Pubkey::default(), 1);
    bank_forks.write().unwrap().insert(bank1);
    let bob = Keypair::new();
    let max_complete_transaction_status_slot = Arc::new(AtomicU64::default());
    let subscriptions = Arc::new(RpcSubscriptions::new_for_tests(
        &exit,
        max_complete_transaction_status_slot,
        bank_forks.clone(),
        Arc::new(RwLock::new(BlockCommitmentCache::default())),
        OptimisticallyConfirmedBank::locked_from_bank_forks_root(&bank_forks),
    ));
    let (trigger, pubsub_service) =
        PubSubService::new(PubSubConfig::default(), &subscriptions, pubsub_addr);
    std::thread::sleep(Duration::from_millis(400));
    let config = Some(RpcAccountInfoConfig {
        commitment: Some(CommitmentConfig::finalized()),
        encoding: None,
        data_slice: None,
    });
    let (mut client, receiver) = PubsubClient::account_subscribe(
        &format!("ws://0.0.0.0:{}/", pubsub_addr.port()),
        &bob.pubkey(),
        config,
    )
    .unwrap();

    // Transfer 100 lamports from alice to bob
    let tx = system_transaction::transfer(&alice, &bob.pubkey(), 100, blockhash);
    bank_forks
        .write()
        .unwrap()
        .get(1)
        .unwrap()
        .process_transaction(&tx)
        .unwrap();
    let commitment_slots = CommitmentSlots {
        slot: 1,
        ..CommitmentSlots::default()
    };
    subscriptions.notify_subscribers(commitment_slots);
    let commitment_slots = CommitmentSlots {
        slot: 2,
        root: 1,
        highest_confirmed_slot: 1,
        highest_confirmed_root: 1,
    };
    subscriptions.notify_subscribers(commitment_slots);

    let expected = json!({
    "context": { "slot": 1 },
        "value": {
            "owner": system_program::id().to_string(),
            "lamports": 100,
            "data": "",
            "executable": false,
            "rentEpoch": 0,
        },
    });

    // Read notification
    let mut errors: Vec<(Value, Value)> = Vec::new();
    let response = receiver.recv();
    match response {
        Ok(response) => {
            let actual = serde_json::to_value(response).unwrap();
            if expected != actual {
                errors.push((expected, actual));
            }
        }
        Err(_) => eprintln!("unexpected websocket receive timeout"),
    }

    exit.store(true, Ordering::Relaxed);
    trigger.cancel();
    client.shutdown().unwrap();
    pubsub_service.close().unwrap();
    assert_eq!(errors, [].to_vec());
}

#[test]
#[serial]
fn test_block_subscription() {
    // setup BankForks
    let exit = Arc::new(AtomicBool::new(false));
    let GenesisConfigInfo {
        genesis_config,
        mint_keypair: alice,
        ..
    } = create_genesis_config(10_000);
    let bank = Bank::new_for_tests(&genesis_config);
    let bank_forks = Arc::new(RwLock::new(BankForks::new(bank)));

    // setup Blockstore
    let ledger_path = get_tmp_ledger_path!();
    let blockstore = Blockstore::open(&ledger_path).unwrap();
    let blockstore = Arc::new(blockstore);

    // populate ledger with test txs
    let bank = bank_forks.read().unwrap().working_bank();
    let keypair1 = Keypair::new();
    let keypair2 = Keypair::new();
    let keypair3 = Keypair::new();
    let max_complete_transaction_status_slot = Arc::new(AtomicU64::new(blockstore.max_root()));
    let _confirmed_block_signatures = create_test_transactions_and_populate_blockstore(
        vec![&alice, &keypair1, &keypair2, &keypair3],
        0,
        bank,
        blockstore.clone(),
        max_complete_transaction_status_slot,
    );
    let max_complete_transaction_status_slot = Arc::new(AtomicU64::default());
    // setup RpcSubscriptions && PubSubService
    let subscriptions = Arc::new(RpcSubscriptions::new_for_tests_with_blockstore(
        &exit,
        max_complete_transaction_status_slot,
        blockstore.clone(),
        bank_forks.clone(),
        Arc::new(RwLock::new(BlockCommitmentCache::default())),
        OptimisticallyConfirmedBank::locked_from_bank_forks_root(&bank_forks),
    ));
    let pubsub_addr = SocketAddr::new(
        IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)),
        rpc_port::DEFAULT_RPC_PUBSUB_PORT,
    );
    let pub_cfg = PubSubConfig {
        enable_block_subscription: true,
        ..PubSubConfig::default()
    };
    let (trigger, pubsub_service) = PubSubService::new(pub_cfg, &subscriptions, pubsub_addr);

    std::thread::sleep(Duration::from_millis(400));

    // setup PubsubClient
    let (mut client, receiver) = PubsubClient::block_subscribe(
        &format!("ws://0.0.0.0:{}/", pubsub_addr.port()),
        RpcBlockSubscribeFilter::All,
        Some(RpcBlockSubscribeConfig {
            commitment: Some(CommitmentConfig {
                commitment: CommitmentLevel::Confirmed,
            }),
            encoding: Some(UiTransactionEncoding::Json),
            transaction_details: Some(TransactionDetails::Signatures),
            show_rewards: None,
        }),
    )
    .unwrap();

    // trigger Gossip notification
    let slot = bank_forks.read().unwrap().highest_slot();
    subscriptions.notify_gossip_subscribers(slot);
    let maybe_actual = receiver.recv_timeout(Duration::from_millis(400));
    match maybe_actual {
        Ok(actual) => {
            let complete_block = blockstore.get_complete_block(slot, false).unwrap();
            let block = complete_block.clone().configure(
                UiTransactionEncoding::Json,
                TransactionDetails::Signatures,
                false,
            );
            let expected = RpcBlockUpdate {
                slot,
                block: Some(block),
                err: None,
            };
            let block = complete_block.configure(
                UiTransactionEncoding::Json,
                TransactionDetails::Signatures,
                false,
            );
            assert_eq!(actual.value.slot, expected.slot);
            assert!(block.eq(&actual.value.block.unwrap()));
        }
        Err(e) => {
            eprintln!("unexpected websocket receive timeout");
            assert_eq!(Some(e), None);
        }
    }

    // cleanup
    exit.store(true, Ordering::Relaxed);
    trigger.cancel();
    client.shutdown().unwrap();
    pubsub_service.close().unwrap();
}

#[test]
#[serial]
fn test_program_subscription() {
    let pubsub_addr = SocketAddr::new(
        IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)),
        rpc_port::DEFAULT_RPC_PUBSUB_PORT,
    );
    let exit = Arc::new(AtomicBool::new(false));

    let GenesisConfigInfo {
        genesis_config,
        mint_keypair: alice,
        ..
    } = create_genesis_config(10_000);
    let bank = Bank::new_for_tests(&genesis_config);
    let blockhash = bank.last_blockhash();
    let bank_forks = Arc::new(RwLock::new(BankForks::new(bank)));
    let bank0 = bank_forks.read().unwrap().get(0).unwrap().clone();
    let bank1 = Bank::new_from_parent(&bank0, &Pubkey::default(), 1);
    bank_forks.write().unwrap().insert(bank1);
    let bob = Keypair::new();
    let max_complete_transaction_status_slot = Arc::new(AtomicU64::default());
    let subscriptions = Arc::new(RpcSubscriptions::new_for_tests(
        &exit,
        max_complete_transaction_status_slot,
        bank_forks.clone(),
        Arc::new(RwLock::new(BlockCommitmentCache::default())),
        OptimisticallyConfirmedBank::locked_from_bank_forks_root(&bank_forks),
    ));
    let (trigger, pubsub_service) =
        PubSubService::new(PubSubConfig::default(), &subscriptions, pubsub_addr);
    std::thread::sleep(Duration::from_millis(400));
    let config = Some(RpcProgramAccountsConfig {
        ..RpcProgramAccountsConfig::default()
    });

    let program_id = Pubkey::new_unique();
    let (mut client, receiver) = PubsubClient::program_subscribe(
        &format!("ws://0.0.0.0:{}/", pubsub_addr.port()),
        &program_id,
        config,
    )
    .unwrap();

    // Create new program account at bob's address
    let tx = system_transaction::create_account(&alice, &bob, blockhash, 100, 0, &program_id);
    bank_forks
        .write()
        .unwrap()
        .get(1)
        .unwrap()
        .process_transaction(&tx)
        .unwrap();
    let commitment_slots = CommitmentSlots {
        slot: 1,
        ..CommitmentSlots::default()
    };
    subscriptions.notify_subscribers(commitment_slots);
    let commitment_slots = CommitmentSlots {
        slot: 2,
        root: 1,
        highest_confirmed_slot: 1,
        highest_confirmed_root: 1,
    };
    subscriptions.notify_subscribers(commitment_slots);

    // Poll notifications generated by the transfer
    let mut notifications = Vec::new();
    let mut pubkeys = HashSet::new();
    loop {
        let response = receiver.recv_timeout(Duration::from_millis(100));
        match response {
            Ok(response) => {
                notifications.push(response.clone());
                pubkeys.insert(response.value.pubkey);
            }
            Err(_) => {
                break;
            }
        }
    }

    // Shutdown
    exit.store(true, Ordering::Relaxed);
    trigger.cancel();
    client.shutdown().unwrap();
    pubsub_service.close().unwrap();

    assert_eq!(notifications.len(), 1);
    assert!(pubkeys.contains(&bob.pubkey().to_string()));
}

#[test]
#[serial]
fn test_root_subscription() {
    let pubsub_addr = SocketAddr::new(
        IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)),
        rpc_port::DEFAULT_RPC_PUBSUB_PORT,
    );
    let exit = Arc::new(AtomicBool::new(false));

    let GenesisConfigInfo { genesis_config, .. } = create_genesis_config(10_000);
    let bank = Bank::new_for_tests(&genesis_config);
    let bank_forks = Arc::new(RwLock::new(BankForks::new(bank)));
    let bank0 = bank_forks.read().unwrap().get(0).unwrap().clone();
    let bank1 = Bank::new_from_parent(&bank0, &Pubkey::default(), 1);
    bank_forks.write().unwrap().insert(bank1);
    let max_complete_transaction_status_slot = Arc::new(AtomicU64::default());
    let subscriptions = Arc::new(RpcSubscriptions::new_for_tests(
        &exit,
        max_complete_transaction_status_slot,
        bank_forks.clone(),
        Arc::new(RwLock::new(BlockCommitmentCache::default())),
        OptimisticallyConfirmedBank::locked_from_bank_forks_root(&bank_forks),
    ));
    let (trigger, pubsub_service) =
        PubSubService::new(PubSubConfig::default(), &subscriptions, pubsub_addr);
    std::thread::sleep(Duration::from_millis(400));
    let (mut client, receiver) =
        PubsubClient::root_subscribe(&format!("ws://0.0.0.0:{}/", pubsub_addr.port())).unwrap();

    let roots = vec![1, 2, 3];
    subscriptions.notify_roots(roots.clone());

    // Read notifications
    let mut errors: Vec<(Slot, Slot)> = Vec::new();
    for expected in roots {
        let response = receiver.recv();
        match response {
            Ok(response) => {
                if expected != response {
                    errors.push((expected, response));
                }
            }
            Err(_) => eprintln!("unexpected websocket receive timeout"),
        }
    }

    exit.store(true, Ordering::Relaxed);
    trigger.cancel();
    client.shutdown().unwrap();
    pubsub_service.close().unwrap();
    assert_eq!(errors, [].to_vec());
}

#[test]
#[serial]
fn test_slot_subscription() {
    let pubsub_addr = SocketAddr::new(
        IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)),
        rpc_port::DEFAULT_RPC_PUBSUB_PORT,
    );
    let exit = Arc::new(AtomicBool::new(false));
    let GenesisConfigInfo { genesis_config, .. } = create_genesis_config(10_000);
    let bank = Bank::new_for_tests(&genesis_config);
    let bank_forks = Arc::new(RwLock::new(BankForks::new(bank)));
    let optimistically_confirmed_bank =
        OptimisticallyConfirmedBank::locked_from_bank_forks_root(&bank_forks);
    let max_complete_transaction_status_slot = Arc::new(AtomicU64::default());
    let subscriptions = Arc::new(RpcSubscriptions::new_for_tests(
        &exit,
        max_complete_transaction_status_slot,
        bank_forks,
        Arc::new(RwLock::new(BlockCommitmentCache::default())),
        optimistically_confirmed_bank,
    ));
    let (trigger, pubsub_service) =
        PubSubService::new(PubSubConfig::default(), &subscriptions, pubsub_addr);
    std::thread::sleep(Duration::from_millis(400));

    let (mut client, receiver) =
        PubsubClient::slot_subscribe(&format!("ws://0.0.0.0:{}/", pubsub_addr.port())).unwrap();

    let mut errors: Vec<(SlotInfo, SlotInfo)> = Vec::new();

    for i in 0..3 {
        subscriptions.notify_slot(i + 1, i, i);

        let maybe_actual = receiver.recv_timeout(Duration::from_millis(400));

        match maybe_actual {
            Ok(actual) => {
                let expected = SlotInfo {
                    slot: i + 1,
                    parent: i,
                    root: i,
                };

                if actual != expected {
                    errors.push((actual, expected));
                }
            }
            Err(_err) => {
                eprintln!("unexpected websocket receive timeout");
                break;
            }
        }
    }

    exit.store(true, Ordering::Relaxed);
    trigger.cancel();
    client.shutdown().unwrap();
    pubsub_service.close().unwrap();

    assert_eq!(errors, [].to_vec());
}