import { LogContext } from '@rocicorp/logger';
import { assert, unreachable } from "../../../../shared/src/asserts.js";
import { deepEqual } from "../../../../shared/src/json.js";
import { must } from "../../../../shared/src/must.js";
import { buildPipeline } from "../../../../zql/src/builder/builder.js";
import { Debug, runtimeDebugFlags, } from "../../../../zql/src/builder/debug-delegate.js";
import { MeasurePushOperator } from "../../../../zql/src/query/measure-push-operator.js";
import { TableSource } from "../../../../zqlite/src/table-source.js";
import { reloadPermissionsIfChanged, } from "../../auth/load-permissions.js";
import { computeZqlSpecs, mustGetTableSpec } from "../../db/lite-tables.js";
import { getOrCreateHistogram } from "../../observability/metrics.js";
import { getSubscriptionState } from "../replicator/schema/replication-state.js";
import { checkClientSchema } from "./client-schema.js";
import { ResetPipelinesSignal, Snapshotter, } from "./snapshotter.js";
/**
 * Manages the state of IVM pipelines for a given ViewSyncer (i.e. client group).
 */
export class PipelineDriver {
    #tables = new Map();
    // We probs need the original query hash
    // so we can decide not to re-transform a custom query
    // that is already hydrated.
    #pipelines = new Map();
    #lc;
    #snapshotter;
    #storage;
    #shardID;
    #logConfig;
    #tableSpecs = new Map();
    #streamer = null;
    #replicaVersion = null;
    #permissions = null;
    #advanceTime = getOrCreateHistogram('sync', 'ivm.advance-time', {
        description: 'Time to advance all queries for a given client group for in response to a single change.',
        unit: 's',
    });
    #inspectorDelegate;
    constructor(lc, logConfig, snapshotter, shardID, storage, clientGroupID, inspectorDelegate) {
        this.#lc = lc.withContext('clientGroupID', clientGroupID);
        this.#snapshotter = snapshotter;
        this.#storage = storage;
        this.#shardID = shardID;
        this.#logConfig = logConfig;
        this.#inspectorDelegate = inspectorDelegate;
    }
    /**
     * Initializes the PipelineDriver to the current head of the database.
     * Queries can then be added (i.e. hydrated) with {@link addQuery()}.
     *
     * Must only be called once.
     */
    init(clientSchema) {
        assert(!this.#snapshotter.initialized(), 'Already initialized');
        const { db } = this.#snapshotter.init().current();
        const fullTables = new Map();
        computeZqlSpecs(this.#lc, db.db, this.#tableSpecs, fullTables);
        if (clientSchema) {
            checkClientSchema(this.#shardID, clientSchema, this.#tableSpecs, fullTables);
        }
        const { replicaVersion } = getSubscriptionState(db);
        this.#replicaVersion = replicaVersion;
    }
    /**
     * @returns Whether the PipelineDriver has been initialized.
     */
    initialized() {
        return this.#snapshotter.initialized();
    }
    /** @returns The replica version. The PipelineDriver must have been initialized. */
    get replicaVersion() {
        return must(this.#replicaVersion, 'Not yet initialized');
    }
    /**
     * Returns the current version of the database. This will reflect the
     * latest version change when calling {@link advance()} once the
     * iteration has begun.
     */
    currentVersion() {
        assert(this.initialized(), 'Not yet initialized');
        return this.#snapshotter.current().version;
    }
    /**
     * Returns the current supported schema version range of the database.  This
     * will reflect changes to supported schema version range when calling
     * {@link advance()} once the iteration has begun.
     */
    currentSchemaVersions() {
        assert(this.initialized(), 'Not yet initialized');
        return this.#snapshotter.current().schemaVersions;
    }
    /**
     * Returns the current upstream {app}.permissions, or `null` if none are defined.
     */
    currentPermissions() {
        assert(this.initialized(), 'Not yet initialized');
        const res = reloadPermissionsIfChanged(this.#lc, this.#snapshotter.current().db, this.#shardID.appID, this.#permissions);
        if (res.changed) {
            this.#permissions = res.permissions;
            this.#lc.debug?.('Reloaded permissions', JSON.stringify(this.#permissions));
        }
        return this.#permissions;
    }
    advanceWithoutDiff() {
        const { db, version } = this.#snapshotter.advanceWithoutDiff().curr;
        for (const table of this.#tables.values()) {
            table.setDB(db.db);
        }
        return version;
    }
    /**
     * Clears the current pipelines and TableSources, returning the PipelineDriver
     * to its initial state. This should be called in response to a schema change,
     * as TableSources need to be recomputed.
     */
    reset(clientSchema) {
        for (const { input } of this.#pipelines.values()) {
            input.destroy();
        }
        this.#pipelines.clear();
        this.#tables.clear();
        const { db } = this.#snapshotter.current();
        const fullTables = new Map();
        computeZqlSpecs(this.#lc, db.db, this.#tableSpecs, fullTables);
        if (clientSchema) {
            checkClientSchema(this.#shardID, clientSchema, this.#tableSpecs, fullTables);
        }
        const { replicaVersion } = getSubscriptionState(db);
        this.#replicaVersion = replicaVersion;
    }
    /**
     * Clears storage used for the pipelines. Call this when the
     * PipelineDriver will no longer be used.
     */
    destroy() {
        this.#storage.destroy();
        this.#snapshotter.destroy();
    }
    /** @return The Set of query hashes for all added queries. */
    addedQueries() {
        const byOriginalHash = new Map();
        for (const pipeline of this.#pipelines.values()) {
            const { originalHash, transformedAst, transformationHash } = pipeline;
            if (!byOriginalHash.has(originalHash)) {
                byOriginalHash.set(originalHash, []);
            }
            byOriginalHash.get(originalHash).push({
                transformationHash,
                transformedAst,
            });
        }
        return [new Set(this.#pipelines.keys()), byOriginalHash];
    }
    totalHydrationTimeMs() {
        let total = 0;
        for (const pipeline of this.#pipelines.values()) {
            total += pipeline.hydrationTimeMs;
        }
        return total;
    }
    /**
     * Adds a pipeline for the query. The method will hydrate the query using the
     * driver's current snapshot of the database and return a stream of results.
     * Henceforth, updates to the query will be returned when the driver is
     * {@link advance}d. The query and its pipeline can be removed with
     * {@link removeQuery()}.
     *
     * If a query with an identical hash has already been added, this method is a
     * no-op and no RowChanges are generated.
     *
     * @param timer The caller-controlled {@link Timer} used to determine the
     *        final hydration time. (The caller may pause and resume the timer
     *        when yielding the thread for time-slicing).
     * @return The rows from the initial hydration of the query.
     */
    *addQuery(transformationHash, queryID, query, timer) {
        assert(this.initialized());
        this.#inspectorDelegate.addQuery(transformationHash, queryID, query);
        if (this.#pipelines.has(transformationHash)) {
            this.#lc.info?.(`query ${transformationHash} already added`, query);
            return;
        }
        const debugDelegate = runtimeDebugFlags.trackRowsVended
            ? new Debug()
            : undefined;
        const input = buildPipeline(query, {
            debug: debugDelegate,
            enableNotExists: true, // Server-side can handle NOT EXISTS
            getSource: name => this.#getSource(name),
            createStorage: () => this.#createStorage(),
            decorateSourceInput: (input, _queryID) => new MeasurePushOperator(input, transformationHash, this.#inspectorDelegate, 'query-update-server'),
            decorateInput: input => input,
            addEdge() { },
            decorateFilterInput: input => input,
        }, queryID);
        const schema = input.getSchema();
        input.setOutput({
            push: change => {
                const streamer = this.#streamer;
                assert(streamer, 'must #startAccumulating() before pushing changes');
                streamer.accumulate(transformationHash, schema, [change]);
            },
        });
        yield* hydrate(input, transformationHash, this.#tableSpecs);
        const hydrationTimeMs = timer.totalElapsed();
        if (runtimeDebugFlags.trackRowCountsVended) {
            if (hydrationTimeMs > this.#logConfig.slowHydrateThreshold) {
                let totalRowsConsidered = 0;
                const lc = this.#lc
                    .withContext('hash', transformationHash)
                    .withContext('hydrationTimeMs', hydrationTimeMs);
                for (const tableName of this.#tables.keys()) {
                    const entries = Object.entries(debugDelegate?.getVendedRowCounts()[tableName] ?? {});
                    totalRowsConsidered += entries.reduce((acc, entry) => acc + entry[1], 0);
                    lc.info?.(tableName + ' VENDED: ', entries);
                }
                lc.info?.(`Total rows considered: ${totalRowsConsidered}`);
            }
        }
        debugDelegate?.reset();
        // Note: This hydrationTime is a wall-clock overestimate, as it does
        // not take time slicing into account. The view-syncer resets this
        // to a more precise processing-time measurement with setHydrationTime().
        this.#pipelines.set(transformationHash, {
            input,
            hydrationTimeMs,
            originalHash: queryID,
            transformedAst: query,
            transformationHash,
        });
    }
    /**
     * Removes the pipeline for the query. This is a no-op if the query
     * was not added.
     */
    removeQuery(hash) {
        const pipeline = this.#pipelines.get(hash);
        if (pipeline) {
            this.#pipelines.delete(hash);
            pipeline.input.destroy();
        }
    }
    /**
     * Returns the value of the row with the given primary key `pk`,
     * or `undefined` if there is no such row. The pipeline must have been
     * initialized.
     */
    getRow(table, pk) {
        assert(this.initialized(), 'Not yet initialized');
        const source = must(this.#tables.get(table));
        return source.getRow(pk);
    }
    /**
     * Advances to the new head of the database.
     *
     * @param timer The caller-controlled {@link Timer} that will be used to
     *        measure the progress of the advancement and abort with a
     *        {@link ResetPipelinesSignal} if it is estimated to take longer
     *        than a hydration.
     * @return The resulting row changes for all added queries. Note that the
     *         `changes` must be iterated over in their entirety in order to
     *         advance the database snapshot.
     */
    advance(timer) {
        assert(this.initialized());
        const diff = this.#snapshotter.advance(this.#tableSpecs);
        const { prev, curr, changes } = diff;
        this.#lc.debug?.(`${prev.version} => ${curr.version}: ${changes} changes`);
        const totalHydrationTimeMs = this.totalHydrationTimeMs();
        // Cancel the advancement processing if it takes longer than half the
        // total hydration time to make it through half of the advancement.
        // This serves as both a circuit breaker for very large transactions,
        // as well as a bound on the amount of time the previous connection locks
        // the inactive WAL file (as the lock prevents WAL2 from switching to the
        // free WAL when the current one is over the size limit, which can make
        // the WAL grow continuously and compound slowness).
        //
        // Note: 1/2 is a conservative estimate policy. A lower proportion would
        // flag slowness sooner, at the expense of larger estimation error.
        function checkProgress(pos) {
            // Check every 10 changes
            if (pos % 10 === 0) {
                const elapsed = timer.totalElapsed();
                if (elapsed > totalHydrationTimeMs / 2 && pos <= changes / 2) {
                    throw new ResetPipelinesSignal(`advancement exceeded timeout at ${pos} of ${changes} changes (${elapsed} ms)`);
                }
            }
        }
        return {
            version: curr.version,
            numChanges: changes,
            changes: this.#advance(diff, 
            // Somewhat arbitrary: only check progress if there are at least 20
            // changes (Note that the first check doesn't happen until 10 changes).
            changes >= 20 ? checkProgress : () => { }),
        };
    }
    *#advance(diff, onChange) {
        let pos = 0;
        for (const { table, prevValue, nextValue, rowKey } of diff) {
            const start = performance.now();
            let type;
            try {
                if (prevValue && nextValue) {
                    // Rows are ultimately referred to by the union key (in #streamNodes())
                    // so an update is represented as an `edit` if and only if the
                    // unionKey-based row keys are the same in prevValue and nextValue.
                    const { unionKey } = must(this.#tableSpecs.get(table)).tableSpec;
                    if (Object.keys(rowKey).length === unionKey.length ||
                        deepEqual(getRowKey(unionKey, prevValue), getRowKey(unionKey, nextValue))) {
                        type = 'edit';
                        yield* this.#push(table, {
                            type: 'edit',
                            row: nextValue,
                            oldRow: prevValue,
                        });
                        continue;
                    }
                    // If the unionKey-based row keys differed, they will be
                    // represented as a remove of the old key and an add of the new key.
                }
                if (prevValue) {
                    type = 'remove';
                    yield* this.#push(table, { type: 'remove', row: prevValue });
                }
                if (nextValue) {
                    type = 'add';
                    yield* this.#push(table, { type: 'add', row: nextValue });
                }
            }
            finally {
                onChange(++pos);
            }
            const elapsed = performance.now() - start;
            this.#advanceTime.record(elapsed / 1000, {
                table,
                type,
            });
        }
        // Set the new snapshot on all TableSources.
        const { curr } = diff;
        for (const table of this.#tables.values()) {
            table.setDB(curr.db.db);
        }
        this.#lc.debug?.(`Advanced to ${curr.version}`);
    }
    /** Implements `BuilderDelegate.getSource()` */
    #getSource(tableName) {
        let source = this.#tables.get(tableName);
        if (source) {
            return source;
        }
        const tableSpec = mustGetTableSpec(this.#tableSpecs, tableName);
        const { primaryKey } = tableSpec.tableSpec;
        const { db } = this.#snapshotter.current();
        source = new TableSource(this.#lc, this.#logConfig, db.db, tableName, tableSpec.zqlSpec, primaryKey);
        this.#tables.set(tableName, source);
        this.#lc.debug?.(`created TableSource for ${tableName}`);
        return source;
    }
    /** Implements `BuilderDelegate.createStorage()` */
    #createStorage() {
        return this.#storage.createStorage();
    }
    *#push(table, change) {
        const source = this.#tables.get(table);
        if (!source) {
            return;
        }
        this.#startAccumulating();
        for (const _ of source.genPush(change)) {
            yield* this.#stopAccumulating().stream();
            this.#startAccumulating();
        }
        this.#stopAccumulating();
    }
    #startAccumulating() {
        assert(this.#streamer === null);
        this.#streamer = new Streamer(this.#tableSpecs);
    }
    #stopAccumulating() {
        const streamer = this.#streamer;
        assert(streamer);
        this.#streamer = null;
        return streamer;
    }
}
class Streamer {
    #tableSpecs;
    constructor(tableSpecs) {
        this.#tableSpecs = tableSpecs;
    }
    #changes = [];
    accumulate(hash, schema, changes) {
        this.#changes.push([hash, schema, changes]);
        return this;
    }
    *stream() {
        for (const [hash, schema, changes] of this.#changes) {
            yield* this.#streamChanges(hash, schema, changes);
        }
    }
    *#streamChanges(queryHash, schema, changes) {
        // We do not sync rows gathered by the permissions
        // system to the client.
        if (schema.system === 'permissions') {
            return;
        }
        for (const change of changes) {
            const { type } = change;
            switch (type) {
                case 'add':
                case 'remove': {
                    yield* this.#streamNodes(queryHash, schema, type, () => [
                        change.node,
                    ]);
                    break;
                }
                case 'child': {
                    const { child } = change;
                    const childSchema = must(schema.relationships[child.relationshipName]);
                    yield* this.#streamChanges(queryHash, childSchema, [child.change]);
                    break;
                }
                case 'edit':
                    yield* this.#streamNodes(queryHash, schema, type, () => [
                        { row: change.node.row, relationships: {} },
                    ]);
                    break;
                default:
                    unreachable(type);
            }
        }
    }
    *#streamNodes(queryHash, schema, op, nodes) {
        const { tableName: table, system } = schema;
        // The primaryKey here is used for referencing rows in CVR and del-row
        // patches sent in pokes. This is the "unionKey", i.e. the union of all
        // columns in unique indexes. This allows clients to migrate from, e.g.
        // pk1 to pk2, as del-patches will be keyed by [...pk1, ...pk2].
        const primaryKey = must(this.#tableSpecs.get(table)).tableSpec.unionKey;
        // We do not sync rows gathered by the permissions
        // system to the client.
        if (system === 'permissions') {
            return;
        }
        for (const node of nodes()) {
            const { relationships, row } = node;
            const rowKey = getRowKey(primaryKey, row);
            yield {
                type: op,
                queryHash,
                table,
                rowKey,
                row: op === 'remove' ? undefined : row,
            };
            for (const [relationship, children] of Object.entries(relationships)) {
                const childSchema = must(schema.relationships[relationship]);
                yield* this.#streamNodes(queryHash, childSchema, op, children);
            }
        }
    }
}
function* toAdds(nodes) {
    for (const node of nodes) {
        yield { type: 'add', node };
    }
}
function getRowKey(cols, row) {
    return Object.fromEntries(cols.map(col => [col, must(row[col])]));
}
/**
 * Core hydration logic used by {@link PipelineDriver#addQuery}, extracted to a
 * function for reuse by bin-analyze so that bin-analyze's hydration logic
 * is as close as possible to zero-cache's real hydration logic.
 */
export function* hydrate(input, hash, tableSpecs) {
    const res = input.fetch({});
    const streamer = new Streamer(tableSpecs).accumulate(hash, input.getSchema(), toAdds(res));
    yield* streamer.stream();
}
//# sourceMappingURL=pipeline-driver.js.map