#include "saltywitch.h"

struct saltywitch_generichash_context {
	uint64_t outlen;
	ErlNifMutex *mutex;
	crypto_generichash_state *state;
};

static ErlNifResourceType *_generichash_context_rtype = NULL;

static void _generichash_dtor(ErlNifEnv *env, void *context)
{
	struct saltywitch_generichash_context *ctx = context;

	if (ctx == NULL) {
		return;
	}

	if (ctx->state != NULL) {
		sodium_free(ctx->state);
		ctx->state = NULL;
	}

	if (ctx->mutex != NULL) {
		enif_mutex_destroy(ctx->mutex);
		ctx->mutex = NULL;
	}
}

int saltywitch_nif_init_generichash(ErlNifEnv *env)
{
	_generichash_context_rtype = enif_open_resource_type(env,
	                                                     NULL,
	                                                     "generichash_context",
	                                                     _generichash_dtor,
	                                                     ERL_NIF_RT_CREATE | ERL_NIF_RT_TAKEOVER,
	                                                     NULL);

	if (_generichash_context_rtype == NULL) {
		return -1;
	}

	return 0;
}

ERL_NIF_TERM saltywitch_generichash_keygen(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
	ERL_NIF_TERM term;

	unsigned char *key = enif_make_new_binary(env, crypto_generichash_KEYBYTES, &term);
	if (key == NULL) {
		return saltywitch_exception(env, atom_error, atom_err_nif_alloc);
	}

	crypto_generichash_keygen(key);

	return term;
}

ERL_NIF_TERM saltywitch_generichash(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
	unsigned char *hash = NULL;
	ERL_NIF_TERM term;
	ErlNifBinary bin;

	ErlNifBinary key = {
	    .size = 0,
	    .data = NULL,
	};

	uint64_t len = crypto_generichash_BYTES;

	if (enif_inspect_binary(env, argv[0], &bin) == false) {
		return enif_make_badarg(env);
	}

	switch (argc) {
	case 3:
		if (enif_get_uint64(env, argv[2], &len) == false) {
			return enif_make_badarg(env);
		}
		/* fallthrough */
	case 2:
		if (enif_inspect_binary(env, argv[1], &key) == false) {
			return enif_make_badarg(env);
		}
	}

	hash = enif_make_new_binary(env, len, &term);
	if (hash == NULL) {
		return enif_raise_exception(env, enif_make_atom(env, "opaque"));
	}

	crypto_generichash(hash, len, bin.data, bin.size, key.data, key.size);

	return term;
}

ERL_NIF_TERM saltywitch_generichash_init(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
	struct saltywitch_generichash_context *ctx;
	ERL_NIF_TERM term;

	ERL_NIF_TERM err    = atom_badarg;
	ERL_NIF_TERM reason = atom_err_opaque;

	ErlNifBinary key = {
	    .size = 0,
	    .data = NULL,
	};

	uint64_t len = crypto_generichash_BYTES;

	switch (argc) {
	case 2:
		if (enif_get_uint64(env, argv[1], &len) == false) {
			reason = atom_err_invalid_type;
			goto badarg_FAULT;
		}

		if (len < crypto_generichash_BYTES_MIN) {
			reason = atom_err_output_too_small;
			goto badarg_FAULT;
		}

		if (len > crypto_generichash_BYTES_MAX) {
			reason = atom_err_output_too_large;
			goto badarg_FAULT;
		}
		/* fallthrough */
	case 1:
		if (enif_inspect_binary(env, argv[0], &key) == false) {
			reason = atom_err_invalid_type;
			goto badarg_FAULT;
		}

		if (key.size < crypto_generichash_KEYBYTES_MIN) {
			reason = atom_err_key_too_small;
			goto badarg_FAULT;
		}

		if (key.size > crypto_generichash_KEYBYTES_MAX) {
			reason = atom_err_key_too_large;
			goto badarg_FAULT;
		}
	}

	ctx = enif_alloc_resource(_generichash_context_rtype, sizeof(struct saltywitch_generichash_context));
	if (ctx == NULL) {
		reason = atom_err_nif_alloc;
		goto resource_alloc_FAULT;
	}

	ctx->outlen = len;

	ctx->mutex = enif_mutex_create("saltywitch.generichash");
	if (ctx->mutex == NULL) {
		reason = atom_err_nif_alloc;
		goto mutex_FAULT;
	}

	ctx->state = sodium_malloc(sizeof(crypto_generichash_state));
	if (ctx->state == NULL) {
		goto state_alloc_FAULT;
	}

	crypto_generichash_init(ctx->state, key.data, key.size, len);

	term = enif_make_resource(env, ctx);

	enif_release_resource(ctx);

	return term;

state_alloc_FAULT:
	enif_mutex_destroy(ctx->mutex);
mutex_FAULT:
	enif_release_resource(ctx);
resource_alloc_FAULT:
	err = atom_error;
badarg_FAULT:
	return saltywitch_exception(env, err, reason);
}

ERL_NIF_TERM saltywitch_generichash_update(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
	struct saltywitch_generichash_context *ctx = NULL;
	ErlNifBinary message;

	ERL_NIF_TERM err    = atom_badarg;
	ERL_NIF_TERM reason = atom_err_opaque;

	if (enif_get_resource(env, argv[0], _generichash_context_rtype, (void **) &ctx) == false) {
		reason = atom_err_invalid_type;
		goto badarg_FAULT;
	}

	if (ctx == NULL) {
		goto badarg_FAULT;
	}

	if (ctx->mutex == NULL || ctx->state == NULL) {
		goto badarg_FAULT;
	}

	if (enif_inspect_iolist_as_binary(env, argv[1], &message) == false) {
		reason = atom_err_invalid_type;
		goto badarg_FAULT;
	}

	enif_mutex_lock(ctx->mutex);
	crypto_generichash_update(ctx->state, message.data, message.size);
	enif_mutex_unlock(ctx->mutex);

	return atom_ok;

badarg_FAULT:
	return saltywitch_exception(env, err, reason);
}

ERL_NIF_TERM saltywitch_generichash_final(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
	struct saltywitch_generichash_context *ctx = NULL;
	unsigned char *hash;
	ERL_NIF_TERM term;

	ERL_NIF_TERM err    = atom_badarg;
	ERL_NIF_TERM reason = atom_err_opaque;

	if (enif_get_resource(env, argv[0], _generichash_context_rtype, (void **) &ctx) == false) {
		reason = atom_err_invalid_type;
		goto badarg_FAULT;
	}

	if (ctx == NULL) {
		goto badarg_FAULT;
	}

	if (ctx->mutex == NULL || ctx->state == NULL) {
		goto badarg_FAULT;
	}

	hash = enif_make_new_binary(env, ctx->outlen, &term);
	if (hash == NULL) {
		reason = atom_err_nif_alloc;
		goto error_FAULT;
	}

	enif_mutex_lock(ctx->mutex);
	crypto_generichash_final(ctx->state, hash, ctx->outlen);
	sodium_free(ctx->state);
	ctx->state  = NULL;
	ctx->outlen = 0;
	enif_mutex_unlock(ctx->mutex);

	return term;

error_FAULT:
	err = atom_error;
badarg_FAULT:
	return saltywitch_exception(env, err, reason);
}