#define SUPERVERBOSE printf
#include "config.h"
/* Needed before including bolt12_merkle.c: */
  #include <stdio.h>
#include "../bolt12_merkle.c"
#include <assert.h>
#include <ccan/array_size/array_size.h>
#include <common/channel_id.h>
#include <common/channel_type.h>
#include <common/features.h>
#include <common/setup.h>
#include <secp256k1_schnorrsig.h>
#include <wire/peer_wire.h>

/* AUTOGENERATED MOCKS START */
/* Generated stub for features_unsupported */
int features_unsupported(const struct feature_set *our_features UNNEEDED,
			 const u8 *their_features UNNEEDED,
			 enum feature_place p UNNEEDED)
{ fprintf(stderr, "features_unsupported called!\n"); abort(); }
/* Generated stub for fromwire_blinded_path */
struct blinded_path *fromwire_blinded_path(const tal_t *ctx UNNEEDED, const u8 **cursor UNNEEDED, size_t *plen UNNEEDED)
{ fprintf(stderr, "fromwire_blinded_path called!\n"); abort(); }
/* Generated stub for fromwire_channel_id */
bool fromwire_channel_id(const u8 **cursor UNNEEDED, size_t *max UNNEEDED,
			 struct channel_id *channel_id UNNEEDED)
{ fprintf(stderr, "fromwire_channel_id called!\n"); abort(); }
/* Generated stub for siphash_seed */
const struct siphash_seed *siphash_seed(void)
{ fprintf(stderr, "siphash_seed called!\n"); abort(); }
/* Generated stub for towire_blinded_path */
void towire_blinded_path(u8 **p UNNEEDED, const struct blinded_path *blinded_path UNNEEDED)
{ fprintf(stderr, "towire_blinded_path called!\n"); abort(); }
/* Generated stub for towire_channel_id */
void towire_channel_id(u8 **pptr UNNEEDED, const struct channel_id *channel_id UNNEEDED)
{ fprintf(stderr, "towire_channel_id called!\n"); abort(); }
/* AUTOGENERATED MOCKS END */

/* Contat several tal objects */
#define concat(p, ...) concat_((p), __VA_ARGS__, NULL)

static LAST_ARG_NULL void *concat_(const void *p, ...)
{
	va_list ap;
	size_t len = 0;
	u8 *ret = tal_arr(tmpctx, u8, len);

	va_start(ap, p);
	do {
		tal_resize(&ret, len + tal_bytelen(p));
		memcpy(ret + len, p, tal_bytelen(p));
		len += tal_bytelen(p);
	} while ((p = va_arg(ap, const void *)) != NULL);
	va_end(ap);
	return ret;
}

/* Just return type field from tlv */
static const u8 *tlv_typenum(u64 type)
{
	u8 *p = tal_arr(tmpctx, u8, 0);

	towire_bigsize(&p, type);
	return p;
}

/* Hashes a tal object */
static struct sha256 *SHA256(const void *obj)
{
	struct sha256 *ret = tal(tmpctx, struct sha256);
	sha256(ret, obj, tal_bytelen(obj));
	return ret;
}

static u8 *tlv(u64 type, const void *contents, size_t len)
{
	u8 *ret = tal_arr(tmpctx, u8, 0);

	towire_bigsize(&ret, type);
	towire_bigsize(&ret, len);
	towire(&ret, contents, len);
	return ret;
}

/* Concatenate these two in lesser, greater order. */
static u8 *ordered(const struct sha256 *a, const struct sha256 *b)
{
	u8 *ret = tal_arr(tmpctx, u8, sizeof(*a) + sizeof(*b));

	if (memcmp(a, b, sizeof(*a)) < 0) {
		memcpy(ret, a, sizeof(*a));
		memcpy(ret + sizeof(*a), b, sizeof(*b));
	} else {
		memcpy(ret, b, sizeof(*b));
		memcpy(ret + sizeof(*b), a, sizeof(*a));
	}
	return ret;
}

/* BOLT-offers #12:
 * Thus we define H(`tag`,`msg`) as SHA256(SHA256(`tag`) ||
 * SHA256(`tag`) || `msg`) */

static struct sha256 *H(const void *tag, const void *msg)
{
	const struct sha256 *taghash = SHA256(tag);
	const u8 *full = concat(taghash, taghash, msg);
	struct sha256 *ret = SHA256(full);

	printf("test: H(tag=%s,msg=%s) -> SHA256(%s|%s|msg) -> %s\n",
	       tal_hex(tmpctx, tag), tal_hex(tmpctx, msg),
	       fmt_sha256(tmpctx, taghash),
	       fmt_sha256(tmpctx, taghash),
	       fmt_sha256(tmpctx, ret));
	return ret;
}

#define json_out(fmt, ...) printf("JSON: " fmt "\n" , ## __VA_ARGS__)

/* AAAA... BBBB... etc */
static struct pubkey *pubkey_for_letter(const tal_t *ctx, char letter)
{
	struct secret secret;
	struct pubkey *pk;

	pk = tal(ctx, struct pubkey);
	memset(&secret, letter, sizeof(secret));
	assert(pubkey_from_secret(&secret, pk));
	return pk;
}

static secp256k1_keypair keypair_for_letter(char letter)
{
	struct secret secret;
	secp256k1_keypair kp;

	memset(&secret, letter, sizeof(secret));

	if (secp256k1_keypair_create(secp256k1_ctx, &kp,
				     secret.data) != 1)
		abort();
	return kp;
}

static struct bip340sig *invoice_signature(const tal_t *ctx, struct tlv_invoice *inv, char letter)
{
	struct sha256 merkle, sha;
	struct bip340sig *sig;
	secp256k1_keypair kp = keypair_for_letter(letter);

	/* Update fields[] array from our settings */
	tlv_update_fields(inv, tlv_invoice, &inv->fields);
	merkle_tlv(inv->fields, &merkle);
	inv->signature = tal(inv, struct bip340sig);
	sighash_from_merkle("invoice", "signature", &merkle, &sha);

	sig = tal(ctx, struct bip340sig);
	assert(secp256k1_schnorrsig_sign32(secp256k1_ctx, sig->u8,
					   sha.u.u8,
					   &kp,
					   NULL) == 1);
	return sig;
}

struct iter {
	const struct tlv_invoice *inv;
	size_t n_field;
	u64 *omitted;
	struct sha256 *hashes;
};

static const struct tlv_field *iter_next_field(bool *omitted, struct iter *iter)
{
	const struct tlv_field *f;
	if (iter->n_field == tal_count(iter->inv->fields))
		return NULL;

	f = &iter->inv->fields[iter->n_field++];
	for (size_t i = 0; i < tal_count(iter->omitted); i++) {
		if (f->numtype == iter->omitted[i]) {
			/* To be sure, return a fake field here. */
			struct tlv_field *tmp_f = tal(tmpctx, struct tlv_field);
			tmp_f->numtype = f->numtype;
			tmp_f->length = 0;
			tmp_f->value = NULL;
			*omitted = true;
			return tmp_f;
		}
	}

	*omitted = false;
	return f;
}

static void iter_calc_nonce(const struct sha256_ctx *lnnonce_ctx,
			    bigsize_t fieldtype,
			    struct sha256 *hash,
			    struct iter *iter)
{
	struct sha256_ctx ctx;

	/* If omit the first field, lnnonce_ctx is NULL, so calc manually */
	bolt12_lnnonce_ctx(&ctx, &iter->inv->fields[0]);

	bolt12_calc_nonce(&ctx, fieldtype, hash, NULL);
}

static void iter_resolve_omitted(struct sha256 *hash, struct iter *iter)
{
	assert(tal_count(iter->hashes) > 0);
	*hash = iter->hashes[0];
	tal_arr_remove(&iter->hashes, 0);
}

int main(int argc, char *argv[])
{
	const char *LnBranch, *LnNonce, *LnLeaf;
	struct tlv_invoice *inv;
	u8 *tlv0;
	struct preimage preimage;
	struct sha256 hash, merkle, part_merkle;
	struct iter iter;

	common_setup(argv[0]);

	/* Note: no nul term */
	LnBranch = tal_dup_arr(tmpctx, char, "LnBranch", strlen("LnBranch"), 0);
	LnLeaf = tal_dup_arr(tmpctx, char, "LnLeaf", strlen("LnLeaf"), 0);
	LnNonce = tal_dup_arr(tmpctx, char, "LnNonce", strlen("LnNonce"), 0);

	memset(&preimage, 0, sizeof(preimage));
	sha256(&hash, &preimage, sizeof(preimage));

	inv = tlv_invoice_new(tmpctx);
	inv->invreq_metadata = tal_arrz(inv, u8, 32);
	inv->offer_issuer = tal_dup_arr(inv, char, "test", 4, 0);
	inv->invreq_amount = tal(inv, u64); *inv->invreq_amount = 1000;
	inv->invreq_payer_id = pubkey_for_letter(inv, 'A');
	inv->invoice_payment_hash = tal_dup(inv, struct sha256, &hash);
	inv->invoice_node_id = pubkey_for_letter(inv, 'B');
	inv->signature = invoice_signature(inv, inv, 'B');
	tlv_update_fields(inv, tlv_invoice, &inv->fields);

	/* Merkle normally. */
	merkle_tlv(inv->fields, &merkle);

	/* This is useful for calculating the rest */
	tlv0 = tlv(inv->fields[0].numtype, inv->fields[0].value, inv->fields[0].length);

	/* Every single omission. */
	for (size_t i = 0; i < tal_count(inv->fields); i++) {
		u8 *omitted_tlv;
		struct sha256 ohash;

		iter.inv = inv;
		iter.omitted = tal_arr(tmpctx, u64, 0);
		iter.n_field = 0;
		iter.hashes = tal_arr(tmpctx, struct sha256, 0);

		omitted_tlv = tlv(inv->fields[i].numtype,
				  inv->fields[i].value,
				  inv->fields[i].length);
		ohash = *H(LnBranch,
			   ordered(H(LnLeaf, omitted_tlv),
				   H(concat(LnNonce, tlv0),
				     tlv_typenum(inv->fields[i].numtype))));
		tal_arr_expand(&iter.omitted, inv->fields[i].numtype);
		tal_arr_expand(&iter.hashes, ohash);
		merkle_tlv_full(&part_merkle,
				iter_next_field,
				iter_calc_nonce,
				iter_resolve_omitted,
				&iter);

		/* Same result */
		assert(sha256_eq(&part_merkle, &merkle));
	}

	/* Now in pairs! */
	for (size_t i = 0; i + 1 < tal_count(inv->fields); i += 2) {
		u8 *omitted_tlv1, *omitted_tlv2;
		struct sha256 ohash1, ohash2, both;

		iter.inv = inv;
		iter.omitted = tal_arr(tmpctx, u64, 0);
		iter.n_field = 0;
		iter.hashes = tal_arr(tmpctx, struct sha256, 0);

		omitted_tlv1 = tlv(inv->fields[i].numtype,
				   inv->fields[i].value,
				   inv->fields[i].length);
		omitted_tlv2 = tlv(inv->fields[i+1].numtype,
				   inv->fields[i+1].value,
				   inv->fields[i+1].length);
		ohash1 = *H(LnBranch,
			    ordered(H(LnLeaf, omitted_tlv1),
				    H(concat(LnNonce, tlv0),
				      tlv_typenum(inv->fields[i].numtype))));
		ohash2 = *H(LnBranch,
			    ordered(H(LnLeaf, omitted_tlv2),
				    H(concat(LnNonce, tlv0),
				      tlv_typenum(inv->fields[i+1].numtype))));
		/* Combine them */
		both = *H(LnBranch, ordered(&ohash1, &ohash2));
		tal_arr_expand(&iter.omitted, inv->fields[i].numtype);
		tal_arr_expand(&iter.omitted, inv->fields[i+1].numtype);
		tal_arr_expand(&iter.hashes, both);
		merkle_tlv_full(&part_merkle,
				iter_next_field,
				iter_calc_nonce,
				iter_resolve_omitted,
				&iter);

		/* Same result */
		assert(sha256_eq(&part_merkle, &merkle));
	}

	common_shutdown();
}
