import { BaseCheckpointSaver, copyCheckpoint, getCheckpointId, WRITES_IDX_MAP, } from "./base.js";
import { TASKS } from "./serde/types.js";
function _generateKey(threadId, checkpointNamespace, checkpointId) {
    return JSON.stringify([threadId, checkpointNamespace, checkpointId]);
}
export class MemorySaver extends BaseCheckpointSaver {
    constructor(serde) {
        super(serde);
        // thread ID ->  checkpoint namespace -> checkpoint ID -> checkpoint mapping
        Object.defineProperty(this, "storage", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: {}
        });
        Object.defineProperty(this, "writes", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: {}
        });
    }
    async _getPendingSends(threadId, checkpointNs, parentCheckpointId) {
        let pendingSends = [];
        if (parentCheckpointId !== undefined) {
            const key = _generateKey(threadId, checkpointNs, parentCheckpointId);
            pendingSends = await Promise.all(Object.values(this.writes[key] || {})
                ?.filter(([_taskId, channel]) => {
                return channel === TASKS;
            })
                .map(([_taskId, _channel, writes]) => {
                return this.serde.loadsTyped("json", writes);
            }) ?? []);
        }
        return pendingSends;
    }
    async getTuple(config) {
        const thread_id = config.configurable?.thread_id;
        const checkpoint_ns = config.configurable?.checkpoint_ns ?? "";
        let checkpoint_id = getCheckpointId(config);
        if (checkpoint_id) {
            const saved = this.storage[thread_id]?.[checkpoint_ns]?.[checkpoint_id];
            if (saved !== undefined) {
                const [checkpoint, metadata, parentCheckpointId] = saved;
                const key = _generateKey(thread_id, checkpoint_ns, checkpoint_id);
                const pending_sends = await this._getPendingSends(thread_id, checkpoint_ns, parentCheckpointId);
                const deserializedCheckpoint = {
                    ...(await this.serde.loadsTyped("json", checkpoint)),
                    pending_sends,
                };
                const pendingWrites = await Promise.all(Object.values(this.writes[key] || {}).map(async ([taskId, channel, value]) => {
                    return [
                        taskId,
                        channel,
                        await this.serde.loadsTyped("json", value),
                    ];
                }));
                const checkpointTuple = {
                    config,
                    checkpoint: deserializedCheckpoint,
                    metadata: (await this.serde.loadsTyped("json", metadata)),
                    pendingWrites,
                };
                if (parentCheckpointId !== undefined) {
                    checkpointTuple.parentConfig = {
                        configurable: {
                            thread_id,
                            checkpoint_ns,
                            checkpoint_id: parentCheckpointId,
                        },
                    };
                }
                return checkpointTuple;
            }
        }
        else {
            const checkpoints = this.storage[thread_id]?.[checkpoint_ns];
            if (checkpoints !== undefined) {
                // eslint-disable-next-line prefer-destructuring
                checkpoint_id = Object.keys(checkpoints).sort((a, b) => b.localeCompare(a))[0];
                const saved = checkpoints[checkpoint_id];
                const [checkpoint, metadata, parentCheckpointId] = saved;
                const key = _generateKey(thread_id, checkpoint_ns, checkpoint_id);
                const pending_sends = await this._getPendingSends(thread_id, checkpoint_ns, parentCheckpointId);
                const deserializedCheckpoint = {
                    ...(await this.serde.loadsTyped("json", checkpoint)),
                    pending_sends,
                };
                const pendingWrites = await Promise.all(Object.values(this.writes[key] || {}).map(async ([taskId, channel, value]) => {
                    return [
                        taskId,
                        channel,
                        await this.serde.loadsTyped("json", value),
                    ];
                }));
                const checkpointTuple = {
                    config: {
                        configurable: {
                            thread_id,
                            checkpoint_id,
                            checkpoint_ns,
                        },
                    },
                    checkpoint: deserializedCheckpoint,
                    metadata: (await this.serde.loadsTyped("json", metadata)),
                    pendingWrites,
                };
                if (parentCheckpointId !== undefined) {
                    checkpointTuple.parentConfig = {
                        configurable: {
                            thread_id,
                            checkpoint_ns,
                            checkpoint_id: parentCheckpointId,
                        },
                    };
                }
                return checkpointTuple;
            }
        }
        return undefined;
    }
    async *list(config, options) {
        // eslint-disable-next-line prefer-const
        let { before, limit, filter } = options ?? {};
        const threadIds = config.configurable?.thread_id
            ? [config.configurable?.thread_id]
            : Object.keys(this.storage);
        const configCheckpointNamespace = config.configurable?.checkpoint_ns;
        const configCheckpointId = config.configurable?.checkpoint_id;
        for (const threadId of threadIds) {
            for (const checkpointNamespace of Object.keys(this.storage[threadId] ?? {})) {
                if (configCheckpointNamespace !== undefined &&
                    checkpointNamespace !== configCheckpointNamespace) {
                    continue;
                }
                const checkpoints = this.storage[threadId]?.[checkpointNamespace] ?? {};
                const sortedCheckpoints = Object.entries(checkpoints).sort((a, b) => b[0].localeCompare(a[0]));
                for (const [checkpointId, [checkpoint, metadataStr, parentCheckpointId],] of sortedCheckpoints) {
                    // Filter by checkpoint ID from config
                    if (configCheckpointId && checkpointId !== configCheckpointId) {
                        continue;
                    }
                    // Filter by checkpoint ID from before config
                    if (before &&
                        before.configurable?.checkpoint_id &&
                        checkpointId >= before.configurable.checkpoint_id) {
                        continue;
                    }
                    // Parse metadata
                    const metadata = (await this.serde.loadsTyped("json", metadataStr));
                    if (filter &&
                        !Object.entries(filter).every(([key, value]) => metadata[key] === value)) {
                        continue;
                    }
                    // Limit search results
                    if (limit !== undefined) {
                        if (limit <= 0)
                            break;
                        limit -= 1;
                    }
                    const key = _generateKey(threadId, checkpointNamespace, checkpointId);
                    const writes = Object.values(this.writes[key] || {});
                    const pending_sends = await this._getPendingSends(threadId, checkpointNamespace, parentCheckpointId);
                    const pendingWrites = await Promise.all(writes.map(async ([taskId, channel, value]) => {
                        return [
                            taskId,
                            channel,
                            await this.serde.loadsTyped("json", value),
                        ];
                    }));
                    const deserializedCheckpoint = {
                        ...(await this.serde.loadsTyped("json", checkpoint)),
                        pending_sends,
                    };
                    const checkpointTuple = {
                        config: {
                            configurable: {
                                thread_id: threadId,
                                checkpoint_ns: checkpointNamespace,
                                checkpoint_id: checkpointId,
                            },
                        },
                        checkpoint: deserializedCheckpoint,
                        metadata,
                        pendingWrites,
                    };
                    if (parentCheckpointId !== undefined) {
                        checkpointTuple.parentConfig = {
                            configurable: {
                                thread_id: threadId,
                                checkpoint_ns: checkpointNamespace,
                                checkpoint_id: parentCheckpointId,
                            },
                        };
                    }
                    yield checkpointTuple;
                }
            }
        }
    }
    async put(config, checkpoint, metadata) {
        const preparedCheckpoint = copyCheckpoint(checkpoint);
        delete preparedCheckpoint.pending_sends;
        const threadId = config.configurable?.thread_id;
        const checkpointNamespace = config.configurable?.checkpoint_ns ?? "";
        if (threadId === undefined) {
            throw new Error(`Failed to put checkpoint. The passed RunnableConfig is missing a required "thread_id" field in its "configurable" property.`);
        }
        if (!this.storage[threadId]) {
            this.storage[threadId] = {};
        }
        if (!this.storage[threadId][checkpointNamespace]) {
            this.storage[threadId][checkpointNamespace] = {};
        }
        const [, serializedCheckpoint] = this.serde.dumpsTyped(preparedCheckpoint);
        const [, serializedMetadata] = this.serde.dumpsTyped(metadata);
        this.storage[threadId][checkpointNamespace][checkpoint.id] = [
            serializedCheckpoint,
            serializedMetadata,
            config.configurable?.checkpoint_id, // parent
        ];
        return {
            configurable: {
                thread_id: threadId,
                checkpoint_ns: checkpointNamespace,
                checkpoint_id: checkpoint.id,
            },
        };
    }
    async putWrites(config, writes, taskId) {
        const threadId = config.configurable?.thread_id;
        const checkpointNamespace = config.configurable?.checkpoint_ns;
        const checkpointId = config.configurable?.checkpoint_id;
        if (threadId === undefined) {
            throw new Error(`Failed to put writes. The passed RunnableConfig is missing a required "thread_id" field in its "configurable" property`);
        }
        if (checkpointId === undefined) {
            throw new Error(`Failed to put writes. The passed RunnableConfig is missing a required "checkpoint_id" field in its "configurable" property.`);
        }
        const outerKey = _generateKey(threadId, checkpointNamespace, checkpointId);
        const outerWrites_ = this.writes[outerKey];
        if (this.writes[outerKey] === undefined) {
            this.writes[outerKey] = {};
        }
        writes.forEach(([channel, value], idx) => {
            const [, serializedValue] = this.serde.dumpsTyped(value);
            const innerKey = [
                taskId,
                WRITES_IDX_MAP[channel] || idx,
            ];
            const innerKeyStr = `${innerKey[0]},${innerKey[1]}`;
            if (innerKey[1] >= 0 && outerWrites_ && innerKeyStr in outerWrites_) {
                return;
            }
            this.writes[outerKey][innerKeyStr] = [taskId, channel, serializedValue];
        });
    }
}
//# sourceMappingURL=memory.js.map