/*-------------------------------------------------------------------------
 * Copyright (c) Microsoft Corporation.  All rights reserved.
 *
 * include/planner/bson_aggregation_pipeline_private.h
 *
 * Private helpers for the bson_aggregation_pipeline definition
 *
 *-------------------------------------------------------------------------
 */

#include <nodes/parsenodes.h>
#include <nodes/makefuncs.h>

#include "metadata/collection.h"
#include <utils/version_utils.h>
#include "collation/collation.h"
#include "utils/docdb_make_funcs.h"

#ifndef BSON_AGGREGATION_PIPELINE_PRIVATE_H
#define BSON_AGGREGATION_PIPELINE_PRIVATE_H

typedef struct AggregationStageDefinition AggregationStageDefinition;

/*
 * for nested stage, use this to record its parent stage name
 */
typedef enum ParentStageName
{
	ParentStageName_NOPARENT = 0,
	ParentStageName_LOOKUP,
	ParentStageName_FACET,
	ParentStageName_UNIONWITH,
	ParentStageName_INVERSEMATCH,
} ParentStageName;


/*
 * Enums to represent all kind of aggregations stages.
 * Please keep the list sorted within their groups for easier readability
 */
typedef enum
{
	Stage_Invalid = 0,

	/* Start internal stages Mongo */
	Stage_Internal_InhibitOptimization = 1,

	/* Start Mongo Public stages */
	Stage_AddFields = 10,
	Stage_Bucket,
	Stage_BucketAuto,
	Stage_ChangeStream,
	Stage_CollStats,
	Stage_Count,
	Stage_CurrentOp,
	Stage_Densify,
	Stage_Documents,
	Stage_Facet,
	Stage_Fill,
	Stage_GeoNear,
	Stage_GraphLookup,
	Stage_Group,
	Stage_IndexStats,
	Stage_Limit,
	Stage_ListLocalSessions,
	Stage_ListSearchIndexes,
	Stage_ListSessions,
	Stage_Lookup,
	Stage_Match,
	Stage_Merge,
	Stage_Out,
	Stage_Project,
	Stage_Redact,
	Stage_ReplaceRoot,
	Stage_ReplaceWith,
	Stage_Sample,
	Stage_Search,
	Stage_SearchMeta,
	Stage_Set,
	Stage_SetWindowFields,
	Stage_Skip,
	Stage_Sort,
	Stage_SortByCount,
	Stage_UnionWith,
	Stage_Unset,
	Stage_Unwind,
	Stage_VectorSearch,

	/* Start of pg_documentdb Custom or internal stages */
	Stage_InverseMatch = 100,
	Stage_LookupUnwind,
	Stage_SortGroup,
} Stage;


typedef enum JoinStageStatus
{
	JoinStageStatus_Unknown = 0,
	JoinStageStatus_NoJoinsOrUnions,
	JoinStageStatus_HasJoinsOrUnions,
} JoinStageStatus;


/*
 * Parsed representation of a $unwind stage spec. Populated by
 * TryParseUnwindStage.
 */
typedef struct UnwindArgs
{
	/* The path value from either the shorthand "$path" form or the
	 * { path: "$path", ... } document form. BSON_TYPE_EOD if absent. */
	bson_value_t pathValue;

	/* Value of the optional "includeArrayIndex" field, zero-length if
	 * absent. Validated to be a non-empty UTF8 with no '$' prefix. */
	StringView includeArrayIndex;

	/* Value of the optional "preserveNullAndEmptyArrays" field. */
	bool preserveNullAndEmptyArrays;

	/* True iff the input spec was a document (not the shorthand string). */
	bool hasOptions;
} UnwindArgs;


/*
 * Callback invoked by TryParseUnwindStage when a spec fails validation.
 */
typedef void (*UnwindParseErrorHandler)(int errCode, const char *errMessage,
										const char *errArg);


/*
 * Shared context during aggregation pipeline build phase.
 */
typedef struct AggregationPipelineBuildContext
{
	/* *********************************************************************
	 * Query Builder Fields
	 * *********************************************************************
	 */

	/* The current stage number (used for tagging stage identifiers) */
	int stageNum;

	/* Whether or not a subquery stage should be injected before the next stage */
	bool requiresSubQuery;

	/* If true, allows 1 project transform, then forces a subquery stage. */
	bool requiresSubQueryAfterProject;

	/* Whether the query should retain an expanded target list*/
	bool expandTargetList;

	/* the level of nested pipeline for stages that have nested pipelines ($facet/$lookup). */
	int nestedPipelineLevel;

	/* The current parameter count (Note: Increment this before use) */
	int currentParamCount;

	/* The number of nested levels (incremented by MigrateSubQuery) */
	int numNestedLevels;

	/*
	 * Whether or not to apply the optimization transformation on the stages
	 */
	bool optimizePipelineStages;

	/* Whether or not the query has join stages (to optimize subquery behavior) */
	JoinStageStatus joinStatus;

	/* *********************************************************************
	 * Namespace Fields
	 * *********************************************************************
	 */

	/* The namespace 'db.coll' associated with this query */
	const char *namespaceName;

	/* The current Mongo collection */
	MongoCollection *mongoCollection;

	/* The database name associated with this request */
	text *databaseNameDatum;

	/* The collection name associated with this request (if applicable) */
	StringView collectionNameView;

	/* Whether or not the aggregation query allows direct shard delegation
	 * This allows queries to go directly against a local shard *iff* it's available.
	 * This can be done for base streaming queries. TODO: Investigate whether or not
	 * this can be extended to other types of queries.
	 */
	bool allowShardBaseTable;

	/* *********************************************************************
	 * Cursor Fields
	 * *********************************************************************
	 */

	/* Whether or not the query requires a persisted cursor */
	bool requiresPersistentCursor;

	/* Whether or not the stage can handle a single batch cursor */
	bool isSingleRowResult;

	/* Whether or not the query requires a tailable cursor */
	bool requiresTailableCursor;

	/* Whether or not it's a point read query */
	bool isPointReadQuery;

	/* *********************************************************************
	 * Pipeline stage fields
	 * *********************************************************************
	 */

	/* The sort specification that precedes it (if available).
	 * If the stage changes the sort order, this is reset.
	 * BSON_TYPE_EOD if not available.
	 */
	bson_value_t sortSpec;

	/* Parent Stage Name For $redact scenarios */
	ParentStageName parentStageName;

	/* *********************************************************************
	 * Global Query Metadata fields
	 * *********************************************************************
	 */

	/*
	 * The variable spec expression that preceds it.
	 */
	Expr *variableSpec;

	/*
	 * String indicating a standard ICU collation. An example string is "und-u-ks-level1-kc-true".
	 * We parse the Mongo collation spec and covert it to an ICU standard collation string.
	 * This string uniquely identify collation-based string comparison logic by postgres.
	 * See: https://www.postgresql.org/docs/current/collation.html
	 */
	const char collationString[MAX_ICU_COLLATION_LENGTH];

	/* *********************************************************************
	 * Feature tracking fields
	 * *********************************************************************
	 */

	/* Tracks shard key filter type applied in HandleMatch. */
	enum
	{
		ShardKeyFilter_None = 0,
		ShardKeyFilter_Equality = 1,
		ShardKeyFilter_In = 2,
	} shardKeyFilterAppliedType;
} AggregationPipelineBuildContext;


typedef struct
{
	/* The bson value of the pipeline spec */
	bson_value_t stageValue;

	/* Definition of internal handlers */
	const AggregationStageDefinition *stageDefinition;
} AggregationStage;


/* Core Infra exports */
Query * MutateQueryWithPipeline(Query *query, List *aggregationStages,
								AggregationPipelineBuildContext *context);
Query * MigrateQueryToSubQuery(Query *parse, AggregationPipelineBuildContext *context);
Aggref * CreateMultiArgAggregate(Oid aggregateFunctionId, List *args, List *argTypes,
								 ParseState *parseState);
List * ExtractAggregationStages(const bson_value_t *pipelineValue,
								AggregationPipelineBuildContext *context);
Query * GenerateBaseTableQuery(text *databaseDatum, const StringView *collectionNameView,
							   pg_uuid_t *collectionUuid, const bson_value_t *indexHint,
							   AggregationPipelineBuildContext *context);
Query * GenerateBaseAgnosticQuery(text *databaseDatum,
								  AggregationPipelineBuildContext *context);
RangeTblEntry * MakeSubQueryRte(Query *subQuery, int stageNum, int pipelineDepth,
								const char *prefix, bool includeAllColumns);

bool CanInlineLookupPipeline(const bson_value_t *pipeline,
							 const StringView *lookupPath,
							 bool hasLet,
							 pgbson **inlinedPipeline,
							 pgbson **nonInlinedPipeline,
							 bool *pipelineIsValid);

void ParseCursorDocument(bson_iter_t *iterator, QueryData *queryData);
const char * CreateNamespaceName(text *databaseName,
								 const StringView *collectionName);

Query * HandleMatchWithIndexFilter(const bson_value_t *existingValue, Query *query,
								   AggregationPipelineBuildContext *context,
								   HTAB *requiredFilterPathNameHashSet);
Query * HandleMatch(const bson_value_t *existingValue, Query *query,
					AggregationPipelineBuildContext *context);
Query * HandleSimpleProjectionStage(const bson_value_t *existingValue, Query *query,
									AggregationPipelineBuildContext *context,
									const char *stageName, Oid functionOid,
									Oid (*functionOidWithLet)(void),
									Oid (*functionOidWithLetAndCollation)(void));
Query * HandleGroup(const bson_value_t *existingValue, Query *query,
					AggregationPipelineBuildContext *context);

/* Sub-Pipeline related aggregation stages */
Query * HandleFacet(const bson_value_t *existingValue, Query *query,
					AggregationPipelineBuildContext *context);

Query * HandleLookup(const bson_value_t *existingValue, Query *query,
					 AggregationPipelineBuildContext *context);

Query * HandleLookupUnwind(const bson_value_t *existingValue, Query *query,
						   AggregationPipelineBuildContext *context);

Query * HandleGraphLookup(const bson_value_t *existingValue, Query *query,
						  AggregationPipelineBuildContext *context);

Query * HandleDocumentsStage(const bson_value_t *existingValue, Query *query,
							 AggregationPipelineBuildContext *context);

Query * HandleUnionWith(const bson_value_t *existingValue, Query *query,
						AggregationPipelineBuildContext *context);

Query * HandleInternalInhibitOptimization(const bson_value_t *existingValue, Query *query,
										  AggregationPipelineBuildContext *context);

Query * HandleInverseMatch(const bson_value_t *existingValue, Query *query,
						   AggregationPipelineBuildContext *context);

/* Metadata based query stages */
Query * HandleCollStats(const bson_value_t *existingValue, Query *query,
						AggregationPipelineBuildContext *context);
Query * HandleIndexStats(const bson_value_t *existingValue, Query *query,
						 AggregationPipelineBuildContext *context);
Query * HandleCurrentOp(const bson_value_t *existingValue, Query *query,
						AggregationPipelineBuildContext *context);
Query * HandleListExtendedIndexes(const bson_value_t *existingValue, Query *query,
								  AggregationPipelineBuildContext *context);
Query * HandleChangeStream(const bson_value_t *existingValue, Query *query,
						   AggregationPipelineBuildContext *context);

bool CanInlineLookupStageLookup(const bson_value_t *lookupStage,
								const StringView *lookupPath,
								bool hasLet);
bool CanInlineLookupWithUnwind(const bson_value_t *lookUpStageValue,
							   const bson_value_t *unwindStageValue,
							   bool *isPreserveNullAndEmptyArrays);

bool TryParseUnwindStage(const bson_value_t *unwindStageValue, UnwindArgs *args,
						 UnwindParseErrorHandler onError);

/* vector search related aggregation stages */
Query * HandleSearch(const bson_value_t *existingValue, Query *query,
					 AggregationPipelineBuildContext *context);

/* output to collection related aggregation pipeline */
Query * HandleMerge(const bson_value_t *existingValue, Query *query,
					AggregationPipelineBuildContext *context);
Query * HandleOut(const bson_value_t *existingValue, Query *query,
				  AggregationPipelineBuildContext *context);

/* Native vector search related aggregation stages */
Query * HandleNativeVectorSearch(const bson_value_t *existingValue, Query *query,
								 AggregationPipelineBuildContext *context);

/* Metadata based query generators */
Query * GenerateConfigDatabaseQuery(AggregationPipelineBuildContext *context);

bool IsPartitionByFieldsOnShardKey(const pgbson *partitionByFields,
								   const MongoCollection *collection);

Expr * GenerateMultiExpressionRepathExpression(List *repathArgs,
											   bool overrideArrayInProjection);
Stage GetAggregationStageAtPosition(const List *aggregationStages, int position);

/* Helper methods */

/* GUC defined in feature_flag_configs.c; consumed by the inline helpers below. */
extern bool EnableAddShardKeyOnlyOnPrimaryKeyFilters;

/*
 * Helper function that creates a UNION ALL Set operation statement
 * that returns a single BSON field.
 */
inline static SetOperationStmt *
MakeBsonSetOpStatement(void)
{
	SetOperationStmt *setOpStatement = makeNode(SetOperationStmt);
	setOpStatement->all = true;
	setOpStatement->op = SETOP_UNION;
	setOpStatement->colCollations = list_make1_oid(InvalidOid);
	setOpStatement->colTypes = list_make1_oid(BsonTypeId());
	setOpStatement->colTypmods = list_make1_int(-1);
	return setOpStatement;
}


inline static bool
ShouldSkipShardKeyFilterOnBaseTable(AggregationPipelineBuildContext *context)
{
	if (!EnableAddShardKeyOnlyOnPrimaryKeyFilters)
	{
		return false;
	}

	return context->mongoCollection != NULL &&
		   context->mongoCollection->shardKey == NULL &&
		   context->mongoCollection->isSingleShardTable;
}


#endif
