// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
import PriorityQueue from "priorityqueuejs";
import semaphore from "semaphore";
import { StatusCodes, SubStatusCodes } from "../common/statusCodes.js";
import { ErrorResponse } from "../request/ErrorResponse.js";
import { QueryRange } from "../routing/QueryRange.js";
import { SmartRoutingMapProvider } from "../routing/smartRoutingMapProvider.js";
import { DocumentProducer } from "./documentProducer.js";
import { getInitialHeader, mergeHeaders } from "./headerUtils.js";
import { RidSkipCountFilter } from "./queryFilteringStrategy/RidSkipCountFilter.js";
import { DiagnosticNodeInternal, DiagnosticNodeType, } from "../diagnostics/DiagnosticNodeInternal.js";
import { createParallelQueryResult } from "./parallelQueryResult.js";
/** @hidden */
export var ParallelQueryExecutionContextBaseStates;
(function (ParallelQueryExecutionContextBaseStates) {
    ParallelQueryExecutionContextBaseStates["started"] = "started";
    ParallelQueryExecutionContextBaseStates["inProgress"] = "inProgress";
    ParallelQueryExecutionContextBaseStates["ended"] = "ended";
})(ParallelQueryExecutionContextBaseStates || (ParallelQueryExecutionContextBaseStates = {}));
/** @hidden */
export class ParallelQueryExecutionContextBase {
    clientContext;
    collectionLink;
    query;
    options;
    partitionedQueryExecutionInfo;
    correlatedActivityId;
    rangeManager;
    queryProcessingStrategy;
    documentProducerComparator;
    err;
    state;
    static STATES = ParallelQueryExecutionContextBaseStates;
    routingProvider;
    requestContinuation;
    respHeaders;
    unfilledDocumentProducersQueue;
    bufferedDocumentProducersQueue;
    // TODO: update type of buffer from any --> generic can be used here
    buffer;
    partitionDataPatchMap = new Map();
    patchCounter = 0;
    updatedContinuationRanges = new Map();
    sem;
    diagnosticNodeWrapper;
    /**
     * Provides the ParallelQueryExecutionContextBase.
     * This is the base class that ParallelQueryExecutionContext and OrderByQueryExecutionContext will derive from.
     *
     * When handling a parallelized query, it instantiates one instance of
     * DocumentProcuder per target partition key range and aggregates the result of each.
     *
     * @param clientContext - The service endpoint to use to create the client.
     * @param collectionLink - The Collection Link
     * @param options - Represents the feed options.
     * @param partitionedQueryExecutionInfo - PartitionedQueryExecutionInfo
     * @hidden
     */
    constructor(clientContext, collectionLink, query, options, partitionedQueryExecutionInfo, correlatedActivityId, rangeManager, queryProcessingStrategy, documentProducerComparator) {
        this.clientContext = clientContext;
        this.collectionLink = collectionLink;
        this.query = query;
        this.options = options;
        this.partitionedQueryExecutionInfo = partitionedQueryExecutionInfo;
        this.correlatedActivityId = correlatedActivityId;
        this.rangeManager = rangeManager;
        this.queryProcessingStrategy = queryProcessingStrategy;
        this.documentProducerComparator = documentProducerComparator;
        this.clientContext = clientContext;
        this.collectionLink = collectionLink;
        this.query = query;
        this.options = options;
        this.partitionedQueryExecutionInfo = partitionedQueryExecutionInfo;
        this.correlatedActivityId = correlatedActivityId;
        this.diagnosticNodeWrapper = {
            consumed: false,
            diagnosticNode: new DiagnosticNodeInternal(clientContext.diagnosticLevel, DiagnosticNodeType.PARALLEL_QUERY_NODE, null),
        };
        this.diagnosticNodeWrapper.diagnosticNode.addData({ stateful: true });
        this.err = undefined;
        this.state = ParallelQueryExecutionContextBase.STATES.started;
        this.routingProvider = new SmartRoutingMapProvider(this.clientContext);
        this.buffer = [];
        this.requestContinuation = options
            ? options.continuationToken || options.continuation
            : undefined;
        // Validate continuation token usage immediately
        if (this.requestContinuation && !this.options.enableQueryControl) {
            throw new Error("Continuation tokens are supported when enableQueryControl is set true in FeedOptions");
        }
        // response headers of undergoing operation
        this.respHeaders = getInitialHeader();
        // Make priority queue for documentProducers
        this.unfilledDocumentProducersQueue = new PriorityQueue((a, b) => this.compareDocumentProducersByRange(a, b));
        this.bufferedDocumentProducersQueue = new PriorityQueue((a, b) => this.documentProducerComparator(b, a));
        // Creating the documentProducers
        this.sem = semaphore(1);
        this.sem.take(() => this._initializeDocumentProducers());
    }
    /**
     * Determine if there are still remaining resources to processs based on the value of the continuation
     * token or the elements remaining on the current batch in the QueryIterator.
     * @returns true if there is other elements to process in the ParallelQueryExecutionContextBase.
     */
    hasMoreResults() {
        return (!this.err &&
            (this.buffer.length > 0 || this.state !== ParallelQueryExecutionContextBase.STATES.ended));
    }
    /**
     * Fetches more results from the query execution context.
     * @param diagnosticNode - Optional diagnostic node for tracing.
     * @returns A promise that resolves to the fetched results.
     * @hidden
     */
    async fetchMore(diagnosticNode) {
        await this.bufferDocumentProducers(diagnosticNode);
        await this.fillBufferFromBufferQueue();
        return this.drainBufferedItems();
    }
    /**
     * Processes buffered document producers
     * @returns A promise that resolves when processing is complete.
     */
    async processBufferedDocumentProducers() {
        while (this.hasBufferedProducers() &&
            this.shouldProcessBufferedProducers(this.isUnfilledQueueEmpty())) {
            const producer = this.getNextBufferedProducer();
            if (!producer)
                break;
            await this.processDocumentProducer(producer);
        }
    }
    /**
     * Processes a single document producer using template method pattern.
     * Common structure with query-specific processing delegated to subclasses.
     */
    async processDocumentProducer(producer) {
        const response = await this.fetchFromProducer(producer);
        this._mergeWithActiveResponseHeaders(response.headers);
        if (response.result) {
            this.addToBuffer(response.result);
            this.handlePartitionMapping(producer, response.result);
        }
        // Handle producer lifecycle
        if (producer.peakNextItem() !== undefined) {
            this.requeueProducer(producer);
        }
        else if (producer.hasMoreResults()) {
            this.moveToUnfilledQueue(producer);
        }
    }
    /**
     * Handles partition mapping updates - implemented in base class using template method pattern.
     * Child classes provide query-specific parameters through abstract methods.
     */
    handlePartitionMapping(producer, result) {
        const itemCount = result?.length || 0;
        const continuationToken = this.getContinuationToken(producer);
        const mapping = {
            itemCount,
            partitionKeyRange: producer.targetPartitionKeyRange,
            continuationToken,
        };
        this.updatePartitionMapping(mapping);
    }
    /**
     * Gets the continuation token to use - implemented by subclasses.
     */
    getContinuationToken(producer) {
        const hasMoreBufferedItems = producer.peakNextItem() !== undefined;
        return hasMoreBufferedItems ? producer.previousContinuationToken : producer.continuationToken;
    }
    /**
     * Updates partition mapping - creates new entry or merges with existing for ORDER BY queries.
     */
    updatePartitionMapping(mapping) {
        const currentPatch = this.partitionDataPatchMap.get(this.patchCounter.toString());
        const isSamePartition = currentPatch?.partitionKeyRange?.id === mapping.partitionKeyRange.id;
        if (isSamePartition && currentPatch) {
            currentPatch.itemCount += mapping.itemCount;
            currentPatch.continuationToken = mapping.continuationToken;
            return;
        }
        // Create new partition mapping entry
        this.partitionDataPatchMap.set((++this.patchCounter).toString(), mapping);
    }
    /**
     * Checks if the unfilled queue is empty (used by ORDER BY for processing control).
     */
    isUnfilledQueueEmpty() {
        return this.unfilledDocumentProducersQueue.size() === 0;
    }
    /**
     * Initializes document producers and fills the priority queue.
     * Handles both continuation token and fresh query scenarios.
     */
    async _initializeDocumentProducers() {
        try {
            const targetPartitionRanges = await this._onTargetPartitionRanges();
            const documentProducers = this.requestContinuation
                ? await this._createDocumentProducersFromContinuation(targetPartitionRanges)
                : this._createDocumentProducersFromFresh(targetPartitionRanges);
            // Fill up our priority queue with documentProducers
            this._enqueueDocumentProducers(documentProducers);
            this.sem.leave();
        }
        catch (err) {
            this.err = err;
            this.sem.leave();
        }
    }
    /**
     * Creates document producers from continuation token scenario.
     */
    async _createDocumentProducersFromContinuation(targetPartitionRanges) {
        // Parse continuation token to get range mappings and check for split/merge scenarios
        const parsedToken = this._parseContinuationToken(this.requestContinuation);
        const continuationRanges = await this._handlePartitionRangeChanges(parsedToken);
        // Use strategy to create additional query info from parsed token
        const additionalQueryInfo = this.queryProcessingStrategy.createAdditionalQueryInfo(parsedToken);
        const filterResult = this.rangeManager.filterPartitionRanges(targetPartitionRanges, continuationRanges, additionalQueryInfo);
        // Extract ranges and tokens from the combined result
        const rangeTokenPairs = filterResult.rangeTokenPairs;
        // Use strategy to create filter context for continuation token processing
        const filterContext = this.queryProcessingStrategy.createFilterContext(parsedToken);
        return rangeTokenPairs.map((rangeTokenPair) => this._createDocumentProducerFromRangeTokenPair(rangeTokenPair, continuationRanges, filterContext));
    }
    /**
     * Creates document producers from fresh query scenario (no continuation token).
     */
    _createDocumentProducersFromFresh(targetPartitionRanges) {
        return targetPartitionRanges.map((partitionTargetRange) => this._createTargetPartitionQueryExecutionContext(partitionTargetRange, undefined));
    }
    /**
     * Creates a document producer from a range token pair (continuation token scenario).
     */
    _createDocumentProducerFromRangeTokenPair(rangeTokenPair, continuationRanges, filterContext) {
        const partitionTargetRange = rangeTokenPair.range;
        const continuationToken = rangeTokenPair.continuationToken;
        const filterCondition = rangeTokenPair.filteringCondition || undefined;
        // Find EPK ranges for this partition range from processed continuation response
        const matchingContinuationRange = continuationRanges.find((cr) => cr.range.id === partitionTargetRange.id);
        const startEpk = matchingContinuationRange?.epkMin;
        const endEpk = matchingContinuationRange?.epkMax;
        // Use strategy to determine partition-specific filter context
        const targetPartitionId = continuationRanges.length > 0 && continuationRanges[continuationRanges.length - 1].range
            ? continuationRanges[continuationRanges.length - 1].range.id
            : undefined;
        const partitionFilterContext = this.queryProcessingStrategy.getPartitionFilterContext(filterContext, targetPartitionId, partitionTargetRange.id);
        return this._createTargetPartitionQueryExecutionContext(partitionTargetRange, continuationToken, startEpk, endEpk, !!(startEpk && endEpk), // populateEpkRangeHeaders - true if both EPK values are present
        filterCondition, partitionFilterContext);
    }
    /**
     * Enqueues document producers into the unfilled queue.
     */
    _enqueueDocumentProducers(documentProducers) {
        documentProducers.forEach((documentProducer) => {
            try {
                this.unfilledDocumentProducersQueue.enq(documentProducer);
            }
            catch (e) {
                this.err = e;
            }
        });
    }
    /**
     * Checks if there are buffered document producers ready for processing.
     * Encapsulates queue size checking.
     */
    hasBufferedProducers() {
        return this.bufferedDocumentProducersQueue.size() > 0;
    }
    /**
     * Gets the next buffered document producer for processing.
     * Encapsulates queue dequeuing logic.
     */
    getNextBufferedProducer() {
        if (this.bufferedDocumentProducersQueue.size() > 0) {
            return this.bufferedDocumentProducersQueue.deq();
        }
        return undefined;
    }
    /**
     * Adds items to the result buffer. Handles both single items and arrays.
     */
    addToBuffer(items) {
        if (Array.isArray(items)) {
            if (items.length > 0) {
                this.buffer.push(...items);
            }
        }
        else if (items) {
            this.buffer.push(items);
        }
    }
    /**
     * Moves a producer to the unfilled queue for later processing.
     */
    moveToUnfilledQueue(producer) {
        this.unfilledDocumentProducersQueue.enq(producer);
    }
    /**
     * Re-queues a producer to the buffered queue for further processing.
     */
    requeueProducer(producer) {
        this.bufferedDocumentProducersQueue.enq(producer);
    }
    /**
     * Compares two document producers based on their partition key ranges and EPK values.
     * Primary comparison: minInclusive values for left-to-right range traversal
     * Secondary comparison: EPK ranges when minInclusive values are identical
     * @param a - First document producer
     * @param b - Second document producer
     * @returns Comparison result for priority queue ordering
     * @hidden
     */
    compareDocumentProducersByRange(a, b) {
        const aMinInclusive = a.targetPartitionKeyRange.minInclusive;
        const bMinInclusive = b.targetPartitionKeyRange.minInclusive;
        const minInclusiveComparison = bMinInclusive.localeCompare(aMinInclusive);
        // If minInclusive values are the same, check minEPK ranges if they exist
        if (minInclusiveComparison === 0) {
            const aMinEpk = a.startEpk;
            const bMinEpk = b.startEpk;
            if (aMinEpk && bMinEpk) {
                return bMinEpk.localeCompare(aMinEpk);
            }
        }
        return minInclusiveComparison;
    }
    /**
     * Detects partition splits/merges by analyzing parsed continuation token ranges and comparing with current topology
     * @param parsed - The continuation token containing range mappings to analyze
     * @returns Array of processed ranges with EPK info
     */
    async _handlePartitionRangeChanges(parsed) {
        const processedRanges = [];
        // Extract range mappings from the already parsed token
        const rangeMappings = parsed.rangeMappings;
        if (!rangeMappings || rangeMappings.length === 0) {
            return [];
        }
        // Check each range mapping for potential splits/merges
        for (const rangeWithToken of rangeMappings) {
            // Create a new QueryRange instance from the simplified range data
            const range = rangeWithToken.queryRange;
            const queryRange = new QueryRange(range.min, range.max, true, // isMinInclusive - assumption: always true
            false);
            const rangeMin = queryRange.min;
            const rangeMax = queryRange.max;
            // Get current overlapping ranges for this continuation token range
            const overlappingRanges = await this.routingProvider.getOverlappingRanges(this.collectionLink, [queryRange], this.getDiagnosticNode());
            // Detect split/merge scenario based on the number of overlapping ranges
            if (overlappingRanges.length === 0) {
                continue;
            }
            else if (overlappingRanges.length === 1) {
                // Check if it's the same range (no change) or a merge scenario
                const currentRange = overlappingRanges[0];
                if (currentRange.minInclusive !== rangeMin || currentRange.maxExclusive !== rangeMax) {
                    // Merge scenario - include EPK ranges from original continuation token range
                    await this._handleContinuationTokenMerge(rangeWithToken, currentRange);
                    processedRanges.push({
                        range: currentRange,
                        continuationToken: rangeWithToken.continuationToken,
                        epkMin: rangeMin, // Original range min becomes EPK min
                        epkMax: rangeMax, // Original range max becomes EPK max
                    });
                }
                else {
                    // Same range - no merge, no EPK ranges needed
                    processedRanges.push({
                        range: currentRange,
                        continuationToken: rangeWithToken.continuationToken,
                    });
                }
            }
            else {
                // Split scenario - one range from continuation token now maps to multiple ranges
                await this._handleContinuationTokenSplit(rangeWithToken, overlappingRanges);
                // Add all overlapping ranges with the same continuation token to processed ranges
                overlappingRanges.forEach((rangeValue) => {
                    processedRanges.push({
                        range: rangeValue,
                        continuationToken: rangeWithToken.continuationToken,
                    });
                });
            }
        }
        return processedRanges;
    }
    /**
     * Parses the continuation token based on query type
     * @param continuationToken - The continuation token string to parse
     * @returns Parsed continuation token object (ORDER BY or Parallel query token)
     * @throws ErrorResponse when continuation token is malformed or cannot be parsed
     */
    _parseContinuationToken(continuationToken) {
        try {
            return this.queryProcessingStrategy.parseContinuationToken(continuationToken);
        }
        catch (e) {
            throw new ErrorResponse(`Invalid continuation token format. Expected token with rangeMappings property. ` +
                `Ensure the continuation token was generated by a compatible query and has not been modified.`);
        }
    }
    /**
     * Handles partition merge scenario for continuation token ranges
     */
    async _handleContinuationTokenMerge(rangeWithToken, _newMergedRange) {
        const rangeKey = `${rangeWithToken.queryRange.min}-${rangeWithToken.queryRange.max}`;
        this.updatedContinuationRanges.set(rangeKey, {
            oldRange: {
                min: rangeWithToken.queryRange.min,
                max: rangeWithToken.queryRange.max,
                isMinInclusive: true, // Assumption: min is always inclusive
                isMaxInclusive: false, // Assumption: max is always exclusive
            },
            newRanges: [
                {
                    min: rangeWithToken.queryRange.min,
                    max: rangeWithToken.queryRange.max,
                    isMinInclusive: true, // Assumption: min is always inclusive
                    isMaxInclusive: false, // Assumption: max is always exclusive
                },
            ],
            continuationToken: rangeWithToken.continuationToken,
        });
    }
    /**
     * Handles partition split scenario for continuation token ranges
     */
    async _handleContinuationTokenSplit(rangeWithToken, overlappingRanges) {
        const rangeKey = `${rangeWithToken.queryRange.min}-${rangeWithToken.queryRange.max}`;
        this.updatedContinuationRanges.set(rangeKey, {
            oldRange: {
                min: rangeWithToken.queryRange.min,
                max: rangeWithToken.queryRange.max,
                isMinInclusive: true, // Assumption: min is always inclusive
                isMaxInclusive: false, // Assumption: max is always exclusive
            },
            newRanges: overlappingRanges.map((range) => ({
                min: range.minInclusive,
                max: range.maxExclusive,
                isMinInclusive: true,
                isMaxInclusive: false,
            })),
            continuationToken: rangeWithToken.continuationToken,
        });
    }
    /**
     * Handles partition merge scenario for continuation token ranges
     */
    _mergeWithActiveResponseHeaders(headers) {
        mergeHeaders(this.respHeaders, headers);
    }
    _getAndResetActiveResponseHeaders() {
        const ret = this.respHeaders;
        this.respHeaders = getInitialHeader();
        return ret;
    }
    getDiagnosticNode() {
        return this.diagnosticNodeWrapper.diagnosticNode;
    }
    async _onTargetPartitionRanges() {
        // invokes the callback when the target partition ranges are ready
        const parsedRanges = this.partitionedQueryExecutionInfo.queryRanges;
        const queryRanges = parsedRanges.map((item) => QueryRange.parseFromDict(item));
        return this.routingProvider.getOverlappingRanges(this.collectionLink, queryRanges, this.getDiagnosticNode());
    }
    /**
     * Gets the replacement ranges for a partitionkeyrange that has been split
     */
    async _getReplacementPartitionKeyRanges(documentProducer, diagnosticNode) {
        const partitionKeyRange = documentProducer.targetPartitionKeyRange;
        // Download the new routing map
        this.routingProvider = new SmartRoutingMapProvider(this.clientContext);
        // Get the queryRange that relates to this partitionKeyRange
        const queryRange = QueryRange.parsePartitionKeyRange(partitionKeyRange);
        return this.routingProvider.getOverlappingRanges(this.collectionLink, [queryRange], diagnosticNode);
    }
    async _enqueueReplacementDocumentProducers(error, diagnosticNode, documentProducer) {
        // Get the replacement ranges
        const replacementPartitionKeyRanges = await this._getReplacementPartitionKeyRanges(documentProducer, diagnosticNode);
        if (replacementPartitionKeyRanges.length === 0) {
            throw error;
        }
        if (this.requestContinuation) {
            // Update composite continuation token to handle partition split
            this._updateContinuationTokenOnPartitionChange(documentProducer, replacementPartitionKeyRanges);
        }
        if (replacementPartitionKeyRanges.length === 1) {
            // Partition is gone due to Merge
            // Create the replacement documentProducer with populateEpkRangeHeaders Flag set to true to set startEpk and endEpk headers
            const replacementDocumentProducer = this._createTargetPartitionQueryExecutionContext(replacementPartitionKeyRanges[0], documentProducer.continuationToken, documentProducer.startEpk, documentProducer.endEpk, true);
            this.unfilledDocumentProducersQueue.enq(replacementDocumentProducer);
        }
        else {
            // Create the replacement documentProducers
            const replacementDocumentProducers = [];
            replacementPartitionKeyRanges.forEach((partitionKeyRange) => {
                const queryRange = QueryRange.parsePartitionKeyRange(partitionKeyRange);
                // Create replacment document producers with the parent's continuationToken
                const replacementDocumentProducer = this._createTargetPartitionQueryExecutionContext(partitionKeyRange, documentProducer.continuationToken, queryRange.min, queryRange.max, false);
                replacementDocumentProducers.push(replacementDocumentProducer);
            });
            // add document producers to the queue
            replacementDocumentProducers.forEach((replacementDocumentProducer) => {
                if (replacementDocumentProducer.hasMoreResults()) {
                    this.unfilledDocumentProducersQueue.enq(replacementDocumentProducer);
                }
            });
        }
    }
    _updateContinuationTokenOnPartitionChange(originalDocumentProducer, replacementPartitionKeyRanges) {
        const rangeWithToken = this._createQueryRangeWithContinuationToken(originalDocumentProducer);
        if (replacementPartitionKeyRanges.length === 1) {
            this._handleContinuationTokenMerge(rangeWithToken, replacementPartitionKeyRanges[0]);
        }
        else {
            this._handleContinuationTokenSplit(rangeWithToken, replacementPartitionKeyRanges);
        }
    }
    /**
     * Creates a QueryRangeWithContinuationToken object from a DocumentProducer.
     * Uses the DocumentProducer's target partition key range and continuation token.
     * @param documentProducer - The DocumentProducer to convert
     * @returns QueryRangeWithContinuationToken object for token operations
     */
    _createQueryRangeWithContinuationToken(documentProducer) {
        const partitionRange = documentProducer.targetPartitionKeyRange;
        // Create a simplified QueryRange using the partition key range boundaries
        const simplifiedQueryRange = {
            min: documentProducer.startEpk || partitionRange.minInclusive,
            max: documentProducer.endEpk || partitionRange.maxExclusive,
        };
        return {
            queryRange: simplifiedQueryRange,
            continuationToken: documentProducer.continuationToken,
        };
    }
    static _needPartitionKeyRangeCacheRefresh(error) {
        // TODO: any error
        return (error.code === StatusCodes.Gone &&
            "substatus" in error &&
            error["substatus"] === SubStatusCodes.PartitionKeyRangeGone);
    }
    /**
     * Replaces the format placeholder in the rewritten query with the provided filter condition.
     * Handles both string queries and SqlQuerySpec objects.
     */
    _replaceFormatPlaceholder(rewrittenQuery, formatPlaceHolder, filterCondition) {
        const replacement = filterCondition ?? "true";
        // If rewrittenQuery has a query property, it's a SqlQuerySpec object
        if (typeof rewrittenQuery === "object" && rewrittenQuery !== null && rewrittenQuery.query) {
            return rewrittenQuery.query.replace(formatPlaceHolder, replacement);
        }
        // Otherwise, it's a string
        return rewrittenQuery.replace(formatPlaceHolder, replacement);
    }
    /**
     * Creates target partition range Query Execution Context
     */
    _createTargetPartitionQueryExecutionContext(partitionKeyTargetRange, continuationToken, startEpk, endEpk, populateEpkRangeHeaders, filterCondition, filterContext) {
        const rewrittenQuery = this.partitionedQueryExecutionInfo.queryInfo?.rewrittenQuery;
        let sqlQuerySpec;
        const query = this.query;
        if (typeof query === "string") {
            sqlQuerySpec = { query };
        }
        else {
            sqlQuerySpec = query;
        }
        const formatPlaceHolder = "{documentdb-formattableorderbyquery-filter}";
        if (rewrittenQuery) {
            sqlQuerySpec = JSON.parse(JSON.stringify(sqlQuerySpec));
            const replacedQuery = this._replaceFormatPlaceholder(rewrittenQuery, formatPlaceHolder, filterCondition);
            sqlQuerySpec["query"] = replacedQuery;
        }
        const options = { ...this.options };
        options.continuationToken = continuationToken;
        let filter;
        if (filterContext) {
            filter = new RidSkipCountFilter(filterContext);
        }
        return new DocumentProducer(this.clientContext, this.collectionLink, sqlQuerySpec, partitionKeyTargetRange, options, this.correlatedActivityId, startEpk, endEpk, populateEpkRangeHeaders, filter);
    }
    async drainBufferedItems() {
        return new Promise((resolve, reject) => {
            this.sem.take(() => {
                if (this.err) {
                    // if there is a prior error return error
                    this.sem.leave();
                    this.err.headers = this._getAndResetActiveResponseHeaders();
                    reject(this.err);
                    return;
                }
                // return undefined if there is no more results
                if (this.buffer.length === 0) {
                    this.sem.leave();
                    const partitionDataPatchMap = this.partitionDataPatchMap;
                    this.partitionDataPatchMap = new Map();
                    this.patchCounter = 0;
                    // Get and reset updated continuation ranges
                    const updatedContinuationRanges = Object.fromEntries(this.updatedContinuationRanges);
                    this.updatedContinuationRanges.clear();
                    const result = createParallelQueryResult([], partitionDataPatchMap, updatedContinuationRanges, undefined);
                    return resolve({
                        result: this.state === ParallelQueryExecutionContextBase.STATES.ended ? undefined : result,
                        headers: this._getAndResetActiveResponseHeaders(),
                    });
                }
                // draing the entire buffer object and return that in result of return object
                const bufferedResults = this.buffer;
                this.buffer = [];
                // reset the patchToRangeMapping
                const partitionDataPatchMap = this.partitionDataPatchMap;
                this.partitionDataPatchMap = new Map();
                this.patchCounter = 0;
                // Get and reset updated continuation ranges
                const updatedContinuationRanges = Object.fromEntries(this.updatedContinuationRanges);
                this.updatedContinuationRanges.clear();
                // release the lock before returning
                this.sem.leave();
                const result = createParallelQueryResult(bufferedResults, partitionDataPatchMap, updatedContinuationRanges, undefined);
                return resolve({
                    result,
                    headers: this._getAndResetActiveResponseHeaders(),
                });
            });
        });
    }
    /**
     * Buffers document producers based on the maximum degree of parallelism.
     * Moves document producers from the unfilled queue to the buffered queue.
     * @param diagnosticNode - The diagnostic node for logging and tracing.
     * @returns A promise that resolves when buffering is complete.
     */
    async bufferDocumentProducers(diagnosticNode) {
        return new Promise((resolve, reject) => {
            this.sem.take(async () => {
                if (this.err) {
                    this.sem.leave();
                    reject(this.err);
                    return;
                }
                this.updateStates(this.err);
                if (this.state === ParallelQueryExecutionContextBase.STATES.ended) {
                    this.sem.leave();
                    resolve();
                    return;
                }
                if (this.unfilledDocumentProducersQueue.size() === 0) {
                    this.sem.leave();
                    resolve();
                    return;
                }
                try {
                    const maxDegreeOfParallelism = this.options.maxDegreeOfParallelism === undefined ||
                        this.options.maxDegreeOfParallelism < 1
                        ? this.unfilledDocumentProducersQueue.size() // number of partitions
                        : Math.min(this.options.maxDegreeOfParallelism, this.unfilledDocumentProducersQueue.size());
                    const documentProducers = [];
                    while (documentProducers.length < maxDegreeOfParallelism &&
                        this.unfilledDocumentProducersQueue.size() > 0) {
                        let documentProducer;
                        try {
                            documentProducer = this.unfilledDocumentProducersQueue.deq();
                        }
                        catch (e) {
                            this.err = e;
                            this.err.headers = this._getAndResetActiveResponseHeaders();
                            reject(this.err);
                            return;
                        }
                        documentProducers.push(documentProducer);
                    }
                    const bufferDocumentProducer = async (documentProducer) => {
                        try {
                            const headers = await documentProducer.bufferMore(diagnosticNode);
                            this._mergeWithActiveResponseHeaders(headers);
                            // Always track this document producer in patchToRangeMapping, even if it has no results
                            // This ensures we maintain a record of all partition ranges that were scanned
                            const nextItem = documentProducer.peakNextItem();
                            if (nextItem !== undefined) {
                                this.bufferedDocumentProducersQueue.enq(documentProducer);
                            }
                            else {
                                // Track document producer with no results in patchToRangeMapping
                                // This represents a scanned partition that yielded no results
                                // IMPORTANT: Only include if continuation token is NOT null/exhausted
                                // Document producers with no data in buffer and no continuation token are exhausted and should not be added to partitionDataPatchMap to prevent infinite loops in order by queries
                                if (documentProducer.continuationToken &&
                                    documentProducer.continuationToken !== "" &&
                                    documentProducer.continuationToken.toLowerCase() !== "null") {
                                    const patchKey = `empty-${documentProducer.targetPartitionKeyRange.id}-${documentProducer.targetPartitionKeyRange.minInclusive}`;
                                    this.partitionDataPatchMap.set(patchKey, {
                                        itemCount: 0, // 0 items for empty result set
                                        partitionKeyRange: documentProducer.targetPartitionKeyRange,
                                        continuationToken: documentProducer.continuationToken,
                                    });
                                }
                                if (documentProducer.hasMoreResults()) {
                                    this.unfilledDocumentProducersQueue.enq(documentProducer);
                                }
                            }
                        }
                        catch (err) {
                            if (ParallelQueryExecutionContextBase._needPartitionKeyRangeCacheRefresh(err)) {
                                // We want the document producer enqueued
                                // So that later parts of the code can repair the execution context
                                // refresh the partition key ranges and ctreate new document producers and add it to the queue
                                await this._enqueueReplacementDocumentProducers(err, diagnosticNode, documentProducer);
                                resolve();
                            }
                            else {
                                this.err = err;
                                this.err.headers = this._getAndResetActiveResponseHeaders();
                                reject(err);
                            }
                        }
                    };
                    try {
                        await Promise.all(documentProducers.map((producer) => bufferDocumentProducer(producer)));
                    }
                    catch (err) {
                        this.err = err;
                        this.err.headers = this._getAndResetActiveResponseHeaders();
                        reject(err);
                        return;
                    }
                    resolve();
                }
                catch (err) {
                    this.err = err;
                    this.err.headers = this._getAndResetActiveResponseHeaders();
                    reject(err);
                }
                finally {
                    this.sem.leave();
                }
            });
        });
    }
    /**
     * Drains the buffer of filled document producers and appends their items to the main buffer.
     * Uses template method pattern - delegates actual processing to subclasses.
     * @returns A promise that resolves when the buffer is filled.
     */
    async fillBufferFromBufferQueue() {
        return new Promise((resolve, reject) => {
            this.sem.take(async () => {
                if (this.err) {
                    // if there is a prior error return error
                    this.sem.leave();
                    this.err.headers = this._getAndResetActiveResponseHeaders();
                    reject(this.err);
                    return;
                }
                if (this.state === ParallelQueryExecutionContextBase.STATES.ended ||
                    this.bufferedDocumentProducersQueue.size() === 0) {
                    this.sem.leave();
                    resolve();
                    return;
                }
                try {
                    await this.processBufferedDocumentProducers();
                    this.updateStates(this.err);
                }
                catch (err) {
                    this.err = err;
                    this.err.headers = this._getAndResetActiveResponseHeaders();
                    reject(this.err);
                    return;
                }
                finally {
                    // release the lock before returning
                    this.sem.leave();
                }
                resolve();
                return;
            });
        });
    }
    updateStates(error) {
        if (error) {
            this.err = error;
            this.state = ParallelQueryExecutionContextBase.STATES.ended;
            return;
        }
        if (this.state === ParallelQueryExecutionContextBase.STATES.started) {
            this.state = ParallelQueryExecutionContextBase.STATES.inProgress;
        }
        const hasNoActiveProducers = this.unfilledDocumentProducersQueue.size() === 0 &&
            this.bufferedDocumentProducersQueue.size() === 0;
        if (hasNoActiveProducers) {
            this.state = ParallelQueryExecutionContextBase.STATES.ended;
        }
    }
}
//# sourceMappingURL=parallelQueryExecutionContextBase.js.map