#include "config.h"
#include "../shutdown_scriptpubkey.c"
#include <assert.h>
#include <ccan/array_size/array_size.h>
#include <ccan/err/err.h>
#include <common/amount.h>
#include <common/randbytes.h>
#include <common/setup.h>
#include <common/utils.h>
#include <stdio.h>
#include <wire/wire.h>

/* AUTOGENERATED MOCKS START */
/* Generated stub for amount_asset_is_main */
bool amount_asset_is_main(struct amount_asset *asset UNNEEDED)
{ fprintf(stderr, "amount_asset_is_main called!\n"); abort(); }
/* Generated stub for amount_asset_to_sat */
struct amount_sat amount_asset_to_sat(struct amount_asset *asset UNNEEDED)
{ fprintf(stderr, "amount_asset_to_sat called!\n"); abort(); }
/* Generated stub for amount_feerate */
 bool amount_feerate(u32 *feerate UNNEEDED, struct amount_sat fee UNNEEDED, size_t weight UNNEEDED)
{ fprintf(stderr, "amount_feerate called!\n"); abort(); }
/* Generated stub for amount_sat */
struct amount_sat amount_sat(u64 satoshis UNNEEDED)
{ fprintf(stderr, "amount_sat called!\n"); abort(); }
/* Generated stub for amount_sat_add */
 bool amount_sat_add(struct amount_sat *val UNNEEDED,
				       struct amount_sat a UNNEEDED,
				       struct amount_sat b UNNEEDED)
{ fprintf(stderr, "amount_sat_add called!\n"); abort(); }
/* Generated stub for amount_sat_eq */
bool amount_sat_eq(struct amount_sat a UNNEEDED, struct amount_sat b UNNEEDED)
{ fprintf(stderr, "amount_sat_eq called!\n"); abort(); }
/* Generated stub for amount_sat_greater_eq */
bool amount_sat_greater_eq(struct amount_sat a UNNEEDED, struct amount_sat b UNNEEDED)
{ fprintf(stderr, "amount_sat_greater_eq called!\n"); abort(); }
/* Generated stub for amount_sat_sub */
 bool amount_sat_sub(struct amount_sat *val UNNEEDED,
				       struct amount_sat a UNNEEDED,
				       struct amount_sat b UNNEEDED)
{ fprintf(stderr, "amount_sat_sub called!\n"); abort(); }
/* Generated stub for amount_sat_to_asset */
struct amount_asset amount_sat_to_asset(struct amount_sat *sat UNNEEDED, const u8 *asset UNNEEDED)
{ fprintf(stderr, "amount_sat_to_asset called!\n"); abort(); }
/* Generated stub for amount_tx_fee */
struct amount_sat amount_tx_fee(u32 fee_per_kw UNNEEDED, size_t weight UNNEEDED)
{ fprintf(stderr, "amount_tx_fee called!\n"); abort(); }
/* AUTOGENERATED MOCKS END */

/* Thanks ChatGPT! */
static const u8 *construct_script(const tal_t *ctx,
				  const u8 *data, size_t data_len)
{
	u8 *script = tal_arr(ctx, u8, 0);
	towire_u8(&script, OP_RETURN);
	if (data_len >= 6 && data_len <= 75) {
		towire_u8(&script, data_len);
		towire(&script, data, data_len);
	} else if (data_len >= 76 && data_len <= 80) {
		towire_u8(&script, 76);
		towire_u8(&script, data_len);
		towire(&script, data, data_len);
	} else {
		return tal_free(script); // Invalid case
	}
	return script;
}

static void test_valid_op_returns(void)
{
	u8 data[80];
	const u8 *script;

	for (size_t i = 6; i <= 80; i++) {
		memset(data, i, sizeof(data));
		script = construct_script(tmpctx, data, i);
		assert(is_valid_op_return(script, tal_bytelen(script)));
	}
}

static void test_invalid_op_return_too_short(void)
{
	u8 data[80];
	const u8 *script;

	for (size_t i = 0; i < 6; i++) {
		memset(data, i, sizeof(data));
		script = construct_script(tmpctx, data, i);
		assert(!is_valid_op_return(script, tal_bytelen(script)));
	}
}

static void test_invalid_op_return_too_long(void)
{
	u8 data[100];
	const u8 *script;

	for (size_t i = 81; i < sizeof(data); i++) {
		memset(data, i, sizeof(data));
		script = construct_script(tmpctx, data, i);
		assert(!is_valid_op_return(script, tal_bytelen(script)));
	}
}

// Test case: Invalid OP_RETURN with incorrect push length (e.g., 77 bytes using wrong prefix)
static void test_invalid_op_return_wrong_push_length(void)
{
	u8 script[90] = {0x6a, 0x76, 0x77}; // Invalid push of 77 bytes
	assert(!is_valid_op_return(script, 3 + 77));
}

// Test case: Invalid OP_RETURN with incorrect OP_RETURN opcode
static void test_invalid_op_return_wrong_opcode(void)
{
	u8 script[10] = {0x00, 0x06, 1, 2, 3, 4, 5, 6}; // 0x00 instead of 0x6a
	assert(!is_valid_op_return(script, 8));
}

// Test case: Invalid OP_RETURN with no data
static void test_invalid_op_return_empty(void)
{
	u8 script[1] = {0x6a}; // Only OP_RETURN, no data
	assert(!is_valid_op_return(script, 1));
}

static const u8 *construct_witness_script(const tal_t *ctx, u8 version,
                                          const u8 *data, size_t data_len)
{
	u8 *script = tal_arr(ctx, u8, 0);

	if (version < OP_1 || version > OP_16 || data_len < 2 || data_len > 40)
		return tal_free(script); // Invalid case

	towire_u8(&script, version);  // OP_1 to OP_16
	towire_u8(&script, data_len); // Push length
	towire(&script, data, data_len); // Data

	return script;
}

static void test_valid_witnessprogs(void)
{
	u8 data[40];
	const u8 *script;

	for (u8 version = OP_1; version <= OP_16; version++) {
		for (size_t i = 2; i <= 40; i++) {
			memset(data, i, sizeof(data));
			script = construct_witness_script(tmpctx, version, data, i);
			assert(is_valid_witnessprog(script, tal_bytelen(script)));
		}
	}
}

static void test_invalid_witnessprogs(void)
{
	u8 data[41];
	const u8 *script;

	// Test: Invalid versions (0 and >16)
	memset(data, 0xAA, sizeof(data));
	script = construct_witness_script(tmpctx, OP_0, data, 20);
	assert(!is_valid_witnessprog(script, tal_bytelen(script)));

	script = construct_witness_script(tmpctx, OP_16 + 1, data, 20);
	assert(!is_valid_witnessprog(script, tal_bytelen(script)));

	// Test: Invalid data lengths (1 byte and >40 bytes)
	script = construct_witness_script(tmpctx, OP_1, data, 1);
	assert(!is_valid_witnessprog(script, tal_bytelen(script)));

	script = construct_witness_script(tmpctx, OP_2, data, 41);
	assert(!is_valid_witnessprog(script, tal_bytelen(script)));

	// Test: Completely empty script (invalid)
	script = tal_arr(tmpctx, u8, 0);
	assert(!is_valid_witnessprog(script, tal_bytelen(script)));
}

// Test runner
int main(int argc, char *argv[])
{
	common_setup(argv[0]);

	test_valid_op_returns();
	test_invalid_op_return_too_short();
	test_invalid_op_return_too_long();
	test_invalid_op_return_wrong_push_length();
	test_invalid_op_return_wrong_opcode();
	test_invalid_op_return_empty();
	test_valid_witnessprogs();
	test_invalid_witnessprogs();
	common_shutdown();
	return 0;
}
