/*--------------------------------------------------------------------------
 *
 * test_copy_ndjson.c
 *		A REAL custom COPY format (newline-delimited JSON), written against
 *		the API proposed in the "Make COPY format extendable" thread, to see
 *		whether an extension can actually implement a format with it.
 *
 * Unlike test_copy_custom_format, this one moves data: COPY TO writes one
 * JSON object per row and COPY FROM reads it back.  That is what exercises
 * the parts of the API an extension really needs.
 *
 * Portions Copyright (c) 2026, PostgreSQL Global Development Group
 *
 * -------------------------------------------------------------------------
 */

#include "postgres.h"

#include "commands/copy.h"
#include "commands/copy_state.h"
#include "commands/copyapi.h"
#include "catalog/pg_proc_d.h"
#include "mb/pg_wchar.h"
#include "utils/builtins.h"
#include "utils/fmgroids.h"
#include "utils/json.h"
#include "utils/jsonfuncs.h"
#include "utils/lsyscache.h"
#include "utils/memutils.h"

PG_MODULE_MAGIC;

typedef struct NdjsonToState
{
	TupleDesc	tupDesc;
	StringInfoData row;
} NdjsonToState;

typedef struct NdjsonFromState
{
	TupleDesc	tupDesc;
	StringInfoData line;
	bool		eof;
} NdjsonFromState;

/* lookup info for json_object_field_text(), filled in at CopyFromStart */
static FmgrInfo ndjson_field_finfo;

/* ---------------------------------------------------------------- COPY TO */

static void
NdjsonToOutFunc(CopyToState cstate, Oid atttypid, FmgrInfo *finfo)
{
	Oid			func_oid;
	bool		is_varlena;

	getTypeOutputInfo(atttypid, &func_oid, &is_varlena);
	fmgr_info(func_oid, finfo);
}

static void
NdjsonToStart(CopyToState cstate, TupleDesc tupDesc)
{
	NdjsonToState *st = palloc0_object(NdjsonToState);

	st->tupDesc = tupDesc;
	initStringInfo(&st->row);
	cstate->format_private = (void *) st;
}

static void
NdjsonToOneRow(CopyToState cstate, TupleTableSlot *slot)
{
	NdjsonToState *st = (NdjsonToState *) cstate->format_private;
	ListCell   *lc;
	bool		first = true;

	resetStringInfo(&st->row);
	appendStringInfoChar(&st->row, '{');

	foreach(lc, cstate->attnumlist)
	{
		int			attnum = lfirst_int(lc);
		Form_pg_attribute attr = TupleDescAttr(st->tupDesc, attnum - 1);
		Datum		value = slot->tts_values[attnum - 1];
		bool		isnull = slot->tts_isnull[attnum - 1];

		if (!first)
			appendStringInfoChar(&st->row, ',');
		first = false;

		escape_json(&st->row, NameStr(attr->attname));
		appendStringInfoChar(&st->row, ':');

		if (isnull)
			appendStringInfoString(&st->row, "null");
		else
		{
			char	   *str = OutputFunctionCall(&cstate->out_functions[attnum - 1],
												 value);

			escape_json(&st->row, str);
		}
	}

	appendStringInfoChar(&st->row, '}');

	/*
	 * CopySendEndOfRow is the raw primitive and does not terminate the row;
	 * the built-in text/CSV formats go through CopySendTextLikeEndOfRow for
	 * that, which is static too.  So a line-oriented custom format has to
	 * emit its own terminator.  Worth documenting in copyapi.h: nothing in
	 * the API says so, and the bundled test module never writes any data, so
	 * there is no example to copy from.
	 */
	appendStringInfoChar(&st->row, '\n');

	/*
	 * Here is the whole point of this module: an extension has to be able to
	 * hand these bytes to whatever destination the COPY is going to (file,
	 * PROGRAM, frontend, callback).  Only CopySendData/CopySendEndOfRow know
	 * how to do that, and in the posted patch both are static in copyto.c.
	 */
	CopySendData(cstate, st->row.data, st->row.len);
	CopySendEndOfRow(cstate);
}

static void
NdjsonToEnd(CopyToState cstate)
{
}

static const CopyToRoutine NdjsonToRoutine = {
	.CopyToOutFunc = NdjsonToOutFunc,
	.CopyToStart = NdjsonToStart,
	.CopyToOneRow = NdjsonToOneRow,
	.CopyToEnd = NdjsonToEnd,
};

/* -------------------------------------------------------------- COPY FROM */

static void
NdjsonFromInFunc(CopyFromState cstate, Oid atttypid, FmgrInfo *finfo,
				 Oid *typioparam)
{
	Oid			func_oid;

	getTypeInputInfo(atttypid, &func_oid, typioparam);
	fmgr_info(func_oid, finfo);
}

static void
NdjsonFromStart(CopyFromState cstate, TupleDesc tupDesc)
{
	NdjsonFromState *st = palloc0_object(NdjsonFromState);

	st->tupDesc = tupDesc;
	initStringInfo(&st->line);
	cstate->format_private = (void *) st;

	fmgr_info(F_JSON_OBJECT_FIELD_TEXT, &ndjson_field_finfo);
}

/*
 * Read one newline-terminated line from the COPY source.
 *
 * Same story as on the TO side: the source can be a file, a program or the
 * frontend, and CopyGetData is the only thing that knows the difference.
 * It is static in copyfromparse.c.
 */
static bool
ndjson_read_line(CopyFromState cstate, NdjsonFromState *st)
{
	char		c;

	resetStringInfo(&st->line);

	for (;;)
	{
		if (CopyGetData(cstate, &c, 1, 1) != 1)
			return st->line.len > 0;	/* last line without newline */

		if (c == '\n')
			return true;
		if (c == '\r')
			continue;
		appendStringInfoChar(&st->line, c);
	}
}

static bool
NdjsonFromOneRow(CopyFromState cstate, ExprContext *econtext,
				 Datum *values, bool *nulls)
{
	NdjsonFromState *st = (NdjsonFromState *) cstate->format_private;
	Datum		json;
	ListCell   *lc;

	if (!ndjson_read_line(cstate, st))
		return false;

	/* an extension is free to use whatever parser it likes; json_in is handy */
	json = DirectFunctionCall1(json_in, CStringGetDatum(st->line.data));

	memset(nulls, true, st->tupDesc->natts * sizeof(bool));

	foreach(lc, cstate->attnumlist)
	{
		int			attnum = lfirst_int(lc);
		Form_pg_attribute attr = TupleDescAttr(st->tupDesc, attnum - 1);
		Datum		field;
		bool		isnull;
		LOCAL_FCINFO(fcinfo, 2);

		/*
		 * json_object_field_text returns NULL for a missing key or a JSON
		 * null, so it cannot be called through DirectFunctionCall.
		 */
		InitFunctionCallInfoData(*fcinfo, &ndjson_field_finfo, 2,
								 InvalidOid, NULL, NULL);
		fcinfo->args[0].value = json;
		fcinfo->args[0].isnull = false;
		fcinfo->args[1].value = CStringGetTextDatum(NameStr(attr->attname));
		fcinfo->args[1].isnull = false;

		field = FunctionCallInvoke(fcinfo);
		isnull = fcinfo->isnull;

		if (isnull)
		{
			nulls[attnum - 1] = true;
			values[attnum - 1] = (Datum) 0;
		}
		else
		{
			char	   *str = TextDatumGetCString(field);

			values[attnum - 1] = InputFunctionCall(&cstate->in_functions[attnum - 1],
												   str,
												   cstate->typioparams[attnum - 1],
												   attr->atttypmod);
			nulls[attnum - 1] = false;
		}
	}

	return true;
}

static void
NdjsonFromEnd(CopyFromState cstate)
{
}

static const CopyFromRoutine NdjsonFromRoutine = {
	.CopyFromInFunc = NdjsonFromInFunc,
	.CopyFromStart = NdjsonFromStart,
	.CopyFromOneRow = NdjsonFromOneRow,
	.CopyFromEnd = NdjsonFromEnd,
};

void
_PG_init(void)
{
	RegisterCopyCustomFormat("ndjson", &NdjsonToRoutine, &NdjsonFromRoutine,
							 NULL, NULL);
}