// Copyright 2026 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 THIRD_PARTY_CREDENTIO_TESTING_LOCAL_KEY_CRYPTO_WRITE_HANDLER_H_
#define THIRD_PARTY_CREDENTIO_TESTING_LOCAL_KEY_CRYPTO_WRITE_HANDLER_H_

#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <vector>

#include "absl/base/nullability.h"
#include "absl/base/thread_annotations.h"
#include "absl/container/flat_hash_map.h"
#include "absl/functional/any_invocable.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/clock_interface.h"
#include "absl/time/time.h"
#include "crypto/algorithms.h"
#include "crypto/crypto_write_handler.h"
#include "crypto/timestamp_write_handler.h"
#include "tink/public_key_sign.h"

namespace credentio {

// Test signer that uses a fixed key and a caller-supplied certificate chain.
//
// To mimic OCSP responses going stale, `ocsp_responses()` returns an empty list
// once a week has passed on `clock` since construction.
//
// This class is thread-safe.
class LocalKeySigner final : public FixedKeySigner {
 public:
  // `certificate_chain` and `ocsp_responses` are DER-encoded, leaf first.
  // `clock` must outlive this object.
  LocalKeySigner(
      std::unique_ptr<crypto::tink::PublicKeySign> absl_nonnull signer,
      std::vector<std::string> certificate_chain,
      std::vector<std::string> ocsp_responses, absl::Clock* absl_nonnull clock,
      SigningAlgorithm algorithm = SigningAlgorithm::kEs256);

  LocalKeySigner(const LocalKeySigner&) = delete;
  LocalKeySigner& operator=(const LocalKeySigner&) = delete;

  SigningAlgorithm algorithm() const override { return algorithm_; }
  const std::vector<std::string>& certificate_chain() const override {
    return certificate_chain_;
  }
  const std::vector<std::string>& ocsp_responses() const override;
  uint64_t max_signature_byte_size() const override;
  absl::StatusOr<SignResponse> Sign(
      absl::string_view to_be_signed) const override;

 private:
  const std::unique_ptr<crypto::tink::PublicKeySign> absl_nonnull signer_;
  const std::vector<std::string> certificate_chain_;
  const std::vector<std::string> ocsp_responses_;
  absl::Clock* absl_nonnull const clock_;
  const absl::Time create_time_;
  const SigningAlgorithm algorithm_;

  mutable absl::Mutex mutex_;
  mutable std::optional<uint64_t> max_signature_byte_size_
      ABSL_GUARDED_BY(mutex_);
};

// Test implementation of CryptoWriteHandler that uses a fixed claim signer and
// a fixed timestamp signer, both supplied by the caller.
//
// For local testing only: keys are held in process memory and timestamps are
// issued locally from the supplied TSA key rather than by an RFC 3161
// time-stamping authority.
//
// This class is thread-safe.
class LocalKeyCryptoWriteHandler final : public CryptoWriteHandler {
 public:
  // `clock` must outlive this object. If `tsa_signer` is null or
  // `is_timestamping_disabled` is true, `CreateTimestamp()` fails with
  // UNAVAILABLE. `tsa_cert_pem` is the PEM-encoded certificate of `tsa_signer`.
  LocalKeyCryptoWriteHandler(
      absl::Clock* absl_nonnull clock,
      std::unique_ptr<crypto::tink::PublicKeySign> absl_nonnull signer,
      std::unique_ptr<crypto::tink::PublicKeySign> absl_nullable tsa_signer,
      bool is_timestamping_disabled, std::vector<std::string> certificate_chain,
      std::vector<std::string> ocsp_responses, SigningAlgorithm algorithm,
      std::string tsa_cert_pem);

  LocalKeyCryptoWriteHandler(const LocalKeyCryptoWriteHandler&) = delete;
  LocalKeyCryptoWriteHandler& operator=(const LocalKeyCryptoWriteHandler&) =
      delete;

  const FixedKeySigner& GetSigner() override { return signer_; }

  void CreateTimestamp(HashAlgorithm hash_algorithm, std::string hashed_message,
                       TimestampWriteHandler::Options options,
                       absl::AnyInvocable<void(absl::StatusOr<std::string>) &&>
                           cb) const override;

  uint64_t max_timestamp_byte_size(HashAlgorithm hash_algorithm) const override;

 private:
  bool is_timestamping_enabled() const {
    return !is_timestamping_disabled_ && tsa_signer_ != nullptr;
  }

  absl::Clock* absl_nonnull const clock_;
  LocalKeySigner signer_;
  const std::unique_ptr<crypto::tink::PublicKeySign> absl_nullable tsa_signer_;
  const bool is_timestamping_disabled_;
  const std::string tsa_cert_pem_;

  mutable absl::Mutex mutex_;
  mutable absl::flat_hash_map<HashAlgorithm, uint64_t>
      max_timestamp_byte_size_map_ ABSL_GUARDED_BY(mutex_);
};

// Credentials for the test claim signer.
struct LocalKeyClaimSignerCredentials {
  // The signing key, as a PEM-encoded EC P-256 private key.
  std::string key_pem;
  // The certificate chain for `key_pem`, DER-encoded, leaf first.
  std::vector<std::string> certificate_chain;
  // DER-encoded OCSP responses covering `certificate_chain`. May be empty.
  std::vector<std::string> ocsp_responses;
};

// Credentials for the test timestamp authority.
struct LocalKeyTsaCredentials {
  // The signing key, as a PEM-encoded EC P-256 private key.
  std::string key_pem;
  // The certificate for `key_pem`, PEM-encoded.
  std::string cert_pem;
};

// Options for `CreateLocalKeyCryptoWriteHandler()`.
struct LocalKeyCryptoWriteHandlerOptions {
  // Credentials to sign claims with. Required.
  LocalKeyClaimSignerCredentials claim_signer;
  // Credentials to sign timestamps with. Required unless
  // `disable_timestamping` is true.
  LocalKeyTsaCredentials tsa;
  // The clock to read the signing and timestamping time from. It must be set
  // to a time within the validity period of the certificates above, and must
  // outlive the returned handler. If null, the real clock is used.
  absl::Clock* absl_nullable clock = nullptr;
  // If true, timestamping is disabled and returns UNAVAILABLE.
  bool disable_timestamping = false;
};

// Creates a CryptoWriteHandler that signs claims with ES256 and timestamps with
// the given test credentials. It performs no validation of its own, so the
// caller is responsible for pairing the keys with certificates that their
// verifier trusts.
absl::StatusOr<std::unique_ptr<CryptoWriteHandler>>
CreateLocalKeyCryptoWriteHandler(LocalKeyCryptoWriteHandlerOptions options);

}  // namespace credentio

#endif  // THIRD_PARTY_CREDENTIO_TESTING_LOCAL_KEY_CRYPTO_WRITE_HANDLER_H_