import {ChainConfig} from "@lodestar/config";
import {
  ForkName,
  KZG_COMMITMENT_INCLUSION_PROOF_DEPTH,
  KZG_COMMITMENT_SUBTREE_INDEX0,
  isForkPostElectra,
} from "@lodestar/params";
import {
  computeEpochAtSlot,
  computeStartSlotAtEpoch,
  getBlockHeaderProposerSignatureSetByHeaderSlot,
  getBlockHeaderProposerSignatureSetByParentStateSlot,
} from "@lodestar/state-transition";
import {BlobIndex, Root, Slot, SubnetID, deneb, ssz} from "@lodestar/types";
import {byteArrayEquals, toRootHex, verifyMerkleBranch} from "@lodestar/utils";
import {kzg} from "../../util/kzg.js";
import {BlobSidecarErrorCode, BlobSidecarGossipError, BlobSidecarValidationError} from "../errors/blobSidecarError.js";
import {GossipAction} from "../errors/gossipValidation.js";
import {IBeaconChain} from "../interface.js";
import {RegenCaller} from "../regen/index.js";

export async function validateGossipBlobSidecar(
  fork: ForkName,
  chain: IBeaconChain,
  blobSidecar: deneb.BlobSidecar,
  subnet: SubnetID
): Promise<void> {
  const blobSlot = blobSidecar.signedBlockHeader.message.slot;

  // [REJECT] The sidecar's index is consistent with `MAX_BLOBS_PER_BLOCK` -- i.e. `blob_sidecar.index < MAX_BLOBS_PER_BLOCK`.
  const maxBlobsPerBlock = chain.config.getMaxBlobsPerBlock(computeEpochAtSlot(blobSlot));
  if (blobSidecar.index >= maxBlobsPerBlock) {
    throw new BlobSidecarGossipError(GossipAction.REJECT, {
      code: BlobSidecarErrorCode.INDEX_TOO_LARGE,
      blobIdx: blobSidecar.index,
      maxBlobsPerBlock,
    });
  }

  // [REJECT] The sidecar is for the correct subnet -- i.e. `compute_subnet_for_blob_sidecar(sidecar.index) == subnet_id`.
  if (computeSubnetForBlobSidecar(fork, chain.config, blobSidecar.index) !== subnet) {
    throw new BlobSidecarGossipError(GossipAction.REJECT, {
      code: BlobSidecarErrorCode.INVALID_INDEX,
      blobIdx: blobSidecar.index,
      subnet,
    });
  }

  // [IGNORE] The sidecar is not from a future slot (with a MAXIMUM_GOSSIP_CLOCK_DISPARITY allowance) --
  // i.e. validate that sidecar.slot <= current_slot (a client MAY queue future blocks for processing at
  // the appropriate slot).
  const currentSlotWithGossipDisparity = chain.clock.currentSlotWithGossipDisparity;
  if (currentSlotWithGossipDisparity < blobSlot) {
    throw new BlobSidecarGossipError(GossipAction.IGNORE, {
      code: BlobSidecarErrorCode.FUTURE_SLOT,
      currentSlot: currentSlotWithGossipDisparity,
      blockSlot: blobSlot,
    });
  }

  // [IGNORE] The sidecar is from a slot greater than the latest finalized slot -- i.e. validate that
  // sidecar.slot > compute_start_slot_at_epoch(state.finalized_checkpoint.epoch)
  const finalizedCheckpoint = chain.forkChoice.getFinalizedCheckpoint();
  const finalizedSlot = computeStartSlotAtEpoch(finalizedCheckpoint.epoch);
  if (blobSlot <= finalizedSlot) {
    throw new BlobSidecarGossipError(GossipAction.IGNORE, {
      code: BlobSidecarErrorCode.WOULD_REVERT_FINALIZED_SLOT,
      blockSlot: blobSlot,
      finalizedSlot,
    });
  }

  // Check if the block is already known. We know it is post-finalization, so it is sufficient to check the fork choice.
  //
  // In normal operation this isn't necessary, however it is useful immediately after a
  // reboot if the `observed_block_producers` cache is empty. In that case, without this
  // check, we will load the parent and state from disk only to find out later that we
  // already know this block.
  const blockRoot = ssz.phase0.BeaconBlockHeader.hashTreeRoot(blobSidecar.signedBlockHeader.message);
  const blockHex = toRootHex(blockRoot);
  if (chain.forkChoice.getBlockHexDefaultStatus(blockHex) !== null) {
    throw new BlobSidecarGossipError(GossipAction.IGNORE, {code: BlobSidecarErrorCode.ALREADY_KNOWN, root: blockHex});
  }

  // TODO: freetheblobs - check for badblock
  // TODO: freetheblobs - check that its first blob with valid signature

  // _[IGNORE]_ The blob's block's parent (defined by `sidecar.block_parent_root`) has been seen (via both
  // gossip and non-gossip sources) (a client MAY queue blocks for processing once the parent block is
  // retrieved).
  const parentRoot = toRootHex(blobSidecar.signedBlockHeader.message.parentRoot);
  const parentBlock = chain.forkChoice.getBlockHexDefaultStatus(parentRoot);
  if (parentBlock === null) {
    // If fork choice does *not* consider the parent to be a descendant of the finalized block,
    // then there are two more cases:
    //
    // 1. We have the parent stored in our database. Because fork-choice has confirmed the
    //    parent is *not* in our post-finalization DAG, all other blocks must be either
    //    pre-finalization or conflicting with finalization.
    // 2. The parent is unknown to us, we probably want to download it since it might actually
    //    descend from the finalized root.
    // (Non-Lighthouse): Since we prune all blocks non-descendant from finalized checking the `db.block` database won't be useful to guard
    // against known bad fork blocks, so we throw PARENT_UNKNOWN for cases (1) and (2)
    throw new BlobSidecarGossipError(GossipAction.IGNORE, {
      code: BlobSidecarErrorCode.PARENT_UNKNOWN,
      parentRoot,
      blockRoot: blockHex,
      slot: blobSlot,
    });
  }

  // [REJECT] The blob is from a higher slot than its parent.
  if (parentBlock.slot >= blobSlot) {
    throw new BlobSidecarGossipError(GossipAction.IGNORE, {
      code: BlobSidecarErrorCode.NOT_LATER_THAN_PARENT,
      parentSlot: parentBlock.slot,
      slot: blobSlot,
    });
  }

  // getBlockSlotState also checks for whether the current finalized checkpoint is an ancestor of the block.
  // As a result, we throw an IGNORE (whereas the spec says we should REJECT for this scenario).
  // this is something we should change this in the future to make the code airtight to the spec.
  // [IGNORE] The block's parent (defined by block.parent_root) has been seen (via both gossip and non-gossip sources) (a client MAY queue blocks for processing once the parent block is retrieved).
  // [REJECT] The block's parent (defined by block.parent_root) passes validation.
  const blockState = await chain.regen
    .getBlockSlotState(parentBlock, blobSlot, {dontTransferCache: true}, RegenCaller.validateGossipBlock)
    .catch(() => {
      throw new BlobSidecarGossipError(GossipAction.IGNORE, {
        code: BlobSidecarErrorCode.PARENT_UNKNOWN,
        parentRoot,
        blockRoot: blockHex,
        slot: blobSlot,
      });
    });

  // [REJECT] The proposer signature, signed_beacon_block.signature, is valid with respect to the proposer_index pubkey.
  const signature = blobSidecar.signedBlockHeader.signature;
  if (!chain.seenBlockInputCache.isVerifiedProposerSignature(blobSlot, blockHex, signature)) {
    const signatureSet = getBlockHeaderProposerSignatureSetByParentStateSlot(
      chain.config,
      blockState.slot,
      blobSidecar.signedBlockHeader
    );
    // Don't batch so verification is not delayed
    if (!(await chain.bls.verifySignatureSets([signatureSet], {verifyOnMainThread: true}))) {
      throw new BlobSidecarGossipError(GossipAction.REJECT, {
        code: BlobSidecarErrorCode.PROPOSAL_SIGNATURE_INVALID,
        blockRoot: blockHex,
        index: blobSidecar.index,
        slot: blobSlot,
      });
    }

    chain.seenBlockInputCache.markVerifiedProposerSignature(blobSlot, blockHex, signature);
  }

  // verify if the blob inclusion proof is correct
  if (!validateBlobSidecarInclusionProof(blobSidecar)) {
    throw new BlobSidecarGossipError(GossipAction.REJECT, {
      code: BlobSidecarErrorCode.INCLUSION_PROOF_INVALID,
      slot: blobSidecar.signedBlockHeader.message.slot,
      blobIdx: blobSidecar.index,
    });
  }

  //  _[IGNORE]_ The sidecar is the only sidecar with valid signature received for the tuple
  // `(sidecar.block_root, sidecar.index)`
  //
  // This is already taken care of by the way we group the blobs in getFullBlockInput helper
  // but may be an error can be thrown there for this

  // _[REJECT]_ The sidecar is proposed by the expected `proposer_index` for the block's slot in the
  // context of the current shuffling (defined by `block_parent_root`/`slot`)
  // If the `proposer_index` cannot immediately be verified against the expected shuffling, the sidecar
  // MAY be queued for later processing while proposers for the block's branch are calculated -- in such
  // a case _do not_ `REJECT`, instead `IGNORE` this message.
  const proposerIndex = blobSidecar.signedBlockHeader.message.proposerIndex;
  if (blockState.getBeaconProposer(blobSlot) !== proposerIndex) {
    throw new BlobSidecarGossipError(GossipAction.REJECT, {
      code: BlobSidecarErrorCode.INCORRECT_PROPOSER,
      proposerIndex,
    });
  }

  // blob, proof and commitment as a valid BLS G1 point gets verified in batch validation
  try {
    await validateBlobsAndBlobProofs([blobSidecar.kzgCommitment], [blobSidecar.blob], [blobSidecar.kzgProof]);
  } catch (_e) {
    throw new BlobSidecarGossipError(GossipAction.REJECT, {
      code: BlobSidecarErrorCode.INVALID_KZG_PROOF,
      blobIdx: blobSidecar.index,
    });
  }
}

/**
 * Validate some blob sidecars in a block
 *
 * Requires the block to be known to the node
 *
 * NOTE: chain is optional to skip signature verification. Helpful for testing purposes and so that can control whether
 * signature gets checked depending on the reqresp method that is being checked
 */
export async function validateBlockBlobSidecars(
  chain: IBeaconChain | null,
  blockSlot: Slot,
  blockRoot: Root,
  blockBlobCount: number,
  blobSidecars: deneb.BlobSidecars
): Promise<void> {
  if (blobSidecars.length === 0) {
    return;
  }

  if (blockBlobCount === 0) {
    throw new BlobSidecarValidationError({
      code: BlobSidecarErrorCode.INCORRECT_SIDECAR_COUNT,
      slot: blockSlot,
      expected: blockBlobCount,
      actual: blobSidecars.length,
    });
  }

  // Hash the first sidecar block header and compare the rest via (cheaper) equality
  const firstSidecarSignedBlockHeader = blobSidecars[0].signedBlockHeader;
  const firstSidecarBlockHeader = firstSidecarSignedBlockHeader.message;
  const firstBlockRoot = ssz.phase0.BeaconBlockHeader.hashTreeRoot(firstSidecarBlockHeader);
  if (!byteArrayEquals(blockRoot, firstBlockRoot)) {
    throw new BlobSidecarValidationError(
      {
        code: BlobSidecarErrorCode.INCORRECT_BLOCK,
        slot: blockSlot,
        blobIdx: 0,
        expected: toRootHex(blockRoot),
        actual: toRootHex(firstBlockRoot),
      },
      "BlobSidecar doesn't match corresponding block"
    );
  }

  if (chain !== null) {
    const blockRootHex = toRootHex(blockRoot);
    const signature = firstSidecarSignedBlockHeader.signature;
    if (!chain.seenBlockInputCache.isVerifiedProposerSignature(blockSlot, blockRootHex, signature)) {
      const signatureSet = getBlockHeaderProposerSignatureSetByHeaderSlot(chain.config, firstSidecarSignedBlockHeader);

      if (
        !(await chain.bls.verifySignatureSets([signatureSet], {
          verifyOnMainThread: true,
        }))
      ) {
        throw new BlobSidecarValidationError({
          code: BlobSidecarErrorCode.PROPOSAL_SIGNATURE_INVALID,
          blockRoot: blockRootHex,
          slot: blockSlot,
          index: blobSidecars[0].index,
        });
      }

      chain.seenBlockInputCache.markVerifiedProposerSignature(blockSlot, blockRootHex, signature);
    }
  }

  const commitments = [];
  const blobs = [];
  const proofs = [];
  for (let i = 0; i < blobSidecars.length; i++) {
    const blobSidecar = blobSidecars[i];
    const blobIndex = blobSidecar.index;

    if (
      i !== 0 &&
      !ssz.phase0.SignedBeaconBlockHeader.equals(blobSidecar.signedBlockHeader, firstSidecarSignedBlockHeader)
    ) {
      throw new BlobSidecarValidationError(
        {
          code: BlobSidecarErrorCode.INCORRECT_BLOCK,
          slot: blockSlot,
          blobIdx: blobIndex,
          expected: toRootHex(blockRoot),
          actual: "unknown - compared via equality",
        },
        "BlobSidecar doesn't match corresponding block"
      );
    }

    if (!validateBlobSidecarInclusionProof(blobSidecar)) {
      throw new BlobSidecarValidationError(
        {
          code: BlobSidecarErrorCode.INCLUSION_PROOF_INVALID,
          slot: blockSlot,
          blobIdx: blobIndex,
        },
        "BlobSidecar inclusion proof invalid"
      );
    }

    commitments.push(blobSidecar.kzgCommitment);
    blobs.push(blobSidecar.blob);
    proofs.push(blobSidecar.kzgProof);
  }

  // Final batch KZG proof verification
  let reason: string | undefined = undefined;
  try {
    if (!(await kzg.asyncVerifyBlobKzgProofBatch(blobs, commitments, proofs))) {
      reason = "Invalid verifyBlobKzgProofBatch";
    }
  } catch (e) {
    reason = (e as Error).message;
  }
  if (reason !== undefined) {
    throw new BlobSidecarValidationError(
      {
        code: BlobSidecarErrorCode.INVALID_KZG_PROOF_BATCH,
        slot: blockSlot,
        reason,
      },
      "BlobSidecar has invalid KZG proof batch"
    );
  }
}

export async function validateBlobsAndBlobProofs(
  expectedKzgCommitments: deneb.BlobKzgCommitments,
  blobs: deneb.Blobs,
  proofs: deneb.KZGProofs
): Promise<void> {
  // assert verify_aggregate_kzg_proof(blobs, expected_kzg_commitments, kzg_aggregated_proof)
  let isProofValid: boolean;
  try {
    isProofValid = await kzg.asyncVerifyBlobKzgProofBatch(blobs, expectedKzgCommitments, proofs);
  } catch (e) {
    (e as Error).message = `Error on verifyBlobKzgProofBatch: ${(e as Error).message}`;
    throw e;
  }
  if (!isProofValid) {
    throw Error("Invalid verifyBlobKzgProofBatch");
  }
}

export function validateBlobSidecarInclusionProof(blobSidecar: deneb.BlobSidecar): boolean {
  return verifyMerkleBranch(
    ssz.deneb.KZGCommitment.hashTreeRoot(blobSidecar.kzgCommitment),
    blobSidecar.kzgCommitmentInclusionProof,
    KZG_COMMITMENT_INCLUSION_PROOF_DEPTH,
    KZG_COMMITMENT_SUBTREE_INDEX0 + blobSidecar.index,
    blobSidecar.signedBlockHeader.message.bodyRoot
  );
}

function computeSubnetForBlobSidecar(fork: ForkName, config: ChainConfig, blobIndex: BlobIndex): SubnetID {
  return (
    blobIndex % (isForkPostElectra(fork) ? config.BLOB_SIDECAR_SUBNET_COUNT_ELECTRA : config.BLOB_SIDECAR_SUBNET_COUNT)
  );
}