import { nanoid } from "nanoid";
import { SimpleMachine } from "../../simplemachine/index.js";
import { StatelogClient } from "../../statelogClient.js";
import { nativeTypeReplacer, nativeTypeReviver } from "../revivers/index.js";
import { CoverageCollector } from "../coverageCollector.js";
import { AgencyCancelledError } from "../errors.js";
import { SmoltalkClient } from "../llmClient.js";
import { MemoryManager } from "../memory/manager.js";
import { MemoryFrame } from "../memory/frame.js";
import { getOrCreateStore } from "../memory/registry.js";
import { GlobalStore } from "../state/globalStore.js";
import { StateStack } from "../state/stateStack.js";
import { TraceWriter } from "../trace/traceWriter.js";
import { CheckpointStore, RESULT_ENTRY_LABEL } from "./checkpointStore.js";
import { PendingPromiseStore } from "./pendingPromiseStore.js";
/**
 * Process-wide singleton CoverageCollector used when AGENCY_COVERAGE is set.
 * One collector per process — multiple RuntimeContext instances created within
 * the same process all share it, and a single `process.on("exit")` listener
 * writes the merged data once. AGENCY_COVERAGE_OUTDIR is set by the CLI from
 * agency.json's coverage.outDir; defaults to ".coverage" otherwise.
 */
let _processCoverageCollector = null;
function getProcessCoverageCollector() {
    if (_processCoverageCollector)
        return _processCoverageCollector;
    const collector = new CoverageCollector();
    _processCoverageCollector = collector;
    const outDir = process.env.AGENCY_COVERAGE_OUTDIR ?? ".coverage";
    process.on("exit", () => {
        // process.on("exit") handlers are sync-only and any throw is unhelpful at
        // shutdown — log a warning and move on so coverage failures never make
        // test runs flaky or noisy.
        try {
            collector.write(outDir);
        }
        catch (err) {
            console.warn(`[coverage] failed to write to ${outDir}: ${err.message}`);
        }
    });
    return collector;
}
/**
 * Round-trip `data` through JSON using the native-type replacer/reviver pair
 * so Sets/Maps/Dates/etc. inside a checkpoint payload come back as their
 * real classes instead of plain `{}` / `[]`.
 */
function reviveNative(data) {
    return JSON.parse(JSON.stringify(data, nativeTypeReplacer), nativeTypeReviver);
}
/* bunch of stuff that every node/function in the runtime needs access to,
that we don't want to pass as individual arguments everywhere */
export class RuntimeContext {
    // this is the part of the runtime context that gets
    // serialized/deserialized to support durable execution
    stateStack;
    globals;
    checkpoints;
    callbacks;
    onStreamLock;
    handlers;
    pendingPromises;
    graph;
    _skipNextCheckpoint;
    _pendingArgOverrides;
    _restoreCount;
    /* Here is why this is needed: When you're stepping through the code,
    every step emits a checkpoint and halts execution. When you execute an
    LLM call that then calls a tool, that tool will also emit a checkpoint
    and halt execution. At that point, that interrupt is going to bubble up
    to run prompt, and it's going to look like the tool call threw an
    interrupt, and then everything's going to get messed up from there.
    Longer term, it would be great for the debugger to show what's happening
    as the tool call is getting executed, but for now, we just use this flag
    so that if we're inside a tool call, we don't halt execution and we don't
    emit a checkpoint.*/
    _toolCallDepth;
    /** Nested-dispatch depth counter for `runHandlerChain` (lib/runtime/interrupts.ts).
     *  Incremented on entry to each handler-chain dispatch, decremented in its
     *  `finally`. The dispatcher throws a `HandlerRecursionError` when this counter
     *  exceeds `MAX_HANDLER_CHAIN_DEPTH` — that means a handler raised an interrupt
     *  whose dispatch re-entered the same chain too many times, almost always a
     *  handler accidentally re-invoking itself (see the recursion case debugged in
     *  https://ampcode.com/threads/T-019e7a80-0a51-75ce-840e-89b5f595da5c).
     *
     *  This is NOT the length of `ctx.handlers`. It only grows when a handler's body
     *  triggers another interrupt while the outer dispatch is still in flight. Normal
     *  push/pop of handlers via `with approve` does not increment it. */
    _handlerChainDepth;
    debuggerState;
    traceWriter;
    // we need a single statelog client instance that can be used across the entire execution of the graph,
    // so that all the logs share the same traceId, so they all show up in the same trace in the Statelog dashboard.
    statelogClient;
    smoltalkDefaults;
    _llmClient;
    _interruptResponses = {};
    get llmClient() { return this._llmClient; }
    setLLMClient(client) { this._llmClient = client; }
    setInterruptResponses(responses) {
        this._interruptResponses = responses;
    }
    getInterruptResponse(interruptId) {
        return this._interruptResponses[interruptId]?.response;
    }
    // this is the directory that the runtime is running in. We need this to be able to read files relative to the runtime.
    dirname;
    /** Callbacks registered at module top-level (via `_callback` during
     *  `__initializeGlobals`, when no real caller frame exists yet). Persist for
     *  the whole run. */
    topLevelCallbacks = [];
    abortController;
    traceConfig;
    runId;
    verbose;
    /**
     * Log threshold used by ad-hoc subsystem loggers (memory, etc.).
     * Plumbed in from `AgencyConfig.logLevel` so users can crank it up to
     * `"debug"` when investigating issues without recompiling. The logger
     * itself is stateless — consumers call `createLogger(execCtx.logLevel)`
     * on demand rather than carrying a long-lived Logger instance, which
     * lets each subsystem add its own prefix without coordinating an
     * instance pool.
     */
    logLevel;
    getStaticVars;
    coverageCollector = null;
    // stored so createExecutionContext can create new StatelogClients
    statelogConfig;
    maxRestores;
    // Memory layer (resolved decisions in
    // docs/superpowers/plans/2026-05-12-memory-layer.md and
    // docs/superpowers/plans/2026-05-29-memory-config-in-code.md).
    //
    // `jsonMemoryConfig` is the immutable seed from `agency.json`'s
    // `memory:` block. Each execCtx pushes it as the bottom frame on
    // its stateStack at construction so the active config is always
    // "top of stack or nothing." Code calls to `enableMemory(...)`
    // push additional frames on top. The "code wins over JSON" rule
    // is structural, not coded.
    //
    // Per-execCtx managers are cached by `configKey` so we don't
    // rebuild a heavy MemoryManager every time stdlib code looks up
    // the active manager. Cache lives on each execCtx (not on this
    // parent context) because manager construction takes the
    // execCtx's own statelog client / log level / llm client.
    jsonMemoryConfig;
    memoryManagerCache = {};
    constructor(args) {
        const statelogConfig = {
            ...args.statelogConfig,
            traceId: args.statelogConfig.traceId || nanoid(),
        };
        this.statelogConfig = statelogConfig;
        this.maxRestores = args.maxRestores ?? 100;
        this.statelogClient = new StatelogClient(statelogConfig);
        this.stateStack = new StateStack();
        this.globals = GlobalStore.withTokenStats();
        this.checkpoints = new CheckpointStore(this.maxRestores);
        this.handlers = [];
        this.callbacks = {};
        this.onStreamLock = false;
        // After a debugger rewind, the first debug step would write a duplicate
        // checkpoint to the trace (the user already saw the rewound checkpoint).
        // rewindFrom sets this flag so the first debugStep trace-write is skipped.
        this._skipNextCheckpoint = false;
        this._restoreCount = 0;
        this._toolCallDepth = 0;
        this._handlerChainDepth = 0;
        this.pendingPromises = new PendingPromiseStore();
        this.debuggerState = null;
        this.traceWriter = null;
        this.traceConfig = args.traceConfig || {};
        this.runId = null;
        this.verbose = args.verbose ?? false;
        this.logLevel = args.logLevel ?? "info";
        this.dirname = args.dirname;
        const graphConfig = {
            debug: {
                log: false,
                logData: false,
            },
            statelog: statelogConfig,
        };
        this.graph = new SimpleMachine(graphConfig);
        this.smoltalkDefaults = args.smoltalkDefaults;
        this._llmClient = new SmoltalkClient();
        this.abortController = new AbortController();
        if (process.env.AGENCY_COVERAGE) {
            this.coverageCollector = getProcessCoverageCollector();
        }
        // JSON-derived memory config is kept as an immutable seed; the
        // actual store + manager are looked up lazily through
        // `getActiveMemoryManager()` on each execCtx so there's a single
        // source of truth (the active stateStack's frame stack).
        if (args.memory) {
            this.jsonMemoryConfig = args.memory;
        }
    }
    getRunId() {
        if (!this.runId) {
            throw new Error("runId not set on RuntimeContext");
        }
        return this.runId;
    }
    async createExecutionContext(runId) {
        const execCtx = Object.create(RuntimeContext.prototype);
        execCtx.graph = this.graph;
        execCtx.smoltalkDefaults = this.smoltalkDefaults;
        execCtx._llmClient = this._llmClient;
        execCtx.dirname = this.dirname;
        execCtx.statelogConfig = this.statelogConfig;
        execCtx.stateStack = new StateStack();
        execCtx.globals = GlobalStore.withTokenStats();
        execCtx.maxRestores = this.maxRestores;
        execCtx.checkpoints = new CheckpointStore(this.maxRestores);
        execCtx.handlers = [];
        execCtx.callbacks = {};
        execCtx.topLevelCallbacks = [];
        execCtx.onStreamLock = false;
        execCtx._skipNextCheckpoint = false;
        execCtx._restoreCount = 0;
        execCtx._toolCallDepth = 0;
        execCtx._handlerChainDepth = 0;
        execCtx._interruptResponses = {};
        execCtx.debuggerState = this.debuggerState;
        execCtx.traceWriter = await TraceWriter.create({
            runId,
            traceConfig: this.traceConfig,
        });
        execCtx.traceConfig = this.traceConfig;
        execCtx.runId = runId;
        execCtx.verbose = this.verbose;
        execCtx.logLevel = this.logLevel;
        execCtx.pendingPromises = new PendingPromiseStore();
        execCtx.abortController = new AbortController();
        execCtx.statelogClient = new StatelogClient({
            ...this.statelogConfig,
            traceId: runId,
        });
        execCtx.coverageCollector = this.coverageCollector;
        // Memory layer: per-execCtx state.
        //   - jsonMemoryConfig is forwarded so `getActiveMemoryManager()`
        //     can re-seed the bottom frame on resumes from old
        //     checkpoints that pre-date frames-on-stateStack.
        //   - memoryManagerCache is per-execCtx (statelog client, llm
        //     client, log level are all per-execCtx).
        //   - JSON config seeds the BOTTOM frame on the new stateStack;
        //     after this point the active stack is the only source of
        //     truth. Code-level `enableMemory(...)` pushes on top.
        execCtx.jsonMemoryConfig = this.jsonMemoryConfig;
        execCtx.memoryManagerCache = {};
        if (this.jsonMemoryConfig) {
            execCtx.stateStack.pushMemoryFrame(new MemoryFrame(this.jsonMemoryConfig));
        }
        return execCtx;
    }
    /**
     * Resolve the active `MemoryManager` for the current stateStack.
     *
     * Single rule: top of `stateStack.other.memoryFrames` (accessed via
     * `activeMemoryFrame()`) wins. No fallback to a "default" manager —
     * the JSON config gets seeded as the bottom frame at execCtx
     * creation, so if JSON was set there will always be a frame and
     * this never returns `undefined` for JSON-only setups.
     *
     * The one back-compat seam: a checkpoint written before
     * `memoryFrames` existed (or one taken right after a user called
     * `disableMemory()` on the JSON-seeded bottom frame) has no frames.
     * If `jsonMemoryConfig` is set we re-seed it as a courtesy so old
     * traces resume as before. If the user explicitly popped, they
     * call `enableMemory(...)` again to turn it back on.
     *
     * Managers are cached per execCtx keyed by `configKey` so
     * push/pop/push of the same dir reuses one instance — important
     * for the "pop back to A returns A's manager" semantics.
     */
    getActiveMemoryManager() {
        if (!this.stateStack)
            return undefined;
        let frame = this.stateStack.activeMemoryFrame();
        if (!frame && this.jsonMemoryConfig && !this.stateStack.hasMemoryFrameStack()) {
            // Old-checkpoint back-compat: stateStack restored from a
            // pre-memoryFrames snapshot (the array key isn't present at
            // all). Re-seed the JSON bottom frame so resume behaves like
            // a fresh run. An empty-array stack (user explicitly called
            // `disableMemory()`) is NOT re-seeded — that would silently
            // resurrect what the user asked to turn off.
            this.stateStack.pushMemoryFrame(new MemoryFrame(this.jsonMemoryConfig));
            frame = this.stateStack.activeMemoryFrame();
        }
        if (!frame)
            return undefined;
        const cached = this.memoryManagerCache[frame.configKey];
        if (cached)
            return cached;
        const manager = new MemoryManager({
            store: getOrCreateStore(frame.configKey, this.logLevel),
            config: frame.config,
            llmClient: this._llmClient,
            smoltalkDefaults: this.smoltalkDefaults,
            source: this.traceConfig?.program ?? "agent",
            // Reuse the per-execCtx StatelogClient so memory's own LLM/embed
            // spans nest under the same trace as the agent's calls.
            statelogClient: this.statelogClient,
            // Threshold for memory's internal logger; promoting this to
            // "debug" in agency.json surfaces every tier/extract/compact
            // step on stderr.
            logLevel: this.logLevel,
            memoryIdRef: {
                // memoryId is orthogonal to which frame is active — it lives
                // on `stateStack.other.memoryId` and persists across frame
                // pushes/pops. We capture `this` (the execCtx) so the ref
                // always sees the current stateStack (which can be replaced
                // by `restoreState`).
                get: () => {
                    const id = this.stateStack?.other?.memoryId;
                    return typeof id === "string" ? id : "default";
                },
                set: (id) => {
                    if (!this.stateStack)
                        return;
                    this.stateStack.other.memoryId = id;
                },
            },
        });
        this.memoryManagerCache[frame.configKey] = manager;
        return manager;
    }
    /** Iterate every cached memory manager. Used at shutdown so a
     *  branch that opened a side store doesn't lose its writes. */
    getAllCachedMemoryManagers() {
        return Object.values(this.memoryManagerCache);
    }
    /* Let's chat through what's going on here. Because since this function
    is called "fork"Stack, you may think that it clone the current stateStack.
    And that's exactly what I had earlier:
  
    ```
    return StateStack.fromJSON(this.stateStack.toJSON());
    ```
  
    This function was created for asynchronous threads, so we could keep track
    of their state. But the way we store the state is for each async thread,
    in branches. It's execution starts at the point it is defined. It doesn't
    have any previous state it's not going to wind up to back to a point *before*
    it was defined.
    Another way to think about it: Suppose node main, calls function foo,
    calls function A, which creates an async thread function A1, and A1 throws an interrupt:
  
    ```
      main -> foo -> A -> A1 (async thread) -> interrupt
    ```
  
    When we resume from the interrupt, we deserialize up to A1 using A's state stack.
    And from that point, we deserialize A1 using A1's state stack. So A1's state stack
    should *only* contain the state created after the A1 thread was initialized.
    Otherwise, there will be a bunch of extra frames on the state stack related
    to calling through our A1, which will be a mismatch.
  
    I'm still leaving this function because I think its name helps explain what
    is happening, and also because it's easier to find the code that creates these
    async threads. But we could just replace all calls to `forkStack`
    with `new StateStack()`.
    */
    pushHandler(fn) {
        this.handlers.push(fn);
    }
    popHandler() {
        this.handlers.pop();
    }
    enterToolCall() {
        this._toolCallDepth++;
    }
    exitToolCall() {
        if (this._toolCallDepth > 0)
            this._toolCallDepth--;
    }
    isInsideToolCall() {
        return this._toolCallDepth > 0;
    }
    get aborted() {
        return this.abortController.signal.aborted;
    }
    /**
     * Branch-aware cancellation check. Returns true if either:
     *   - the global ctx is aborted (e.g., user pressed Ctrl-C), OR
     *   - the given branch stack's per-branch abort signal has fired (e.g.,
     *     this branch is a race loser).
     *
     * Pass the local stateStack at any call site that lives inside a fork/race
     * branch so the check sees that branch's signal. Without a stack arg, this
     * is equivalent to `ctx.aborted`.
     */
    isCancelled(stack) {
        if (this.abortController.signal.aborted)
            return true;
        return !!stack?.abortSignal?.aborted;
    }
    /**
     * Branch-aware AbortSignal for HTTP/fetch/streaming calls. Returns a
     * composite signal that fires on either global ctx abort OR the given
     * branch stack's abort. Pass to smoltalk's `abortSignal` so per-branch
     * cancellation actually tears down in-flight network requests.
     *
     * If no stack is given (or the stack has no branch signal), returns the
     * global ctx signal — same behavior as before.
     */
    getAbortSignal(stack) {
        if (!stack?.abortSignal)
            return this.abortController.signal;
        return AbortSignal.any([this.abortController.signal, stack.abortSignal]);
    }
    cancel(reason) {
        if (!this.abortController.signal.aborted) {
            this.abortController.abort(new AgencyCancelledError(reason));
        }
    }
    forkStack() {
        return new StateStack();
    }
    /** Sever references held by an execution context so GC can reclaim them. */
    cleanup() {
        if (!this.abortController.signal.aborted) {
            this.abortController.abort(new AgencyCancelledError("cleanup"));
        }
        this.pendingPromises.clear();
        this.stateStack = null;
        this.globals = null;
        this.checkpoints = null;
        this.statelogClient = null;
        this.callbacks = null;
        this.handlers = null;
        this.traceWriter = null;
        this.memoryManagerCache = {};
    }
    /** Get the most recent result-entry checkpoint for the current function. */
    getResultCheckpoint() {
        const sorted = this.checkpoints.getSorted();
        for (let i = sorted.length - 1; i >= 0; i--) {
            if (sorted[i].pinned && sorted[i].label === RESULT_ENTRY_LABEL) {
                return sorted[i];
            }
        }
        return undefined;
    }
    restoreState(checkpoint) {
        const currentTokenStats = this.globals.getTokenStats();
        const stack = reviveNative(checkpoint.stack);
        const globals = reviveNative(checkpoint.globals);
        this.stateStack = StateStack.fromJSON(stack);
        this.stateStack.deserializeMode();
        this.globals = GlobalStore.fromJSON(globals);
        this.globals.restoreTokenStats(currentTokenStats);
        this.pendingPromises.clear();
    }
    /** @deprecated Use checkpoints.create() instead */
    stateToJSON() {
        return {
            stack: this.stateStack.toJSON(),
            globals: this.globals.toJSON(),
        };
    }
    toJSON() {
        return {
            stateStack: this.stateStack.toJSON(),
            callbacks: Object.keys(this.callbacks),
            onStreamLock: this.onStreamLock,
            graph: this.graph.toJSON(),
            statelogClient: "redacted",
            smoltalkDefaults: "redacted",
            dirname: this.dirname,
        };
    }
    /* Get smoltalk config with missing keys populated with defaults */
    getSmoltalkConfig(config = {}) {
        return { ...this.smoltalkDefaults, ...config };
    }
    async pauseTraceWriter() {
        if (!this.traceWriter) {
            //throw new Error("No trace writer to pause");
        }
        await this.traceWriter?.pause();
        this.traceWriter = null;
    }
    async closeTraceWriter() {
        if (!this.traceWriter) {
            //throw new Error("No trace writer to close");
        }
        await this.traceWriter?.close();
        this.traceWriter = null;
    }
    async writeCheckpointToTraceWriter(checkpoint) {
        if (!this.traceWriter) {
            return;
        }
        await this.traceWriter.writeCheckpoint(checkpoint);
    }
    async writeStaticStateToTrace(values) {
        if (!this.traceWriter) {
            return;
        }
        await this.traceWriter.writeStaticState(values);
    }
    hasDebugger() {
        return this.debuggerState !== null;
    }
    hasTraceWriter() {
        return this.traceWriter !== null;
    }
}
