#ifndef LIB_TARGET_OPENFHEPKE_OPENFHEPKETEMPLATES_H_
#define LIB_TARGET_OPENFHEPKE_OPENFHEPKETEMPLATES_H_

#include <string_view>

namespace mlir {
namespace heir {
namespace openfhe {

constexpr std::string_view kSourceRelativeOpenfheImport = R"cpp(
#include <cassert>
#include <cmath>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <vector>
#include "src/pke/include/openfhe.h"  // from @openfhe
)cpp";
constexpr std::string_view kInstallationRelativeOpenfheImport = R"cpp(
#include <cassert>
#include <cmath>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <vector>
#include "openfhe/pke/openfhe.h"  // from @openfhe
)cpp";
constexpr std::string_view kEmbeddedOpenfheImport = R"cpp(
#include <cassert>
#include <cmath>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <vector>
#include "openfhe.h"
)cpp";

// clang-format off
constexpr std::string_view kModulePreludeTemplate = R"cpp(
using namespace lbcrypto;
using CiphertextT = Ciphertext<DCRTPoly>;
using ConstCiphertextT = ConstCiphertext<DCRTPoly>;
using CCParamsT = CCParams<CryptoContext{0}RNS>;
using CryptoContextT = CryptoContext<DCRTPoly>;
using EvalKeyT = EvalKey<DCRTPoly>;
using PlaintextT = Plaintext;
using PrivateKeyT = PrivateKey<DCRTPoly>;
using PublicKeyT = PublicKey<DCRTPoly>;
)cpp";

constexpr std::string_view kLoadResourceTemplate = R"cpp(
// Assumes a little-endian target.
template <typename T>
std::vector<T> load_resource(const std::string& path, size_t size) {
  std::ifstream file(path, std::ios::binary);
  if (!file.is_open()) {
    std::cerr << "Failed to open file: " << path << std::endl;
    std::abort();
  }
  std::vector<T> data(size);
  file.read(reinterpret_cast<char*>(data.data()), size * sizeof(T));
  if (!file) {
    std::cerr << "Failed to read expected number of bytes from: " << path << std::endl;
    std::abort();
  }
  return data;
}

// Specialization for bool to avoid std::vector<bool>::data() compilation failure.
template <>
inline std::vector<bool> load_resource<bool>(const std::string& path, size_t size) {
  std::ifstream file(path, std::ios::binary);
  if (!file.is_open()) {
    std::cerr << "Failed to open file: " << path << std::endl;
    std::abort();
  }
  std::vector<char> temp(size);
  file.read(temp.data(), size);
  if (!file) {
    std::cerr << "Failed to read expected number of bytes from: " << path << std::endl;
    std::abort();
  }
  std::vector<bool> data(size);
  for (size_t i = 0; i < size; ++i) {
    data[i] = temp[i] != 0;
  }
  return data;
}
)cpp";
// clang-format on

// clang-format off
constexpr std::string_view kPybindImports = R"cpp(
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
)cpp";
// clang-format on

// clang-format off
constexpr std::string_view kPybindCommon = R"cpp(
using namespace lbcrypto;
namespace py = pybind11;

// Minimal bindings required for generated functions to run.
// Cf. https://pybind11.readthedocs.io/en/stable/advanced/classes.html#module-local-class-bindings
// which is a temporary workaround to allow us to have multiple compilations in
// the same python program. Better would be to cache the pybind11 module across
// calls.
void bind_common(py::module &m)
{
    py::class_<PublicKeyImpl<DCRTPoly>, std::shared_ptr<PublicKeyImpl<DCRTPoly>>>(m, "PublicKey", py::module_local())
        .def(py::init<>());
    py::class_<PrivateKeyImpl<DCRTPoly>, std::shared_ptr<PrivateKeyImpl<DCRTPoly>>>(m, "PrivateKey", py::module_local())
        .def(py::init<>());
    py::class_<KeyPair<DCRTPoly>>(m, "KeyPair", py::module_local())
        .def_readwrite("publicKey", &KeyPair<DCRTPoly>::publicKey)
        .def_readwrite("secretKey", &KeyPair<DCRTPoly>::secretKey);
    py::class_<CiphertextImpl<DCRTPoly>, std::shared_ptr<CiphertextImpl<DCRTPoly>>>(m, "Ciphertext", py::module_local())
        .def(py::init<>());
    py::class_<PlaintextImpl, std::shared_ptr<PlaintextImpl>>(m, "Plaintext", py::module_local());
    py::class_<CryptoContextImpl<DCRTPoly>, std::shared_ptr<CryptoContextImpl<DCRTPoly>>>(m, "CryptoContext", py::module_local())
        .def(py::init<>())
        .def("KeyGen", &CryptoContextImpl<DCRTPoly>::KeyGen);
}
)cpp";
// clang-format on

// clang-format off
constexpr std::string_view kPybindModuleTemplate = R"cpp(
PYBIND11_MODULE({0}, m) {{
  bind_common(m);
)cpp";
// clang-format on

// Emit a pybind11 binding that releases the GIL for the duration of the C++
// function call.  This enables multi-threaded C++ code (e.g. OpenMP parallel
// regions inside OpenFHE) to run concurrently with the Python interpreter.
// The `py::call_guard<py::gil_scoped_release>()` helper ensures the GIL is
// relinquished on entry and re-acquired on exit.
constexpr std::string_view kPybindFunctionTemplate =
    "m.def(\"{0}\", &{0}, py::call_guard<py::gil_scoped_release>());";

// Emit a pybind11 class binding for a struct, used for functions with multiple
// return values. {0} is the struct name.
constexpr std::string_view kPybindStructClassTemplate =
    R"cpp(py::class_<{0}>(m, "{0}", py::module_local()).def(py::init<>())
    )cpp";

// Emit a pybind11 field binding for a struct member.
// {0} is the field index (e.g., 0, 1).
// {1} is the struct name.
constexpr std::string_view kPybindStructFieldTemplate =
    "    .def_readwrite(\"arg{0}\", &{1}::arg{0})\n";

// clang-format off
constexpr std::string_view KdebugHeaderImports = R"cpp(
#include <map>
#include <string>
)cpp";
// clang-format on

// clang-format off
constexpr std::string_view kDebugBlockArgCheckTemplate = R"cpp(
auto {0} = {1}.at("asm.is_block_arg");
if ({0} == "1") {{
  std::cout << "Input" << std::endl;
}} else {{
  std::cout << {1}.at("asm.op_name") << std::endl;
}}
)cpp";

constexpr std::string_view kDebugDecryptTemplate = R"cpp(
PlaintextT {0};
{1}->Decrypt({2}, {3}, &{0});
{0}->SetLength(std::stoul({4}.at("message.size")));
std::cout << "  " << {0} << std::endl;
)cpp";
// clang-format on

}  // namespace openfhe
}  // namespace heir
}  // namespace mlir

#endif  // LIB_TARGET_OPENFHEPKE_OPENFHEPKETEMPLATES_H_
