#include <openssl/ecdsa.h>

#include "ec_keys.hh"

#include "loggers.hh" // To be remove to build a security shared library

ec_keys::ec_keys(const ec_elliptic_curves p_elliptic_curve): _elliptic_curve(p_elliptic_curve), _ec_key(nullptr), _ec_group(nullptr), _bn_ctx(nullptr), _pr_key(), _pu_key_x(), _pu_key_y() {
  loggers::get_instance().log(">>> ec_keys::ec_keys: %d", static_cast<int>(p_elliptic_curve));
  
  const int result = init();
  if (result == -1) {
    loggers::get_instance().error("ec_keys::ec_keys: Unsupported elliptic_curve %d", _elliptic_curve);
  }
} // End of constructor

ec_keys::ec_keys(const ec_elliptic_curves p_elliptic_curve, const std::vector<unsigned char>& p_private_key): _elliptic_curve(p_elliptic_curve), _ec_key(nullptr), _ec_group(nullptr), _bn_ctx(nullptr), _pr_key(p_private_key), _pu_key_x(), _pu_key_y() {
  loggers::get_instance().log(">>> ec_keys::ec_keys (1): %d", static_cast<int>(p_elliptic_curve));
  
  // Sanity checks
  if ((_elliptic_curve == ec_elliptic_curves::nist_p_256) || (_elliptic_curve == ec_elliptic_curves::brainpool_p_256_r1)) {
    if (p_private_key.size() != 32) {
      loggers::get_instance().error("ec_keys::ec_keys: Invalid public keys size");
    }
  } else if (_elliptic_curve == ec_elliptic_curves::brainpool_p_384_r1) {
    if ((p_private_key.size() != 48)) {
      loggers::get_instance().error("ec_keys::ec_keys: Invalid public keys size");
    }
  }
  
  int result = init();
  if (result == -1) {
    loggers::get_instance().error("ec_keys::ec_keys: Unsupported elliptic_curve %d", _elliptic_curve);
  }
  ::EC_KEY_set_conv_form(_ec_key, POINT_CONVERSION_COMPRESSED);
  
  // Build private key
  BIGNUM p;
  ::BN_init(&p);
  ::BN_bin2bn(_pr_key.data(), _pr_key.size(), &p);
  // Build public keys
  EC_POINT* ec_point = ::EC_POINT_new(_ec_group);
  ::EC_POINT_mul(_ec_group, ec_point, &p, NULL, NULL, _bn_ctx);
  // Set private key
  ::EC_KEY_set_private_key(_ec_key, &p);
  if (::EC_KEY_check_key(_ec_key) != 0) {
    loggers::get_instance().error("ec_keys::ec_keys (1): Invalid private key");
  }
  // Private key is correct, set public keys
  ::EC_KEY_set_public_key(_ec_key, ec_point);

  BIGNUM xy;
  ::BN_init(&xy);
  ::EC_POINT_point2bn(_ec_group, ec_point, POINT_CONVERSION_COMPRESSED, &xy, _bn_ctx);
  std::vector<unsigned char> v(BN_num_bytes(&xy));
  ::BN_bn2bin(&xy, v.data());
  const int l = v.size() / 2;
  _pu_key_x.resize(l);
  std::copy(v.cbegin(), v.cbegin() + l - 1, _pu_key_x.begin());
  _pu_key_y.resize(l);
  std::copy(v.cbegin() + l, v.cend(), _pu_key_y.begin());
  ::EC_POINT_free(ec_point);
  
} // End of constructor

ec_keys::ec_keys(const ec_elliptic_curves p_elliptic_curve, const std::vector<unsigned char>& p_public_key_x, const std::vector<unsigned char>& p_public_key_y): _elliptic_curve(p_elliptic_curve), _ec_key(nullptr), _ec_group(nullptr), _bn_ctx(nullptr), _pr_key(), _pu_key_x(p_public_key_x), _pu_key_y(p_public_key_y) {
  loggers::get_instance().log(">>> ec_keys::ec_keys (2): %d", static_cast<int>(p_elliptic_curve));
  
  // Sanity checks
  if ((_elliptic_curve == ec_elliptic_curves::nist_p_256) || (_elliptic_curve == ec_elliptic_curves::brainpool_p_256_r1)) {
    if ((p_public_key_x.size() != 32) || (p_public_key_y.size() != 32)) {
      loggers::get_instance().error("ec_keys::ec_keys: Invalid public keys size");
    }
  } else if (_elliptic_curve == ec_elliptic_curves::brainpool_p_384_r1) {
    if ((p_public_key_x.size() != 48) || (p_public_key_y.size() != 48)) {
      loggers::get_instance().error("ec_keys::ec_keys: Invalid public keys size");
    }
  }
  
  int result = init();
  if (result == -1) {
    loggers::get_instance().error("ec_keys::ec_keys: Unsupported elliptic_curve %d", _elliptic_curve);
  }
  ::EC_KEY_set_conv_form(_ec_key, POINT_CONVERSION_COMPRESSED);
  
  // Set public key
  BIGNUM x;
  ::BN_init(&x);
  ::BN_bin2bn(_pu_key_x.data(), _pu_key_x.size(), &x);
  BIGNUM y;
  ::BN_init(&y);
  ::BN_bin2bn(_pu_key_y.data(), _pu_key_y.size(), &y);
  EC_POINT* ec_point = ::EC_POINT_new(_ec_group);
  result = 0;
  switch (_elliptic_curve) {
  case ec_elliptic_curves::nist_p_256: // Use primary
    // No break;
  case ec_elliptic_curves::brainpool_p_256_r1:
    // No break;
  case ec_elliptic_curves::brainpool_p_384_r1:
    result = ::EC_POINT_set_affine_coordinates_GFp(_ec_group, ec_point, &x, &y, _bn_ctx); // Use primary elliptic curve
    break;
  default: // Use Binary
    result = ::EC_POINT_set_affine_coordinates_GF2m(_ec_group, ec_point, &x, &y, _bn_ctx);
  } // End of 'switch' statement
  if (result == 0) {
    loggers::get_instance().error("ec_keys::ec_keys (1): Failed to get coordinates");
  }
  ::EC_KEY_set_public_key(_ec_key, ec_point);
  ::EC_POINT_free(ec_point);
} // End of constructor

ec_keys::~ec_keys() {
  loggers::get_instance().log(">>> ec_keys::~ec_keys");
  
  if(_ec_key != nullptr) {
    ::EC_KEY_free(_ec_key);
  }
  if (_bn_ctx != nullptr) {
    ::BN_CTX_free(_bn_ctx);
  }
  
  loggers::get_instance().log("<<< ec_keys::~ec_keys");
} // End of Destructor

int ec_keys::generate() {
  loggers::get_instance().log(">>> ec_keys::generate");
  
  // Sanity check
  if (!::EC_KEY_generate_key(_ec_key)) { // Generate the private and public key 
    loggers::get_instance().error("ec_keys::generate: Wrong parameters");
    return -1;
  }

  BIGNUM x, y;
  ::BN_init(&x);
  ::BN_init(&y);
  const EC_POINT* ec_point = EC_KEY_get0_public_key(_ec_key);
  int result = 0;
  switch (_elliptic_curve) {
  case ec_elliptic_curves::nist_p_256: // Use primary
    // No break;
  case ec_elliptic_curves::brainpool_p_256_r1:
    // No break;
  case ec_elliptic_curves::brainpool_p_384_r1:
    result = ::EC_POINT_get_affine_coordinates_GFp(_ec_group, ec_point, &x, &y, _bn_ctx); // Use primer on elliptic curve
    break;
  default: // Use binary
    result = ::EC_POINT_get_affine_coordinates_GF2m(_ec_group, ec_point, &x, &y, _bn_ctx);
  } // End of 'switch' statement
  if (result == 0) {
    loggers::get_instance().error("ec_keys::generate: Failed to get coordinates");
  }
  const BIGNUM* p = ::EC_KEY_get0_private_key(_ec_key);
  
  _pr_key.resize(BN_num_bytes(p));
  ::BN_bn2bin(p, _pr_key.data());
  _pu_key_x.resize(BN_num_bytes(&x));
  ::BN_bn2bin(&x, _pu_key_x.data());
  _pu_key_y.resize(BN_num_bytes(&y));
  ::BN_bn2bin(&y, _pu_key_y.data());

  // TODO Compressed
  // int len = ::EC_POINT_point2oct(_ec_group, ec_point, POINT_CONVERSION_COMPRESSED, NULL, 0, _bn_ctx); 
  // std::vector<unsigned char> cy;
  // cy.resize(len);
  // ::EC_POINT_point2oct(_ec_group, ec_point, POINT_CONVERSION_COMPRESSED, (unsigned char *)cy.data(), len, _bn_ctx); 
  
  return 0;
}

int ec_keys::sign(const std::vector<unsigned char>& p_data, std::vector<unsigned char>& p_r_sig, std::vector<unsigned char>& p_s_sig) {
  loggers::get_instance().log(">>> ec_keys::sign");

  // Sanity checks
  if(_pr_key.size() == 0) { // No private key
    return -1;
  }
  if (p_data.size() == 0) {
    return -1;
  }
  
  ECDSA_SIG *signature = ::ECDSA_do_sign(p_data.data(), p_data.size(), _ec_key);
  if (signature == nullptr) {
    loggers::get_instance().warning("ec_keys::sign: Signature failed");
    return -1;
  }
  loggers::get_instance().log("ec_keys::sign: succeed");

  if (::ECDSA_do_verify(p_data.data(), p_data.size(), signature, _ec_key) != 1) {
    loggers::get_instance().warning("ec_keys::sign: Signature not verified");
    return -1;
  }
  
  p_r_sig.resize(BN_num_bytes(signature->r));
  ::BN_bn2bin(signature->r, p_r_sig.data());
  loggers::get_instance().log_to_hexa("ec_keys::sign: r=", p_r_sig.data(), p_r_sig.size());
  p_s_sig.resize(BN_num_bytes(signature->r));
  ::BN_bn2bin(signature->s, p_s_sig.data());
  loggers::get_instance().log_to_hexa("ec_keys::sign: s=", p_s_sig.data(), p_s_sig.size());

  ::ECDSA_SIG_free(signature);
  
  return 0;
}

int ec_keys::sign_verif(const std::vector<unsigned char>& p_data, const std::vector<unsigned char>& p_signature) {
  loggers::get_instance().log(">>> ec_keys::sign_verif");

  // Sanity checks
  if (p_data.size() == 0) {
    return false;
  }

  // Build the signature
  BIGNUM r, s;
  ::BN_init(&r);
  ::BN_init(&s);
  ::BN_bin2bn(p_signature.data(), p_signature.size() / 2, &r);
  ::BN_bin2bn(p_signature.data() + p_signature.size() / 2, p_signature.size() / 2, &s);
  ECDSA_SIG *signature = ECDSA_SIG_new();
  signature->r = &r;
  signature->s = &s;
  // Check the signature
  int result = ::ECDSA_do_verify(p_data.data(), p_data.size(), signature, _ec_key);
  ::ECDSA_SIG_free(signature);
  loggers::get_instance().log("ec_keys::sign_verif: %s", (result == 1) ? "succeed": "failed");
  return result != 1;
}
const int ec_keys::init() {
  loggers::get_instance().log(">>> ec_keys::init: %d", static_cast<int>(_elliptic_curve));

  ::OpenSSL_add_all_algorithms();
  ::ERR_load_crypto_strings();

  int result = -1;
  switch (_elliptic_curve) { // TODO Group this cde into a private  method
  case ec_elliptic_curves::nist_p_256: // Use the ANSI X9.62 Prime 256v1 curve 
    result = ::OBJ_txt2nid("prime256v1");
    break;
  case ec_elliptic_curves::brainpool_p_256_r1:
    result = ::OBJ_txt2nid("brainpoolP256r1");
    break;
  case ec_elliptic_curves::brainpool_p_384_r1:
    result = ::OBJ_txt2nid("brainpoolP384r1");
    break;
  default:
    loggers::get_instance().error("ec_keys::ec_keys: Unsupported EC elliptic_curve");
  } // End of 'switch' statement
  if (result < 0) {
    loggers::get_instance().warning("ec_keys::ec_keys: Unaible to set EC elliptic_curve");
    return -1;
  }

  _ec_key = ::EC_KEY_new_by_curve_name(result); // Set the elliptic curve
  ::EC_KEY_set_asn1_flag(_ec_key, OPENSSL_EC_NAMED_CURVE); // Used to save and retrieve keys
  _ec_group = ::EC_KEY_get0_group(_ec_key); // Get pointer to the EC_GROUP
  _bn_ctx = ::BN_CTX_new();

  return 0;
} // End of init

