/*
 * 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.evm.frame.MessageFrame;
import org.hyperledger.besu.evm.gascalculator.GasCalculator;

import jakarta.validation.constraints.NotNull;
import org.apache.tuweni.bytes.Bytes;
import org.bouncycastle.crypto.digests.SHAKEDigest;

/**
 * AERE PQC precompile: Falcon HashToPoint (FIPS 206 / NIST Falcon round-3) at 0x0AE7.
 *
 * <p>HashToPoint is the SHAKE256-driven map from a (nonce, message) pair to a challenge polynomial
 * {@code c} in Z_q[x]/(x^n+1), q = 12289. It is the single most expensive step of an on-chain
 * Falcon verification: a hand-rolled Solidity Falcon-512 verify spends the bulk of its ~10.5M gas
 * inside the in-EVM Keccak-f[1600] permutations that drive this rejection sampler. Exposing it
 * natively lets a Solidity Falcon verifier replace that whole loop with one ~500-gas staticcall,
 * collapsing per-auth Falcon cost.
 *
 * <p>Input layout: {@code logn(1) || nonce(40) || message(rest)} where {@code logn} is 9
 * (Falcon-512, n=512) or 10 (Falcon-1024, n=1024). Output: {@code n} coefficients, each a
 * big-endian uint16 in [0, q), i.e. {@code 2*n} bytes. Malformed input (length &lt; 41, or logn not
 * in {9,10}) returns EMPTY (0x).
 *
 * <p>Algorithm (matches the reference {@code hash_to_point_vartime} exactly): absorb
 * {@code nonce || message} into a SHAKE256 sponge, then repeatedly squeeze two bytes, interpret
 * them as a big-endian 16-bit value {@code w}, and keep {@code w mod q} whenever {@code w < 5q =
 * 61445}, until n coefficients are collected. Uses the audited Bouncy Castle SHAKE256 XOF.
 */
public class HashToPointPrecompiledContract extends AbstractPrecompiledContract {

  private static final int Q = 12289;
  private static final int REJECT_BOUND = 5 * Q; // 61445
  private static final int NONCE_LEN = 40;
  private static final int MIN_INPUT = 1 + NONCE_LEN; // logn byte + 40-byte nonce

  private static final int BASE_GAS = 60;
  private static final int GAS_PER_WORD = 12;

  /**
   * Instantiates a new HashToPoint precompiled contract.
   *
   * @param gasCalculator the gas calculator
   */
  HashToPointPrecompiledContract(final GasCalculator gasCalculator) {
    super("AereHashToPoint", gasCalculator);
  }

  /** Ring degree n from the logn selector byte, or 0 if the selector is invalid. */
  private static int degree(final Bytes input) {
    if (input.size() < MIN_INPUT) {
      return 0;
    }
    final int logn = input.get(0) & 0xff;
    if (logn == 9) {
      return 512;
    }
    if (logn == 10) {
      return 1024;
    }
    return 0;
  }

  @Override
  public long gasRequirement(final Bytes input) {
    final int n = degree(input);
    if (n == 0) {
      // Malformed: charge only for the bytes actually presented for hashing.
      final long words = ((long) input.size() + 31) / 32;
      return BASE_GAS + GAS_PER_WORD * words;
    }
    // Absorbed bytes (everything after the logn selector) + expected squeeze. The sampler keeps a
    // sample with probability 61445/65536, so it squeezes ~2*n / 0.9375 bytes on average; charge a
    // conservative fixed 70/64 (~1.094x) expansion so gas is a pure function of the input.
    final long absorbBytes = input.size() - 1L;
    final long squeezeBytes = (2L * n * 70L) / 64L;
    final long words = (absorbBytes + 31) / 32 + (squeezeBytes + 31) / 32;
    return BASE_GAS + GAS_PER_WORD * words;
  }

  @NotNull
  @Override
  public PrecompileContractResult computePrecompile(
      final Bytes input, @NotNull final MessageFrame messageFrame) {
    final int n = degree(input);
    if (n == 0) {
      return PrecompileContractResult.success(Bytes.EMPTY);
    }
    try {
      // Absorb nonce || message (everything after the 1-byte logn selector).
      final byte[] absorbed = input.slice(1).toArrayUnsafe();
      final SHAKEDigest shake = new SHAKEDigest(256);
      shake.update(absorbed, 0, absorbed.length);

      final byte[] out = new byte[2 * n];
      final byte[] two = new byte[2];
      int filled = 0;
      while (filled < n) {
        shake.doOutput(two, 0, 2); // incremental squeeze, keeps the sponge in squeezing phase
        final int w = ((two[0] & 0xff) << 8) | (two[1] & 0xff);
        if (w < REJECT_BOUND) {
          final int coeff = w % Q;
          out[2 * filled] = (byte) (coeff >>> 8);
          out[2 * filled + 1] = (byte) (coeff & 0xff);
          filled++;
        }
      }
      return PrecompileContractResult.success(Bytes.wrap(out));
    } catch (final Throwable t) {
      return PrecompileContractResult.success(Bytes.EMPTY);
    }
  }
}