/*
{{copyright}} 
*/

/*
{{version}} 
*/

/*
{{license}} 
*/

#pragma once

#include "data.h"

#include <filesystem>
#include <fstream>
#include <memory>
#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <span>
#include <string>
#include <vector>

namespace sdi_toolBox::crypto
{
//--------------------------------------------------------------
class DigitalSign
{
  // Custom deleter for EVP_PKEY to ensure it's automatically freed
  struct EVP_PKEY_Deleter
  {
    void operator()(EVP_PKEY *p) const
    {
      if (p)
      {
        EVP_PKEY_free(p);
      }
    }
  };
  using pEVP_PKEY = std::unique_ptr<EVP_PKEY, EVP_PKEY_Deleter>;

  struct BIO_Deleter
  {
    void operator()(BIO *b) const
    {
      if (b)
      {
        BIO_free(b);
      }
    }
  };
  using pBIO = std::unique_ptr<BIO, BIO_Deleter>;

  struct EVP_MD_CTX_Deleter
  {
    void operator()(EVP_MD_CTX *ctx) const
    {
      if (ctx)
      {
        EVP_MD_CTX_free(ctx);
      }
    }
  };
  using pEVP_MD_CTX = std::unique_ptr<EVP_MD_CTX, EVP_MD_CTX_Deleter>;

  public:
  DigitalSign() = default;                                          // Default constructor
  virtual ~DigitalSign()                             = default;     // Default destructor
  DigitalSign(const DigitalSign &obj)                = delete;      // Copy constructor
  DigitalSign(DigitalSign &&obj) noexcept            = delete;      // Move constructor
  DigitalSign &operator=(const DigitalSign &obj)     = delete;      // Copy assignment operator
  DigitalSign &operator=(DigitalSign &&obj) noexcept = delete;      // Move assignment operator

  void loadPEMKeyFromFile(const std::filesystem::path &filepath, bool is_private);     // Load key from PEM file
  void loadPEMKeyFromMemory(const std::string &key, bool is_private);                  // Load key from memory

  std::vector<uint8_t> sign(const Data8 &data, const EVP_MD *digest_type = EVP_sha256()) const;        // Sign data with private key
  std::string          signHex(const Data8 &data, const EVP_MD *digest_type = EVP_sha256()) const;     // Sign data with private key

  bool verify(const Data8 &data, const std::span<const uint8_t> &signature, const EVP_MD *digest_type = EVP_sha256()) const;     // Verify signature with public key
  bool verify(const Data8 &data, const std::string &signature, const EVP_MD *digest_type = EVP_sha256()) const;                  // Verify signature with public key

  protected:
  pEVP_PKEY m_privateKey;
  pEVP_PKEY m_publicKey;

  private:
  static std::vector<uint8_t> hexToBytes(const std::string &hexString);              // Converts a hexadecimal string into a vector of bytes
  static std::string          bytesToHex(const std::span<const uint8_t> &bytes);     // Converts a span of bytes into a hexadecimal string representation
};
//--------------------------------------------------------------
/* Load key from PEM file */
inline void DigitalSign::loadPEMKeyFromFile(const std::filesystem::path &filepath, const bool is_private)
{
  std::ifstream file(filepath, std::ifstream::binary);
  if (!file.is_open())
    throw std::runtime_error("Error: Could not open file " + filepath.string());

  const std::string key((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
  loadPEMKeyFromMemory(key, is_private);
}
//--------------------------------------------------------------
/* Load key from memory */
inline void DigitalSign::loadPEMKeyFromMemory(const std::string &key, const bool is_private)
{
  if (key.empty())
    throw std::runtime_error("Error: Input stream is empty or could not be read");

  // Create a memory BIO (Basic Input/Output) for reading the key
  const auto bio = static_cast<pBIO>(BIO_new_mem_buf(key.data(), static_cast<int>(key.size())));
  if (!bio)
    throw std::runtime_error("Error: Failed to create memory BIO");

  // Use the BIO to read the key
  if (is_private)
  {
    m_privateKey = static_cast<pEVP_PKEY>(PEM_read_bio_PrivateKey(bio.get(), nullptr, nullptr, nullptr));
    if (!m_privateKey)
      throw std::runtime_error("Error: Failed to read private key from memory buffer");
  }
  else
  {
    m_publicKey = static_cast<pEVP_PKEY>(PEM_read_bio_PUBKEY(bio.get(), nullptr, nullptr, nullptr));
    if (!m_publicKey)
      throw std::runtime_error("Error: Failed to read private key from memory buffer");
  }
}
//--------------------------------------------------------------
/* Sign data with private key */
inline std::vector<uint8_t> DigitalSign::sign(const Data8 &data, const EVP_MD *digest_type) const
{
  if (!m_privateKey)
    throw std::runtime_error("Error: Private key is not loaded for signing");

  // Initialize the context for signing
  const auto md_ctx = static_cast<pEVP_MD_CTX>(EVP_MD_CTX_new());
  if (!md_ctx)
    throw std::runtime_error("Error: Failed to create EVP_MD_CTX");

  // Set the context for signing with the private key and digest
  if (EVP_DigestSignInit(md_ctx.get(), nullptr, digest_type, nullptr, m_privateKey.get()) != 1)
    throw std::runtime_error("Error: EVP_DigestSignInit failed");

  // Provide the data to be signed
  if (EVP_DigestSignUpdate(md_ctx.get(), data.getBytes().data(), data.getSize()) != 1)
    throw std::runtime_error("Error: EVP_DigestSignUpdate failed");

  // Determine the required signature size
  size_t sigLen = 0;
  if (EVP_DigestSignFinal(md_ctx.get(), nullptr, &sigLen) != 1)
    throw std::runtime_error("Error: EVP_DigestSignFinal (size) failed");

  // Allocate space and generate the signature
  std::vector<uint8_t> signature(sigLen);
  if (EVP_DigestSignFinal(md_ctx.get(), signature.data(), &sigLen) != 1)
    throw std::runtime_error("Error: EVP_DigestSignFinal (signature) failed");

  // Resize signature to the actual length used
  signature.resize(sigLen);

  return signature;
}
//--------------------------------------------------------------
/* Sign data */
inline std::string DigitalSign::signHex(const Data8 &data, const EVP_MD *digest_type) const
{
  return bytesToHex(sign(data, digest_type));
}
//--------------------------------------------------------------
/* Verify signature with public key  */
inline bool DigitalSign::verify(const Data8 &data, const std::span<const uint8_t> &signature, const EVP_MD *digest_type) const
{
  if (!m_publicKey)
    throw std::runtime_error("Error: Public key is not loaded for verifying");

  // Initialize the context for verifying
  const auto md_ctx = static_cast<pEVP_MD_CTX>(EVP_MD_CTX_new());
  if (!md_ctx)
    throw std::runtime_error("Error: Failed to create EVP_MD_CTX");

  // Set the context for verifying with the public key and digest
  if (EVP_DigestVerifyInit(md_ctx.get(), nullptr, digest_type, nullptr, m_publicKey.get()) != 1)
    throw std::runtime_error("Error: EVP_DigestVerifyInit failed");

  // Provide the original data
  if (EVP_DigestVerifyUpdate(md_ctx.get(), data.getBytes().data(), data.getSize()) != 1)
    throw std::runtime_error("Error: EVP_DigestVerifyUpdate failed");

  // Verify the signature
  // Returns 1 for a good signature, 0 for a bad signature, and -1 on error
  const auto result = EVP_DigestVerifyFinal(md_ctx.get(), signature.data(), signature.size());
  if (result == 1)
    return true;     // Signature is valid
  if (result == 0)
    return false;     // Signature is NOT valid

  throw std::runtime_error("Error: EVP_DigestVerifyFinal failed during execution");
}
//--------------------------------------------------------------
/* Verify signature with public key  */
inline bool DigitalSign::verify(const Data8 &data, const std::string &signature, const EVP_MD *digest_type) const
{
  return verify(data, hexToBytes(signature), digest_type);
}
//--------------------------------------------------------------
/* Converts a hexadecimal string into a vector of bytes */
inline std::vector<uint8_t> DigitalSign::hexToBytes(const std::string &hexString)
{
  // A valid hex string must have an even length
  if (hexString.length() % 2 != 0)
    throw std::invalid_argument("Hex string must have an even length");

  std::vector<uint8_t> bytes;
  bytes.reserve(hexString.length() / 2);

  for (size_t i = 0; i < hexString.length(); i += 2)
  {
    std::string   byteString = hexString.substr(i, 2);
    unsigned long byteVal;

    try
    {
      byteVal = std::stoul(byteString, nullptr, 16);
    }
    catch (const std::exception & /*e*/)
    {
      throw std::invalid_argument("Invalid hexadecimal character sequence: " + byteString);     // Non-hex characters ?
    }

    // Check if the value fits in a uint8_t (should always be true for 2 chars)
    if (byteVal > 0xFF)
      throw std::invalid_argument("Internal error: Converted value exceeds 0xFF");

    bytes.push_back(static_cast<uint8_t>(byteVal));
  }

  return bytes;
}
//--------------------------------------------------------------
/* Converts a span of bytes into a hexadecimal string representation */
inline std::string DigitalSign::bytesToHex(const std::span<const uint8_t> &bytes)
{
  std::stringstream ss;
  ss << std::hex << std::uppercase << std::setfill('0');

  for (const uint8_t byte : bytes)
    ss << std::format("{0:02X}", byte);

  return ss.str();
}
//--------------------------------------------------------------
}     // namespace sdi_toolBox::crypto
