/*
 * The MIT License (MIT)

 * Copyright (c) 2025 libsig Developers

 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:

 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.

 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

#ifndef LIBSIG_ECC_H_
#define LIBSIG_ECC_H_

#include <iostream>
#include <string>
#include <memory>
#include <fstream>
#include <filesystem>

#include <openssl/evp.h>
#include <openssl/err.h>
#include <openssl/pem.h>
#include <openssl/bio.h>
#include <openssl/ec.h>

namespace fs = std::filesystem;

namespace libsig {

struct EVP_PKEY_Deleter {
  void operator()(EVP_PKEY* p) const { EVP_PKEY_free(p); }
};
using EVP_PKEY_ptr = std::unique_ptr<EVP_PKEY, EVP_PKEY_Deleter>;

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

class ECCrypto {
 public:
  ECCrypto() {
    OpenSSL_add_all_algorithms();
    ERR_load_crypto_strings();
  }

  ~ECCrypto() {
    EVP_cleanup();
    ERR_free_strings();
  }

  // Generate key-pairs using secp256k1
  EVP_PKEY_ptr generateKeyPair(int curve_nid = NID_secp256k1) {
    EVP_PKEY_ptr pkey(nullptr);
        
    EVP_PKEY_CTX* ctx = EVP_PKEY_CTX_new_id(EVP_PKEY_EC, nullptr);
    if (!ctx) {
      std::cerr << "Failed to create EVP_PKEY_CTX" << std::endl;
      return pkey;
    }

    if (EVP_PKEY_keygen_init(ctx) <= 0) {
      std::cerr << "Failed to initialize key generation" << std::endl;
      EVP_PKEY_CTX_free(ctx);
      return pkey;
    }

    if (EVP_PKEY_CTX_set_ec_paramgen_curve_nid(ctx, curve_nid) <= 0) {
      std::cerr << "Failed to set EC curve" << std::endl;
      EVP_PKEY_CTX_free(ctx);
      return pkey;
    }
        
    EVP_PKEY* key = nullptr;
    if (EVP_PKEY_keygen(ctx, &key) <= 0) {
      std::cerr << "Failed to generate EC key pair" << std::endl;
      EVP_PKEY_CTX_free(ctx);
      return pkey;
    }
        
    EVP_PKEY_CTX_free(ctx);
    pkey.reset(key);
    return pkey;
  }

  // store public key as PEM
  std::string publicKeyToPEM(const EVP_PKEY_ptr& pkey) {
    if (!pkey) return "";

    BIO_ptr bio(BIO_new(BIO_s_mem()));
    if (!PEM_write_bio_PUBKEY(bio.get(), pkey.get())) {
      std::cerr << "Failed to write public key to PEM" << std::endl;
      return "";
    }

    BUF_MEM* mem = nullptr;
    BIO_get_mem_ptr(bio.get(), &mem);
    std::string result(mem->data, mem->length);
    return result;
  }

  // store private key as PEM
  std::string privateKeyToPEM(const EVP_PKEY_ptr& pkey) {
    if (!pkey) return "";

    BIO_ptr bio(BIO_new(BIO_s_mem()));
    if (!PEM_write_bio_PrivateKey(bio.get(), pkey.get(), 
        nullptr, nullptr, 0, nullptr, nullptr)) {
      std::cerr << "Failed to write private key to PEM" << std::endl;
      return "";
    }

    BUF_MEM* mem = nullptr;
    BIO_get_mem_ptr(bio.get(), &mem);
    std::string result(mem->data, mem->length);
    return result;
  }

  // Save keys to file
  void saveKeyTofile(const std::string& filepath, 
  	                 const std::string& key) {
    try {
      fs::path path(filepath);
      if (path.has_parent_path() && 
          !fs::exists(path.parent_path())) {
        fs::create_directories(path.parent_path());
      }
        
      std::ofstream file(filepath);
      if (!file) {
        exit(0);
      }
      file << key;
    } catch (const std::exception& e) {
      std::cerr << "Error: " << e.what() << std::endl;
    }
  }

  // Load public key from file
  EVP_PKEY_ptr LoadPublicKey(const std::string& file_path) {
    std::cout << file_path << std::endl;
    std::ifstream file(file_path);
    std::string pem_key = {std::istreambuf_iterator<char>(file), 
            std::istreambuf_iterator<char>()};

    BIO_ptr bio(BIO_new_mem_buf(pem_key.data(), pem_key.size()));
    if (!bio) {
        throw std::runtime_error("Failed to create BIO");
    }

    EVP_PKEY* pkey = PEM_read_bio_PUBKEY(bio.get(), nullptr, nullptr, nullptr);
    if (!pkey) {
        ERR_print_errors_fp(stderr);
        throw std::runtime_error("Failed to parse public key");
    }

    return EVP_PKEY_ptr(pkey);
  }

  // Load private key from file
  EVP_PKEY_ptr LoadPrivateKey(const std::string& file_path) {
    std::ifstream file(file_path);
    std::string pem_key = {std::istreambuf_iterator<char>(file), 
            std::istreambuf_iterator<char>()};
    BIO_ptr bio(BIO_new_mem_buf(pem_key.data(), pem_key.size()));
    if (!bio) {
        throw std::runtime_error("Failed to create BIO");
    }
    EVP_PKEY* pkey = PEM_read_bio_PrivateKey(
        bio.get(), 
        nullptr, 
        [](char* buf, int size, int rwflag, void* u) -> int { // 密码回调
            if (!u) return 0;
            const char* pass = static_cast<const char*>(u);
            int len = strlen(pass);
            if (len > size) len = size;
            memcpy(buf, pass, len);
            return len;
        }, 
        (void*)nullptr
    );
    if (!pkey) {
        ERR_print_errors_fp(stderr);
        throw std::runtime_error("Failed to parse private key");
    }

    return EVP_PKEY_ptr(pkey);
  }

  // sign a message using peivate key and sha256 hash
  bool signMessage(const EVP_PKEY_ptr& pkey, 
                   const std::string& message, 
                   std::string& signature) {
    if (!pkey || message.empty()) return false;

    EVP_MD_CTX* ctx = EVP_MD_CTX_new();
    if (!ctx) {
      std::cerr << "Failed to create EVP_MD_CTX" << std::endl;
      return false;
    }

    if (EVP_DigestSignInit(ctx, nullptr, 
        EVP_sha256(), nullptr, pkey.get()) <= 0) {
      std::cerr << "Failed to initialize signing" << std::endl;
      EVP_MD_CTX_free(ctx);
      return false;
    }

    if (EVP_DigestSignUpdate(ctx, message.data(), message.size()) <= 0) {
      std::cerr << "Failed to update signing" << std::endl;
      EVP_MD_CTX_free(ctx);
      return false;
    }

    size_t sig_len = 0;
    if (EVP_DigestSignFinal(ctx, nullptr, &sig_len) <= 0) {
      std::cerr << "Failed to get signature length" << std::endl;
      EVP_MD_CTX_free(ctx);
      return false;
    }

    signature.resize(sig_len);
    if (EVP_DigestSignFinal(ctx, 
        reinterpret_cast<unsigned char*>(&signature[0]), 
        &sig_len) <= 0) {
      std::cerr << "Failed to create signature" << std::endl;
      EVP_MD_CTX_free(ctx);
      return false;
    }

    signature.resize(sig_len);
    EVP_MD_CTX_free(ctx);
    return true;
  }

  // Verify signature by public key
  bool verifySignature(const EVP_PKEY_ptr& pkey, 
                       const std::string& message, 
                       const std::string& signature) {
    if (!pkey || message.empty() || signature.empty()) return false;

    EVP_MD_CTX* ctx = EVP_MD_CTX_new();
    if (!ctx) {
      std::cerr << "Failed to create EVP_MD_CTX" << std::endl;
      return false;
    }

    if (EVP_DigestVerifyInit(ctx, nullptr, EVP_sha256(), nullptr, pkey.get()) <= 0) {
      std::cerr << "Failed to initialize verification" << std::endl;
      EVP_MD_CTX_free(ctx);
      return false;
    }

    if (EVP_DigestVerifyUpdate(ctx, message.data(), message.size()) <= 0) {
      std::cerr << "Failed to update verification" << std::endl;
      EVP_MD_CTX_free(ctx);
      return false;
    }

    int result = EVP_DigestVerifyFinal(ctx, 
      reinterpret_cast<const unsigned char*>(signature.data()), 
      signature.size());

    EVP_MD_CTX_free(ctx);
    return result == 1;
  }

 private:
  void handleErrors() {
    ERR_print_errors_fp(stderr);
    abort();
  }
};

} // namespace libsig

#endif // LIBSIG_ECC_H_