import {ForkName, ForkSeq} from "@lodestar/params";
import {
  computeEpochAtSlot,
  createAggregateSignatureSetFromComponents,
  isAggregatorFromCommitteeLength,
} from "@lodestar/state-transition";
import {IndexedAttestation, RootHex, SignedAggregateAndProof, electra, ssz} from "@lodestar/types";
import {toRootHex} from "@lodestar/utils";
import {AttestationError, AttestationErrorCode, GossipAction} from "../errors/index.js";
import {IBeaconChain} from "../index.js";
import {RegenCaller} from "../regen/index.js";
import {
  getAttestationDataSigningRoot,
  getCommitteeValidatorIndices,
  getSeenAttDataKeyFromSignedAggregateAndProof,
  getShufflingForAttestationVerification,
  verifyHeadBlockAndTargetRoot,
  verifyPropagationSlotRange,
} from "./attestation.js";
import {getAggregateAndProofSignatureSet, getSelectionProofSignatureSet} from "./signatureSets/index.js";

export type AggregateAndProofValidationResult = {
  indexedAttestation: IndexedAttestation;
  committeeValidatorIndices: Uint32Array;
  attDataRootHex: RootHex;
};

export async function validateApiAggregateAndProof(
  fork: ForkName,
  chain: IBeaconChain,
  signedAggregateAndProof: SignedAggregateAndProof
): Promise<AggregateAndProofValidationResult> {
  const skipValidationKnownAttesters = true;
  const prioritizeBls = true;
  return validateAggregateAndProof(fork, chain, signedAggregateAndProof, null, {
    skipValidationKnownAttesters,
    prioritizeBls,
  });
}

export async function validateGossipAggregateAndProof(
  fork: ForkName,
  chain: IBeaconChain,
  signedAggregateAndProof: SignedAggregateAndProof,
  serializedData: Uint8Array
): Promise<AggregateAndProofValidationResult> {
  return validateAggregateAndProof(fork, chain, signedAggregateAndProof, serializedData);
}

async function validateAggregateAndProof(
  fork: ForkName,
  chain: IBeaconChain,
  signedAggregateAndProof: SignedAggregateAndProof,
  serializedData: Uint8Array | null = null,
  opts: {skipValidationKnownAttesters: boolean; prioritizeBls: boolean} = {
    skipValidationKnownAttesters: false,
    prioritizeBls: false,
  }
): Promise<AggregateAndProofValidationResult> {
  const {skipValidationKnownAttesters, prioritizeBls} = opts;
  // Do checks in this order:
  // - do early checks (w/o indexed attestation)
  // - > obtain indexed attestation and committes per slot
  // - do middle checks w/ indexed attestation
  // - > verify signature
  // - do late checks w/ a valid signature

  const aggregateAndProof = signedAggregateAndProof.message;
  const aggregate = aggregateAndProof.aggregate;
  const {aggregationBits} = aggregate;
  const attData = aggregate.data;
  const attSlot = attData.slot;

  let committeeIndex: number | null;
  if (ForkSeq[fork] >= ForkSeq.gloas) {
    // [REJECT] `aggregate.data.index < 2`.
    if (attData.index >= 2) {
      throw new AttestationError(GossipAction.REJECT, {
        code: AttestationErrorCode.INVALID_PAYLOAD_STATUS_VALUE,
        attDataIndex: attData.index,
      });
    }
    // [REJECT] `aggregate.data.index == 0` if `block.slot == aggregate.data.slot`.
    const block = chain.forkChoice.getBlockDefaultStatus(attData.beaconBlockRoot);

    // If block is unknown, we don't handle it here. It will throw error later on at `verifyHeadBlockAndTargetRoot()`
    if (block !== null && block.slot === attData.slot && attData.index !== 0) {
      throw new AttestationError(GossipAction.REJECT, {
        code: AttestationErrorCode.PREMATURELY_INDICATED_PAYLOAD_PRESENT,
      });
    }

    // [REJECT] If `aggregate.data.index == 1` (payload present for a past
    //   block), the execution payload for `block` passes validation.
    // [IGNORE] When `aggregate.data.index == 1` (payload present for a past block),
    // the corresponding execution payload for `block` has been seen (a client MAY queue
    // attestations for processing once the payload is retrieved and SHOULD request the
    // payload envelope via `ExecutionPayloadEnvelopesByRoot`).
    if (block !== null && attData.index === 1 && !chain.seenPayloadEnvelope(toRootHex(attData.beaconBlockRoot))) {
      throw new AttestationError(GossipAction.IGNORE, {
        code: AttestationErrorCode.EXECUTION_PAYLOAD_NOT_SEEN,
        beaconBlockRoot: toRootHex(attData.beaconBlockRoot),
      });
    }

    // [REJECT] len(committee_indices) == 1, where committee_indices = get_committee_indices(aggregate)
    committeeIndex = (aggregate as electra.Attestation).committeeBits.getSingleTrueBit();
    if (committeeIndex === null) {
      throw new AttestationError(GossipAction.REJECT, {code: AttestationErrorCode.NOT_EXACTLY_ONE_COMMITTEE_BIT_SET});
    }
  } else if (ForkSeq[fork] >= ForkSeq.electra) {
    committeeIndex = (aggregate as electra.Attestation).committeeBits.getSingleTrueBit();
    // [REJECT] len(committee_indices) == 1, where committee_indices = get_committee_indices(aggregate)
    if (committeeIndex === null) {
      throw new AttestationError(GossipAction.REJECT, {code: AttestationErrorCode.NOT_EXACTLY_ONE_COMMITTEE_BIT_SET});
    }
    // [REJECT] aggregate.data.index == 0
    if (attData.index !== 0) {
      throw new AttestationError(GossipAction.REJECT, {code: AttestationErrorCode.NON_ZERO_ATTESTATION_DATA_INDEX});
    }
  } else {
    committeeIndex = attData.index;
  }

  const seenAttDataKey = serializedData ? getSeenAttDataKeyFromSignedAggregateAndProof(fork, serializedData) : null;
  const cachedAttData = seenAttDataKey ? chain.seenAttestationDatas.get(attSlot, committeeIndex, seenAttDataKey) : null;

  const attEpoch = computeEpochAtSlot(attSlot);
  const attTarget = attData.target;
  const targetEpoch = attTarget.epoch;

  chain.metrics?.gossipAttestation.attestationSlotToClockSlot.observe(
    {caller: RegenCaller.validateGossipAggregateAndProof},
    chain.clock.currentSlot - attSlot
  );

  if (!cachedAttData) {
    // [REJECT] The attestation's epoch matches its target -- i.e. attestation.data.target.epoch == compute_epoch_at_slot(attestation.data.slot)
    if (targetEpoch !== attEpoch) {
      throw new AttestationError(GossipAction.REJECT, {code: AttestationErrorCode.BAD_TARGET_EPOCH});
    }

    // Pre-deneb:
    // [IGNORE] aggregate.data.slot is within the last ATTESTATION_PROPAGATION_SLOT_RANGE slots (with a MAXIMUM_GOSSIP_CLOCK_DISPARITY allowance)
    // -- i.e. aggregate.data.slot + ATTESTATION_PROPAGATION_SLOT_RANGE >= current_slot >= aggregate.data.slot
    // (a client MAY queue future aggregates for processing at the appropriate slot).
    // Post-deneb:
    // [IGNORE] `aggregate.data.slot` is equal to or earlier than the `current_slot` (with a `MAXIMUM_GOSSIP_CLOCK_DISPARITY` allowance)
    // -- i.e. `aggregate.data.slot <= current_slot`
    //   (a client MAY queue future aggregates for processing at the appropriate slot).
    // [IGNORE] the epoch of `aggregate.data.slot` is either the current or previous epoch
    //   (with a `MAXIMUM_GOSSIP_CLOCK_DISPARITY` allowance)
    // -- i.e. `compute_epoch_at_slot(aggregate.data.slot) in (get_previous_epoch(state), get_current_epoch(state))`
    verifyPropagationSlotRange(fork, chain, attSlot);
  }

  // [IGNORE] The aggregate is the first valid aggregate received for the aggregator with
  // index aggregate_and_proof.aggregator_index for the epoch aggregate.data.target.epoch.
  const aggregatorIndex = aggregateAndProof.aggregatorIndex;
  if (chain.seenAggregators.isKnown(targetEpoch, aggregatorIndex)) {
    throw new AttestationError(GossipAction.IGNORE, {
      code: AttestationErrorCode.AGGREGATOR_ALREADY_KNOWN,
      targetEpoch,
      aggregatorIndex,
    });
  }

  // _[IGNORE]_ A valid aggregate attestation defined by `hash_tree_root(aggregate.data)` whose `aggregation_bits`
  // is a non-strict superset has _not_ already been seen.
  const attDataRootHex = cachedAttData
    ? cachedAttData.attDataRootHex
    : toRootHex(ssz.phase0.AttestationData.hashTreeRoot(attData));
  if (
    !skipValidationKnownAttesters &&
    chain.seenAggregatedAttestations.isKnown(targetEpoch, committeeIndex, attDataRootHex, aggregationBits)
  ) {
    throw new AttestationError(GossipAction.IGNORE, {
      code: AttestationErrorCode.ATTESTERS_ALREADY_KNOWN,
      targetEpoch,
      aggregateRoot: attDataRootHex,
    });
  }

  // [IGNORE] The block being voted for (attestation.data.beacon_block_root) has been seen (via both gossip
  // and non-gossip sources) (a client MAY queue attestations for processing once block is retrieved).
  // Lighthouse doesn't check maxSkipSlots option here but Lodestar wants to be more strict
  // to be more DOS protection

  // [REJECT] The aggregate attestation's target block is an ancestor of the block named in the LMD vote
  // -- i.e. `get_checkpoint_block(store, aggregate.data.beacon_block_root, aggregate.data.target.epoch) == aggregate.data.target.root`
  const attHeadBlock = verifyHeadBlockAndTargetRoot(
    chain,
    attData.beaconBlockRoot,
    attTarget.root,
    attSlot,
    attEpoch,
    RegenCaller.validateGossipAggregateAndProof,
    chain.opts.maxSkipSlots
  );

  // [IGNORE] The current finalized_checkpoint is an ancestor of the block defined by aggregate.data.beacon_block_root
  // -- i.e. get_ancestor(store, aggregate.data.beacon_block_root, compute_start_slot_at_epoch(store.finalized_checkpoint.epoch)) == store.finalized_checkpoint.root
  // > Altready check in `chain.forkChoice.hasBlock(attestation.data.beaconBlockRoot)`

  const shuffling = await getShufflingForAttestationVerification(
    chain,
    attEpoch,
    attHeadBlock,
    RegenCaller.validateGossipAttestation
  );

  // [REJECT] The committee index is within the expected range
  // -- i.e. data.index < get_committee_count_per_slot(state, data.target.epoch)
  const committeeValidatorIndices = cachedAttData
    ? cachedAttData.committeeValidatorIndices
    : getCommitteeValidatorIndices(shuffling, attSlot, committeeIndex);

  // [REJECT] The number of aggregation bits matches the committee size
  // -- i.e. `len(aggregation_bits) == len(get_beacon_committee(state, aggregate.data.slot, index))`.
  if (aggregate.aggregationBits.bitLen !== committeeValidatorIndices.length) {
    throw new AttestationError(GossipAction.REJECT, {code: AttestationErrorCode.WRONG_NUMBER_OF_AGGREGATION_BITS});
  }
  const attestingIndices = aggregate.aggregationBits.intersectValues(committeeValidatorIndices);

  const indexedAttestation: IndexedAttestation = {
    attestingIndices,
    data: attData,
    signature: aggregate.signature,
  };

  // TODO: Check this before regen
  // [REJECT] The attestation has participants -- that is,
  // len(get_attesting_indices(state, aggregate.data, aggregate.aggregation_bits)) >= 1.
  if (attestingIndices.length < 1) {
    // missing attestation participants
    throw new AttestationError(GossipAction.REJECT, {code: AttestationErrorCode.EMPTY_AGGREGATION_BITFIELD});
  }

  // [REJECT] aggregate_and_proof.selection_proof selects the validator as an aggregator for the slot
  // -- i.e. is_aggregator(state, aggregate.data.slot, aggregate.data.index, aggregate_and_proof.selection_proof) returns True.
  if (!isAggregatorFromCommitteeLength(committeeValidatorIndices.length, aggregateAndProof.selectionProof)) {
    throw new AttestationError(GossipAction.REJECT, {code: AttestationErrorCode.INVALID_AGGREGATOR});
  }

  // [REJECT] The aggregator's validator index is within the committee
  // -- i.e. aggregate_and_proof.aggregator_index in get_beacon_committee(state, aggregate.data.slot, aggregate.data.index).
  if (!committeeValidatorIndices.includes(aggregateAndProof.aggregatorIndex)) {
    throw new AttestationError(GossipAction.REJECT, {code: AttestationErrorCode.AGGREGATOR_NOT_IN_COMMITTEE});
  }

  // [REJECT] The aggregate_and_proof.selection_proof is a valid signature of the aggregate.data.slot
  // by the validator with index aggregate_and_proof.aggregator_index.
  // [REJECT] The aggregator signature, signed_aggregate_and_proof.signature, is valid.
  // [REJECT] The signature of aggregate is valid.
  const signingRoot = cachedAttData ? cachedAttData.signingRoot : getAttestationDataSigningRoot(chain.config, attData);
  const indexedAttestationSignatureSet = createAggregateSignatureSetFromComponents(
    indexedAttestation.attestingIndices,
    signingRoot,
    indexedAttestation.signature
  );
  const signatureSets = [
    getSelectionProofSignatureSet(chain.config, attSlot, aggregatorIndex, signedAggregateAndProof),
    getAggregateAndProofSignatureSet(chain.config, attEpoch, aggregatorIndex, signedAggregateAndProof),
    indexedAttestationSignatureSet,
  ];
  // no need to write to SeenAttestationDatas

  if (!(await chain.bls.verifySignatureSets(signatureSets, {batchable: true, priority: prioritizeBls}))) {
    throw new AttestationError(GossipAction.REJECT, {code: AttestationErrorCode.INVALID_SIGNATURE});
  }

  // It's important to double check that the attestation still hasn't been observed, since
  // there can be a race-condition if we receive two attestations at the same time and
  // process them in different threads.
  if (chain.seenAggregators.isKnown(targetEpoch, aggregatorIndex)) {
    throw new AttestationError(GossipAction.IGNORE, {
      code: AttestationErrorCode.AGGREGATOR_ALREADY_KNOWN,
      targetEpoch,
      aggregatorIndex,
    });
  }

  // Same race-condition check as above for seen aggregators
  if (
    !skipValidationKnownAttesters &&
    chain.seenAggregatedAttestations.isKnown(targetEpoch, committeeIndex, attDataRootHex, aggregationBits)
  ) {
    throw new AttestationError(GossipAction.IGNORE, {
      code: AttestationErrorCode.ATTESTERS_ALREADY_KNOWN,
      targetEpoch,
      aggregateRoot: attDataRootHex,
    });
  }

  chain.seenAggregators.add(targetEpoch, aggregatorIndex);
  chain.seenAggregatedAttestations.add(
    targetEpoch,
    committeeIndex,
    attDataRootHex,
    {aggregationBits, trueBitCount: attestingIndices.length},
    false
  );

  return {indexedAttestation, committeeValidatorIndices, attDataRootHex};
}