#include <stdint.h>
//#include <stdio.h>
//#include <unistd.h>
#include <errno.h>
#include <alloca.h>

#include <ucode/module.h>

#include <mbedtls/error.h>
#include <mbedtls/md.h>
#include <mbedtls/pk.h>

#define TRUE ucv_boolean_new(true)
#define FALSE ucv_boolean_new(false)

static uc_resource_type_t *pk_type;

static uc_value_t *
md_digest(uc_vm_t *vm, size_t nargs)
{
	uc_value_t *alg = uc_fn_arg(0);
	uc_value_t *input = uc_fn_arg(1);

	if (ucv_type(alg) != UC_STRING)
		uc_vm_raise_exception(vm, EXCEPTION_TYPE, "alg not a string");

	if (ucv_type(input) != UC_STRING)
		uc_vm_raise_exception(vm, EXCEPTION_TYPE, "input not a string");

	const mbedtls_md_info_t *info = mbedtls_md_info_from_string(ucv_string_get(alg));

	if (!info)
		uc_vm_raise_exception(vm, EXCEPTION_TYPE, "unknown MD algorithm");

	unsigned char size = mbedtls_md_get_size(info);
	char *output = alloca(size);

	if (mbedtls_md(info, (const unsigned char*)ucv_string_get(input), ucv_string_length(input), (unsigned char*)output))
		uc_vm_raise_exception(vm, EXCEPTION_RUNTIME, "bad input data");

	uc_value_t *rv = ucv_string_new_length(output, size);
	return rv;
}

#if 0

static uc_value_t *
ecsa_verify(uc_vm_t *vm, size_t nargs)
{
	// Read key from 0x04 | x | y
	int mbedtls_ecp_point_read_binary(const mbedtls_ecp_group *grp,
					  mbedtls_ecp_point *pt,
					  const unsigned char *buf, size_t ilen);

	// Verify
	int mbedtls_ecdsa_verify(mbedtls_ecp_group *grp,
				 const unsigned char *buf, size_t blen,
				 const mbedtls_ecp_point *Q,
				 const mbedtls_mpi *r,
				 const mbedtls_mpi *s);
}
#endif

static uc_value_t *
pk_init(uc_vm_t *vm, size_t nargs)
{
	const mbedtls_pk_info_t *info = mbedtls_pk_info_from_type(MBEDTLS_PK_ECDSA);
	mbedtls_pk_context *ctx = malloc(sizeof(mbedtls_pk_context));

	mbedtls_pk_init(ctx);
	/* if (mbedtls_pk_setup(ctx, info)) */
	/* 	uc_vm_raise_exception(vm, EXCEPTION_RUNTIME, "ECDSA not available"); */
		
	return uc_resource_new(pk_type, ctx);
}

static void
pk_free(void *ptr)
{
	mbedtls_pk_context *ctx = ptr;

	mbedtls_pk_free(ctx);
	free(ctx);
}

static uc_value_t *
pk_set_public_key(uc_vm_t *vm, size_t nargs)
{
	mbedtls_pk_context *ctx = uc_fn_thisval("mbedtls.pk");

        if (!ctx)
		uc_vm_raise_exception(vm, EXCEPTION_TYPE, "invalid mbedtls.pk object");

	uc_value_t *key = uc_fn_arg(0);

	if (ucv_type(key) != UC_STRING)
		uc_vm_raise_exception(vm, EXCEPTION_TYPE, "key is not a string");

	int err = mbedtls_pk_parse_public_key(ctx, (const unsigned char*)ucv_string_get(key), ucv_string_length(key));
	if (err)
		uc_vm_raise_exception(vm, EXCEPTION_RUNTIME, "not a valid PEM/DER key %s: %s",
				      mbedtls_high_level_strerr(err),
				      mbedtls_low_level_strerr(err));
	return NULL;
}

static uc_value_t *
pk_verify(uc_vm_t *vm, size_t nargs)
{
	mbedtls_pk_context *ctx = uc_fn_thisval("mbedtls.pk");

        if (!ctx)
		uc_vm_raise_exception(vm, EXCEPTION_TYPE, "invalid mbedtls.pk object");

	uc_value_t *md_alg = uc_fn_arg(0);
	uc_value_t *hash = uc_fn_arg(1);
	uc_value_t *sig = uc_fn_arg(2);

	if (ucv_type(md_alg) != UC_STRING)
		uc_vm_raise_exception(vm, EXCEPTION_TYPE, "md_alg is not a string");

	if (ucv_type(hash) != UC_STRING)
		uc_vm_raise_exception(vm, EXCEPTION_TYPE, "hash is not a string");

	if (ucv_type(sig) != UC_STRING)
		uc_vm_raise_exception(vm, EXCEPTION_TYPE, "sig is not a string");

	const mbedtls_md_info_t *md_info = mbedtls_md_info_from_string(ucv_string_get(md_alg));

	if (!md_info)
		uc_vm_raise_exception(vm, EXCEPTION_TYPE, "unknown MD algorithm");

	const mbedtls_md_type_t md_type = mbedtls_md_get_type(md_info);
	
	int err = mbedtls_pk_verify(ctx, md_type,
				    (const unsigned char*)ucv_string_get(hash), ucv_string_length(hash),
				    (const unsigned char*)ucv_string_get(sig), ucv_string_length(sig));
	if (err)
		uc_vm_raise_exception(vm, EXCEPTION_RUNTIME, "validation failed: %s: %s",
				      mbedtls_high_level_strerr(err),
				      mbedtls_low_level_strerr(err));

	return TRUE;
}		    

static const uc_function_list_t global_fns[] = {
	{ "md_digest",			md_digest },
	{ "pk",				pk_init },
};

static const uc_function_list_t pk_fns[] = {
	{ "set_public_key",		pk_set_public_key },
	{ "verify",			pk_verify },
};

void uc_module_init(uc_vm_t *vm, uc_value_t *scope)
{
	uc_function_list_register(scope, global_fns);

	pk_type = uc_type_declare(vm, "mbedtls.pk", pk_fns, pk_free);	
}