// Copyright 2022 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine License
// Agreement.

#ifndef WVCDM_CORE_TEST_LICENSE_HOLDER_H_
#define WVCDM_CORE_TEST_LICENSE_HOLDER_H_

#include <string>
#include <vector>

#include <gtest/gtest.h>

#include "cdm_engine.h"
#include "config_test_env.h"
#include "url_request.h"
#include "wv_attributes.h"
#include "wv_cdm_event_listener.h"

#include "log.h"

namespace wvcdm {

// An event listener to tell when a renewal is needed.
class SimpleEventListener : public wvcdm::WvCdmEventListener {
 public:
  SimpleEventListener() { renewal_needed_ = false; }
  // We will want to know when a renewal is needed.
  void OnSessionRenewalNeeded(const CdmSessionId& /*session*/) override {
    renewal_needed_ = true;
  }
  void OnSessionKeysChange(const CdmSessionId&, const CdmKeyStatusMap&,
                           bool) override {}
  void OnExpirationUpdate(const CdmSessionId&,
                          int64_t expiry_time UNUSED) override {}
  bool renewal_needed() { return renewal_needed_; }
  void set_renewal_needed(bool renewal_needed) {
    renewal_needed_ = renewal_needed;
  }

 private:
  bool renewal_needed_ = false;
};

class LicenseHolder {
 public:
  // Set up a new license holder with the specified content id.
  LicenseHolder(const std::string& content_id, CdmEngine* cdm_engine,
                const ConfigTestEnv& config)
      : content_id_(content_id), cdm_engine_(cdm_engine), config_(config) {}
  ~LicenseHolder();
  // Open a session for the license. Must be called before FetchLicense or
  // ReloadLicense.
  void OpenSession();
  // Close the session and release resources.
  void CloseSession();
  // Generate a license request, send it to the license server, and wait for the
  // response. The response is *not* loaded into the cdm engine.
  void FetchLicense();
  // Load the license response into the CDM engine. A call to FetchLicense()
  // must be made first.
  void LoadLicense();
  // Attempt to load the license response into the CDM engine, but expect a
  // failure.
  void FailLoadLicense();
  // Reload the license. Call OpenSession() before calling
  // ReloadLicense(). Also, the key_set_id must have been set previously. The
  // key_set_id is set by calling LoadLicense(), or by calling set_key_set_id().
  void ReloadLicense();
  // Attempt to reload a license, but expect a failure.
  void FailReloadLicense();
  // Generate the renewal request, and send it to the server.
  void GenerateAndPostRenewalRequest(const std::string& renewal_policy_id);
  // Fetch the renewal response. This can add a few seconds of latency.
  void FetchRenewal();
  // Load the renewal response that was fetched in FetchRenewal().
  void LoadRenewal();
  // Generate the release request, and send it to the server.
  void GenerateAndPostReleaseRequest(const std::string& release_policy_id);
  // Fetch the release response. This can add a few seconds of latency.
  void FetchRelease();
  // Load the release response that was fetched in FetchRelease().
  void LoadRelease();
  // Releases the license and frees up entry in usage table.
  void RemoveLicense();

  // Try to decrypt some random data. This does not verify that the data is
  // decrypted correctly. Returns the result of the decrypt operation.
  CdmResponseType Decrypt(const std::string& key_id);
  // Try to copy the clear lead to a secure buffer. Returns the result of the
  // copy buffer operation.
  CdmResponseType DecryptClearLead(const std::string& key_id,
                                   size_t num_samples);
  // Try to decrypt some random data to a secure buffer. If the test harness
  // does not allow creating a secure buffer, then this function fails
  // immediately. Otherwise, a secure buffer is created and used for a
  // decryption operation.
  void DecryptSecure(const KeyId& key_id);
  // Try to decrypt some random data, but expect failure. The failure may
  // be either the expected_status, or NEED_KEY. We allow NEED_KEY in case
  // the server recognized that we cannot support the given key.
  void FailDecrypt(const KeyId& key_id, CdmResponseEnum expected_status);

  const std::string& content_id() const { return content_id_; }
  void set_content_id(const std::string& content_id) {
    content_id_ = content_id;
  }
  // The session id.  This is only valid after a call to OpenSession.
  const std::string& session_id() { return session_id_; }
  // Returns true if the license is offline.
  bool can_persist() const { return can_persist_; }
  // Sets whether the license is offline or not.
  void set_can_persist(bool can_persist) { can_persist_ = can_persist; }
  uint64_t start_of_rental_clock() const { return start_of_rental_clock_; }
  const std::string& key_set_id() const { return key_set_id_; }
  void set_key_set_id(const std::string& key_set_id) {
    key_set_id_ = key_set_id;
  }
  SimpleEventListener& event_listener() { return event_listener_; }
  void set_log_core_message(bool log_core_message) {
    log_core_message_ = log_core_message;
  }
  void set_pst_length(uint64_t pst_length) { pst_length_ = pst_length; }
  void set_renewal_info_check(bool renewal_info_check) {
    renewal_info_check_ = renewal_info_check;
  }
  void set_client_id_check(bool client_id_check) {
    client_id_check_ = client_id_check;
  }

 private:
  std::string content_id_;
  CdmSessionId session_id_;
  CdmKeySetId key_set_id_;
  std::string key_response_;
  bool can_persist_ = false;
  uint64_t start_of_rental_clock_ = 0u;
  CdmEngine* cdm_engine_ = nullptr;
  const ConfigTestEnv& config_;
  SimpleEventListener event_listener_;
  std::unique_ptr<UrlRequest> request_in_flight_;
  std::string request_message_;
  std::string request_response_;
  bool log_core_message_ = false;
  uint64_t pst_length_ = 0u;
  bool renewal_info_check_ = false;
  bool client_id_check_ = false;

  // Generate the license request.
  void GenerateKeyRequest(const InitializationData& init_data,
                          CdmKeyRequest* key_request);

  // Generate a URL for the specified policy. The license request should be sent
  // to this url.
  std::string MakeUrl(const std::string& server_url,
                      const std::string& renewal_policy_id);
  // Fetch the key response from the server.
  void GetKeyResponse(const CdmKeyRequest& key_request);
};

}  // namespace wvcdm

#endif  // WVCDM_CORE_TEST_LICENSE_HOLDER_H_