import { assert, unreachable } from "../../../shared/src/asserts.js";
import { hasOwn } from "../../../shared/src/has-own.js";
import { throwOutput } from "./operator.js";
import { compareValues } from "./data.js";
import { assertOrderingIncludesPK } from "../query/complete-ordering.js";
//#region ../zql/src/ivm/take.ts
var 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.
*/
var Take = class {
	#input;
	#storage;
	#limit;
	#partitionKey;
	#partitionKeyComparator;
	#rowHiddenFromFetch;
	#output = throwOutput;
	constructor(input, storage, limit, partitionKey) {
		assert(limit >= 0, "Limit must be non-negative");
		const { sort } = input.getSchema();
		assert(sort !== void 0, "Take requires sorted input");
		assertOrderingIncludesPK(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 === void 0) return;
			for (const inputNode of this.#input.fetch(req)) {
				if (inputNode === "yield") {
					yield inputNode;
					continue;
				}
				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;
		}
		const maxBound = this.#storage.get(MAX_BOUND_KEY);
		if (maxBound === void 0) return;
		for (const inputNode of this.#input.fetch(req)) {
			if (inputNode === "yield") {
				yield inputNode;
				continue;
			}
			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 !== void 0 && this.getSchema().compareRows(takeState.bound, inputNode.row) >= 0) yield inputNode;
		}
	}
	*#initialFetch(req) {
		assert(req.start === void 0, "Start should be undefined");
		assert(!req.reverse, "Reverse should be false");
		if (this.#limit === 0) return;
		assert(constraintMatchesPartitionKey(req.constraint, this.#partitionKey), "Constraint should match partition key");
		const takeStateKey = getTakeStateKey(this.#partitionKey, req.constraint);
		assert(this.#storage.get(takeStateKey) === void 0, "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)) {
				if (inputNode === "yield") {
					yield "yield";
					continue;
				}
				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));
				assert(!downstreamEarlyReturn, "Unexpected early return prevented full hydration");
			}
		}
	}
	#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") {
			yield* 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 === void 0 || compareRows(takeState.bound, change.node.row) < 0 ? change.node.row : takeState.bound, maxBound);
				yield* this.#output.push(change, this);
				return;
			}
			if (takeState.bound === void 0 || compareRows(change.node.row, takeState.bound) >= 0) return;
			let beforeBoundNode;
			let boundNode;
			if (this.#limit === 1) for (const node of this.#input.fetch({
				start: {
					row: takeState.bound,
					basis: "at"
				},
				constraint
			})) {
				if (node === "yield") {
					yield node;
					continue;
				}
				boundNode = node;
				break;
			}
			else for (const node of this.#input.fetch({
				start: {
					row: takeState.bound,
					basis: "at"
				},
				constraint,
				reverse: true
			})) if (node === "yield") {
				yield node;
				continue;
			} else if (boundNode === void 0) boundNode = node;
			else {
				beforeBoundNode = node;
				break;
			}
			assert(boundNode !== void 0, "Take: boundNode must be found during fetch");
			const removeChange = {
				type: "remove",
				node: boundNode
			};
			this.#setTakeState(takeStateKey, takeState.size, beforeBoundNode === void 0 || compareRows(change.node.row, beforeBoundNode.row) > 0 ? change.node.row : beforeBoundNode.row, maxBound);
			yield* this.#pushWithRowHiddenFromFetch(change.node.row, removeChange);
			yield* this.#output.push(change, this);
		} else if (change.type === "remove") {
			if (takeState.bound === void 0) return;
			if (compareRows(change.node.row, takeState.bound) > 0) return;
			let beforeBoundNode;
			for (const node of this.#input.fetch({
				start: {
					row: takeState.bound,
					basis: "after"
				},
				constraint,
				reverse: true
			})) {
				if (node === "yield") {
					yield node;
					continue;
				}
				beforeBoundNode = node;
				break;
			}
			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
			})) {
				if (node === "yield") {
					yield node;
					continue;
				}
				const push = compareRows(node.row, takeState.bound) > 0;
				newBound = {
					node,
					push
				};
				if (push) break;
			}
			if (newBound?.push) {
				yield* this.#output.push(change, this);
				this.#setTakeState(takeStateKey, takeState.size, newBound.node.row, maxBound);
				yield* this.#output.push({
					type: "add",
					node: newBound.node
				}, this);
				return;
			}
			this.#setTakeState(takeStateKey, takeState.size - 1, newBound?.node.row, maxBound);
			yield* this.#output.push(change, this);
		} else if (change.type === "child") {
			if (takeState.bound && compareRows(change.node.row, takeState.bound) <= 0) yield* 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 that = this;
		const replaceBoundAndForwardChange = function* () {
			that.#setTakeState(takeStateKey, takeState.size, change.node.row, maxBound);
			yield* that.#output.push(change, that);
		};
		if (oldCmp === 0) {
			if (newCmp === 0) {
				yield* this.#output.push(change, this);
				return;
			}
			if (newCmp < 0) {
				if (this.#limit === 1) {
					yield* replaceBoundAndForwardChange();
					return;
				}
				let beforeBoundNode;
				for (const node of this.#input.fetch({
					start: {
						row: takeState.bound,
						basis: "after"
					},
					constraint,
					reverse: true
				})) {
					if (node === "yield") {
						yield node;
						continue;
					}
					beforeBoundNode = node;
					break;
				}
				assert(beforeBoundNode !== void 0, "Take: beforeBoundNode must be found during fetch");
				this.#setTakeState(takeStateKey, takeState.size, beforeBoundNode.row, maxBound);
				yield* this.#output.push(change, this);
				return;
			}
			assert(newCmp > 0, "New comparison must be greater than 0");
			let newBoundNode;
			for (const node of this.#input.fetch({
				start: {
					row: takeState.bound,
					basis: "at"
				},
				constraint
			})) {
				if (node === "yield") {
					yield node;
					continue;
				}
				newBoundNode = node;
				break;
			}
			assert(newBoundNode !== void 0, "Take: newBoundNode must be found during fetch");
			if (compareRows(newBoundNode.row, change.node.row) === 0) {
				yield* replaceBoundAndForwardChange();
				return;
			}
			this.#setTakeState(takeStateKey, takeState.size, newBoundNode.row, maxBound);
			yield* this.#pushWithRowHiddenFromFetch(newBoundNode.row, {
				type: "remove",
				node: change.oldNode
			});
			yield* this.#output.push({
				type: "add",
				node: newBoundNode
			}, this);
			return;
		}
		if (oldCmp > 0) {
			assert(newCmp !== 0, "Invalid state. Row has duplicate primary key");
			if (newCmp > 0) return;
			assert(newCmp < 0, "New comparison must be less than 0");
			let oldBoundNode;
			let newBoundNode;
			for (const node of this.#input.fetch({
				start: {
					row: takeState.bound,
					basis: "at"
				},
				constraint,
				reverse: true
			})) if (node === "yield") {
				yield node;
				continue;
			} else if (oldBoundNode === void 0) oldBoundNode = node;
			else {
				newBoundNode = node;
				break;
			}
			assert(oldBoundNode !== void 0, "Take: oldBoundNode must be found during fetch");
			assert(newBoundNode !== void 0, "Take: newBoundNode must be found during fetch");
			this.#setTakeState(takeStateKey, takeState.size, newBoundNode.row, maxBound);
			yield* this.#pushWithRowHiddenFromFetch(change.node.row, {
				type: "remove",
				node: oldBoundNode
			});
			yield* this.#output.push({
				type: "add",
				node: change.node
			}, this);
			return;
		}
		if (oldCmp < 0) {
			assert(newCmp !== 0, "Invalid state. Row has duplicate primary key");
			if (newCmp < 0) {
				yield* this.#output.push(change, this);
				return;
			}
			assert(newCmp > 0, "New comparison must be greater than 0");
			let afterBoundNode;
			for (const node of this.#input.fetch({
				start: {
					row: takeState.bound,
					basis: "after"
				},
				constraint
			})) {
				if (node === "yield") {
					yield node;
					continue;
				}
				afterBoundNode = node;
				break;
			}
			assert(afterBoundNode !== void 0, "Take: afterBoundNode must be found during fetch");
			if (compareRows(afterBoundNode.row, change.node.row) === 0) {
				yield* replaceBoundAndForwardChange();
				return;
			}
			yield* this.#output.push({
				type: "remove",
				node: change.oldNode
			}, this);
			this.#setTakeState(takeStateKey, takeState.size, afterBoundNode.row, maxBound);
			yield* this.#output.push({
				type: "add",
				node: afterBoundNode
			}, this);
			return;
		}
		unreachable();
	}
	*#pushWithRowHiddenFromFetch(row, change) {
		this.#rowHiddenFromFetch = row;
		try {
			yield* this.#output.push(change, this);
		} finally {
			this.#rowHiddenFromFetch = void 0;
		}
	}
	#setTakeState(takeStateKey, size, bound, maxBound) {
		this.#storage.set(takeStateKey, {
			size,
			bound
		});
		if (bound !== void 0 && (maxBound === void 0 || this.getSchema().compareRows(bound, maxBound) > 0)) this.#storage.set(MAX_BOUND_KEY, bound);
	}
	destroy() {
		this.#input.destroy();
	}
};
function getTakeStateKey(partitionKey, rowOrConstraint) {
	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 === void 0 || partitionKey === void 0) 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;
	};
}
//#endregion
export { Take };

//# sourceMappingURL=take.js.map