// Part of the Concrete Compiler Project, under the BSD3 License with Zama
// Exceptions. See
// https://github.com/zama-ai/concrete-compiler-internal/blob/main/LICENSE.txt
// for license information.

#ifndef CONCRETELANG_CLIENTLIB_EVALUATION_KEYS_H_
#define CONCRETELANG_CLIENTLIB_EVALUATION_KEYS_H_

#include <memory>

#include "concrete-core-ffi.h"
#include "concretelang/Common/Error.h"

typedef struct LweCircuitBootstrapPrivateFunctionalPackingKeyswitchKeys64
    LweCircuitBootstrapPrivateFunctionalPackingKeyswitchKeys64;

int destroy_lwe_circuit_bootstrap_private_functional_packing_keyswitch_keys_u64(
    LweCircuitBootstrapPrivateFunctionalPackingKeyswitchKeys64 *);

namespace concretelang {
namespace clientlib {

// =============================================

/// Wrapper for `LweKeyswitchKey64` so that it cleans up properly.
class LweKeyswitchKey {
private:
  LweKeyswitchKey64 *ksk;

protected:
  friend std::ostream &operator<<(std::ostream &ostream,
                                  const LweKeyswitchKey &wrappedKsk);
  friend std::istream &operator>>(std::istream &istream,
                                  LweKeyswitchKey &wrappedKsk);

public:
  LweKeyswitchKey(LweKeyswitchKey64 *ksk) : ksk{ksk} {}
  LweKeyswitchKey(LweKeyswitchKey &other) = delete;
  LweKeyswitchKey(LweKeyswitchKey &&other) : ksk{other.ksk} {
    other.ksk = nullptr;
  }
  ~LweKeyswitchKey() {
    if (this->ksk != nullptr) {
      CAPI_ASSERT_ERROR(destroy_lwe_keyswitch_key_u64(this->ksk));

      this->ksk = nullptr;
    }
  }

  LweKeyswitchKey64 *get() { return this->ksk; }
};

// =============================================

/// Wrapper for `LweBootstrapKey64` so that it cleans up properly.
class LweBootstrapKey {
private:
  LweBootstrapKey64 *bsk;

protected:
  friend std::ostream &operator<<(std::ostream &ostream,
                                  const LweBootstrapKey &wrappedBsk);
  friend std::istream &operator>>(std::istream &istream,
                                  LweBootstrapKey &wrappedBsk);

public:
  LweBootstrapKey(LweBootstrapKey64 *bsk) : bsk{bsk} {}
  LweBootstrapKey(LweBootstrapKey &other) = delete;
  LweBootstrapKey(LweBootstrapKey &&other) : bsk{other.bsk} {
    other.bsk = nullptr;
  }
  ~LweBootstrapKey() {
    if (this->bsk != nullptr) {
      CAPI_ASSERT_ERROR(destroy_lwe_bootstrap_key_u64(this->bsk));
      this->bsk = nullptr;
    }
  }

  LweBootstrapKey64 *get() { return this->bsk; }
};

// =============================================

/// Wrapper for `LweCircuitBootstrapPrivateFunctionalPackingKeyswitchKeys64` so
/// that it cleans up properly.
class PackingKeyswitchKey {
private:
  LweCircuitBootstrapPrivateFunctionalPackingKeyswitchKeys64 *key;

protected:
  friend std::ostream &operator<<(std::ostream &ostream,
                                  const PackingKeyswitchKey &wrappedFpksk);
  friend std::istream &operator>>(std::istream &istream,
                                  PackingKeyswitchKey &wrappedFpksk);

public:
  PackingKeyswitchKey(
      LweCircuitBootstrapPrivateFunctionalPackingKeyswitchKeys64 *key)
      : key{key} {}
  PackingKeyswitchKey(PackingKeyswitchKey &other) = delete;
  PackingKeyswitchKey(PackingKeyswitchKey &&other) : key{other.key} {
    other.key = nullptr;
  }
  ~PackingKeyswitchKey() {
    if (this->key != nullptr) {
      CAPI_ASSERT_ERROR(
          destroy_lwe_circuit_bootstrap_private_functional_packing_keyswitch_keys_u64(
              this->key));
      this->key = nullptr;
    }
  }

  LweCircuitBootstrapPrivateFunctionalPackingKeyswitchKeys64 *get() {
    return this->key;
  }
};

// =============================================

/// Evalution keys required for execution.
class EvaluationKeys {
private:
  std::shared_ptr<LweKeyswitchKey> sharedKsk;
  std::shared_ptr<LweBootstrapKey> sharedBsk;
  std::shared_ptr<PackingKeyswitchKey> sharedFpksk;

protected:
  friend std::ostream &operator<<(std::ostream &ostream,
                                  const EvaluationKeys &evaluationKeys);
  friend std::istream &operator>>(std::istream &istream,
                                  EvaluationKeys &evaluationKeys);

public:
  EvaluationKeys()
      : sharedKsk{std::shared_ptr<LweKeyswitchKey>(nullptr)},
        sharedBsk{std::shared_ptr<LweBootstrapKey>(nullptr)} {}

  EvaluationKeys(std::shared_ptr<LweKeyswitchKey> sharedKsk,
                 std::shared_ptr<LweBootstrapKey> sharedBsk,
                 std::shared_ptr<PackingKeyswitchKey> sharedFpksk)
      : sharedKsk{sharedKsk}, sharedBsk{sharedBsk}, sharedFpksk{sharedFpksk} {}

  LweKeyswitchKey64 *getKsk() { return this->sharedKsk->get(); }
  LweBootstrapKey64 *getBsk() { return this->sharedBsk->get(); }
  LweCircuitBootstrapPrivateFunctionalPackingKeyswitchKeys64 *getFpksk() {
    return this->sharedFpksk->get();
  };
};

// =============================================

} // namespace clientlib
} // namespace concretelang

#endif