import { assert, unreachable } from "../../../shared/src/asserts.js";
import { hasOwn } from "../../../shared/src/has-own.js";
import { must } from "../../../shared/src/must.js";
import { assertOrderingIncludesPK } from "../builder/builder.js";
import {} from "./change.js";
import { compareValues } from "./data.js";
import { throwOutput, } from "./operator.js";
import { first, take } from "./stream.js";
const MAX_BOUND_KEY = 'maxBound';
/**
 * The Take operator is for implementing limit queries. It takes the first n
 * nodes of its input as determined by the input’s comparator. It then keeps
 * a *bound* of the last item it has accepted so that it can evaluate whether
 * new incoming pushes should be accepted or rejected.
 *
 * Take can count rows globally or by unique value of some field.
 *
 * Maintains the invariant that its output size is always <= limit, even
 * mid processing of a push.
 */
export class Take {
    #input;
    #storage;
    #limit;
    #partitionKey;
    #partitionKeyComparator;
    // Fetch overlay needed for some split push cases.
    #rowHiddenFromFetch;
    #output = throwOutput;
    constructor(input, storage, limit, partitionKey) {
        assert(limit >= 0, 'Limit must be non-negative');
        assertOrderingIncludesPK(input.getSchema().sort, input.getSchema().primaryKey);
        input.setOutput(this);
        this.#input = input;
        this.#storage = storage;
        this.#limit = limit;
        this.#partitionKey = partitionKey;
        this.#partitionKeyComparator =
            partitionKey && makePartitionKeyComparator(partitionKey);
    }
    setOutput(output) {
        this.#output = output;
    }
    getSchema() {
        return this.#input.getSchema();
    }
    *fetch(req) {
        if (!this.#partitionKey ||
            (req.constraint &&
                constraintMatchesPartitionKey(req.constraint, this.#partitionKey))) {
            const takeStateKey = getTakeStateKey(this.#partitionKey, req.constraint);
            const takeState = this.#storage.get(takeStateKey);
            if (!takeState) {
                yield* this.#initialFetch(req);
                return;
            }
            if (takeState.bound === undefined) {
                return;
            }
            for (const inputNode of this.#input.fetch(req)) {
                if (this.getSchema().compareRows(takeState.bound, inputNode.row) < 0) {
                    return;
                }
                if (this.#rowHiddenFromFetch &&
                    this.getSchema().compareRows(this.#rowHiddenFromFetch, inputNode.row) === 0) {
                    continue;
                }
                yield inputNode;
            }
            return;
        }
        // There is a partition key, but the fetch is not constrained or constrained
        // on a different key.  Thus we don't have a single take state to bound by.
        // This currently only happens with nested sub-queries
        // e.g. issues include issuelabels include label.  We could remove this
        // case if we added a translation layer (powered by some state) in join.
        // Specifically we need joinKeyValue => parent constraint key
        const maxBound = this.#storage.get(MAX_BOUND_KEY);
        if (maxBound === undefined) {
            return;
        }
        for (const inputNode of this.#input.fetch(req)) {
            if (this.getSchema().compareRows(inputNode.row, maxBound) > 0) {
                return;
            }
            const takeStateKey = getTakeStateKey(this.#partitionKey, inputNode.row);
            const takeState = this.#storage.get(takeStateKey);
            if (takeState?.bound !== undefined &&
                this.getSchema().compareRows(takeState.bound, inputNode.row) >= 0) {
                yield inputNode;
            }
        }
    }
    *#initialFetch(req) {
        assert(req.start === undefined, 'Start should be undefined');
        assert(!req.reverse, 'Reverse should be false');
        assert(constraintMatchesPartitionKey(req.constraint, this.#partitionKey), 'Constraint should match partition key');
        if (this.#limit === 0) {
            return;
        }
        const takeStateKey = getTakeStateKey(this.#partitionKey, req.constraint);
        assert(this.#storage.get(takeStateKey) === undefined, 'Take state should be undefined');
        let size = 0;
        let bound;
        let downstreamEarlyReturn = true;
        let exceptionThrown = false;
        try {
            for (const inputNode of this.#input.fetch(req)) {
                yield inputNode;
                bound = inputNode.row;
                size++;
                if (size === this.#limit) {
                    break;
                }
            }
            downstreamEarlyReturn = false;
        }
        catch (e) {
            exceptionThrown = true;
            throw e;
        }
        finally {
            if (!exceptionThrown) {
                this.#setTakeState(takeStateKey, size, bound, this.#storage.get(MAX_BOUND_KEY));
                // If it becomes necessary to support downstream early return, this
                // assert should be removed, and replaced with code that consumes
                // the input stream until limit is reached or the input stream is
                // exhausted so that takeState is properly hydrated.
                assert(!downstreamEarlyReturn, 'Unexpected early return prevented full hydration');
            }
        }
    }
    *cleanup(req) {
        assert(req.start === undefined, 'Start should be undefined');
        assert(constraintMatchesPartitionKey(req.constraint, this.#partitionKey), 'Constraint should match partition key');
        const takeStateKey = getTakeStateKey(this.#partitionKey, req.constraint);
        this.#storage.del(takeStateKey);
        let size = 0;
        for (const inputNode of this.#input.cleanup(req)) {
            if (size === this.#limit) {
                return;
            }
            size++;
            yield inputNode;
        }
    }
    #getStateAndConstraint(row) {
        const takeStateKey = getTakeStateKey(this.#partitionKey, row);
        const takeState = this.#storage.get(takeStateKey);
        let maxBound;
        let constraint;
        if (takeState) {
            maxBound = this.#storage.get(MAX_BOUND_KEY);
            constraint =
                this.#partitionKey &&
                    Object.fromEntries(this.#partitionKey.map(key => [key, row[key]]));
        }
        return { takeState, takeStateKey, maxBound, constraint };
    }
    push(change) {
        if (change.type === 'edit') {
            this.#pushEditChange(change);
            return;
        }
        const { takeState, takeStateKey, maxBound, constraint } = this.#getStateAndConstraint(change.node.row);
        if (!takeState) {
            return;
        }
        const { compareRows } = this.getSchema();
        if (change.type === 'add') {
            if (takeState.size < this.#limit) {
                this.#setTakeState(takeStateKey, takeState.size + 1, takeState.bound === undefined ||
                    compareRows(takeState.bound, change.node.row) < 0
                    ? change.node.row
                    : takeState.bound, maxBound);
                this.#output.push(change, this);
                return;
            }
            // size === limit
            if (takeState.bound === undefined ||
                compareRows(change.node.row, takeState.bound) >= 0) {
                return;
            }
            // added row < bound
            let beforeBoundNode;
            let boundNode;
            if (this.#limit === 1) {
                boundNode = must(first(this.#input.fetch({
                    start: {
                        row: takeState.bound,
                        basis: 'at',
                    },
                    constraint,
                })));
            }
            else {
                [boundNode, beforeBoundNode] = take(this.#input.fetch({
                    start: {
                        row: takeState.bound,
                        basis: 'at',
                    },
                    constraint,
                    reverse: true,
                }), 2);
            }
            const removeChange = {
                type: 'remove',
                node: boundNode,
            };
            // Remove before add to maintain invariant that
            // output size <= limit.
            this.#setTakeState(takeStateKey, takeState.size, beforeBoundNode === undefined ||
                compareRows(change.node.row, beforeBoundNode.row) > 0
                ? change.node.row
                : beforeBoundNode.row, maxBound);
            this.#withRowHiddenFromFetch(change.node.row, () => {
                this.#output.push(removeChange, this);
            });
            this.#output.push(change, this);
        }
        else if (change.type === 'remove') {
            if (takeState.bound === undefined) {
                // change is after bound
                return;
            }
            const compToBound = compareRows(change.node.row, takeState.bound);
            if (compToBound > 0) {
                // change is after bound
                return;
            }
            const [beforeBoundNode] = take(this.#input.fetch({
                start: {
                    row: takeState.bound,
                    basis: 'after',
                },
                constraint,
                reverse: true,
            }), 1);
            let newBound;
            if (beforeBoundNode) {
                const push = compareRows(beforeBoundNode.row, takeState.bound) > 0;
                newBound = {
                    node: beforeBoundNode,
                    push,
                };
            }
            if (!newBound?.push) {
                for (const node of this.#input.fetch({
                    start: {
                        row: takeState.bound,
                        basis: 'at',
                    },
                    constraint,
                })) {
                    const push = compareRows(node.row, takeState.bound) > 0;
                    newBound = {
                        node,
                        push,
                    };
                    if (push) {
                        break;
                    }
                }
            }
            if (newBound?.push) {
                this.#output.push(change, this);
                this.#setTakeState(takeStateKey, takeState.size, newBound.node.row, maxBound);
                this.#output.push({
                    type: 'add',
                    node: newBound.node,
                }, this);
                return;
            }
            this.#setTakeState(takeStateKey, takeState.size - 1, newBound?.node.row, maxBound);
            this.#output.push(change, this);
        }
        else if (change.type === 'child') {
            // A 'child' change should be pushed to output if its row
            // is <= bound.
            if (takeState.bound &&
                compareRows(change.node.row, takeState.bound) <= 0) {
                this.#output.push(change, this);
            }
        }
    }
    #pushEditChange(change) {
        assert(!this.#partitionKeyComparator ||
            this.#partitionKeyComparator(change.oldNode.row, change.node.row) === 0, 'Unexpected change of partition key');
        const { takeState, takeStateKey, maxBound, constraint } = this.#getStateAndConstraint(change.oldNode.row);
        if (!takeState) {
            return;
        }
        assert(takeState.bound, 'Bound should be set');
        const { compareRows } = this.getSchema();
        const oldCmp = compareRows(change.oldNode.row, takeState.bound);
        const newCmp = compareRows(change.node.row, takeState.bound);
        const replaceBoundAndForwardChange = () => {
            this.#setTakeState(takeStateKey, takeState.size, change.node.row, maxBound);
            this.#output.push(change, this);
        };
        // The bounds row was changed.
        if (oldCmp === 0) {
            // The new row is the new bound.
            if (newCmp === 0) {
                // no need to update the state since we are keeping the bounds
                this.#output.push(change, this);
                return;
            }
            if (newCmp < 0) {
                if (this.#limit === 1) {
                    replaceBoundAndForwardChange();
                    return;
                }
                // New row will be in the result but it might not be the bounds any
                // more. We need to find the row before the bounds to determine the new
                // bounds.
                const beforeBoundNode = must(first(this.#input.fetch({
                    start: {
                        row: takeState.bound,
                        basis: 'after',
                    },
                    constraint,
                    reverse: true,
                })));
                this.#setTakeState(takeStateKey, takeState.size, beforeBoundNode.row, maxBound);
                this.#output.push(change, this);
                return;
            }
            assert(newCmp > 0, 'New comparison must be greater than 0');
            // Find the first item at the old bounds. This will be the new bounds.
            const newBoundNode = must(first(this.#input.fetch({
                start: {
                    row: takeState.bound,
                    basis: 'at',
                },
                constraint,
            })));
            // The next row is the new row. We can replace the bounds and keep the
            // edit change.
            if (compareRows(newBoundNode.row, change.node.row) === 0) {
                replaceBoundAndForwardChange();
                return;
            }
            // The new row is now outside the bounds, so we need to remove the old
            // row and add the new bounds row.
            this.#setTakeState(takeStateKey, takeState.size, newBoundNode.row, maxBound);
            this.#withRowHiddenFromFetch(newBoundNode.row, () => {
                this.#output.push({
                    type: 'remove',
                    node: change.oldNode,
                }, this);
            });
            this.#output.push({
                type: 'add',
                node: newBoundNode,
            }, this);
            return;
        }
        if (oldCmp > 0) {
            assert(newCmp !== 0, 'Invalid state. Row has duplicate primary key');
            // Both old and new outside of bounds
            if (newCmp > 0) {
                return;
            }
            // old was outside, new is inside. Pushing out the old bounds
            assert(newCmp < 0, 'New comparison must be less than 0');
            const [oldBoundNode, newBoundNode] = take(this.#input.fetch({
                start: {
                    row: takeState.bound,
                    basis: 'at',
                },
                constraint,
                reverse: true,
            }), 2);
            // Remove before add to maintain invariant that
            // output size <= limit.
            this.#setTakeState(takeStateKey, takeState.size, newBoundNode.row, maxBound);
            this.#withRowHiddenFromFetch(change.node.row, () => {
                this.#output.push({
                    type: 'remove',
                    node: oldBoundNode,
                }, this);
            });
            this.#output.push({
                type: 'add',
                node: change.node,
            }, this);
            return;
        }
        if (oldCmp < 0) {
            assert(newCmp !== 0, 'Invalid state. Row has duplicate primary key');
            // Both old and new inside of bounds
            if (newCmp < 0) {
                this.#output.push(change, this);
                return;
            }
            // old was inside, new is larger than old bound
            assert(newCmp > 0, 'New comparison must be greater than 0');
            // at this point we need to find the row after the bound and use that or
            // the newRow as the new bound.
            const afterBoundNode = must(first(this.#input.fetch({
                start: {
                    row: takeState.bound,
                    basis: 'after',
                },
                constraint,
            })));
            // The new row is the new bound. Use an edit change.
            if (compareRows(afterBoundNode.row, change.node.row) === 0) {
                replaceBoundAndForwardChange();
                return;
            }
            this.#output.push({
                type: 'remove',
                node: change.oldNode,
            }, this);
            this.#setTakeState(takeStateKey, takeState.size, afterBoundNode.row, maxBound);
            this.#output.push({
                type: 'add',
                node: afterBoundNode,
            }, this);
            return;
        }
        unreachable();
    }
    #withRowHiddenFromFetch(row, fn) {
        this.#rowHiddenFromFetch = row;
        try {
            fn();
        }
        finally {
            this.#rowHiddenFromFetch = undefined;
        }
    }
    #setTakeState(takeStateKey, size, bound, maxBound) {
        this.#storage.set(takeStateKey, {
            size,
            bound,
        });
        if (bound !== undefined &&
            (maxBound === undefined ||
                this.getSchema().compareRows(bound, maxBound) > 0)) {
            this.#storage.set(MAX_BOUND_KEY, bound);
        }
    }
    destroy() {
        this.#input.destroy();
    }
}
function getTakeStateKey(partitionKey, rowOrConstraint) {
    // The order must be consistent. We always use the order as defined by the
    // partition key.
    const partitionValues = [];
    if (partitionKey && rowOrConstraint) {
        for (const key of partitionKey) {
            partitionValues.push(rowOrConstraint[key]);
        }
    }
    return JSON.stringify(['take', ...partitionValues]);
}
function constraintMatchesPartitionKey(constraint, partitionKey) {
    if (constraint === undefined || partitionKey === undefined) {
        return constraint === partitionKey;
    }
    if (partitionKey.length !== Object.keys(constraint).length) {
        return false;
    }
    for (const key of partitionKey) {
        if (!hasOwn(constraint, key)) {
            return false;
        }
    }
    return true;
}
function makePartitionKeyComparator(partitionKey) {
    return (a, b) => {
        for (const key of partitionKey) {
            const cmp = compareValues(a[key], b[key]);
            if (cmp !== 0) {
                return cmp;
            }
        }
        return 0;
    };
}
//# sourceMappingURL=take.js.map