#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include "argon2.h"
#include "random.h"
#include "util.h"

static void
die(int check, char *why)
{
	if (check) {
		fprintf(stderr, "Error: %s\n", why);
		exit(1);
	}
}

int
main(int argc, char *argv[])
{
	if (argc != 2) {
		fprintf(stderr, "Usage: %s <challenge>\n", argv[0]);
	}

	int iters, memory, digits;
	unsigned char salt[16];

	char *challenge = argv[1];
	char *algo = strtok(challenge, ":");
	if (strcmp(algo, "argon2id") != 0) {
		fprintf(stderr, "Error: unknown challenge type %s\n", algo);
		return 1;
	}
	char *endptr;

	char *iterstr = strtok(NULL, ":");
	iters = strtoul(iterstr, &endptr, 10);
	die(*endptr, "Invalid challenge");

	char *memorystr = strtok(NULL, ":");
	memory = strtoul(memorystr, &endptr, 10);
	die(*endptr, "Invalid challenge");

	char *digitsstr = strtok(NULL, ":");
	digits = strtoul(digitsstr, &endptr, 10);
	die(*endptr, "Invalid challenge");

	char *saltstr = strtok(NULL, ":");
	die(strlen(saltstr) != 32, "Invalid challenge");
	int r = dechex(saltstr, strlen(saltstr), salt, sizeof(salt));
	die(r == -1, "Invalid challenge");

	unsigned char password[16];
	unsigned char hash[32];
	argon2_context context = {
		.out = hash,
		.outlen = sizeof(hash),
		.salt = salt,
		.saltlen = sizeof(salt),
		.pwd = password,
		.pwdlen = sizeof(password),
		.t_cost = iters,
		.m_cost = memory,
		.lanes = 1,
		.threads = 1,
		.flags = ARGON2_DEFAULT_FLAGS,
	};

	if (isatty(STDERR_FILENO)) {
		// TODO: Provide a better estimate based on the difficulty
		fprintf(stderr, "Generating a proof of work.\n");
		fprintf(stderr, "This may take anywhere from tens of seconds to a few hours on slow computers.\n");
	}

	unsigned long nattempts = 0;
	int valid = 0;
	while (!valid) {
		if (isatty(STDERR_FILENO)) {
			fprintf(stderr, "\rAttempt %lu", ++nattempts);
		}

		ssize_t b = get_random_bytes(password, sizeof(password));
		assert(b == sizeof(password));

		r = argon2id_ctx(&context);
		die(r != 0, "argon2id failed\n");

		valid = 1;
		for (int i = 0; i < digits / 2; ++i) {
			if (hash[i] != 0) {
				valid = 0;
				break;
			}
		}
	}

	if (isatty(STDERR_FILENO)) {
		fprintf(stderr, ": found.\nHere is your proof:\n");
	}

	char proof[33];
	enchex(password, sizeof(password), proof, sizeof(proof));
	printf("%s\n", proof);

	return 0;
}
