"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.PipelinedQueryExecutionContext = void 0;
const ErrorResponse_js_1 = require("../request/ErrorResponse.js");
const OffsetLimitEndpointComponent_js_1 = require("./EndpointComponent/OffsetLimitEndpointComponent.js");
const OrderByEndpointComponent_js_1 = require("./EndpointComponent/OrderByEndpointComponent.js");
const OrderedDistinctEndpointComponent_js_1 = require("./EndpointComponent/OrderedDistinctEndpointComponent.js");
const UnorderedDistinctEndpointComponent_js_1 = require("./EndpointComponent/UnorderedDistinctEndpointComponent.js");
const GroupByEndpointComponent_js_1 = require("./EndpointComponent/GroupByEndpointComponent.js");
const orderByQueryExecutionContext_js_1 = require("./orderByQueryExecutionContext.js");
const parallelQueryExecutionContext_js_1 = require("./parallelQueryExecutionContext.js");
const GroupByValueEndpointComponent_js_1 = require("./EndpointComponent/GroupByValueEndpointComponent.js");
const NonStreamingOrderByDistinctEndpointComponent_js_1 = require("./EndpointComponent/NonStreamingOrderByDistinctEndpointComponent.js");
const NonStreamingOrderByEndpointComponent_js_1 = require("./EndpointComponent/NonStreamingOrderByEndpointComponent.js");
const QueryValidationHelper_js_1 = require("./QueryValidationHelper.js");
const ContinuationTokenParser_js_1 = require("./ContinuationTokenParser.js");
const LegacyFetchImplementation_js_1 = require("./LegacyFetchImplementation.js");
const QueryControlFetchImplementation_js_1 = require("./QueryControlFetchImplementation.js");
const constants_js_1 = require("../common/constants.js");
/** @hidden */
class PipelinedQueryExecutionContext {
    clientContext;
    collectionLink;
    query;
    options;
    partitionedQueryExecutionInfo;
    emitRawOrderByPayload;
    supportsContinuationTokens;
    fetchBuffer;
    endpoint;
    fetchImplementation;
    constructor(clientContext, collectionLink, query, options, partitionedQueryExecutionInfo, correlatedActivityId, emitRawOrderByPayload = false, supportsContinuationTokens = true) {
        this.clientContext = clientContext;
        this.collectionLink = collectionLink;
        this.query = query;
        this.options = options;
        this.partitionedQueryExecutionInfo = partitionedQueryExecutionInfo;
        this.emitRawOrderByPayload = emitRawOrderByPayload;
        this.supportsContinuationTokens = supportsContinuationTokens;
        // Validate that queryInfo is present in partitioned query execution info
        if (!partitionedQueryExecutionInfo.queryInfo) {
            throw new ErrorResponse_js_1.ErrorResponse("Query execution requires valid query plan information. " +
                "The partitioned query execution info is missing queryInfo. " +
                "This may indicate an invalid query or a problem with query planning.");
        }
        if (!this.options.maxItemCount) {
            this.options.maxItemCount = constants_js_1.QueryExecution.DEFAULT_PAGE_SIZE;
        }
        const pageSize = this.options.maxItemCount;
        // Extract query information and characteristics
        const analyzedQueryInfo = this.analyzeQueryInfo(partitionedQueryExecutionInfo.queryInfo);
        const { sortOrders, nonStreamingOrderBy, isOrderByQuery, isGroupByQuery, isUnorderedDistinctQuery, querySupportsTokens, } = analyzedQueryInfo;
        // Reject continuation token usage for unsupported query types
        if (!querySupportsTokens) {
            (0, QueryValidationHelper_js_1.rejectContinuationTokenForUnsupportedQueries)(this.options.continuationToken, [
                QueryValidationHelper_js_1.QueryTypes.nonStreamingOrderBy(nonStreamingOrderBy),
                QueryValidationHelper_js_1.QueryTypes.groupBy(isGroupByQuery),
                QueryValidationHelper_js_1.QueryTypes.unorderedDistinct(isUnorderedDistinctQuery),
            ]);
        }
        // Parse continuation token fields once for reuse in pipeline construction
        const queryContinuationFields = this.options.continuationToken
            ? (0, ContinuationTokenParser_js_1.parseContinuationTokenFields)(this.options.continuationToken)
            : undefined;
        // Pick between non-streaming vs streaming execution context
        this.endpoint = nonStreamingOrderBy
            ? this.createNonStreamingEndpoint(partitionedQueryExecutionInfo, sortOrders, correlatedActivityId, options)
            : this.createStreamingEndpoint(partitionedQueryExecutionInfo, sortOrders, correlatedActivityId, isGroupByQuery, queryContinuationFields);
        this.fetchBuffer = [];
        // Initialize the appropriate fetch implementation based on enableQueryControl
        if (this.options.enableQueryControl) {
            const querySupportsContinuationTokens = this.supportsContinuationTokens && querySupportsTokens;
            this.fetchImplementation = new QueryControlFetchImplementation_js_1.QueryControlFetchImplementation(this.endpoint, pageSize, this.collectionLink, this.options.continuationToken, isOrderByQuery, querySupportsContinuationTokens);
        }
        else {
            this.fetchImplementation = new LegacyFetchImplementation_js_1.LegacyFetchImplementation(this.endpoint, pageSize);
        }
    }
    hasMoreResults() {
        return this.fetchBuffer.length !== 0 || this.endpoint.hasMoreResults();
    }
    async fetchMore(diagnosticNode) {
        return this.fetchImplementation.fetchMore(diagnosticNode, this.fetchBuffer);
    }
    /**
     * Creates a non-streaming endpoint for vector search and similar queries that require buffering
     */
    createNonStreamingEndpoint(partitionedQueryExecutionInfo, sortOrders, correlatedActivityId, options) {
        const queryInfo = partitionedQueryExecutionInfo.queryInfo; // Safe to use ! after validation in constructor
        if (!options.allowUnboundedNonStreamingQueries) {
            this.checkQueryConstraints(queryInfo);
        }
        const vectorSearchBufferSize = this.calculateVectorSearchBufferSize(queryInfo, options);
        this.validateVectorSearchBufferSize(vectorSearchBufferSize, options);
        const baseContext = new parallelQueryExecutionContext_js_1.ParallelQueryExecutionContext(this.clientContext, this.collectionLink, this.query, this.options, this.partitionedQueryExecutionInfo, correlatedActivityId);
        return this.wrapWithNonStreamingComponent(baseContext, queryInfo, sortOrders, vectorSearchBufferSize);
    }
    /**
     * Creates a streaming endpoint with proper pipeline components
     */
    createStreamingEndpoint(partitionedQueryExecutionInfo, sortOrders, correlatedActivityId, isGroupByQuery, queryContinuationFields) {
        const queryInfo = partitionedQueryExecutionInfo.queryInfo; // Safe to use ! after validation in constructor
        // Create base execution context
        let endpoint = this.createBaseExecutionContext(partitionedQueryExecutionInfo, sortOrders, correlatedActivityId);
        // Apply pipeline transformations
        endpoint = this.applyGroupByComponents(endpoint, queryInfo, isGroupByQuery);
        endpoint = this.applyDistinctComponents(endpoint, queryInfo, queryContinuationFields);
        endpoint = this.applyLimitComponents(endpoint, queryInfo, queryContinuationFields);
        return endpoint;
    }
    /**
     * Creates the base execution context (OrderBy or Parallel)
     */
    createBaseExecutionContext(partitionedQueryExecutionInfo, sortOrders, correlatedActivityId) {
        if (Array.isArray(sortOrders) && sortOrders.length > 0) {
            // OrderBy queries need special wrapping for payload structure
            return new OrderByEndpointComponent_js_1.OrderByEndpointComponent(new orderByQueryExecutionContext_js_1.OrderByQueryExecutionContext(this.clientContext, this.collectionLink, this.query, this.options, partitionedQueryExecutionInfo, correlatedActivityId), this.emitRawOrderByPayload);
        }
        // Parallel queries
        return new parallelQueryExecutionContext_js_1.ParallelQueryExecutionContext(this.clientContext, this.collectionLink, this.query, this.options, partitionedQueryExecutionInfo, correlatedActivityId);
    }
    /**
     * Wraps base context with appropriate non-streaming component
     */
    wrapWithNonStreamingComponent(baseContext, queryInfo, sortOrders, vectorSearchBufferSize) {
        const distinctType = queryInfo.distinctType;
        if (distinctType === "None") {
            return new NonStreamingOrderByEndpointComponent_js_1.NonStreamingOrderByEndpointComponent(baseContext, sortOrders, vectorSearchBufferSize, queryInfo.offset, this.emitRawOrderByPayload);
        }
        return new NonStreamingOrderByDistinctEndpointComponent_js_1.NonStreamingOrderByDistinctEndpointComponent(baseContext, queryInfo, vectorSearchBufferSize, this.emitRawOrderByPayload);
    }
    /**
     * Applies GROUP BY components to the pipeline if needed
     */
    applyGroupByComponents(endpoint, queryInfo, isGroupByQuery) {
        if (!isGroupByQuery) {
            return endpoint;
        }
        return queryInfo.hasSelectValue
            ? new GroupByValueEndpointComponent_js_1.GroupByValueEndpointComponent(endpoint, queryInfo)
            : new GroupByEndpointComponent_js_1.GroupByEndpointComponent(endpoint, queryInfo);
    }
    /**
     * Applies DISTINCT components to the pipeline if needed
     */
    applyDistinctComponents(endpoint, queryInfo, queryContinuationFields) {
        const distinctType = queryInfo.distinctType;
        if (distinctType === "Ordered") {
            const lastHash = queryContinuationFields?.hashedLastResult;
            return new OrderedDistinctEndpointComponent_js_1.OrderedDistinctEndpointComponent(endpoint, lastHash);
        }
        if (distinctType === "Unordered") {
            return new UnorderedDistinctEndpointComponent_js_1.UnorderedDistinctEndpointComponent(endpoint);
        }
        return endpoint;
    }
    /**
     * Applies TOP and OFFSET+LIMIT components to the pipeline if needed
     */
    applyLimitComponents(endpoint, queryInfo, queryContinuationFields) {
        // Apply TOP component (TOP N is effectively OFFSET 0 LIMIT N)
        let top = queryInfo.top;
        if (typeof top === "number") {
            if (queryContinuationFields?.limit !== undefined) {
                top = queryContinuationFields.limit;
            }
            endpoint = new OffsetLimitEndpointComponent_js_1.OffsetLimitEndpointComponent(endpoint, 0, top);
        }
        // Apply OFFSET+LIMIT component
        let limit = queryInfo.limit;
        let offset = queryInfo.offset;
        if (queryContinuationFields) {
            if (queryContinuationFields.limit !== undefined) {
                limit = queryContinuationFields.limit;
            }
            if (queryContinuationFields.offset !== undefined) {
                offset = queryContinuationFields.offset;
            }
        }
        if (typeof limit === "number" && typeof offset === "number") {
            endpoint = new OffsetLimitEndpointComponent_js_1.OffsetLimitEndpointComponent(endpoint, offset, limit);
        }
        return endpoint;
    }
    /**
     * Validates vector search buffer size constraints
     */
    validateVectorSearchBufferSize(vectorSearchBufferSize, options) {
        const maxBufferSize = options["vectorSearchBufferSize"]
            ? options["vectorSearchBufferSize"]
            : constants_js_1.QueryExecution.DEFAULT_MAX_VECTOR_SEARCH_BUFFER_SIZE;
        if (vectorSearchBufferSize > maxBufferSize) {
            throw new ErrorResponse_js_1.ErrorResponse(`Executing a vector search query with TOP or OFFSET + LIMIT value ${vectorSearchBufferSize} larger than the vectorSearchBufferSize ${maxBufferSize} ` +
                `is not allowed`);
        }
    }
    calculateVectorSearchBufferSize(queryInfo, options) {
        if (queryInfo.top === 0 || queryInfo.limit === 0)
            return 0;
        return queryInfo.top
            ? queryInfo.top
            : queryInfo.limit
                ? queryInfo.offset + queryInfo.limit
                : options["vectorSearchBufferSize"] && options["vectorSearchBufferSize"] > 0
                    ? options["vectorSearchBufferSize"]
                    : constants_js_1.QueryExecution.DEFAULT_MAX_VECTOR_SEARCH_BUFFER_SIZE;
    }
    checkQueryConstraints(queryInfo) {
        const hasTop = queryInfo.top || queryInfo.top === 0;
        const hasLimit = queryInfo.limit || queryInfo.limit === 0;
        if (!hasTop && !hasLimit) {
            throw new ErrorResponse_js_1.ErrorResponse("Executing a non-streaming search query without TOP or LIMIT can consume a large number of RUs " +
                "very fast and have long runtimes. Please ensure you are using one of the above two filters " +
                "with your vector search query.");
        }
        return;
    }
    /**
     * Analyzes query information and extracts key characteristics for query execution planning
     */
    analyzeQueryInfo(queryInfo) {
        const sortOrders = queryInfo.orderBy;
        const nonStreamingOrderBy = queryInfo.hasNonStreamingOrderBy;
        const isOrderByQuery = Array.isArray(sortOrders) && sortOrders.length > 0;
        // Check if this is a GROUP BY query
        const isGroupByQuery = Object.keys(queryInfo.groupByAliasToAggregateType || {}).length > 0 ||
            (queryInfo.aggregates?.length || 0) > 0 ||
            (queryInfo.groupByExpressions?.length || 0) > 0;
        // Check if this is an unordered DISTINCT query
        const isUnorderedDistinctQuery = queryInfo.distinctType === "Unordered";
        // Determine if this query type supports continuation tokens
        const querySupportsTokens = !isUnorderedDistinctQuery && !isGroupByQuery && !nonStreamingOrderBy;
        return {
            sortOrders,
            nonStreamingOrderBy,
            isOrderByQuery,
            isGroupByQuery,
            isUnorderedDistinctQuery,
            querySupportsTokens,
        };
    }
}
exports.PipelinedQueryExecutionContext = PipelinedQueryExecutionContext;
//# sourceMappingURL=pipelinedQueryExecutionContext.js.map