/*
 * Copyright contributors to the AERE Network.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 */
package org.hyperledger.besu.evm.precompile;

import org.hyperledger.besu.crypto.SecureRandomProvider;
import org.hyperledger.besu.evm.frame.MessageFrame;
import org.hyperledger.besu.evm.gascalculator.GasCalculator;

import java.security.SecureRandom;

import jakarta.validation.constraints.NotNull;
import org.apache.tuweni.bytes.Bytes;
import org.bouncycastle.pqc.crypto.mlkem.MLKEMGenerator;
import org.bouncycastle.pqc.crypto.mlkem.MLKEMParameters;
import org.bouncycastle.pqc.crypto.mlkem.MLKEMPublicKeyParameters;

/**
 * AERE PQC precompile: ML-KEM-768 (FIPS 203) DETERMINISTIC encapsulation at 0x0AE6.
 *
 * <p>This is AERE's first post-quantum CONFIDENTIALITY primitive on-chain. Every other native PQC
 * precompile (0x0AE1-0x0AE5) is a signature or hash and gives post-quantum AUTHENTICITY only. This
 * precompile makes a Module-Lattice KEM key-agreement transcript verifiable on-chain: given an
 * encapsulation key {@code ek} and the 32-byte encapsulation randomness {@code m} ("coins"), it
 * recomputes the ciphertext {@code c} and shared secret {@code K} that ML-KEM.Encaps(ek, m)
 * produces. A verifier compares the recomputed {@code (c, K)} against a claimed transcript; equality
 * proves the KEM step was performed honestly with the stated coins. This serves UMBRA's PQXDH
 * handshake settlement and the AERE PQC key-registry.
 *
 * <p>Input layout: {@code ek(1184) || m(32)} = 1216 bytes exactly.
 * Output layout: {@code c(1088) || K(32)} = 1120 bytes, or EMPTY (0x) on any malformed input.
 *
 * <p>Determinism: FIPS-203 Encaps normally draws {@code m} from a CSPRNG, which cannot run inside a
 * consensus-critical precompile. We take {@code m} from calldata and drive Bouncy Castle's
 * ML-KEM.Encaps_internal (K-PKE.Encrypt with explicit coins), so every node computes the identical
 * {@code (c, K)}. No cryptography is reimplemented here; the audited Bouncy Castle BCPQC ML-KEM
 * implementation on the classpath does the work.
 */
public class MLKEM768PrecompiledContract extends AbstractPrecompiledContract {

  /** ML-KEM-768 encapsulation-key (public key) length, FIPS 203. */
  static final int EK_LEN = 1184;

  /** Encapsulation randomness ("coins" m) length. */
  static final int M_LEN = 32;

  /** Expected total calldata length. */
  static final int INPUT_LEN = EK_LEN + M_LEN; // 1216

  /** ML-KEM-768 ciphertext length. */
  static final int CT_LEN = 1088;

  /** ML-KEM shared-secret length. */
  static final int SS_LEN = 32;

  /**
   * Fixed gas. ML-KEM-768 encapsulation is dominated by one A*r matrix-vector product in the NTT
   * domain (k=3) plus SHA3/SHAKE hashing; measured on the AERE Besu scratch fork it sits between
   * ML-DSA-44 verify (55k) and Falcon-1024 verify (75k). Priced fixed like the other lattice
   * precompiles.
   */
  private static final long GAS = 60_000L;

  // The generator constructor requires a SecureRandom, but the DETERMINISTIC encapsulation path
  // (internalGenerateEncapsulated with caller-supplied coins m) never draws from it: the output
  // depends only on (ek, m). Uses Besu's approved provider rather than constructing one directly.
  private static final SecureRandom RNG = SecureRandomProvider.publicSecureRandom();

  /**
   * Instantiates a new ML-KEM-768 precompiled contract.
   *
   * @param gasCalculator the gas calculator
   */
  MLKEM768PrecompiledContract(final GasCalculator gasCalculator) {
    super("AereMLKEM768", gasCalculator);
  }

  @Override
  public long gasRequirement(final Bytes input) {
    return GAS;
  }

  @NotNull
  @Override
  public PrecompileContractResult computePrecompile(
      final Bytes input, @NotNull final MessageFrame messageFrame) {
    if (input.size() != INPUT_LEN) {
      return PrecompileContractResult.success(Bytes.EMPTY);
    }
    try {
      final byte[] ek = input.slice(0, EK_LEN).toArrayUnsafe();
      final byte[] m = input.slice(EK_LEN, M_LEN).toArrayUnsafe();

      final MLKEMPublicKeyParameters pub =
          new MLKEMPublicKeyParameters(MLKEMParameters.ml_kem_768, ek);

      // Deterministic Encaps: feed the caller-supplied coins m as the encapsulation randomness.
      final MLKEMGenerator gen = new MLKEMGenerator(RNG);
      final org.bouncycastle.crypto.SecretWithEncapsulation enc =
          gen.internalGenerateEncapsulated(pub, m);

      final byte[] ss = enc.getSecret();
      final byte[] ct = enc.getEncapsulation();
      if (ct.length != CT_LEN || ss.length != SS_LEN) {
        return PrecompileContractResult.success(Bytes.EMPTY);
      }

      final byte[] out = new byte[CT_LEN + SS_LEN];
      System.arraycopy(ct, 0, out, 0, CT_LEN);
      System.arraycopy(ss, 0, out, CT_LEN, SS_LEN);
      return PrecompileContractResult.success(Bytes.wrap(out));
    } catch (final Throwable t) {
      // Consensus rule for the non-signature PQC precompiles: malformed input -> EMPTY, never fault.
      return PrecompileContractResult.success(Bytes.EMPTY);
    }
  }
}