/*-------------------------------------------------------------------------
 *
 * json_io.c
 *	  Primary input/output and conversion procedures
 *	  for JSON data type.
 *
 * Copyright (c) 2010, PostgreSQL Global Development Group
 * Written by Joey Adams <joeyadams3.14159@gmail.com>.
 *
 *-------------------------------------------------------------------------
 */

#include "json.h"
#include "util.h"

#include "utils/array.h"
#include "mb/pg_wchar.h"

PG_MODULE_MAGIC;

static json_type decide_json_type(Oid type, char category);
static const char *datum_to_json(Datum datum, TypeInfo *typeInfo,
			  json_type target_type);
static const char *array_to_json(Datum datum);
static const char *build_array_string(const char **values,
				   const int *dim, int ndim);
static void build_array_string_recurse(StringInfo string, const char ***values,
						   const int *dim, int ndim);

void
report_corrupt_json(void)
{
	ereport(ERROR,
			(errcode(ERRCODE_DATA_CORRUPTED),
			 errmsg("corrupt JSON content")));
}

PG_FUNCTION_INFO_V1(json_in);
Datum		json_in(PG_FUNCTION_ARGS);
Datum
json_in(PG_FUNCTION_ARGS)
{
	char	   *string = PG_GETARG_CSTRING(0);
	jsontype   *vardata;

	if (!json_validate_server_encoded(string))
		ereport(ERROR,
				(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
				 errmsg("could not parse JSON value")));

	vardata = cstring_to_text(string);
	PG_RETURN_JSON_P(vardata);
}

PG_FUNCTION_INFO_V1(json_out);
Datum		json_out(PG_FUNCTION_ARGS);
Datum
json_out(PG_FUNCTION_ARGS)
{
	jsontype   *vardata = PG_GETARG_JSON_P(0);
	char	   *string = text_to_cstring((text *) vardata);

	Assert(json_validate_server_encoded(string));

	PG_RETURN_CSTRING(string);
}

/*
 * Performs JSON value extraction in UTF-8 C-String land.
 *
 * If the input was JSON-encoded NULL, this function returns NULL
 * (indicating that we want from_json to yield actual SQL NULL).
 */
static const char *
from_json_cstring(const char *input)
{
	size_t		len;
	JSON	   *json = json_decode(input);
	const char *cstring_out = NULL;

	if (json == NULL)
		report_corrupt_json();

	switch (json->type)
	{
		case JSON_NULL:
			cstring_out = NULL;
			break;
		case JSON_STRING:
			cstring_out = json_get_string(json, &len);
			if (strlen(cstring_out) != len)
				ereport(ERROR,
						(errcode(ERRCODE_UNTRANSLATABLE_CHARACTER),
						 errmsg("could not convert JSON-encoded string to TEXT: contains \\u0000"),
						 errdetail("the string was: %s", utf8_to_server(input, strlen(input)))));

			/*
			 * todo: suggest using from_json_as_bytea instead when that is
			 * available.
			 */
			break;
		case JSON_NUMBER:
			cstring_out = json_get_number(json);
			break;
		case JSON_BOOL:
			cstring_out = json_get_bool(json) ? "true" : "false";
			break;
		case JSON_OBJECT:
			ereport(ERROR,
					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
					 errmsg("from_json cannot be applied to JSON objects")));
			break;
		case JSON_ARRAY:
			ereport(ERROR,
					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
					 errmsg("from_json cannot be applied to JSON arrays")));
			break;
		default:
			Assert(false);
	}

	return cstring_out;
}

PG_FUNCTION_INFO_V1(from_json);
Datum		from_json(PG_FUNCTION_ARGS);
Datum
from_json(PG_FUNCTION_ARGS)
{
	char	   *cstring_in = text_to_utf8_cstring(PG_GETARG_JSON_P(0));
	const char *cstring_out = from_json_cstring(cstring_in);
	text	   *vardata_out;

	pfree(cstring_in);

	if (cstring_out != NULL)
	{
		vardata_out = utf8_cstring_to_text(cstring_out);
		PG_RETURN_TEXT_P(vardata_out);
	}
	else
	{
		PG_RETURN_NULL();
	}
}

PG_FUNCTION_INFO_V1(to_json);
Datum		to_json(PG_FUNCTION_ARGS);
Datum
to_json(PG_FUNCTION_ARGS)
{
	Oid			argtype = get_fn_expr_argtype(fcinfo->flinfo, 0);
	TypeInfo   *typeInfo;
	json_type	target_type;

	typeInfo = fcinfo->flinfo->fn_extra;
	if (typeInfo == NULL)
	{
		typeInfo = fcinfo->flinfo->fn_extra =
			MemoryContextAlloc(fcinfo->flinfo->fn_mcxt, sizeof(TypeInfo));
		getTypeInfo(typeInfo, argtype, IOFunc_output, fcinfo->flinfo->fn_mcxt);
	}
	else if (typeInfo->type != argtype)
	{
		getTypeInfo(typeInfo, argtype, IOFunc_output, fcinfo->flinfo->fn_mcxt);
	}

	target_type = decide_json_type(argtype, typeInfo->typcategory);

	/*
	 * If NULL argument is given, return the string "null", not actual NULL.
	 * The null check is done after the type check so users can identify type
	 * bugs earlier.
	 */
	if (PG_ARGISNULL(0))
		PG_RETURN_JSON_P(cstring_to_text("null"));

	PG_RETURN_JSON_P(utf8_cstring_to_text(
				  datum_to_json(PG_GETARG_DATUM(0), typeInfo, target_type)));
}

/*
 * decide_json_type
 *
 * Given a type and its corresponding typcategory, determine what
 * JSON type to encode it to.  If it's not possible to encode the type
 * (or if support for the type isn't implemented yet), report an error.
 */
static json_type
decide_json_type(Oid type, char category)
{
	switch (category)
	{
		case 'S':
			return JSON_STRING;
		case 'P':
			if (type == CSTRINGOID)
				return JSON_STRING;
			goto invalid;
		case 'N':
			if (type == CASHOID)
				goto invalid;
			return JSON_NUMBER;
		case 'B':
			if (type == BOOLOID)
				return JSON_BOOL;
			goto invalid;
		case 'A':
			return JSON_ARRAY;
		default:
			goto invalid;
	}

invalid:
	ereport(ERROR,
			(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
		 errmsg("to_json cannot convert %s to JSON", format_type_be(type))));
	return JSON_INVALID;
}

/*
 * datum_to_json
 *	  Converts a datum to a UTF-8-encoded JSON string.
 *
 *	  typeInfo comes from the getTypeInfo function, and
 *	  target_type comes from decide_json_type .
 *
 *	  See to_json and array_to_json for examples of
 *	  how to invoke this function.
 */
static const char *
datum_to_json(Datum datum, TypeInfo *typeInfo, json_type target_type)
{
	char	   *cstring;
	char	   *cstring_utf8;
	JSON	   *node;
	char	   *encoded_utf8;

	switch (target_type)
	{
		case JSON_STRING:
		case JSON_NUMBER:
			cstring = OutputFunctionCall(&typeInfo->proc, datum);
			cstring_utf8 = server_to_utf8(cstring, strlen(cstring));

			if (cstring != cstring_utf8)
				pfree(cstring);

			if (target_type == JSON_STRING)
				node = json_mkstring(cstring_utf8, strlen(cstring_utf8));
			else
				node = json_mknumber(cstring_utf8, strlen(cstring_utf8));

			encoded_utf8 = json_encode(node, 0);
			if (encoded_utf8 == NULL)
			{
				/*
				 * Currently, we assume that output conversion for all types
				 * in category 'N' (excluding CASH) will, by coincidence,
				 * always give a valid JSON-encoded number.
				 *
				 * If this error message appears, it means the assumption is
				 * incorrect, and we need to either blacklist the type in
				 * decide_json_type (like we did with CASH), or add code that
				 * can parse the number type.
				 */
				ereport(ERROR,
						(errcode(ERRCODE_INTERNAL_ERROR),
						 errmsg("to_json failed to convert %s \"%s\" to a valid JSON %s.\n",
								format_type_be(typeInfo->type),
								cstring,
						 target_type == JSON_STRING ? "string" : "number")));
			}

			return encoded_utf8;

		case JSON_BOOL:
			return DatumGetBool(datum) ? "true" : "false";

		case JSON_ARRAY:
			return array_to_json(datum);

		default:
			Assert(false);
			return NULL;
	}
}

/*
 * array_to_json
 *	  Converts a PostgreSQL array datum to a UTF-8-encoded JSON string.
 *
 *	  Note: We assume that any type with typcategory = 'A'
 *			is compatible with array_out.
 */
static const char *
array_to_json(Datum datum)
{
	ArrayType  *v = DatumGetArrayTypeP(datum);
	Oid			element_type = ARR_ELEMTYPE(v);
	TypeInfo	element_typeinfo;
	json_type	target_type;
	int		   *dim;
	int			ndim;
	int			nitems;
	int			i;
	char	   *s;
	bits8	   *bitmap;
	int			bitmask;
	Datum		elt;
	const char **values;
	const char *ret;

	/*
	 * todo: Consider caching the TypeInfo of array items (which we're
	 * computing now) in the fcinfo->flinfo->fn_mcxt of to_json like we do
	 * with the TypeInfo of the array itself.
	 */
	getTypeInfo(&element_typeinfo, element_type, IOFunc_output, CurrentMemoryContext);

	target_type = decide_json_type(element_type, element_typeinfo.typcategory);

	ndim = ARR_NDIM(v);
	dim = ARR_DIMS(v);
	nitems = ArrayGetNItems(ndim, dim);

	/*
	 * Unfortunately, we can't make SELECT
	 * to_json(ARRAY[[],[],[]]::integer[][]); yield '[[],[],[]]' because ndim
	 * is set to 0 when the array has no items.
	 */
	if (nitems <= 0)
		return "[]";

	values = palloc(nitems * sizeof(const char *));

	s = ARR_DATA_PTR(v);
	bitmap = ARR_NULLBITMAP(v);
	bitmask = 1;

	for (i = 0; i < nitems; i++)
	{
		if (bitmap != NULL && (*bitmap & bitmask) == 0)
		{
			values[i] = NULL;
		}
		else
		{
			elt = fetch_att(s, element_typeinfo.typbyval, element_typeinfo.typlen);
			s = att_addlength_datum(s, element_typeinfo.typlen, elt);
			s = (char *) att_align_nominal(s, element_typeinfo.typalign);

			values[i] = datum_to_json(elt, &element_typeinfo, target_type);
		}

		if (bitmap != NULL)
		{
			bitmask <<= 1;
			if (bitmask == 0x100)
			{
				bitmap++;
				bitmask = 1;
			}
		}
	}

	ret = build_array_string(values, dim, ndim);

	pfree(values);

	return ret;
}

/*
 * build_array_string
 *	  Takes a multidimensional array of (JSON-formatted) strings and
 *	  converts it to JSON by wrapping the strings in array brackets and commas.
 */
static const char *
build_array_string(const char **values, const int *dim, int ndim)
{
	StringInfoData string[1];

	initStringInfo(string);
	build_array_string_recurse(string, &values, dim, ndim);

	return string->data;
}

static void
build_array_string_recurse(StringInfo string, const char ***values,
						   const int *dim, int ndim)
{
	int			i,
				count = dim[0];

	appendStringInfoChar(string, '[');

	for (i = 0; i < count; i++)
	{
		if (ndim > 1)
		{
			build_array_string_recurse(string, values, dim + 1, ndim - 1);
		}
		else
		{
			const char *value = *(*values)++;

			if (value != NULL)
				appendStringInfoString(string, value);
			else
				appendStringInfoString(string, "null");
		}

		if (i + 1 < count)
			appendStringInfoChar(string, ',');
	}

	appendStringInfoChar(string, ']');
}