/*
 * SPDX-License-Identifier: CC0-1.0
 *
 * Copyright (C) 2025 W. Kosior <koszko@koszko.org>
 */

#define _POSIX_C_SOURCE 200809L /* for getline() */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "pqcrypto_blind_sig.h"

static const char sample_message[] = "This is a message to blind-sign.";

#define RETRIES_MAX 1000U

static void signing_example(blind_sig_ctx_t const ctx) {
	unsigned long retries_left = RETRIES_MAX;

	blind_sig_priv_key_t priv_key;
	blind_sig_pub_key_t pub_key;
	blind_sig_signer_state_t signer_state;
	blind_sig_user_state_t user_state;

	blind_sig_proto_p1_t p1_result;
	blind_sig_proto_p2_t p2_result;
	blind_sig_proto_p3_t p3_result;
	blind_sig_proto_p4_t p4_result;
	blind_sig_proto_p5_t p5_result;

	blind_sig_t sig;

	blind_sig_key_init_gen(priv_key, pub_key, ctx);

	blind_sig_signer_state_init(signer_state, ctx);

	blind_sig_user_state_init(user_state, ctx);

p1:
	blind_sig_proto_p1_init_do(p1_result, signer_state, ctx);

/* p2: */
	blind_sig_proto_p2_init_do(p2_result, user_state, pub_key,
				   sample_message, sizeof(sample_message),
				   p1_result, &retries_left, ctx);

	if (!p2_result->ok) {
		blind_sig_proto_p1_clear(p1_result, ctx);
		blind_sig_proto_p2_clear(p2_result, ctx);

		goto too_many_retries;
	}

/* p3: */
	blind_sig_proto_p3_init_do(p3_result, signer_state, priv_key,
				   p2_result, ctx);

	if (!p3_result->ok) {
		blind_sig_proto_p1_clear(p1_result, ctx);
		blind_sig_proto_p2_clear(p2_result, ctx);
		blind_sig_proto_p3_clear(p3_result, ctx);

		if (!(retries_left && retries_left--))
			goto too_many_retries;

		goto p1;
	}

/* p4: */
	blind_sig_proto_p4_init_do(p4_result, sig, user_state, p3_result, ctx);

/* p5: */
	blind_sig_proto_p5_init_do(p5_result, signer_state, pub_key, p4_result,
				   ctx);

	if (!p5_result->ok) {
		blind_sig_proto_p1_clear(p1_result, ctx);
		blind_sig_proto_p2_clear(p2_result, ctx);
		blind_sig_proto_p3_clear(p3_result, ctx);
		blind_sig_proto_p4_clear(p4_result, ctx);
		blind_sig_proto_p5_clear(p5_result, ctx);

		if (!(retries_left && retries_left--))
			goto too_many_retries;

		goto p1;
	}

/* verification: */
	if (!p4_result->needs_restart) {
		if (blind_sig_verify(sig, pub_key, ctx))
			fprintf(stderr, "Signature verified successfully.\n");
		else
			fprintf(stderr, "Signature not valid.\n");
	}

/* out: */
	if (!p4_result->needs_restart)
		blind_sig_clear(sig, ctx);

	blind_sig_priv_key_clear(priv_key, ctx);

	blind_sig_pub_key_clear(pub_key, ctx);

	blind_sig_signer_state_clear(signer_state, ctx);

	blind_sig_user_state_clear(user_state, ctx);

	blind_sig_proto_p1_clear(p1_result, ctx);
	blind_sig_proto_p2_clear(p2_result, ctx);
	blind_sig_proto_p3_clear(p3_result, ctx);
	blind_sig_proto_p4_clear(p4_result, ctx);
	blind_sig_proto_p5_clear(p5_result, ctx);

	return;

too_many_retries:
	fprintf(stderr, "Too many retries.\n");
};

static int perform_key_gen(blind_sig_ctx_t ctx) {
	blind_sig_priv_key_t priv_key;
	blind_sig_pub_key_t pub_key;

	blind_sig_key_init_gen(priv_key, pub_key, ctx);

	blind_sig_priv_key_file_serialize("priv_key.example", priv_key, ctx);
	blind_sig_pub_key_file_serialize("pub_key.example", pub_key, ctx);

	blind_sig_priv_key_clear(priv_key, ctx);
	blind_sig_pub_key_clear(pub_key, ctx);

	return EXIT_SUCCESS;
}

static int perform_signing_p1(blind_sig_ctx_t ctx) {
	blind_sig_signer_state_t signer_state;
	blind_sig_proto_p1_t p1_result;

	blind_sig_signer_state_init(signer_state, ctx);

	blind_sig_proto_p1_init_do(p1_result, signer_state, ctx);

	blind_sig_signer_state_file_serialize("signer_state.example",
					      signer_state, ctx);
	blind_sig_proto_p1_file_serialize("p1_result.example", p1_result, ctx);

	blind_sig_signer_state_clear(signer_state, ctx);
	blind_sig_proto_p1_clear(p1_result, ctx);

	return EXIT_SUCCESS;
}

static int perform_signing_p2(blind_sig_ctx_t ctx) {
	blind_sig_user_state_t user_state;
	blind_sig_pub_key_t pub_key;
	ssize_t message_bytes;
	char * message_buf = NULL;
	size_t message_buf_bytes;
	blind_sig_proto_p1_t p1_result;
	unsigned long retries_left = RETRIES_MAX;
	blind_sig_proto_p2_t p2_result;
	int result = EXIT_SUCCESS;

	blind_sig_user_state_init(user_state, ctx);

	blind_sig_pub_key_file_deserialize("pub_key.example", pub_key, ctx);

	printf("Write a line-one message to blind-sign.\n> ");
	fflush(stdout);

	message_bytes = getline(&message_buf, &message_buf_bytes, stdin);
	if (message_bytes <= 0)
		abort();

	putchar('\n');

	blind_sig_proto_p1_file_deserialize("p1_result.example", p1_result,
					    ctx);

	blind_sig_proto_p2_init_do(p2_result, user_state, pub_key, message_buf,
				   message_bytes, p1_result, &retries_left,
				   ctx);

	if (!p2_result->ok) {
		fprintf(stderr, "Too many retries.\n");
		result = EXIT_FAILURE;
	}

	blind_sig_user_state_file_serialize("user_state.example", user_state,
					    ctx);
	blind_sig_proto_p2_file_serialize("p2_result.example", p2_result, ctx);

	blind_sig_user_state_clear(user_state, ctx);
	blind_sig_pub_key_clear(pub_key, ctx);
	free(message_buf);
	blind_sig_proto_p1_clear(p1_result, ctx);
	blind_sig_proto_p2_clear(p2_result, ctx);

	return result;
}

static int perform_signing_p3(blind_sig_ctx_t ctx) {
	blind_sig_signer_state_t signer_state;
	blind_sig_priv_key_t priv_key;
	blind_sig_proto_p2_t p2_result;
	blind_sig_proto_p3_t p3_result;
	int result;

	blind_sig_signer_state_file_deserialize("signer_state.example",
						signer_state, ctx);
	blind_sig_priv_key_file_deserialize("priv_key.example", priv_key, ctx);
	blind_sig_proto_p2_file_deserialize("p2_result.example", p2_result,
					    ctx);

	blind_sig_proto_p3_init_do(p3_result, signer_state, priv_key, p2_result,
				   ctx);

	result = p3_result->ok ? EXIT_SUCCESS : EXIT_FAILURE;

	blind_sig_signer_state_file_serialize("signer_state.example",
					      signer_state, ctx);
	blind_sig_proto_p3_file_serialize("p3_result.example", p3_result, ctx);

	blind_sig_signer_state_clear(signer_state, ctx);
	blind_sig_priv_key_clear(priv_key, ctx);
	blind_sig_proto_p2_clear(p2_result, ctx);
	blind_sig_proto_p3_clear(p3_result, ctx);

	return result;
}

static int perform_signing_p4(blind_sig_ctx_t ctx) {
	blind_sig_t sig;
	blind_sig_user_state_t user_state;
	blind_sig_proto_p3_t p3_result;
	blind_sig_proto_p4_t p4_result;
	int result;

	blind_sig_user_state_file_deserialize("user_state.example", user_state,
					      ctx);
	blind_sig_proto_p3_file_deserialize("p3_result.example", p3_result,
					    ctx);

	blind_sig_proto_p4_init_do(p4_result, sig, user_state, p3_result, ctx);

	result = p4_result->needs_restart? EXIT_FAILURE : EXIT_SUCCESS;

	/* No serialization of state — it doesn't get modified here. */
	if (!p4_result->needs_restart)
		blind_sig_file_serialize("sig.example", sig, ctx);
	blind_sig_proto_p4_file_serialize("p4_result.example", p4_result, ctx);

	if (!p4_result->needs_restart)
		blind_sig_clear(sig, ctx);
	blind_sig_user_state_clear(user_state, ctx);
	blind_sig_proto_p3_clear(p3_result, ctx);
	blind_sig_proto_p4_clear(p4_result, ctx);

	return result;
}

static int perform_signing_p5(blind_sig_ctx_t ctx) {
	blind_sig_signer_state_t signer_state;
	blind_sig_pub_key_t pub_key;
	blind_sig_proto_p4_t p4_result;
	blind_sig_proto_p5_t p5_result;
	int result;

	blind_sig_signer_state_file_deserialize("signer_state.example",
						signer_state, ctx);
	blind_sig_pub_key_file_deserialize("pub_key.example", pub_key, ctx);
	blind_sig_proto_p4_file_deserialize("p4_result.example", p4_result,
					    ctx);

	blind_sig_proto_p5_init_do(p5_result, signer_state, pub_key, p4_result,
				   ctx);

	result = p5_result->ok ? EXIT_SUCCESS : EXIT_FAILURE;

	/* Again, no serialization of state — it doesn't get modified here. */
	blind_sig_proto_p5_file_serialize("p5_result.example", p5_result, ctx);

	blind_sig_signer_state_clear(signer_state, ctx);
	blind_sig_pub_key_clear(pub_key, ctx);
	blind_sig_proto_p4_clear(p4_result, ctx);
	blind_sig_proto_p5_clear(p5_result, ctx);

	return result;
}

static int perform_verification(blind_sig_ctx_t ctx) {
	blind_sig_t sig;
	blind_sig_pub_key_t pub_key;
	int result;

	blind_sig_file_deserialize("sig.example", sig, ctx);
	blind_sig_pub_key_file_deserialize("pub_key.example", pub_key, ctx);

	if (blind_sig_verify(sig, pub_key, ctx)) {
		fprintf(stderr, "Got valid signature for message.\n");
		fwrite(sig->message->buf, sig->message->bytes, 1, stdout);
		result = EXIT_SUCCESS;
	} else {
		fprintf(stderr, "Got invalid signature.\n");
		result = EXIT_FAILURE;
	}

	blind_sig_clear(sig, ctx);
	blind_sig_pub_key_clear(pub_key, ctx);

	return result;
}

static void perform_example(void) {
	blind_sig_params_t params;
	blind_sig_ctx_t ctx;

	fprintf(stderr, "Preparing blind signature context (\"Current III\" set"
		" from Rückert's paper).\n");

	*params = *blind_sig_params_current_III;
	params->log_level = BLIND_SIG_LOG_INFO_2_PARAMETERS;

	/* fprintf(stderr, "Preparing blind signature context (toy set of" */
	/* 	" parameters).\n"); */

	/* *params = *blind_sig_params_toy; */

	blind_sig_ctx_init(ctx, params);

	fprintf(stderr, "Making a test signature for ASCII buffer: \"%s\".\n",
		sample_message);

	signing_example(ctx);

	fprintf(stderr, "Done, cleaning up.\n");

	blind_sig_ctx_clear(ctx);
}

int main(const int argc, const char * const * argv) {
	blind_sig_params_t params;
	blind_sig_ctx_t ctx;
	char const * action = argc > 1 ? argv[1] : "example";
	int result = EXIT_SUCCESS;

	if (!strcmp(action, "example")) {
		perform_example();
		goto out;
	}

	*params = *blind_sig_params_current_III;
	params->log_level = BLIND_SIG_LOG_INFO_1_CONTROL;

	blind_sig_ctx_init(ctx, params);

#define X(action_name)							\
	!strcmp(action, #action_name) ? perform_##action_name(ctx) :

	result =
		X(key_gen)
		X(signing_p1)
		X(signing_p2)
		X(signing_p3)
		X(signing_p4)
		X(signing_p5)
		X(verification)
		(fprintf(stderr, "Bad action `%s'.\n", action), EXIT_FAILURE);

#undef X

out:
	flint_cleanup();

	return result;
}