/*-------------------------------------------------------------------------
 * Copyright (c) Microsoft Corporation.  All rights reserved.
 *
 * src/commands/bson_get_rewritten_sql.c
 *
 * Implementation of bson_get_rewritten_sql which takes a database name
 * and a BSON command spec (aggregate/find/count/distinct) and returns
 * the SQL that would be generated by the query planner.
 *
 *-------------------------------------------------------------------------
 */
#include <postgres.h>
#include <fmgr.h>
#include <utils/builtins.h>
#include <utils/ruleutils.h>

#include "io/bson_core.h"
#include "io/bsonvalue_utils.h"
#include "aggregation/bson_aggregation_pipeline.h"

/* Export the function to PostgreSQL */
PG_FUNCTION_INFO_V1(command_bson_get_rewritten_sql);

/*
 * command_bson_get_rewritten_sql takes a database name (text) and
 * a BSON command spec, determines the command type from the first
 * key of the spec, generates the corresponding Query tree, and
 * deparses it to SQL text.
 */
Datum
command_bson_get_rewritten_sql(PG_FUNCTION_ARGS)
{
	text *database = PG_GETARG_TEXT_PP(0);
	pgbson *spec = PG_GETARG_PGBSON(1);

	/* Extract the first key to determine the command type */
	bson_iter_t iter;
	PgbsonInitIterator(spec, &iter);

	if (!bson_iter_next(&iter))
	{
		ereport(ERROR, (errmsg("empty BSON spec provided")));
	}

	StringView commandName = bson_iter_key_string_view(&iter);
	Query *query;

	CursorParamKind cursorParams = CursorParamKind_Persistent;
	bool setStatementTimeout = false;

	if (StringViewEqualsCString(&commandName, "aggregate"))
	{
		QueryData queryData = GenerateFirstPageQueryData();
		query = GenerateAggregationQuery(database, spec, &queryData,
										 cursorParams, setStatementTimeout);
	}
	else if (StringViewEqualsCString(&commandName, "find"))
	{
		QueryData queryData = GenerateFirstPageQueryData();
		query = GenerateFindQuery(database, spec, &queryData,
								  cursorParams, setStatementTimeout);
	}
	else if (StringViewEqualsCString(&commandName, "count"))
	{
		query = GenerateCountQuery(database, spec, setStatementTimeout);
	}
	else if (StringViewEqualsCString(&commandName, "distinct"))
	{
		query = GenerateDistinctQuery(database, spec, setStatementTimeout);
	}
	else
	{
		ereport(ERROR, (errmsg("unsupported command type: %.*s",
							   commandName.length, commandName.string)));
	}

	bool pretty = false;
	char *sqlText = pg_get_querydef(query, pretty);

	PG_RETURN_TEXT_P(cstring_to_text(sqlText));
}
