import { Firestore } from '@google-cloud/firestore';
import { SessionStore, SessionStoreOptions, SessionSnapshot, SnapshotMutator } from 'genkit/beta';

/**
 * Copyright 2026 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/**
 * Options for {@link FirestoreSessionStore}.
 */
interface FirestoreSessionStoreOptions {
    /**
     * An explicit Firestore instance. Defaults to a new {@link Firestore}
     * instance (which picks up Application Default Credentials and the
     * `FIRESTORE_EMULATOR_HOST` environment variable).
     */
    db?: Firestore;
    /**
     * The collection where snapshot documents are stored (keyed by
     * `snapshotId`). Defaults to `"genkit-sessions"`. Two companion collections
     * are derived from it: `"<collection>-pointers"` holds one pointer document
     * per session and `"<collection>-shards"` holds the sharded checkpoint
     * state.
     */
    collection?: string;
    /**
     * Number of turns between full-state checkpoints. A larger value stores
     * fewer (but reconstructs over more) diffs; a smaller value reconstructs
     * faster at the cost of more frequent full-state writes. Defaults to
     * {@link DEFAULT_CHECKPOINT_INTERVAL}.
     *
     * Cost tuning: per-save reconstruction reads grow ~linearly with this value
     * (so per-interval read work is ~quadratic in it), while checkpoint write and
     * storage cost shrink with it. Lower it (e.g. 10) for small-state, read-heavy
     * sessions; raise it (e.g. 50-100) for large per-turn state retained for a
     * long time, where checkpoint write/storage amplification dominates.
     */
    checkpointInterval?: number;
    /**
     * Maximum size in bytes of a single shard / diff document. Checkpoint state
     * is split into chunks of this size, and any diff exceeding it is promoted to
     * a (sharded) checkpoint so that no document approaches Firestore's 1 MiB
     * limit. Defaults to {@link DEFAULT_SHARD_SIZE}.
     */
    shardSize?: number;
    /**
     * Returns a per-tenant prefix derived from the call's
     * {@link SessionStoreOptions} (e.g. the authenticated user id from
     * `options.context`). When set, all snapshot, pointer and shard documents are
     * nested under a tenant-scoped subcollection keyed by this prefix, so reads
     * and writes are isolated per tenant: one tenant can never see (or even
     * address) another's snapshots, even if they get hold of a `snapshotId` -
     * resolving it still requires the matching, auth-derived prefix. Defaults to
     * `"global"`.
     *
     * The prefix is used only to build document paths; the stored ids
     * (`snapshotId`, `parentId`, `checkpointId`, ...) remain prefix-agnostic.
     */
    snapshotPathPrefix?: (options?: SessionStoreOptions) => string;
}
/**
 * A Firestore-backed {@link SessionStore} that persists session snapshots as
 * incremental JSON Patch diffs anchored to periodic, sharded full-state
 * checkpoints.
 *
 * Storage layout (the `<prefix>` segment is the per-tenant prefix returned by
 * {@link FirestoreSessionStoreOptions.snapshotPathPrefix}, or `"global"` when
 * none is configured):
 *
 * - `<collection>/<prefix>/snapshots/<snapshotId>` - one document per snapshot.
 *   A `diff` document holds the patch from its parent (`statePatch`); a
 *   `checkpoint` document holds a full-state materialization (sharded out of
 *   band).
 * - `<collection>-shards/<prefix>/shards/<checkpointId>_<index>` - the sharded
 *   full state for a checkpoint.
 * - `<collection>-pointers/<prefix>/pointers/<sessionId>` - one document per
 *   session pointing at the latest leaf snapshot and the metadata needed to
 *   reconstruct it.
 *
 * Reconstruction uses only document-ID lookups (`getAll`), so it needs no
 * secondary indexes and is strongly consistent. No single document approaches
 * the 1 MiB limit (state is sharded by `shardSize`), and the number of *diff*
 * documents touched per read/write is bounded by `checkpointInterval` rather
 * than total session length - so the store scales to arbitrarily long sessions
 * (e.g. coding agents, long-lived chatbots). Note that checkpoints still store
 * the full accumulated state, so checkpoint shard count (and the bytes written
 * per checkpoint) grow with the state's size; tune `checkpointInterval` to
 * trade per-save diff reads against checkpoint write amplification.
 */
declare class FirestoreSessionStore<S = unknown> implements SessionStore<S> {
    protected db: Firestore;
    protected collection: string;
    protected checkpointInterval: number;
    protected shardSize: number;
    protected snapshotPathPrefix?: (options?: SessionStoreOptions) => string;
    constructor(opts?: FirestoreSessionStoreOptions);
    /** Resolves the (per-tenant) prefix for the given call options. */
    private prefixFor;
    /** The (per-tenant) snapshots subcollection. */
    private snapshotsCol;
    /** The (per-tenant) pointers subcollection. */
    private pointersCol;
    /** The (per-tenant) shards subcollection. */
    private shardsCol;
    getSnapshot(opts: {
        snapshotId?: string;
        sessionId?: string;
        context?: SessionStoreOptions['context'];
    }): Promise<SessionSnapshot<S> | undefined>;
    saveSnapshot(snapshotId: string | undefined, mutator: SnapshotMutator<S>, options?: SessionStoreOptions): Promise<string | null>;
    onSnapshotStateChange(snapshotId: string, callback: (snapshot: SessionSnapshot<S>) => void, options?: SessionStoreOptions): void | (() => void);
    /**
     * Validates that exactly one of `snapshotId` / `sessionId` is provided, and
     * that the provided one is a non-blank string. Blank / whitespace-only ids
     * are rejected up front (rather than silently treated as "absent") so callers
     * get a clear `INVALID_ARGUMENT` instead of an unusable document key.
     */
    private normalize;
    /** Builds a {@link Reader} bound to a transaction or the bare instance. */
    private reader;
    /**
     * Resolves a parent's chain metadata (nearest checkpoint, shard count and
     * segment path) *without* materializing its - potentially large - state.
     *
     * In the common linear case the parent is the session's current leaf, so the
     * metadata is read straight off the pointer and this performs *zero*
     * document reads. Otherwise it reads the single parent document. Crucially,
     * resolving only the metadata lets `saveSnapshot` decide diff-vs-checkpoint
     * (which only needs `segmentPath.length`) before paying for a full state
     * reconstruction - so checkpoint-boundary turns, which would rewrite the
     * whole state anyway, skip reconstruction entirely.
     */
    private loadParentChainMeta;
    /**
     * Reconstructs the state of `id` by reading its document to learn its
     * checkpoint and segment path, then materializing from that checkpoint.
     * Returns `undefined` when the snapshot does not exist.
     */
    private reconstruct;
    /**
     * Materializes the state of `targetId` from a known checkpoint using a single
     * batched, strongly-consistent `getAll`: the checkpoint's shards, the
     * (bounded) segment of diff documents along `segmentPath`, and - only when
     * the target *is* the checkpoint - the checkpoint document itself. The diffs
     * are then applied in order onto the checkpoint's state. Cost is bounded by
     * `checkpointInterval` + shard count, independent of total session length.
     *
     * Note: when `segmentPath` is non-empty the state comes entirely from the
     * shards and the target's metadata from the last segment document, so the
     * checkpoint *document* is not read - saving one read on the common path.
     */
    private reconstructFrom;
    /**
     * Serializes `state` to UTF-8, splits it into `shardSize`-byte chunks, and
     * writes them at `<checkpointId>_<index>`. When `oldShardCount` exceeds the
     * new count (a shrinking re-checkpoint), the now-stale trailing shards are
     * deleted. Returns the number of shards written.
     */
    private writeShards;
    /**
     * Writes a full-state checkpoint at `id` (sharding the state via
     * {@link writeShards}) and returns the snapshot metadata describing it: a
     * checkpoint anchors itself (`checkpointId === id`), has an empty
     * `segmentPath`, and carries no `statePatch`.
     *
     * This is the shared promotion path used whenever a snapshot must be a
     * checkpoint rather than a diff - the session root, an orphaned parent, a
     * checkpoint-interval boundary, an in-place checkpoint rewrite, and the
     * promotion of an oversized diff (whether new turn or leaf upsert) so that no
     * single document approaches Firestore's 1 MiB limit. Pass `oldShardCount`
     * when re-checkpointing an existing checkpoint so stale trailing shards are
     * pruned.
     */
    private writeCheckpoint;
    /** Concatenates ordered shard documents and parses the materialized state. */
    private stitch;
    /** UTF-8 byte length of a JSON-serializable value. */
    private byteLength;
    /** Assembles a {@link SessionSnapshot} from a document and its state. */
    private toSnapshot;
}

export { FirestoreSessionStore, type FirestoreSessionStoreOptions };