import { CheckpointError } from "../errors.js";
import { deepClone } from "../utils.js";
import { checkpointSchema } from "./schemas.js";
/** Label used for pinned checkpoints created at function entry for Result error handling. */
export const RESULT_ENTRY_LABEL = "result-entry";
let globalCheckpointCounter = 0;
/** Reset the global checkpoint ID counter. For use in tests only. */
export function resetGlobalCheckpointCounter() {
    globalCheckpointCounter = 0;
}
export class Checkpoint {
    id;
    stack;
    globals;
    nodeId;
    moduleId;
    scopeName;
    stepPath;
    label;
    pinned;
    constructor(args) {
        this.id = args.id ?? globalCheckpointCounter++;
        this.stack = args.stack;
        this.globals = args.globals;
        this.nodeId = args.nodeId;
        this.moduleId = args.moduleId ?? "";
        this.scopeName = args.scopeName ?? "";
        this.stepPath = args.stepPath ?? "";
        this.label = args.label ?? null;
        this.pinned = args.pinned ?? false;
    }
    getScopeKey() {
        return `${this.moduleId}:${this.scopeName}`;
    }
    get location() {
        return {
            nodeId: this.nodeId,
            moduleId: this.moduleId,
            scopeName: this.scopeName,
            stepPath: this.stepPath,
        };
    }
    getCurrentFrame() {
        return this.stack.stack?.at(-1);
    }
    getThreadMessages(displayIndex) {
        const frame = this.getCurrentFrame();
        if (!frame?.threads)
            return null;
        const { threads, activeStack } = frame.threads;
        const threadIds = Object.keys(threads);
        if (threadIds.length === 0)
            return null;
        let selectedId;
        let selectedIndex;
        if (displayIndex !== undefined) {
            // Explicit index requested (for thread cycling)
            selectedIndex = ((displayIndex % threadIds.length) + threadIds.length) % threadIds.length;
            selectedId = threadIds[selectedIndex];
        }
        else {
            // Default: prefer the active thread at the top of the stack
            selectedId =
                activeStack.findLast((id) => threads[id] != null) ?? threadIds[0];
            selectedIndex = threadIds.indexOf(selectedId);
        }
        const selectedThread = threads[selectedId];
        if (!selectedThread)
            return null;
        const messages = selectedThread.messages.map((m) => ({
            role: m.role ?? "unknown",
            content: this.getContentFromMessage(m),
        }));
        return { threadId: selectedId, threadIndex: selectedIndex, threadCount: threadIds.length, messages };
    }
    getContentFromMessage(message) {
        if (typeof message.content === "string") {
            return message.content;
        }
        else if (Array.isArray(message.content)) {
            return message.content.map((part) => part.text ?? "").join("");
        }
        else if (message.content == null) {
            if (message.role === "assistant" && message.toolCalls) {
                return message.toolCalls
                    .map((toolCall) => `Tool call: ${toolCall.name}(${JSON.stringify(toolCall.arguments)})`)
                    .join("\n");
            }
            return "(no content)";
        }
        return JSON.stringify(message.content);
    }
    getGlobalsForModule() {
        return this.globals.store?.[this.moduleId] ?? null;
    }
    getFilename() {
        return this.moduleId.split("/").pop() || this.moduleId;
    }
    pathEquals(other) {
        return (this.moduleId === other.moduleId &&
            this.scopeName === other.scopeName &&
            this.stepPath === other.stepPath);
    }
    equals(other) {
        return JSON.stringify(this.toJSON()) === JSON.stringify(other.toJSON());
    }
    toJSON() {
        return {
            id: this.id,
            stack: this.stack,
            globals: this.globals,
            nodeId: this.nodeId,
            moduleId: this.moduleId,
            scopeName: this.scopeName,
            stepPath: this.stepPath,
            label: this.label,
            pinned: this.pinned,
        };
    }
    clone(opts = {}) {
        return Checkpoint.fromJSON({ ...this.toJSON(), ...opts });
    }
    getLocation() {
        return `${this.moduleId}:${this.scopeName}#${this.stepPath}`;
    }
    static fromStateStack(stateStack, ctx, opts) {
        const nodeId = ctx.stateStack.currentNodeId();
        if (!nodeId) {
            throw new CheckpointError("Cannot create checkpoint: no current node id in state stack. This error can happen if you call a function that throws an interrupt from the global namespace. Please use `const foo = funcName() with approve` syntax.");
        }
        return new Checkpoint({
            stack: stateStack.toJSON(),
            globals: ctx.globals.toJSON(),
            nodeId,
            ...opts,
        });
    }
    static fromContext(ctx, opts) {
        return Checkpoint.fromStateStack(ctx.stateStack, ctx, opts);
    }
    static fromJSON(json) {
        if (json instanceof Checkpoint) {
            return json;
        }
        const parsed = checkpointSchema.safeParse(json);
        if (!parsed.success) {
            return null;
        }
        return new Checkpoint(parsed.data);
    }
}
export class CheckpointStore {
    checkpoints = {};
    restoreCounts = {};
    locationRestoreCounts = {};
    maxRestores;
    maxSize;
    constructor(maxRestores = 100, maxSize = -1) {
        this.maxRestores = maxRestores;
        this.maxSize = maxSize;
    }
    create(stateStack, ctx, opts) {
        if (opts.removeDuplicate) {
            this.removeDuplicate(opts);
        }
        const checkpoint = Checkpoint.fromStateStack(stateStack, ctx, opts);
        this.checkpoints[checkpoint.id] = checkpoint;
        return checkpoint.id;
    }
    add(checkpoint) {
        this.checkpoints[checkpoint.id] = checkpoint;
        if (checkpoint.id >= globalCheckpointCounter) {
            globalCheckpointCounter = checkpoint.id + 1;
        }
    }
    cloneCheckpoint(_checkpoint, opts = {}) {
        const checkpoint = Checkpoint.fromJSON(_checkpoint);
        if (!checkpoint) {
            throw new CheckpointError("Invalid checkpoint provided for cloning.");
        }
        const id = globalCheckpointCounter++;
        const newCheckpoint = checkpoint.clone({ ...opts, id });
        this.checkpoints[id] = newCheckpoint;
        return id;
    }
    findCheckpoint(location) {
        const sorted = this.getSorted(); // sorted by id ascending
        const currentIdx = sorted.findIndex((cp) => cp.moduleId === location.moduleId &&
            cp.scopeName === location.scopeName &&
            cp.stepPath === location.stepPath);
        if (currentIdx <= 0)
            return null;
        return sorted[currentIdx - 1];
    }
    findBefore(checkpoint) {
        const sorted = this.getSorted(); // sorted by id ascending
        const checkpointId = checkpoint.id;
        sorted.reverse();
        for (const cp of sorted) {
            if (cp.id < checkpointId && !cp.pathEquals(checkpoint) && !cp.pinned) {
                return cp;
            }
        }
        return null;
    }
    createRolling(ctx, opts) {
        // Remove existing unpinned checkpoint at the same location to avoid duplicates
        const id = this.create(ctx.stateStack, ctx, {
            ...opts,
            label: null,
            pinned: false,
            removeDuplicate: true,
        });
        this.removeDebugFlagsFor(id);
        this.evictIfNeeded();
        return id;
    }
    removeDebugFlagsFor(id) {
        const cp = this.checkpoints[id];
        if (!cp)
            return;
        const frame = cp.getCurrentFrame();
        if (!frame || !frame.locals)
            return;
        for (const key of Object.keys(frame.locals)) {
            if (key.startsWith("__dbg_")) {
                delete frame.locals[key];
            }
        }
    }
    removeDuplicate(location) {
        for (const [id, cp] of Object.entries(this.checkpoints)) {
            if (!cp.pinned &&
                cp.moduleId === location.moduleId &&
                cp.scopeName === location.scopeName &&
                cp.stepPath === location.stepPath) {
                delete this.checkpoints[Number(id)];
                return;
            }
        }
    }
    createPinned(stateStack, ctx, opts) {
        return this.create(stateStack, ctx, { ...opts, pinned: true });
    }
    pin(id, label) {
        const cp = this.checkpoints[id];
        if (!cp)
            return;
        cp.pinned = true;
        if (label !== undefined)
            cp.label = label;
    }
    get(id) {
        const cp = this.checkpoints[id];
        return cp;
    }
    delete(id) {
        delete this.checkpoints[id];
        delete this.restoreCounts[id];
    }
    getCheckpoints() {
        return Object.values(this.checkpoints).filter(Boolean);
    }
    // Invalidate all checkpoints with id > the given id
    deleteAfterCheckpoint(id, deletePinned = false) {
        for (const key of Object.keys(this.checkpoints)) {
            const numericKey = Number(key);
            if (numericKey > id) {
                if (!deletePinned && this.checkpoints[numericKey].pinned) {
                    continue;
                }
                delete this.checkpoints[numericKey];
                delete this.restoreCounts[numericKey];
            }
        }
    }
    getLocationRestoreCount(location) {
        return this.locationRestoreCounts[location] || 0;
    }
    trackLocationRestore(location) {
        this.locationRestoreCounts[location] = (this.locationRestoreCounts[location] || 0) + 1;
    }
    trackRestore(id) {
        this.restoreCounts[id] = (this.restoreCounts[id] || 0) + 1;
        if (this.restoreCounts[id] > this.maxRestores) {
            throw new CheckpointError(`Checkpoint ${id} has been restored ${this.maxRestores} times. Possible infinite loop.`);
        }
    }
    getSorted() {
        return Object.values(this.checkpoints).sort((a, b) => a.id - b.id);
    }
    prettyPrint() {
        const cps = this.getSorted()
            .map((cp) => `Checkpoint ${cp.id}: nodeId=${cp.nodeId}, moduleId=${cp.moduleId}, scopeName=${cp.scopeName}, stepPath=${cp.stepPath}, label=${cp.label}, pinned=${cp.pinned}`)
            .join("\n");
        return cps || "No checkpoints";
    }
    toJSON() {
        return {
            checkpoints: deepClone(this.checkpoints),
            counter: globalCheckpointCounter,
        };
    }
    static fromJSON(json, maxRestores = 100) {
        const store = new CheckpointStore(maxRestores);
        const rehydrated = {};
        for (const [id, cpJson] of Object.entries(json.checkpoints)) {
            const checkpoint = Checkpoint.fromJSON(cpJson);
            if (checkpoint) {
                rehydrated[Number(id)] = checkpoint;
            }
            else {
                console.warn(`Failed to rehydrate checkpoint with id ${id}, skipping. Data: ${JSON.stringify(cpJson)}`);
            }
        }
        store.checkpoints = rehydrated;
        globalCheckpointCounter = Math.max(globalCheckpointCounter, json.counter);
        return store;
    }
    evictIfNeeded() {
        if (this.maxSize < 0)
            return; // no max size limit
        const unpinned = Object.values(this.checkpoints)
            .filter((cp) => !cp.pinned)
            .sort((a, b) => a.id - b.id);
        while (unpinned.length > this.maxSize) {
            const oldest = unpinned.shift();
            delete this.checkpoints[oldest.id];
        }
    }
}
