// g++ -o sign sign.cpp -l mbedcrypto -lsecp256k1

#include <mbedtls/ecdsa.h>
#include <mbedtls/sha256.h>
#include <secp256k1.h>
#include <string>
#include <vector>
#include <sstream>
#include <iomanip>

std::vector<uint8_t> hex_decode(const std::string &s) {
  std::vector<uint8_t> bytes;
  for (size_t i = 0; i < s.length(); i += 2) {
    const std::string hex = s.substr(i, 2);
    const uint8_t decimal = std::strtol(hex.c_str(), 0, 16);
    bytes.push_back(decimal);
  }
  return bytes;
}

std::string hex_encode(const unsigned char *data, size_t len) {
  std::stringstream ss;
  for (size_t i = 0; i < len; i++) ss << std::hex << std::setw(2) << std::setfill('0') << (int)data[i];
  return ss.str();
}

static inline int mbd_rand(void *rng_state, unsigned char *output, size_t len) {
  for (size_t i = 0; i < len; ++i) output[i] = rand();
  return 0;
}

int main() {
  std::string device_id = "cbf97fb4ff5c592ab387fd303ac9069b75930cf6e10d58da67f8426dd0f82ff2";
  std::string user_id = "71b9ac79010a448b95836a2be0dde671";
  std::string msg = "5dde4e1bdf9e4966b387ba58f4b3fdc3:" + device_id + ":" + user_id + ":0";

  std::vector<uint8_t> pub(MBEDTLS_ECP_MAX_BYTES, 0);
  std::vector<uint8_t> sigdata(65, 0);

  unsigned char msg_hash[32];
  int ret = 0;
  mbedtls_sha256_ret((uint8_t *)msg.c_str(), msg.size(), msg_hash, 0);

  {
    mbedtls_ecdsa_context ctx_sign;
    mbedtls_ecdsa_init(&ctx_sign);
    // load private key
    mbedtls_ecp_group_load(&ctx_sign.grp, MBEDTLS_ECP_DP_SECP256K1);
    mbedtls_mpi_read_string(&ctx_sign.d, 16, device_id.c_str());
    mbedtls_ecp_mul(&ctx_sign.grp, &ctx_sign.Q, &ctx_sign.d, &ctx_sign.grp.G, nullptr, nullptr);
    // dump public key
    size_t pub_len = 0;
    mbedtls_ecp_point_write_binary(&ctx_sign.grp, &ctx_sign.Q, MBEDTLS_ECP_PF_UNCOMPRESSED, &pub_len, pub.data(), pub.size());
    pub.resize(pub_len);
    printf("pubkey %s\n", hex_encode(pub.data(), pub_len).c_str());
    // sign message
    mbedtls_mpi r, s;
    mbedtls_mpi_init(&r);
    mbedtls_mpi_init(&s);
    mbedtls_ecdsa_sign(&ctx_sign.grp, &r, &s, &ctx_sign.d, msg_hash, sizeof(msg_hash), mbd_rand, nullptr);
    mbedtls_mpi_write_binary(&r, &sigdata[0], 32);
    mbedtls_mpi_write_binary(&s, &sigdata[32], 32);
    sigdata[64] = 1;
    mbedtls_mpi_free(&r);
    mbedtls_mpi_free(&s);
    printf("sigdata %s\n", hex_encode(sigdata.data(), sigdata.size()).c_str());
    mbedtls_ecdsa_free(&ctx_sign);
  }

  {
    mbedtls_ecdsa_context ctx_verify;
    mbedtls_ecdsa_init(&ctx_verify);
    // load public key
    mbedtls_ecp_group_load(&ctx_verify.grp, MBEDTLS_ECP_DP_SECP256K1);
    ret = mbedtls_ecp_point_read_binary(&ctx_verify.grp, &ctx_verify.Q, pub.data(), pub.size());
    if (ret < 0) printf("mbedtls_ecp_point_read_string -0x%x\n", -ret);
    // verify signature
    mbedtls_mpi r, s;
    mbedtls_mpi_init(&r);
    mbedtls_mpi_init(&s);
    mbedtls_mpi_read_binary(&r, &sigdata[0], 32);
    mbedtls_mpi_read_binary(&s, &sigdata[32], 32);
    ret = mbedtls_ecdsa_verify(&ctx_verify.grp, msg_hash, sizeof(msg_hash), &ctx_verify.Q, &r, &s);
    if (ret < 0)
      printf("mbedtls_ecdsa_verify -0x%x\n", -ret);
    else
      printf("mbedtls_ecdsa_verify ok\n");
    mbedtls_mpi_free(&r);
    mbedtls_mpi_free(&s);
    mbedtls_ecdsa_free(&ctx_verify);
  }

  {
    secp256k1_ecdsa_signature sig;
    secp256k1_pubkey pubkey;
    secp256k1_context *ctx = secp256k1_context_create(SECP256K1_CONTEXT_SIGN | SECP256K1_CONTEXT_VERIFY);
    ret = secp256k1_ecdsa_signature_parse_compact(ctx, &sig, sigdata.data());
    if (!ret) printf("secp256k1_ecdsa_signature_parse_compact failed\n");
    ret = secp256k1_ec_pubkey_parse(ctx, &pubkey, pub.data(), pub.size());
    if (!ret) printf("secp256k1_ec_pubkey_parse failed\n");
    ret = secp256k1_ecdsa_verify(ctx, &sig, msg_hash, &pubkey);
    if (!ret)
      printf("secp256k1_ecdsa_verify failed\n");
    else
      printf("secp256k1_ecdsa_verify ok\n");
    secp256k1_context_destroy(ctx);
  }
  return 0;
}