/*-------------------------------------------------------------------------
 * Copyright (c) Microsoft Corporation.  All rights reserved.
 *
 * src/planner/documentdb_custom_planner.c
 *
 * Implementation of the cursor based operations.
 *
 *-------------------------------------------------------------------------
 */
#include <postgres.h>
#include <fmgr.h>
#include <miscadmin.h>
#include <funcapi.h>
#include <utils/portal.h>
#include <utils/varlena.h>
#include <utils/rel.h>
#include <nodes/makefuncs.h>
#include <utils/lsyscache.h>
#include <metadata/metadata_cache.h>
#include <commands/portalcmds.h>
#include "metadata/collection.h"
#include "planner/documents_custom_planner.h"
#include "utils/version_utils.h"

static bool SetPointReadQualsOnIndexScan(IndexScan *indexScan, Expr *queryQuals);
static List * FormatProjections(List *targetEntries);

/*
 * Does a best effort to create a point read plan on the primary key index.
 * If it's unable to, returns NULL - this will fall back to the default postgres planner.
 */
PlannedStmt *
TryCreatePointReadPlan(Query *query)
{
	if (query->limitOffset != NULL)
	{
		if (IsA(query->limitOffset, Const))
		{
			Const *limitValue = (Const *) query->limitOffset;
			if (!limitValue->constisnull && DatumGetInt64(limitValue->constvalue) != 0)
			{
				/* non-zero offset is not valid for this planner */
				return NULL;
			}
		}
		else
		{
			return NULL;
		}
	}

	if (query->limitCount != NULL)
	{
		if (IsA(query->limitCount, Const))
		{
			Const *limitValue = (Const *) query->limitCount;
			if (!limitValue->constisnull &&
				DatumGetInt64(limitValue->constvalue) < 1)
			{
				/* Limit 0 or less is not valid for this planner */
				return NULL;
			}
		}
		else
		{
			/* Fall back */
			return NULL;
		}
	}

	/* Multi table queries not supported */
	if (list_length(query->rtable) != 1)
	{
		return NULL;
	}

	RangeTblEntry *firstEntry = linitial(query->rtable);

	/* Querying anything but tables not supported */
	if (firstEntry->rtekind != RTE_RELATION)
	{
		return NULL;
	}

	PlannedStmt *stmt = makeNode(PlannedStmt);
	stmt->stmt_len = 0;
	stmt->stmt_location = -1;
	stmt->queryId = query->queryId;
	stmt->commandType = query->commandType;
	stmt->canSetTag = true;

#if PG_VERSION_NUM >= 180000
	Bitmapset *unprunableRelIds = NULL;
	unprunableRelIds = bms_add_member(unprunableRelIds, 1);
	stmt->unprunableRelids = unprunableRelIds;
#endif

	stmt->rtable = query->rtable;
	stmt->jitFlags = 0;

	/* Create an IndexScan plan on the BTree index */
	IndexScan *indexScan = makeNode(IndexScan);
	indexScan->scan.plan.plan_rows = 1;

	/* Use custom values that hint that it's our plan */
	indexScan->scan.plan.startup_cost = 1.23123e5;
	indexScan->scan.plan.total_cost = 1.23123e6;
	indexScan->scan.plan.lefttree = NULL;
	indexScan->scan.plan.righttree = NULL;
	indexScan->scan.scanrelid = 1;

	/* We don't need order by since it's a point read */
	indexScan->indexorderby = NIL;
	indexScan->indexorderbyorig = NIL;
	indexScan->indexorderbyops = NIL;
	indexScan->indexorderdir = ForwardScanDirection;

	/* Do not add returns before the RelationClose */
	Relation relation = RelationIdGetRelation(firstEntry->relid);
	indexScan->indexid = relation->rd_pkindex;
	if (indexScan->indexid == InvalidOid)
	{
		/* Load the index */
#if PG_VERSION_NUM >= 180000
		bool deferrableOk = false;
		indexScan->indexid = RelationGetPrimaryKeyIndex(relation, deferrableOk);
#else
		indexScan->indexid = RelationGetPrimaryKeyIndex(relation);
#endif
	}

	RelationClose(relation);

	if (indexScan->indexid == InvalidOid)
	{
		/* We can't detect a PK index - fall back */
		return NULL;
	}

	if (!SetPointReadQualsOnIndexScan(indexScan, (Expr *) query->jointree->quals))
	{
		return NULL;
	}

	/* Finally, filter and set the projections if successful */
	indexScan->scan.plan.targetlist = FormatProjections(query->targetList);
	stmt->planTree = (Plan *) indexScan;

#if (PG_VERSION_NUM >= 160000)

	/* Add the permsInfo on the planned statement */
	stmt->permInfos = query->rteperminfos;
#endif

	return stmt;
}


/*
 * Removes unnecessary projections as per the point read plan.
 */
static List *
FormatProjections(List *targetEntries)
{
	if (list_length(targetEntries) == 1)
	{
		return targetEntries;
	}

	ListCell *cell;
	foreach(cell, targetEntries)
	{
		TargetEntry *currEntry = lfirst(cell);
		if (currEntry->resjunk)
		{
			/* unused - remove */
			targetEntries = foreach_delete_current(targetEntries, cell);
		}
	}

	return targetEntries;
}


static bool
TraverseQualsAndExtract(List *quals, List **queryRuntimeClauses,
						Expr **objectIdExp, Expr **objectIdOrigExp,
						Expr **shardKeyExp, Expr **shardKeyOrigExp)
{
	ListCell *cell;
	foreach(cell, quals)
	{
		Expr *expr = (Expr *) lfirst(cell);
		switch (expr->type)
		{
			case T_OpExpr:
			{
				OpExpr *opExpr = (OpExpr *) expr;

				if (opExpr->opfuncid == 0)
				{
					/* Set the opFuncId for runtime execution */
					opExpr->opfuncid = get_opcode(opExpr->opno);
				}
				if (list_length(opExpr->args) != 2)
				{
					/* Unable to push data into the specified index */
					*queryRuntimeClauses = lappend(*queryRuntimeClauses, expr);
					continue;
				}

				Expr *firstArg = linitial(opExpr->args);
				if (!IsA(firstArg, Var))
				{
					*queryRuntimeClauses = lappend(*queryRuntimeClauses, expr);
					continue;
				}

				Var *firstVar = (Var *) firstArg;
				if (firstVar->varattno == DOCUMENT_DATA_TABLE_OBJECT_ID_VAR_ATTR_NUMBER &&
					opExpr->opno == BsonEqualOperatorId())
				{
					if (*objectIdExp != NULL)
					{
						return false;
					}

					OpExpr *newOpExpr = copyObject(opExpr);
					Var *indexVar = makeVar(INDEX_VAR, 2, BsonTypeId(), -1,
											InvalidOid, 0);
					newOpExpr->args = list_make2(indexVar, lsecond(opExpr->args));
					*objectIdExp = (Expr *) newOpExpr;
					*objectIdOrigExp = (Expr *) opExpr;
				}
				else if (firstVar->varattno ==
						 DOCUMENT_DATA_TABLE_SHARD_KEY_VALUE_VAR_ATTR_NUMBER &&
						 opExpr->opno == BigintEqualOperatorId())
				{
					if (*shardKeyExp != NULL)
					{
						return false;
					}

					OpExpr *newOpExpr = copyObject(opExpr);
					Var *indexVar = makeVar(INDEX_VAR, 1, BsonTypeId(), -1,
											InvalidOid, 0);
					newOpExpr->args = list_make2(indexVar, lsecond(opExpr->args));
					*shardKeyExp = (Expr *) newOpExpr;
					*shardKeyOrigExp = (Expr *) opExpr;
				}
				else
				{
					if (opExpr->opfuncid == BsonEqualMatchRuntimeFunctionId())
					{
						Expr *secondArg = lsecond(opExpr->args);
						pgbsonelement secondElement;
						if (IsA(secondArg, Const) &&
							TryGetSinglePgbsonElementFromPgbson(DatumGetPgBsonPacked(
																	((Const *)
																	 secondArg)->
																	constvalue),
																&secondElement) &&
							strcmp(secondElement.path, "_id") == 0)
						{
							/* Skip the _id runtime filter: We know we will have
							 * an _id filter by ensuring it below.
							 */
							continue;
						}
					}
					else if (opExpr->opfuncid == BsonTextFunctionId())
					{
						/* Text search does not qualify here */
						return false;
					}

					*queryRuntimeClauses = lappend(*queryRuntimeClauses, expr);
				}

				continue;
			}

			case T_FuncExpr:
			{
				FuncExpr *funcExpr = (FuncExpr *) expr;
				if (funcExpr->funcid == BsonTextFunctionId())
				{
					/* Text search does not qualify here */
					return false;
				}
				else if (funcExpr->funcid == BsonIndexHintFunctionOid())
				{
					return false;
				}
				else if (funcExpr->funcid == BsonFullScanFunctionOid())
				{
					/* Strip full scan */
					continue;
				}
				else if (IsClusterVersionAtleast(DocDB_V0, 112, 1) &&
						 funcExpr->funcid == ApiCursorTrackerFunctionId())
				{
					/* Strip the continuation function */
					continue;
				}

				*queryRuntimeClauses = lappend(*queryRuntimeClauses, expr);
				continue;
			}

			case T_BoolExpr:
			{
				BoolExpr *boolExpr = (BoolExpr *) expr;
				if (boolExpr->boolop == AND_EXPR)
				{
					if (!TraverseQualsAndExtract(boolExpr->args, queryRuntimeClauses,
												 objectIdExp, objectIdOrigExp,
												 shardKeyExp, shardKeyOrigExp))
					{
						return false;
					}
				}
				else
				{
					*queryRuntimeClauses = lappend(*queryRuntimeClauses, expr);
				}

				continue;
			}

			default:
			{
				*queryRuntimeClauses = lappend(*queryRuntimeClauses, expr);
				continue;
			}
		}
	}

	return true;
}


/*
 * Scan quals and ensure that there's a point read in there.
 * returns false if it couldn't form a point read plan.
 */
static bool
SetPointReadQualsOnIndexScan(IndexScan *indexScan, Expr *queryQuals)
{
	List *runtimeClauses = NIL;
	List *quals = make_ands_implicit(queryQuals);

	Expr *objectIdExpr = NULL;
	Expr *objectIdOriginalExpr = NULL;
	Expr *shardKeyExpr = NULL;
	Expr *shardKeyOriginalExpr = NULL;

	if (!TraverseQualsAndExtract(quals, &runtimeClauses, &objectIdExpr,
								 &objectIdOriginalExpr, &shardKeyExpr,
								 &shardKeyOriginalExpr))
	{
		return false;
	}

	/* Not a point read */
	if (objectIdExpr == NULL || shardKeyExpr == NULL)
	{
		return false;
	}

	indexScan->scan.plan.qual = runtimeClauses;
	indexScan->indexqual = list_make2(shardKeyExpr, objectIdExpr);
	indexScan->indexqualorig = list_make2(shardKeyOriginalExpr, objectIdOriginalExpr);
	return true;
}
