/*
 * Copyright 2019 Google LLC.
 * 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
 *
 *     https://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.
 */

#ifndef PRIVATE_JOIN_AND_COMPUTE_CRYPTO_DODIS_YAMPOLSKIY_PRF_DY_VERIFIABLE_RANDOM_FUNCTION_H_
#define PRIVATE_JOIN_AND_COMPUTE_CRYPTO_DODIS_YAMPOLSKIY_PRF_DY_VERIFIABLE_RANDOM_FUNCTION_H_

#include <stdint.h>

#include <memory>
#include <optional>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

#include "private_join_and_compute/crypto/big_num.h"
#include "private_join_and_compute/crypto/dodis_yampolskiy_prf/dy_verifiable_random_function.pb.h"
#include "private_join_and_compute/crypto/ec_point.h"
#include "private_join_and_compute/crypto/pedersen_over_zn.h"

namespace private_join_and_compute {

// Implements a Verifiable Random Function that allows provable evaluations of
// the Dodis Yampolskiy (DY) PRF with key k on a point x,
// where k and x are both committed to using a Pedersen commitment. The
// Dodis-Yampolskiy PRF is defined as F_k(x) = g^(1/(k+x)). This class assumes
// the DY PRF is implemented over an elliptic curve group, and that commitments
// are over Zn.
//
// The verification protocol is achieved by proving knowledge of the exponent
// k+x. Specifically, the prover commits to k+x and provides a sigma-protocol
// proving that F_k(x)^(k+x) = g. This sigma-protocol can be made
// non-interactive using the Fiat-Shamir heuristic (namely generating the
// verifier's random challenge by hashing the prover's first message).
class DyVerifiableRandomFunction {
 public:
  struct GenerateKeysProof {};

  // Creates a VRF object from the supplied parameters.
  //
  // "pedersen" should be created externally using the PedersenParameters from
  // parameters_proto.
  //
  // Does not take ownership of context, ec_group or pedersen.
  static StatusOr<std::unique_ptr<DyVerifiableRandomFunction>> Create(
      proto::DyVrfParameters parameters_proto, Context* context,
      ECGroup* ec_group, PedersenOverZn* pedersen);

  // Generates a new public/private keypair for the DY VRF together with a proof
  // that each entry of the commitment is the same key.
  StatusOr<std::tuple<proto::DyVrfPublicKey, proto::DyVrfPrivateKey,
                      proto::DyVrfGenerateKeysProof>>
  GenerateKeyPair();

  // Verifies that the public key has a bounded key, and commits to the same key
  // in each component of the Pedersen batch commitment.
  Status VerifyGenerateKeysProof(const proto::DyVrfPublicKey& public_key,
                                 const proto::DyVrfGenerateKeysProof& proof);

  // Applies the DY VRF to a given batch of messages.
  StatusOr<std::vector<ECPoint>> Apply(
      absl::Span<const BigNum> messages,
      const proto::DyVrfPrivateKey& private_key);

  // Generates a proof that prf_evaluations were generated by applying the VRF
  // to the supplied messages. A commitment and opening to the messages can
  // optionally be provided.
  StatusOr<proto::DyVrfApplyProof> GenerateApplyProof(
      absl::Span<const BigNum> messages,
      absl::Span<const ECPoint> prf_evaluations,
      const proto::DyVrfPublicKey& public_key,
      const proto::DyVrfPrivateKey& private_key,
      const PedersenOverZn::CommitmentAndOpening& commit_and_open_messages);

  // Verifies that prf_evaluations was produced by applying a DY VRF with the
  // supplied public key on the supplied committed messages. "Ok" status
  // corresponds to a passing proof. "Non-ok" status contains an error
  // specifying which check failed.
  //
  // Assumes that the Pedersen commitment parameters and the prover's Public Key
  // are already verified or trustworthy.
  //
  // The challenge can optionally be injected manually into the proof struct. If
  // not, the challenge is generated using the Fiat-Shamir heuristic.
  Status VerifyApplyProof(absl::Span<const ECPoint> prf_evaluations,
                          const proto::DyVrfPublicKey& public_key,
                          const PedersenOverZn::Commitment& commit_messages,
                          const proto::DyVrfApplyProof& proof);

  // Generates the challenge for the ApplyProof using the Fiat-Shamir heuristic.
  // Exposed for testing and special cases such as enclosing proofs.
  StatusOr<BigNum> GenerateApplyProofChallenge(
      absl::Span<const ECPoint> prf_evaluations,
      const proto::DyVrfPublicKey& public_key,
      const PedersenOverZn::Commitment& commit_messages,
      const proto::DyVrfApplyProof::Message1& message_1);

 private:
  struct PublicKey {
    PedersenOverZn::Commitment commit_key;
  };

  struct PrivateKey {
    BigNum key;
    PedersenOverZn::Opening commit_key_opening;
  };

  // Container for proof elements showing that a particular value is the result
  // of applying the DY VRF on committed messages with a particular key. Also
  // has structs for various helpful intermediates.
  struct ApplyProof {
    struct Message1 {
      PedersenOverZn::Commitment commit_dummy_messages_plus_key;
      std::vector<ECPoint> dummy_dy_prf_base_gs;
    };

    struct Message1PrivateState {
      std::vector<BigNum> dummy_messages_plus_key;
      BigNum dummy_key;
      PedersenOverZn::Opening dummy_opening;
    };

    struct Message2 {
      std::vector<BigNum> masked_dummy_messages_plus_key;
      PedersenOverZn::Opening masked_dummy_opening;
    };
  };

  DyVerifiableRandomFunction(proto::DyVrfParameters parameters_proto,
                             Context* context, ECGroup* ec_group,
                             ECPoint dy_prf_base_g, PedersenOverZn* pedersen)
      : parameters_proto_(std::move(parameters_proto)),
        context_(context),
        ec_group_(ec_group),
        dy_prf_base_g_(std::move(dy_prf_base_g)),
        pedersen_(pedersen) {}

  // Produces the first message of the proof that prf_evaluations is the result
  // of a VRF applied to a given set of messages.
  StatusOr<std::pair<std::unique_ptr<ApplyProof::Message1>,
                     std::unique_ptr<ApplyProof::Message1PrivateState>>>
  GenerateApplyProofMessage1(
      absl::Span<const BigNum> messages,
      absl::Span<const ECPoint> prf_evaluations,
      const PedersenOverZn::CommitmentAndOpening& commit_and_open_messages,
      const PublicKey& public_key, const PrivateKey& private_key);

  // Produces the second message of the proof that prf_evaluations is the result
  // of a VRF applied to a given set of messages.
  //
  // The challenge can be optionally generated using the Fiat-Shamir heuristic.
  StatusOr<std::unique_ptr<ApplyProof::Message2>> GenerateApplyProofMessage2(
      absl::Span<const BigNum> messages,
      absl::Span<const ECPoint> prf_evaluations,
      const PedersenOverZn::CommitmentAndOpening& commit_and_open_messages,
      const PublicKey& public_key, const PrivateKey& private_key,
      const ApplyProof::Message1& message_1,
      const ApplyProof::Message1PrivateState& private_state,
      const BigNum& challenge);

  proto::DyVrfPublicKey DyVrfPublicKeyToProto(
      const DyVerifiableRandomFunction::PublicKey& public_key);
  proto::DyVrfPrivateKey DyVrfPrivateKeyToProto(
      const DyVerifiableRandomFunction::PrivateKey& private_key);
  StatusOr<proto::DyVrfApplyProof::Message1> DyVrfApplyProofMessage1ToProto(
      const DyVerifiableRandomFunction::ApplyProof::Message1& message_1);
  proto::DyVrfApplyProof::Message2 DyVrfApplyProofMessage2ToProto(
      const DyVerifiableRandomFunction::ApplyProof::Message2& message_2);

  StatusOr<DyVerifiableRandomFunction::PublicKey> ParseDyVrfPublicKeyProto(
      Context* ctx, const proto::DyVrfPublicKey& public_key_proto);
  StatusOr<DyVerifiableRandomFunction::PrivateKey> ParseDyVrfPrivateKeyProto(
      Context* ctx, const proto::DyVrfPrivateKey& private_key_proto);
  StatusOr<DyVerifiableRandomFunction::ApplyProof::Message1>
  ParseDyVrfApplyProofMessage1Proto(
      Context* ctx, ECGroup* ec_group,
      const proto::DyVrfApplyProof::Message1& message_1_proto);
  StatusOr<DyVerifiableRandomFunction::ApplyProof::Message2>
  ParseDyVrfApplyProofMessage2Proto(
      Context* ctx, const proto::DyVrfApplyProof::Message2& message_2_proto);

  // Generates the challenge for the GenerateKeysProof using the Fiat-Shamir
  // heuristic.
  StatusOr<BigNum> GenerateChallengeForGenerateKeysProof(
      const proto::DyVrfGenerateKeysProof::Statement& statement,
      const proto::DyVrfGenerateKeysProof::Message1& message_1);

  proto::DyVrfParameters parameters_proto_;
  Context* context_;
  ECGroup* ec_group_;
  ECPoint dy_prf_base_g_;
  PedersenOverZn* pedersen_;
};

}  // namespace private_join_and_compute

#endif  // PRIVATE_JOIN_AND_COMPUTE_CRYPTO_DODIS_YAMPOLSKIY_PRF_DY_VERIFIABLE_RANDOM_FUNCTION_H_