#import "KoinsSecp256k1.h"

#import <openssl/bn.h>
#import <openssl/ec.h>
#import <openssl/ecdsa.h>
#import <openssl/obj_mac.h>
#import <openssl/sha.h>

// secp256k1 group order.
static const char *kCurveOrderHex =
    "FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141";

static BIGNUM *KoinsCurveOrder(void) {
  BIGNUM *order = NULL;
  BN_hex2bn(&order, kCurveOrderHex);
  return order;
}

static NSData *KoinsSHA512Half(NSData *input) {
  unsigned char digest[SHA512_DIGEST_LENGTH];
  SHA512((const unsigned char *)input.bytes, (size_t)input.length, digest);
  return [NSData dataWithBytes:digest length:32];
}

// Mirror of ripple-keypairs deriveScalar: hash the bytes (optionally with a
// 32-bit discriminator) until the result is a valid non-zero curve scalar.
static BIGNUM *KoinsDeriveScalar(NSData *bytes, BOOL hasDiscriminator,
                                 uint32_t discriminator, const BIGNUM *order) {
  for (uint32_t i = 0; i <= 4294967295u; i++) {
    NSMutableData *input = [NSMutableData dataWithData:bytes];
    if (hasDiscriminator) {
      uint32_t d = htonl(discriminator);
      [input appendBytes:&d length:4];
    }
    uint32_t idx = htonl(i);
    [input appendBytes:&idx length:4];

    NSData *half = KoinsSHA512Half(input);
    BIGNUM *scalar = BN_bin2bn((const unsigned char *)half.bytes, (int)half.length, NULL);
    if (scalar && !BN_is_zero(scalar) && BN_cmp(scalar, order) < 0) {
      return scalar;
    }
    if (scalar) BN_free(scalar);
  }
  return NULL;
}

@implementation KoinsSecp256k1

+ (nullable NSData *)privateKeyFromEntropy:(NSData *)entropy {
  if (entropy.length != 16) return nil;

  BN_CTX *ctx = BN_CTX_new();
  BIGNUM *order = KoinsCurveOrder();
  EC_GROUP *group = EC_GROUP_new_by_curve_name(NID_secp256k1);
  EC_POINT *publicGen = EC_POINT_new(group);
  NSData *result = nil;

  BIGNUM *privateGen = KoinsDeriveScalar(entropy, NO, 0, order);
  if (privateGen && EC_POINT_mul(group, publicGen, privateGen, NULL, NULL, ctx)) {
    unsigned char compressed[33];
    size_t written = EC_POINT_point2oct(group, publicGen,
                                        POINT_CONVERSION_COMPRESSED,
                                        compressed, sizeof(compressed), ctx);
    if (written == sizeof(compressed)) {
      NSData *publicGenBytes = [NSData dataWithBytes:compressed length:written];
      BIGNUM *offset = KoinsDeriveScalar(publicGenBytes, YES, 0, order);
      if (offset) {
        BIGNUM *derived = BN_new();
        if (BN_mod_add(derived, offset, privateGen, order, ctx)) {
          unsigned char out[32] = {0};
          BN_bn2binpad(derived, out, sizeof(out));
          result = [NSData dataWithBytes:out length:sizeof(out)];
        }
        BN_free(derived);
        BN_free(offset);
      }
    }
  }

  if (privateGen) BN_free(privateGen);
  if (publicGen) EC_POINT_free(publicGen);
  if (group) EC_GROUP_free(group);
  if (order) BN_free(order);
  if (ctx) BN_CTX_free(ctx);
  return result;
}

+ (nullable NSData *)publicKeyFromPrivateKey:(NSData *)privateKey {
  if (privateKey.length != 32) return nil;

  BN_CTX *ctx = BN_CTX_new();
  EC_GROUP *group = EC_GROUP_new_by_curve_name(NID_secp256k1);
  EC_POINT *point = EC_POINT_new(group);
  BIGNUM *scalar = BN_bin2bn((const unsigned char *)privateKey.bytes, (int)privateKey.length, NULL);
  NSData *result = nil;

  if (scalar && EC_POINT_mul(group, point, scalar, NULL, NULL, ctx)) {
    unsigned char compressed[33];
    size_t written = EC_POINT_point2oct(group, point,
                                        POINT_CONVERSION_COMPRESSED,
                                        compressed, sizeof(compressed), ctx);
    if (written == sizeof(compressed)) {
      result = [NSData dataWithBytes:compressed length:written];
    }
  }

  if (scalar) BN_free(scalar);
  if (point) EC_POINT_free(point);
  if (group) EC_GROUP_free(group);
  if (ctx) BN_CTX_free(ctx);
  return result;
}

+ (nullable NSData *)addScalar:(NSData *)left toScalar:(NSData *)right {
  if (left.length != 32 || right.length != 32) return nil;

  BN_CTX *ctx = BN_CTX_new();
  BIGNUM *order = KoinsCurveOrder();
  BIGNUM *a = BN_bin2bn((const unsigned char *)left.bytes, (int)left.length, NULL);
  BIGNUM *b = BN_bin2bn((const unsigned char *)right.bytes, (int)right.length, NULL);
  BIGNUM *sum = BN_new();
  NSData *result = nil;

  if (a && b && sum && BN_mod_add(sum, a, b, order, ctx)) {
    unsigned char out[32] = {0};
    BN_bn2binpad(sum, out, sizeof(out));
    result = [NSData dataWithBytes:out length:sizeof(out)];
  }

  if (sum) BN_free(sum);
  if (b) BN_free(b);
  if (a) BN_free(a);
  if (order) BN_free(order);
  if (ctx) BN_CTX_free(ctx);
  return result;
}

+ (nullable NSData *)signDigest:(NSData *)digest
                 withPrivateKey:(NSData *)privateKey {
  if (digest.length != 32 || privateKey.length != 32) return nil;

  BN_CTX *ctx = BN_CTX_new();
  EC_KEY *key = EC_KEY_new_by_curve_name(NID_secp256k1);
  BIGNUM *scalar = BN_bin2bn((const unsigned char *)privateKey.bytes, (int)privateKey.length, NULL);
  NSData *result = nil;

  if (key && scalar && EC_KEY_set_private_key(key, scalar)) {
    const EC_GROUP *group = EC_KEY_get0_group(key);
    EC_POINT *pub = EC_POINT_new(group);
    if (pub && EC_POINT_mul(group, pub, scalar, NULL, NULL, ctx) &&
        EC_KEY_set_public_key(key, pub)) {
      ECDSA_SIG *sig = ECDSA_do_sign((const unsigned char *)digest.bytes, (int)digest.length, key);
      if (sig) {
        // Ripple only accepts canonical (low-S) signatures.
        BIGNUM *order = KoinsCurveOrder();
        BIGNUM *halfOrder = BN_dup(order);
        BN_rshift1(halfOrder, halfOrder);
        const BIGNUM *r = NULL;
        const BIGNUM *s = NULL;
        ECDSA_SIG_get0(sig, &r, &s);
        if (BN_cmp(s, halfOrder) > 0) {
          BIGNUM *lowS = BN_new();
          BN_sub(lowS, order, s);
          ECDSA_SIG_set0(sig, BN_dup(r), lowS);
        }
        unsigned char *der = NULL;
        int derLen = i2d_ECDSA_SIG(sig, &der);
        if (derLen > 0) {
          result = [NSData dataWithBytes:der length:(NSUInteger)derLen];
          OPENSSL_free(der);
        }
        BN_free(halfOrder);
        BN_free(order);
        ECDSA_SIG_free(sig);
      }
    }
    if (pub) EC_POINT_free(pub);
  }

  if (scalar) BN_free(scalar);
  if (key) EC_KEY_free(key);
  if (ctx) BN_CTX_free(ctx);
  return result;
}

@end