//! Functions for constructing various containers used in the Ethereum consensus layer.
//!
//! Currently only suitable for use in tests.
//! All containers are signed with keys generated by [`interop::secret_key`].

use core::ops::Range;
use std::sync::Arc;

use anyhow::{bail, ensure, Result};
use bls::AggregateSignature;
use deposit_tree::DepositTree;
use helper_functions::{
    accessors, misc,
    signing::{RandaoEpoch, SignForSingleFork as _, SignForSingleForkAtSlot as _},
    slot_report::NullSlotReport,
};
use itertools::{Either, Itertools as _};
use ssz::{BitList, BitVector, ContiguousList, SszHash as _};
use std_ext::ArcExt as _;
use transition_functions::{capella, combined};
use typenum::Unsigned as _;
use types::{
    altair::containers::{
        BeaconBlock as AltairBeaconBlock, BeaconBlockBody as AltairBeaconBlockBody, SyncAggregate,
    },
    bellatrix::containers::{
        BeaconBlock as BellatrixBeaconBlock, BeaconBlockBody as BellatrixBeaconBlockBody,
        ExecutionPayload as BellatrixExecutionPayload,
    },
    capella::containers::{
        BeaconBlock as CapellaBeaconBlock, BeaconBlockBody as CapellaBeaconBlockBody,
        ExecutionPayload as CapellaExecutionPayload,
    },
    combined::{Attestation, BeaconBlock, BeaconState, ExecutionPayload, SignedBeaconBlock},
    config::Config,
    deneb::containers::{
        BeaconBlock as DenebBeaconBlock, BeaconBlockBody as DenebBeaconBlockBody,
        ExecutionPayload as DenebExecutionPayload,
    },
    electra::containers::{
        Attestation as ElectraAttestation, BeaconBlock as ElectraBeaconBlock,
        BeaconBlockBody as ElectraBeaconBlockBody,
    },
    nonstandard::{AttestationEpoch, Phase, RelativeEpoch},
    phase0::{
        consts::GENESIS_SLOT,
        containers::{
            Attestation as Phase0Attestation, AttestationData, BeaconBlock as Phase0BeaconBlock,
            BeaconBlockBody as Phase0BeaconBlockBody, Checkpoint, Deposit, Eth1Data,
        },
        primitives::{Epoch, ExecutionBlockHash, Slot, SubnetId, ValidatorIndex, H256},
    },
    preset::Preset,
    traits::BeaconState as _,
};

type BlockWithState<P> = (Arc<SignedBeaconBlock<P>>, Arc<BeaconState<P>>);

pub fn min_genesis_state<P: Preset>(config: &Config) -> Result<(Arc<BeaconState<P>>, DepositTree)> {
    let (genesis_state, deposit_tree) = interop::quick_start_beacon_state(
        config,
        config.min_genesis_time,
        config.min_genesis_active_validator_count,
    )?;

    Ok((Arc::new(genesis_state), deposit_tree))
}

pub fn empty_block<P: Preset>(
    config: &Config,
    pre_state: Arc<BeaconState<P>>,
    slot: Slot,
    graffiti: H256,
) -> Result<BlockWithState<P>> {
    let advanced_state = advance_state(config, pre_state, slot)?;
    let eth1_data = advanced_state.eth1_data();
    let attestations = core::iter::empty();
    let deposits = ContiguousList::default();
    let sync_aggregate = SyncAggregate::empty();
    let execution_payload = None;

    block(
        config,
        advanced_state,
        eth1_data,
        graffiti,
        attestations,
        deposits,
        sync_aggregate,
        execution_payload,
    )
}

pub fn block_justifying_previous_epoch<P: Preset>(
    config: &Config,
    pre_state: Arc<BeaconState<P>>,
    epoch: Epoch,
    graffiti: H256,
) -> Result<BlockWithState<P>> {
    let block_slot = misc::compute_start_slot_at_epoch::<P>(epoch);
    let advanced_state = advance_state(config, pre_state, block_slot)?;
    let eth1_data = advanced_state.eth1_data();
    let attestation_slots = misc::slots_in_epoch::<P>(epoch - 1);
    let attestations = full_block_attestations(config, &advanced_state, attestation_slots)?;
    let deposits = ContiguousList::default();
    let sync_aggregate = SyncAggregate::empty();
    let execution_payload = None;

    block(
        config,
        advanced_state,
        eth1_data,
        graffiti,
        attestations,
        deposits,
        sync_aggregate,
        execution_payload,
    )
}

pub fn block_justifying_current_epoch<P: Preset>(
    config: &Config,
    pre_state: Arc<BeaconState<P>>,
    epoch: Epoch,
    graffiti: H256,
    execution_payload: Option<ExecutionPayload<P>>,
) -> Result<BlockWithState<P>> {
    let block_slot = misc::compute_start_slot_at_epoch::<P>(epoch + 1) - 1;
    let advanced_state = advance_state(config, pre_state, block_slot)?;
    let eth1_data = advanced_state.eth1_data();
    let attestation_slots = misc::compute_start_slot_at_epoch::<P>(epoch)..block_slot;
    let attestations = full_block_attestations(config, &advanced_state, attestation_slots)?;
    let deposits = ContiguousList::default();
    let sync_aggregate = SyncAggregate::empty();

    block(
        config,
        advanced_state,
        eth1_data,
        graffiti,
        attestations,
        deposits,
        sync_aggregate,
        execution_payload,
    )
}

pub fn block_with_deposits<P: Preset>(
    config: &Config,
    pre_state: Arc<BeaconState<P>>,
    slot: Slot,
    deposits: ContiguousList<Deposit, P::MaxDeposits>,
) -> Result<BlockWithState<P>> {
    let advanced_state = advance_state(config, pre_state, slot)?;
    let eth1_data = advanced_state.eth1_data();
    let graffiti = H256::zero();
    let attestations = core::iter::empty();
    let sync_aggregate = SyncAggregate::empty();
    let execution_payload = None;

    block(
        config,
        advanced_state,
        eth1_data,
        graffiti,
        attestations,
        deposits,
        sync_aggregate,
        execution_payload,
    )
}

pub fn block_with_eth1_vote_and_deposits<P: Preset>(
    config: &Config,
    pre_state: Arc<BeaconState<P>>,
    slot: Slot,
    eth1_vote: Eth1Data,
    deposits: ContiguousList<Deposit, P::MaxDeposits>,
) -> Result<BlockWithState<P>> {
    let advanced_state = advance_state(config, pre_state, slot)?;
    let graffiti = H256::zero();
    let attestations = core::iter::empty();
    let sync_aggregate = SyncAggregate::empty();
    let execution_payload = None;

    block(
        config,
        advanced_state,
        eth1_vote,
        graffiti,
        attestations,
        deposits,
        sync_aggregate,
        execution_payload,
    )
}

pub fn block_with_payload<P: Preset>(
    config: &Config,
    pre_state: Arc<BeaconState<P>>,
    slot: Slot,
    graffiti: H256,
    execution_payload: ExecutionPayload<P>,
) -> Result<BlockWithState<P>> {
    let advanced_state = advance_state(config, pre_state, slot)?;
    let eth1_data = advanced_state.eth1_data();
    let attestations = core::iter::empty();
    let deposits = ContiguousList::default();
    let sync_aggregate = SyncAggregate::empty();
    let execution_payload = Some(execution_payload);

    block(
        config,
        advanced_state,
        eth1_data,
        graffiti,
        attestations,
        deposits,
        sync_aggregate,
        execution_payload,
    )
}

pub fn full_blocks_up_to_epoch<P: Preset>(
    config: &Config,
    genesis_state: Arc<BeaconState<P>>,
    end_epoch: Epoch,
) -> Result<Vec<Arc<SignedBeaconBlock<P>>>> {
    let start_slot = GENESIS_SLOT + 1;
    let end_slot = misc::compute_start_slot_at_epoch::<P>(end_epoch);

    let mut pre_state = genesis_state;
    let mut blocks = vec![];

    // This has to be a loop because it moves out of `pre_state`.
    for slot in start_slot..=end_slot {
        let advanced_state = advance_state(config, pre_state, slot)?;
        let eth1_data = advanced_state.eth1_data();
        let graffiti = H256::zero();
        let attestations = full_block_attestations(config, &advanced_state, (slot - 1)..slot)?;
        let deposits = ContiguousList::default();
        let sync_aggregate = full_sync_aggregate(config, &advanced_state);
        let execution_payload = None;

        let (block, post_state) = block(
            config,
            advanced_state,
            eth1_data,
            graffiti,
            attestations,
            deposits,
            sync_aggregate,
            execution_payload,
        )?;

        pre_state = post_state;
        blocks.push(block);
    }

    Ok(blocks)
}

pub fn singular_attestation<P: Preset>(
    config: &Config,
    state: Arc<BeaconState<P>>,
    epoch: Epoch,
    validator_index: ValidatorIndex,
) -> Result<(Attestation<P>, SubnetId)> {
    let start_slot = misc::compute_start_slot_at_epoch::<P>(epoch);
    let state_in_epoch = advance_state(config, state, start_slot)?;

    for slot in misc::slots_in_epoch::<P>(epoch) {
        let committees = accessors::beacon_committees(&state_in_epoch, slot)?;

        for (committee, committee_index) in committees.zip(0..) {
            let committees_per_slot =
                accessors::get_committee_count_per_slot(&state_in_epoch, RelativeEpoch::Current);

            let subnet_id = misc::compute_subnet_for_attestation::<P>(
                committees_per_slot,
                slot,
                committee_index,
            )?;

            if let Some(position) = committee
                .into_iter()
                .position(|index| index == validator_index)
            {
                let pre_electra = state_in_epoch.phase() < Phase::Electra;
                let index = pre_electra.then_some(committee_index).unwrap_or_default();
                let beacon_block_root = accessors::latest_block_root(&state_in_epoch);
                let root = accessors::epoch_boundary_block_root(&state_in_epoch, beacon_block_root);

                let data = AttestationData {
                    slot,
                    index,
                    beacon_block_root,
                    source: state_in_epoch.current_justified_checkpoint(),
                    target: Checkpoint { epoch, root },
                };

                let secret_key = interop::secret_key(validator_index);
                let signature = data.sign(config, &state_in_epoch, &secret_key).into();

                let attestation = if pre_electra {
                    let mut aggregation_bits = BitList::with_length(committee.len());
                    aggregation_bits.set(position, true);

                    Attestation::from(Phase0Attestation {
                        aggregation_bits,
                        data,
                        signature,
                    })
                } else {
                    let mut aggregation_bits = BitList::with_length(committee.len());
                    aggregation_bits.set(position, true);

                    let mut committee_bits = BitVector::default();
                    committee_bits.set(committee_index.try_into()?, true);

                    Attestation::from(ElectraAttestation {
                        aggregation_bits,
                        data,
                        signature,
                        committee_bits,
                    })
                };

                return Ok((attestation, subnet_id));
            }
        }
    }

    bail!("validator should belong to some committee")
}

pub fn execution_payload<P: Preset>(
    config: &Config,
    state: &Arc<BeaconState<P>>,
    slot: Slot,
    block_hash: ExecutionBlockHash,
) -> Result<ExecutionPayload<P>> {
    let parent_hash = state
        .post_bellatrix()
        .map(|state| state.latest_execution_payload_header().block_hash())
        .unwrap_or_default();

    // `accessors::get_current_epoch(state)` is correct even if `slot` is in a later epoch.
    // Using `misc::compute_epoch_at_slot::<P>(slot)` makes the payload invalid unless `state` is
    // advanced to `slot`. See how RANDAO mixes are copied in `process_randao_mixes_reset`.
    let prev_randao = accessors::get_randao_mix(state, accessors::get_current_epoch(state));
    let timestamp = misc::compute_timestamp_at_slot(config, state, slot);

    let withdrawals = state
        .post_capella()
        .map(capella::get_expected_withdrawals)
        .transpose()?
        .unwrap_or_default()
        .try_into()?;

    let execution_payload = match config.phase_at_slot::<P>(slot) {
        Phase::Phase0 | Phase::Altair => bail!("slot should be post-Bellatrix"),
        Phase::Bellatrix => BellatrixExecutionPayload {
            parent_hash,
            prev_randao,
            timestamp,
            block_hash,
            ..BellatrixExecutionPayload::default()
        }
        .into(),
        Phase::Capella => CapellaExecutionPayload {
            parent_hash,
            prev_randao,
            timestamp,
            block_hash,
            withdrawals,
            ..CapellaExecutionPayload::default()
        }
        .into(),
        Phase::Deneb | Phase::Electra => DenebExecutionPayload {
            parent_hash,
            prev_randao,
            timestamp,
            block_hash,
            withdrawals,
            ..DenebExecutionPayload::default()
        }
        .into(),
    };

    Ok(execution_payload)
}

#[allow(clippy::too_many_arguments)]
#[allow(clippy::too_many_lines)]
fn block<P: Preset>(
    config: &Config,
    advanced_state: Arc<BeaconState<P>>,
    eth1_data: Eth1Data,
    graffiti: H256,
    attestations: impl IntoIterator<Item = Attestation<P>>,
    deposits: ContiguousList<Deposit, P::MaxDeposits>,
    sync_aggregate: SyncAggregate<P>,
    mut execution_payload: Option<ExecutionPayload<P>>,
) -> Result<BlockWithState<P>> {
    let slot = advanced_state.slot();
    let proposer_index = accessors::get_beacon_proposer_index(&advanced_state)?;
    let secret_key = interop::secret_key(proposer_index);
    let parent_root = accessors::latest_block_root(&advanced_state);

    let randao_reveal = RandaoEpoch::from(accessors::get_current_epoch(&advanced_state))
        .sign(config, &advanced_state, &secret_key)
        .into();

    let (phase0_attestations, electra_attestations): (Vec<_>, Vec<_>) = attestations
        .into_iter()
        .partition_map(|attestation| match attestation {
            Attestation::Phase0(attestation) => Either::Left(attestation),
            Attestation::Electra(attestation) => Either::Right(attestation),
        });

    ensure!(
        phase0_attestations.is_empty() || advanced_state.phase() < Phase::Electra,
        "post-Electra block cannot contain Phase 0 attestations",
    );
    ensure!(
        electra_attestations.is_empty() || advanced_state.phase() >= Phase::Electra,
        "pre-Electra block cannot contain Electra attestations",
    );

    // Starting with `consensus-specs` v1.4.0-alpha.0, all Capella blocks must be post-Merge.
    if advanced_state.phase() >= Phase::Capella && execution_payload.is_none() {
        execution_payload = Some(self::execution_payload(
            config,
            &advanced_state,
            advanced_state.slot(),
            ExecutionBlockHash::zero(),
        )?);
    }

    let without_state_root = match advanced_state.phase() {
        Phase::Phase0 => BeaconBlock::from(Phase0BeaconBlock {
            slot,
            proposer_index,
            parent_root,
            state_root: H256::zero(),
            body: Phase0BeaconBlockBody {
                randao_reveal,
                eth1_data,
                graffiti,
                attestations: phase0_attestations.try_into()?,
                deposits,
                ..Phase0BeaconBlockBody::default()
            },
        }),
        Phase::Altair => BeaconBlock::from(AltairBeaconBlock {
            slot,
            proposer_index,
            parent_root,
            state_root: H256::zero(),
            body: AltairBeaconBlockBody {
                randao_reveal,
                eth1_data,
                graffiti,
                attestations: phase0_attestations.try_into()?,
                deposits,
                sync_aggregate,
                ..AltairBeaconBlockBody::default()
            },
        }),
        Phase::Bellatrix => BeaconBlock::from(BellatrixBeaconBlock {
            slot,
            proposer_index,
            parent_root,
            state_root: H256::zero(),
            body: BellatrixBeaconBlockBody {
                randao_reveal,
                eth1_data,
                graffiti,
                attestations: phase0_attestations.try_into()?,
                deposits,
                sync_aggregate,
                ..BellatrixBeaconBlockBody::default()
            },
        }),
        Phase::Capella => BeaconBlock::from(CapellaBeaconBlock {
            slot,
            proposer_index,
            parent_root,
            state_root: H256::zero(),
            body: CapellaBeaconBlockBody {
                randao_reveal,
                eth1_data,
                graffiti,
                attestations: phase0_attestations.try_into()?,
                deposits,
                sync_aggregate,
                ..CapellaBeaconBlockBody::default()
            },
        }),
        Phase::Deneb => BeaconBlock::from(DenebBeaconBlock {
            slot,
            proposer_index,
            parent_root,
            state_root: H256::zero(),
            body: DenebBeaconBlockBody {
                randao_reveal,
                eth1_data,
                graffiti,
                attestations: phase0_attestations.try_into()?,
                deposits,
                sync_aggregate,
                ..DenebBeaconBlockBody::default()
            },
        }),
        Phase::Electra => BeaconBlock::from(ElectraBeaconBlock {
            slot,
            proposer_index,
            parent_root,
            state_root: H256::zero(),
            body: ElectraBeaconBlockBody {
                randao_reveal,
                eth1_data,
                graffiti,
                attestations: electra_attestations.try_into()?,
                deposits,
                sync_aggregate,
                ..ElectraBeaconBlockBody::default()
            },
        }),
    }
    .with_execution_payload(execution_payload)?;

    let mut post_state = advanced_state;

    combined::process_untrusted_block(
        config,
        post_state.make_mut(),
        &without_state_root,
        NullSlotReport,
        true,
    )?;

    let message = without_state_root.with_state_root(post_state.hash_tree_root());
    let signature = message.sign(config, &post_state, &secret_key).into();
    let block = Arc::new(message.with_signature(signature));

    Ok((block, post_state))
}

// `advanced_state` is the one for the block being constructed,
// not the one that the attestations would be constructed with.
//
// Starting with Electra, attestations in blocks contain bits for all committees in the slot.
// Despite that they are represented by the same type as attestations outside blocks.
//
// See [`compute_on_chain_aggregate`] in the Electra Honest Validator specification.
//
// [`compute_on_chain_aggregate`]: https://github.com/ethereum/consensus-specs/blob/v1.5.0-alpha.0/specs/electra/validator.md#attestations
fn full_block_attestations<P: Preset>(
    config: &Config,
    advanced_state: &BeaconState<P>,
    slots: Range<Slot>,
) -> Result<Vec<Attestation<P>>> {
    let mut attestations = vec![];

    for slot in slots {
        // `accessors::latest_block_root` would be incorrect in some cases,
        // though we currently don't have any tests that can demonstrate that.
        // `rapid_upgrade_blocks` includes the attestations as early as possible.
        // `fork_choice_control::extra_tests` don't need validators to attest perfectly.
        let beacon_block_root = accessors::get_block_root_at_slot(advanced_state, slot)?;
        let epoch = misc::compute_epoch_at_slot::<P>(slot);
        let attestation_epoch = accessors::attestation_epoch(advanced_state, epoch)?;

        let source = match attestation_epoch {
            AttestationEpoch::Previous => advanced_state.previous_justified_checkpoint(),
            AttestationEpoch::Current => advanced_state.current_justified_checkpoint(),
        };

        let target = Checkpoint {
            epoch,
            // `accessors::epoch_boundary_block_root` would be incorrect because
            // `advanced_state` may be in a different epoch than the attestations.
            root: accessors::get_block_root(advanced_state, attestation_epoch)?,
        };

        let committees = accessors::beacon_committees(advanced_state, slot)?;

        if advanced_state.phase() < Phase::Electra {
            for (committee, index) in committees.zip(0..) {
                let data = AttestationData {
                    slot,
                    index,
                    beacon_block_root,
                    source,
                    target,
                };

                let signing_root = data.signing_root(config, advanced_state);

                let signature = committee
                    .into_iter()
                    .map(|validator_index| interop::secret_key(validator_index).sign(signing_root))
                    .reduce(AggregateSignature::aggregate)
                    .unwrap_or_default()
                    .into();

                attestations.push(Attestation::from(Phase0Attestation {
                    aggregation_bits: BitList::new(true, committee.len()),
                    data,
                    signature,
                }));
            }
        } else {
            let relative_epoch = attestation_epoch.into();

            let committees_per_slot =
                accessors::get_committee_count_per_slot(advanced_state, relative_epoch);

            let validator_count =
                accessors::active_validator_count_u64(advanced_state, relative_epoch);

            let committees_in_epoch = committees_per_slot * P::SlotsPerEpoch::U64;
            let slots_since_epoch_start = misc::slots_since_epoch_start::<P>(slot);
            let committees_before_slot = slots_since_epoch_start * committees_per_slot;
            let committees_including_slot = committees_before_slot + committees_per_slot - 1;
            let start = validator_count * committees_before_slot / committees_in_epoch;
            let end = validator_count * (committees_including_slot + 1) / committees_in_epoch;
            let active_validators_in_slot = (end - start).try_into()?;

            let aggregation_bits = BitList::new(true, active_validators_in_slot);

            let data = AttestationData {
                slot,
                index: 0,
                beacon_block_root,
                source,
                target,
            };

            let mut committee_bits = BitVector::default();

            for committee_index in 0..committees_per_slot {
                let index = committee_index.try_into()?;
                committee_bits.set(index, true);
            }

            let signing_root = data.signing_root(config, advanced_state);

            let signature = committees
                .flatten()
                .map(|validator_index| interop::secret_key(validator_index).sign(signing_root))
                .reduce(AggregateSignature::aggregate)
                .unwrap_or_default()
                .into();

            attestations.push(Attestation::from(ElectraAttestation {
                aggregation_bits,
                data,
                signature,
                committee_bits,
            }));
        }
    }

    Ok(attestations)
}

fn full_sync_aggregate<P: Preset>(
    config: &Config,
    advanced_state: &BeaconState<P>,
) -> SyncAggregate<P> {
    let Some(advanced_state) = advanced_state.post_altair() else {
        return SyncAggregate::empty();
    };

    let parent_root = accessors::latest_block_root(advanced_state);
    let signing_root = parent_root.signing_root(config, advanced_state, advanced_state.slot() - 1);

    let sync_committee_signature = advanced_state
        .current_sync_committee()
        .pubkeys
        .iter()
        .map(|pubkey| {
            accessors::index_of_public_key(advanced_state, pubkey.to_bytes()).expect(
                "public keys in advanced_state.current_sync_committee \
                 are taken from advanced_state.validators",
            )
        })
        .map(|validator_index| interop::secret_key(validator_index).sign(signing_root))
        .reduce(AggregateSignature::aggregate)
        .unwrap_or_default()
        .into();

    SyncAggregate {
        sync_committee_bits: BitVector::new(true),
        sync_committee_signature,
    }
}

fn advance_state<P: Preset>(
    config: &Config,
    mut state: Arc<BeaconState<P>>,
    slot: Slot,
) -> Result<Arc<BeaconState<P>>> {
    if state.slot() < slot {
        combined::process_slots(config, state.make_mut(), slot)?;
    }

    Ok(state)
}
