use criterion::Criterion;
use criterion::{black_box, Benchmark};
use log::debug;
use ssz::TreeHash;
use state_processing::{
    per_block_processing,
    per_block_processing::{
        process_attestations, process_attester_slashings, process_deposits, process_eth1_data,
        process_exits, process_proposer_slashings, process_randao, process_transfers,
        verify_block_signature,
    },
};
use types::test_utils::{TestingBeaconBlockBuilder, TestingBeaconStateBuilder};
use types::*;

/// Run the benchmarking suite on a foundation spec with 16,384 validators.
pub fn bench_block_processing_n_validators(c: &mut Criterion, validator_count: usize) {
    let spec = ChainSpec::foundation();

    let bench_builder = BlockBenchingBuilder::new(validator_count, &spec);

    let (mut state, keypairs) = build_state(validator_count, &spec);
    let block = build_block(&mut state, &keypairs, &spec);

    assert_eq!(
        block.body.proposer_slashings.len(),
        spec.max_proposer_slashings as usize,
        "The block should have the maximum possible proposer slashings"
    );

    assert_eq!(
        block.body.attester_slashings.len(),
        spec.max_attester_slashings as usize,
        "The block should have the maximum possible attester slashings"
    );

    for attester_slashing in &block.body.attester_slashings {
        let len_1 = attester_slashing
            .slashable_attestation_1
            .validator_indices
            .len();
        let len_2 = attester_slashing
            .slashable_attestation_1
            .validator_indices
            .len();
        assert!(
            (len_1 == len_2) && (len_2 == spec.max_indices_per_slashable_vote as usize),
            "Each attester slashing should have the maximum possible validator indices"
        );
    }

    assert_eq!(
        block.body.attestations.len(),
        spec.max_attestations as usize,
        "The block should have the maximum possible attestations."
    );

    assert_eq!(
        block.body.deposits.len(),
        spec.max_deposits as usize,
        "The block should have the maximum possible deposits."
    );

    assert_eq!(
        block.body.voluntary_exits.len(),
        spec.max_voluntary_exits as usize,
        "The block should have the maximum possible voluntary exits."
    );

    assert_eq!(
        block.body.transfers.len(),
        spec.max_transfers as usize,
        "The block should have the maximum possible transfers."
    );

    bench_block_processing(
        c,
        &block,
        &state,
        &spec,
        &format!("{}_validators", validator_count),
    );
}

fn build_state(validator_count: usize, spec: &ChainSpec) -> (BeaconState, Vec<Keypair>) {
    let mut builder =
        TestingBeaconStateBuilder::from_default_keypairs_file_if_exists(validator_count, &spec);

    // Set the state to be just before an epoch transition.
    let target_slot = (spec.genesis_epoch + 4).end_slot(spec.slots_per_epoch);
    builder.teleport_to_slot(target_slot, &spec);

    // Builds all caches; benches will not contain shuffling/committee building times.
    builder.build_caches(&spec).unwrap();

    builder.build()
}

pub struct BlockBenchingBuilder {
    pub state_builder: TestingBeaconStateBuilder,
    pub block_builder: TestingBeaconBlockBuilder,

    pub num_validators: usize,
    pub num_proposer_slashings: usize,
    pub num_attester_slashings: usize,
    pub num_indices_per_slashable_vote: usize,
    pub num_attestations: usize,
    pub num_deposits: usize,
    pub num_exits: usize,
    pub num_transfers: usize,
}

impl BlockBenchingBuilder {
    pub fn new(num_validators: usize, spec: &ChainSpec) -> Self {
        let mut state_builder =
            TestingBeaconStateBuilder::from_default_keypairs_file_if_exists(num_validators, &spec);
        let mut block_builder = TestingBeaconBlockBuilder::new(spec);

        Self {
            state_builder,
            block_builder,
            num_validators: 0,
            num_proposer_slashings: 0,
            num_attester_slashings: 0,
            num_indices_per_slashable_vote: spec.max_indices_per_slashable_vote as usize,
            num_attestations: 0,
            num_deposits: 0,
            num_exits: 0,
            num_transfers: 0
        }
    }

    pub fn maximize_block_operations(&mut self, spec: &ChainSpec) {
        self.num_proposer_slashings = spec.max_proposer_slashings as usize;
        self.num_attester_slashings = spec.max_attester_slashings as usize;
        self.num_indices_per_slashable_vote = spec.max_indices_per_slashable_vote as usize;
        self.num_attestations = spec.max_attestations as usize;
        self.num_deposits = spec.max_deposits as usize;
        self.num_exits = spec.max_voluntary_exits as usize;
        self.num_transfers = spec.max_transfers as usize;
    }

    pub fn set_slot(&mut self, slot: Slot, spec: &ChainSpec) {
        self.state_builder.teleport_to_slot(slot, &spec);
    }

    pub fn build_caches(&mut self, spec: &ChainSpec) {
        // Builds all caches; benches will not contain shuffling/committee building times.
        self.state_builder.build_caches(&spec).unwrap();
    }

    pub fn build(self, spec: &ChainSpec) -> (BeaconBlock, BeaconState) {
        let (mut state, keypairs) = self.state_builder.build();
        let builder = &mut self.block_builder;

        builder.set_slot(state.slot);

        let proposer_index = state.get_beacon_proposer_index(state.slot, spec).unwrap();
        let keypair = &keypairs[proposer_index];

        builder.set_randao_reveal(&keypair.sk, &state.fork, spec);

        // Used as a stream of validator indices for use in slashings, exits, etc.
        let mut validators_iter = (0..keypairs.len() as u64).into_iter();

        // Insert `ProposerSlashing` objects.
        debug!(
            "Inserting {} proposer slashings...",
            self.num_proposer_slashings
        );
        for _ in 0..self.num_proposer_slashings {
            let validator_index = validators_iter.next().expect("Insufficient validators.");

            builder.insert_proposer_slashing(
                validator_index,
                &keypairs[validator_index as usize].sk,
                &state.fork,
                spec,
            );
        }

        // Insert `AttesterSlashing` objects
        debug!(
            "Inserting {} attester slashings...",
            self.num_attester_slashings
        );
        for _ in 0..self.num_attester_slashings {
            let mut attesters: Vec<u64> = vec![];
            let mut secret_keys: Vec<&SecretKey> = vec![];

            for _ in 0..self.num_indices_per_slashable_vote {
                let validator_index = validators_iter.next().expect("Insufficient validators.");

                attesters.push(validator_index);
                secret_keys.push(&keypairs[validator_index as usize].sk);
            }

            builder.insert_attester_slashing(&attesters, &secret_keys, &state.fork, spec);
        }

        // Insert `Attestation` objects.
        debug!("Inserting {} attestations...", self.num_attestations);
        let all_secret_keys: Vec<&SecretKey> = keypairs.iter().map(|keypair| &keypair.sk).collect();
        builder
            .insert_attestations(&state, &all_secret_keys, self.num_attestations as usize, spec)
            .unwrap();

        // Insert `Deposit` objects.
        debug!("Inserting {} deposits...", self.num_deposits);
        for i in 0..self.num_deposits {
            builder.insert_deposit(32_000_000_000, state.deposit_index + (i as u64), &state, spec);
        }

        // Insert the maximum possible number of `Exit` objects.
        debug!("Inserting {} exits...", self.num_exits);
        for _ in 0..self.num_exits {
            let validator_index = validators_iter.next().expect("Insufficient validators.");

            builder.insert_exit(
                &state,
                validator_index,
                &keypairs[validator_index as usize].sk,
                spec,
            );
        }

        // Insert the maximum possible number of `Transfer` objects.
        debug!("Inserting {} transfers...", self.num_transfers);
        for _ in 0..self.num_transfers {
            let validator_index = validators_iter.next().expect("Insufficient validators.");

            // Manually set the validator to be withdrawn.
            state.validator_registry[validator_index as usize].withdrawable_epoch =
                state.previous_epoch(spec);

            builder.insert_transfer(
                &state,
                validator_index,
                validator_index,
                1,
                keypairs[validator_index as usize].clone(),
                spec,
            );
        }

        let mut block = builder.build(&keypair.sk, &state.fork, spec);

        // Set the eth1 data to be different from the state.
        block.eth1_data.block_hash = Hash256::from_slice(&vec![42; 32]);

        (block, state)
    }
}

/// Run the detailed benchmarking suite on the given `BeaconState`.
///
/// `desc` will be added to the title of each bench.
fn bench_block_processing(
    c: &mut Criterion,
    initial_block: &BeaconBlock,
    initial_state: &BeaconState,
    initial_spec: &ChainSpec,
    desc: &str,
) {
    let state = initial_state.clone();
    let block = initial_block.clone();
    let spec = initial_spec.clone();
    c.bench(
        &format!("{}/block_processing", desc),
        Benchmark::new("verify_block_signature", move |b| {
            b.iter_batched(
                || state.clone(),
                |mut state| {
                    verify_block_signature(&mut state, &block, &spec).unwrap();
                    state
                },
                criterion::BatchSize::SmallInput,
            )
        })
        .sample_size(10),
    );

    let state = initial_state.clone();
    let block = initial_block.clone();
    let spec = initial_spec.clone();
    c.bench(
        &format!("{}/block_processing", desc),
        Benchmark::new("process_randao", move |b| {
            b.iter_batched(
                || state.clone(),
                |mut state| {
                    process_randao(&mut state, &block, &spec).unwrap();
                    state
                },
                criterion::BatchSize::SmallInput,
            )
        })
        .sample_size(10),
    );

    let state = initial_state.clone();
    let block = initial_block.clone();
    c.bench(
        &format!("{}/block_processing", desc),
        Benchmark::new("process_eth1_data", move |b| {
            b.iter_batched(
                || state.clone(),
                |mut state| {
                    process_eth1_data(&mut state, &block.eth1_data).unwrap();
                    state
                },
                criterion::BatchSize::SmallInput,
            )
        })
        .sample_size(10),
    );

    let state = initial_state.clone();
    let block = initial_block.clone();
    let spec = initial_spec.clone();
    c.bench(
        &format!("{}/block_processing", desc),
        Benchmark::new("process_proposer_slashings", move |b| {
            b.iter_batched(
                || state.clone(),
                |mut state| {
                    process_proposer_slashings(&mut state, &block.body.proposer_slashings, &spec)
                        .unwrap();
                    state
                },
                criterion::BatchSize::SmallInput,
            )
        })
        .sample_size(10),
    );

    let state = initial_state.clone();
    let block = initial_block.clone();
    let spec = initial_spec.clone();
    c.bench(
        &format!("{}/block_processing", desc),
        Benchmark::new("process_attester_slashings", move |b| {
            b.iter_batched(
                || state.clone(),
                |mut state| {
                    process_attester_slashings(&mut state, &block.body.attester_slashings, &spec)
                        .unwrap();
                    state
                },
                criterion::BatchSize::SmallInput,
            )
        })
        .sample_size(10),
    );

    let state = initial_state.clone();
    let block = initial_block.clone();
    let spec = initial_spec.clone();
    c.bench(
        &format!("{}/block_processing", desc),
        Benchmark::new("process_attestations", move |b| {
            b.iter_batched(
                || state.clone(),
                |mut state| {
                    process_attestations(&mut state, &block.body.attestations, &spec).unwrap();
                    state
                },
                criterion::BatchSize::SmallInput,
            )
        })
        .sample_size(10),
    );

    let state = initial_state.clone();
    let block = initial_block.clone();
    let spec = initial_spec.clone();
    c.bench(
        &format!("{}/block_processing", desc),
        Benchmark::new("process_deposits", move |b| {
            b.iter_batched(
                || state.clone(),
                |mut state| {
                    process_deposits(&mut state, &block.body.deposits, &spec).unwrap();
                    state
                },
                criterion::BatchSize::SmallInput,
            )
        })
        .sample_size(10),
    );

    let state = initial_state.clone();
    let block = initial_block.clone();
    let spec = initial_spec.clone();
    c.bench(
        &format!("{}/block_processing", desc),
        Benchmark::new("process_exits", move |b| {
            b.iter_batched(
                || state.clone(),
                |mut state| {
                    process_exits(&mut state, &block.body.voluntary_exits, &spec).unwrap();
                    state
                },
                criterion::BatchSize::SmallInput,
            )
        })
        .sample_size(10),
    );

    let state = initial_state.clone();
    let block = initial_block.clone();
    let spec = initial_spec.clone();
    c.bench(
        &format!("{}/block_processing", desc),
        Benchmark::new("process_transfers", move |b| {
            b.iter_batched(
                || state.clone(),
                |mut state| {
                    process_transfers(&mut state, &block.body.transfers, &spec).unwrap();
                    state
                },
                criterion::BatchSize::SmallInput,
            )
        })
        .sample_size(10),
    );

    let state = initial_state.clone();
    let block = initial_block.clone();
    let spec = initial_spec.clone();
    c.bench(
        &format!("{}/block_processing", desc),
        Benchmark::new("per_block_processing", move |b| {
            b.iter_batched(
                || state.clone(),
                |mut state| {
                    per_block_processing(&mut state, &block, &spec).unwrap();
                    state
                },
                criterion::BatchSize::SmallInput,
            )
        })
        .sample_size(10),
    );

    let mut state = initial_state.clone();
    state.drop_cache(RelativeEpoch::Previous);
    let spec = initial_spec.clone();
    c.bench(
        &format!("{}/block_processing", desc),
        Benchmark::new("build_previous_state_epoch_cache", move |b| {
            b.iter_batched(
                || state.clone(),
                |mut state| {
                    state
                        .build_epoch_cache(RelativeEpoch::Previous, &spec)
                        .unwrap();
                    state
                },
                criterion::BatchSize::SmallInput,
            )
        })
        .sample_size(10),
    );

    let mut state = initial_state.clone();
    state.drop_cache(RelativeEpoch::Current);
    let spec = initial_spec.clone();
    c.bench(
        &format!("{}/block_processing", desc),
        Benchmark::new("build_current_state_epoch_cache", move |b| {
            b.iter_batched(
                || state.clone(),
                |mut state| {
                    state
                        .build_epoch_cache(RelativeEpoch::Current, &spec)
                        .unwrap();
                    state
                },
                criterion::BatchSize::SmallInput,
            )
        })
        .sample_size(10),
    );

    let mut state = initial_state.clone();
    state.drop_pubkey_cache();
    c.bench(
        &format!("{}/block_processing", desc),
        Benchmark::new("build_pubkey_cache", move |b| {
            b.iter_batched(
                || state.clone(),
                |mut state| {
                    state.update_pubkey_cache().unwrap();
                    state
                },
                criterion::BatchSize::SmallInput,
            )
        })
        .sample_size(10),
    );

    let block = initial_block.clone();
    c.bench(
        &format!("{}/block_processing", desc),
        Benchmark::new("tree_hash_block", move |b| {
            b.iter(|| black_box(block.hash_tree_root()))
        })
        .sample_size(10),
    );
}