// Copyright (c) The Dijets Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

use crate::{
    cluster::Cluster,
    experiments::{
        compatibility_test::update_batch_instance, Context, Experiment, ExperimentParam,
    },
    instance,
    instance::Instance,
    tx_emitter::{execute_and_wait_transactions, EmitJobRequest},
};
use anyhow::format_err;
use async_trait::async_trait;
use dijets_logger::prelude::*;
use dijets_sdk::{transaction_builder::TransactionFactory, types::LocalAccount};
use dijets_transaction_builder::stdlib::encode_update_dijets_version_script;
use dijets_types::{chain_id::ChainId, transaction::TransactionPayload};
use language_e2e_tests::common_transactions::multi_agent_p2p_script_function;
use std::{collections::HashSet, fmt, time::Duration};
use structopt::StructOpt;

#[derive(StructOpt, Debug)]
pub struct ValidatorVersioningParams {
    #[structopt(
        long,
        default_value = "15",
        help = "Number of nodes to update in the first batch"
    )]
    pub count: usize,
    #[structopt(long, help = "Image tag of newer validator software")]
    pub updated_image_tag: String,
}

pub struct ValidatorVersioning {
    first_batch: Vec<Instance>,
    first_batch_lsr: Vec<Instance>,
    second_batch: Vec<Instance>,
    second_batch_lsr: Vec<Instance>,
    _full_nodes: Vec<Instance>,
    updated_image_tag: String,
}

impl ExperimentParam for ValidatorVersioningParams {
    type E = ValidatorVersioning;
    fn build(self, cluster: &Cluster) -> Self::E {
        if self.count > cluster.validator_instances().len() {
            panic!(
                "Can not reboot {} validators in cluster with {} instances",
                self.count,
                cluster.validator_instances().len()
            );
        }
        let (first_batch, second_batch) = cluster.split_n_validators_random(self.count);
        let first_batch = first_batch.into_validator_instances();
        let second_batch = second_batch.into_validator_instances();
        let mut first_batch_lsr = vec![];
        let mut second_batch_lsr = vec![];
        if !cluster.lsr_instances().is_empty() {
            first_batch_lsr = cluster.lsr_instances_for_validators(&first_batch);
            second_batch_lsr = cluster.lsr_instances_for_validators(&second_batch);
        }

        Self::E {
            first_batch,
            first_batch_lsr,
            second_batch,
            second_batch_lsr,
            _full_nodes: cluster.fullnode_instances().to_vec(),
            updated_image_tag: self.updated_image_tag,
        }
    }
}

#[async_trait]
impl Experiment for ValidatorVersioning {
    fn affected_validators(&self) -> HashSet<String> {
        instance::instancelist_to_set(&self.first_batch)
            .union(&instance::instancelist_to_set(&self.second_batch))
            .cloned()
            .collect()
    }

    async fn run(&mut self, context: &mut Context<'_>) -> anyhow::Result<()> {
        // Mint a number of accounts
        context
            .tx_emitter
            .mint_accounts(
                &EmitJobRequest::for_instances(
                    context.cluster.validator_instances().to_vec(),
                    context.global_emit_job_request,
                    0,
                    0,
                ),
                150,
            )
            .await?;
        let mut account = context.tx_emitter.take_account();
        let secondary_signer_account = context.tx_emitter.take_account();

        // Define the transaction generator
        //
        // TODO: In the future we may want to pass this functor as an argument to the experiment
        // to make versioning test extensible.
        // Define a multi-agent p2p transaction.
        let txn_payload = multi_agent_p2p_script_function(10);

        let tx_factory =
            TransactionFactory::new(ChainId::test()).with_transaction_expiration_time(420);
        let txn_gen = |account: &mut LocalAccount, secondary_signer_account: &LocalAccount| {
            account.sign_multi_agent_with_transaction_builder(
                vec![secondary_signer_account],
                tx_factory.payload(txn_payload.clone()),
            )
        };

        // grab a validator node
        let old_validator_node = context.cluster.random_validator_instance();
        let mut old_client = old_validator_node.json_rpc_client();

        info!("1. Send a transaction using the new feature to a validator node");
        let txn1 = txn_gen(&mut account, &secondary_signer_account);
        if execute_and_wait_transactions(&mut old_client, &mut account, vec![txn1])
            .await
            .is_ok()
        {
            return Err(format_err!(
                "The transaction should be rejected as the new feature is not yet recognized \
                by any of the validator nodes"
            ));
        };
        info!("-- [Expected] The transaction is rejected by the validator node");

        info!("2. Update the first batch of validator nodes");
        update_batch_instance(
            context,
            &self.first_batch,
            &self.first_batch_lsr,
            self.updated_image_tag.clone(),
        )
        .await?;

        // choose an updated validator
        let new_validator_node = self
            .first_batch
            .get(0)
            .expect("getting an updated validator instance requires a non-empty list");
        let mut new_client = new_validator_node.json_rpc_client();

        info!("3. Send the transaction using the new feature to an updated validator node");
        let txn3 = txn_gen(&mut account, &secondary_signer_account);
        if execute_and_wait_transactions(&mut new_client, &mut account, vec![txn3])
            .await
            .is_ok()
        {
            return Err(format_err!(
                "The transaction should be rejected as the feature is under gating",
            ));
        }
        info!("-- The transaction is rejected as expected");

        info!("4. Update the rest of the validator nodes");
        update_batch_instance(
            context,
            &self.second_batch,
            &self.second_batch_lsr,
            self.updated_image_tag.clone(),
        )
        .await?;

        info!("5. Send the transaction using the new feature to an updated validator node again");
        let txn4 = txn_gen(&mut account, &secondary_signer_account);
        if execute_and_wait_transactions(&mut new_client, &mut account, vec![txn4])
            .await
            .is_ok()
        {
            return Err(format_err!(
                "The transaction should be rejected as the feature is still gated",
            ));
        }
        info!("-- The transaction is still rejected as expected, because the new feature is gated");

        info!("6. Activate the new feature multi agent");
        let mut dijets_root_account = context
            .tx_emitter
            .load_dijets_root_account(&new_client)
            .await?;
        let allowed_nonce = 0;
        let update_txn = dijets_root_account.sign_with_transaction_builder(tx_factory.payload(
            TransactionPayload::Script(encode_update_dijets_version_script(allowed_nonce, 3)),
        ));
        execute_and_wait_transactions(&mut new_client, &mut dijets_root_account, vec![update_txn])
            .await?;

        info!("7. Send the transaction using the new feature after Dijets version update");
        let txn5 = txn_gen(&mut account, &secondary_signer_account);
        execute_and_wait_transactions(&mut new_client, &mut account, vec![txn5]).await?;
        info!("-- [Expected] The transaction goes through");

        Ok(())
    }

    fn deadline(&self) -> Duration {
        Duration::from_secs(15 * 60)
    }
}

impl fmt::Display for ValidatorVersioning {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Updating [")?;
        for instance in self.first_batch.iter() {
            write!(f, "{}, ", instance)?;
        }
        for instance in self.second_batch.iter() {
            write!(f, "{}, ", instance)?;
        }
        write!(f, "]")?;
        writeln!(f, "Updated Config: {:?}", self.updated_image_tag)
    }
}
