import * as fs from "fs";
import * as path from "path";
import { AsyncLocalStorage } from "node:async_hooks";
import { nanoid } from "nanoid";
// Bump this when the wire format changes in a way the viewer needs
// to notice. The viewer rejects files with a higher version.
export const STATELOG_FORMAT_VERSION = 1;
const DEFAULT_REQUEST_TIMEOUT_MS = 1500;
// === Client ===
// Shared empty stack returned by `snapshotStack()` when the client is
// disabled. Reusing one frozen array avoids per-fork allocations in the
// default no-op mode. The runner only ever reads from this snapshot
// (passes it back into `runInBranchContext`), never mutates it.
const EMPTY_STACK = Object.freeze([]);
export class StatelogClient {
    host;
    debugMode;
    traceId;
    apiKey;
    projectId;
    logFile;
    enabled = true;
    // Whether the remote http sink is usable. A configured host with no
    // apiKey disables ONLY the remote send; local sinks (logFile, stdout)
    // still receive events.
    remoteEnabled = false;
    // The "root" span stack — used by the outer agent run thread. Code
    // running inside `runInBranchContext` sees a branch-local stack
    // delivered via AsyncLocalStorage instead.
    rootStack = [];
    // Per-branch span stacks live in this AsyncLocalStorage. Each concurrent
    // fork/race branch gets its own stack so its `startSpan`/`endSpan`
    // calls never bleed into the parent or siblings — even though they all
    // share this single StatelogClient instance.
    spanStorage = new AsyncLocalStorage();
    metadata;
    requestTimeoutMs;
    constructor(config) {
        const { host, apiKey, projectId, traceId, debugMode } = config;
        this.host = host;
        this.apiKey = apiKey;
        this.projectId = projectId;
        this.debugMode = debugMode || false;
        this.traceId = traceId || nanoid();
        this.logFile = config.logFile;
        this.requestTimeoutMs =
            config.requestTimeoutMs ?? DEFAULT_REQUEST_TIMEOUT_MS;
        this.metadata = config.metadata;
        // Observability must be explicitly enabled. When false (the default),
        // the entire client is a no-op — no events emitted, no network calls,
        // no file writes.
        if (!config.observability) {
            this.enabled = false;
            return;
        }
        if (this.debugMode) {
            console.log(`Statelog client initialized with host: ${host} and traceId: ${this.traceId}`, { config });
        }
        // Decide whether the remote sink is usable. The remote sink (any
        // host that isn't "stdout") requires an apiKey. If the host is set
        // but the key is missing, we keep the client enabled (so local
        // sinks still work) but skip the http POST inside `post()`.
        const hostLower = this.host.toLowerCase();
        const isRemoteHost = !!this.host && hostLower !== "stdout";
        if (isRemoteHost) {
            if (this.apiKey) {
                this.remoteEnabled = true;
            }
            else if (this.debugMode) {
                console.warn("StatelogClient: remote host configured without apiKey — remote sink disabled. Local sinks (stdout/logFile) will still receive events.");
            }
        }
        // If a logFile is configured, ensure the parent directory exists.
        // We don't truncate here — each run uses its own runId as traceId,
        // so multiple runs writing to the same file are still distinguishable.
        if (this.logFile) {
            try {
                fs.mkdirSync(path.dirname(this.logFile), { recursive: true });
            }
            catch (err) {
                if (this.debugMode)
                    console.warn(`StatelogClient: failed to ensure logFile dir: ${err}`);
            }
        }
    }
    // === Span management ===
    //
    // Spans are tracked per-async-context. Outside fork/race branches the
    // active stack is `rootStack`. Inside `runInBranchContext` it is the
    // per-branch stack delivered by `spanStorage`. This means concurrent
    // branches each get a private stack — their `startSpan`/`endSpan`
    // calls cannot interleave or pop each other's spans, and events
    // emitted inside a branch are attributed to that branch's current
    // span (which inherits the parent fork span at branch entry).
    currentStack() {
        return this.spanStorage.getStore() ?? this.rootStack;
    }
    // Returns a shallow snapshot of the active span stack — used by the
    // runner to seed each branch's stack with the parent's spans so
    // events inside the branch nest under the fork/race span.
    //
    // Short-circuits to a shared empty array when observability is off
    // so the runner can call this unconditionally without paying for an
    // allocation on every fork/race in the no-op default mode.
    snapshotStack() {
        if (!this.enabled)
            return EMPTY_STACK;
        return [...this.currentStack()];
    }
    // Run `fn` with a fresh, branch-local span stack seeded from
    // `parentStack`. Each call to this method creates an independent ALS
    // context; sibling calls (e.g. concurrent fork branches) see
    // independent stacks even though they share this StatelogClient.
    //
    // Spans pushed inside `fn` are popped against this private stack only,
    // never the parent. Defensively copies `parentStack` so the caller's
    // array is never mutated.
    //
    // When observability is disabled the ALS plumbing is skipped entirely
    // — we just invoke `fn()` directly. The runner can therefore always
    // wrap branches in this call without paying ALS overhead in no-op
    // mode.
    runInBranchContext(parentStack, fn) {
        if (!this.enabled)
            return fn();
        return this.spanStorage.run([...parentStack], fn);
    }
    startSpan(type) {
        if (!this.enabled)
            return undefined;
        const stack = this.currentStack();
        const spanId = nanoid(12);
        const parentSpanId = stack.length > 0 ? stack[stack.length - 1].spanId : null;
        stack.push({
            spanId,
            parentSpanId,
            spanType: type,
            startTime: performance.now(),
        });
        return spanId;
    }
    // Pop the span identified by `spanId` from the active stack. The id
    // MUST be the one returned by the matching `startSpan` call:
    //
    //   const id = client.startSpan("agentRun");
    //   try { ... } finally { client.endSpan(id); }
    //
    // - If `spanId` is undefined, this is a no-op.
    // - If the span is at the top of the active stack, it pops it directly.
    // - Otherwise we drop everything above the matched span — this
    //   defends against an inner span that forgot to call `endSpan`.
    // - If the span isn't found in the active stack, this is a no-op.
    //   That can legitimately happen when a branch tries to end a span
    //   that lives on the parent's stack, or vice versa.
    endSpan(spanId) {
        if (!this.enabled || !spanId)
            return undefined;
        const stack = this.currentStack();
        const top = stack[stack.length - 1];
        if (top && top.spanId === spanId) {
            return stack.pop();
        }
        const idx = stack.findIndex((s) => s.spanId === spanId);
        if (idx < 0)
            return undefined;
        const popped = stack[idx];
        stack.length = idx;
        return popped;
    }
    get currentSpan() {
        const stack = this.currentStack();
        return stack.length > 0 ? stack[stack.length - 1] : undefined;
    }
    toJSON() {
        return {
            traceId: this.traceId,
            projectId: this.projectId,
            host: this.host,
            debugMode: this.debugMode,
        };
    }
    // === Existing event methods ===
    async debug(message, data) {
        await this.post({
            type: "debug",
            message: message,
            data,
        });
    }
    async graph({ nodes, edges, startNode, }) {
        await this.post({
            type: "graph",
            nodes,
            edges,
            startNode,
        });
    }
    async enterNode({ nodeId, data, }) {
        await this.post({
            type: "enterNode",
            nodeId,
            data,
        });
    }
    async exitNode({ nodeId, data, timeTaken, }) {
        await this.post({
            type: "exitNode",
            nodeId,
            data,
            timeTaken,
        });
    }
    async beforeHook({ nodeId, startData, endData, timeTaken, }) {
        await this.post({
            type: "beforeHook",
            nodeId,
            startData,
            endData,
            timeTaken,
        });
    }
    async afterHook({ nodeId, startData, endData, timeTaken, }) {
        await this.post({
            type: "afterHook",
            nodeId,
            startData,
            endData,
            timeTaken,
        });
    }
    async followEdge({ fromNodeId, toNodeId, isConditionalEdge, data, }) {
        await this.post({
            type: "followEdge",
            edgeId: `${fromNodeId}->${toNodeId}`,
            fromNodeId,
            toNodeId,
            isConditionalEdge,
            data,
        });
    }
    async promptCompletion({ messages, completion, model, timeTaken, tools, responseFormat, usage, cost, finishReason, stream, }) {
        await this.post({
            type: "promptCompletion",
            messages,
            completion,
            model,
            timeTaken,
            tools,
            responseFormat,
            usage,
            cost,
            finishReason,
            stream,
        });
    }
    /**
     * Emit an `embedCompletion` event. Mirrors `promptCompletion` so the
     * viewer can render the two with a shared formatter, but keeps the
     * embedding-specific bits (dimensions, phase tag) separate.
     *
     * `inputPreview` is intentionally a short slice of the source text —
     * the full text is kept off the wire so a transcript-with-PII never
     * ends up in a remote sink by accident. Vectors are NEVER logged
     * (only their length, via `dimensions`).
     */
    async embedCompletion({ inputPreview, inputCount, model, dimensions, timeTaken, usage, cost, phase, }) {
        await this.post({
            type: "embedCompletion",
            inputPreview,
            inputCount,
            model,
            dimensions,
            timeTaken,
            usage,
            cost,
            phase,
        });
    }
    /**
     * Memory umbrella-span marker events.
     *
     * Why these exist: the logs viewer infers span types from EVENT
     * types (see `lib/logsViewer/tree.ts#inferSpanLabel`). A span that
     * never has an event posted with its own `span_id` is invisible to
     * the viewer — its children get re-parented to the trace root.
     *
     * Calling `startSpan("memoryRemember")` only sets up the AsyncLocal-
     * Storage stack; on its own it produces no event. These methods
     * post a tiny marker event right after `startSpan` so the span
     * materializes in the tree with the right label, and the inner
     * `llmCall` / `embedding` spans nest under it. The payload is also
     * useful in its own right (content/query preview, memoryId).
     */
    async memoryRemember({ contentPreview, memoryId, }) {
        await this.post({
            type: "memoryRemember",
            contentPreview,
            memoryId,
        });
    }
    async memoryRecall({ queryPreview, memoryId, phase, }) {
        await this.post({
            type: "memoryRecall",
            queryPreview,
            memoryId,
            phase,
        });
    }
    async memoryForget({ queryPreview, memoryId, }) {
        await this.post({
            type: "memoryForget",
            queryPreview,
            memoryId,
        });
    }
    async memoryCompaction({ memoryId, messageCount, threshold, }) {
        await this.post({
            type: "memoryCompaction",
            memoryId,
            messageCount,
            threshold,
        });
    }
    async toolCall({ toolName, args, output, model, timeTaken, }) {
        await this.post({
            type: "toolCall",
            toolName,
            args,
            output,
            model,
            timeTaken,
        });
    }
    async diff({ itemA, itemB, message, }) {
        await this.post({
            type: "diff",
            itemA,
            itemB,
            message,
        });
    }
    // === New event methods ===
    async runMetadata(metadata) {
        await this.post({
            type: "runMetadata",
            ...metadata,
        });
    }
    async agentStart({ entryNode, args, }) {
        await this.post({
            type: "agentStart",
            entryNode,
            args,
        });
        if (this.metadata) {
            await this.runMetadata({ ...this.metadata, entryNode });
        }
    }
    async agentEnd({ entryNode, result, timeTaken, tokenStats, }) {
        // Fire-and-forget the remote send: this is the very last event of
        // a run, and waiting for the http round trip here directly delays
        // process exit. The synchronous file/stdout sinks still run inline
        // inside `post()`, so local observability is unaffected.
        await this.post({
            type: "agentEnd",
            entryNode,
            result,
            timeTaken,
            tokenStats,
        }, { noWait: true });
    }
    // --- Interrupt lifecycle ---
    async interruptThrown({ interruptId, interruptData, }) {
        await this.post({
            type: "interruptThrown",
            interruptId,
            interruptData,
        });
    }
    async handlerDecision({ interruptId, handlerIndex, decision, value, }) {
        await this.post({
            type: "handlerDecision",
            interruptId,
            handlerIndex,
            decision,
            value,
        });
    }
    async interruptResolved({ interruptId, outcome, resolvedBy, timeTaken, }) {
        await this.post({
            type: "interruptResolved",
            interruptId,
            outcome,
            resolvedBy,
            timeTaken,
        });
    }
    // --- Checkpoint lifecycle ---
    async checkpointCreated({ checkpointId, reason, sourceLocation, }) {
        await this.post({
            type: "checkpointCreated",
            checkpointId,
            reason,
            sourceLocation,
        });
    }
    async checkpointRestored({ checkpointId, restoreCount, maxRestores, overrides, }) {
        await this.post({
            type: "checkpointRestored",
            checkpointId,
            restoreCount,
            maxRestores,
            overrides,
        });
    }
    // --- Fork/Race lifecycle ---
    async forkStart({ forkId, mode, branchCount, }) {
        await this.post({
            type: "forkStart",
            forkId,
            mode,
            branchCount,
        });
    }
    async forkBranchEnd({ forkId, branchIndex, outcome, timeTaken, }) {
        await this.post({
            type: "forkBranchEnd",
            forkId,
            branchIndex,
            outcome,
            timeTaken,
        });
    }
    async forkEnd({ forkId, mode, timeTaken, winnerIndex, }) {
        await this.post({
            type: "forkEnd",
            forkId,
            mode,
            timeTaken,
            winnerIndex,
        });
    }
    // --- Thread lifecycle ---
    async threadCreated({ threadId, threadType, parentThreadId, }) {
        await this.post({
            type: "threadCreated",
            threadId,
            threadType,
            parentThreadId,
        });
    }
    /** Fired when a previously-closed thread is re-activated via
     *  `ThreadStore.resumeExisting()` (i.e. by `thread(continue: id)`
     *  or `thread(session: name)` on second+ entry). Mirrors
     *  `threadCreated` so trace consumers can distinguish first-entry
     *  from resumption. */
    async threadResumed({ threadId, }) {
        await this.post({
            type: "threadResumed",
            threadId,
        });
    }
    /** Fired when the `onThreadEnd` callback dispatcher itself throws
     *  inside the `Runner.thread` finally block. Individual callback
     *  errors are caught and logged by `fireWithGuard`; this event
     *  covers the rarer case where the dispatcher loop itself blew up
     *  (e.g. a malformed registration). Emitted in lieu of the
     *  previous `console.error` so the failure is observable in
     *  traces. */
    async threadEndHookError({ threadId, error, }) {
        await this.post({
            type: "threadEndHookError",
            threadId,
            error,
        });
    }
    // --- Structured errors ---
    async error({ errorType, message, functionName, retryable, sourceLocation, }) {
        await this.post({
            type: "error",
            errorType,
            message,
            functionName,
            retryable,
            sourceLocation,
        });
    }
    // === Post (wire format) ===
    async post(body, options) {
        if (!this.enabled) {
            return;
        }
        // We need either a host (remote / stdout) or a logFile to do anything.
        if (!this.host && !this.logFile) {
            return;
        }
        const span = this.currentSpan;
        const postBody = JSON.stringify({
            format_version: STATELOG_FORMAT_VERSION,
            trace_id: this.traceId,
            project_id: this.projectId,
            span_id: span?.spanId ?? null,
            parent_span_id: span?.parentSpanId ?? null,
            data: { ...body, timestamp: new Date().toISOString() },
        });
        // File sink: append one JSON object per line. Done synchronously
        // so tests can read the file immediately after an awaited event.
        if (this.logFile) {
            try {
                fs.appendFileSync(this.logFile, postBody + "\n");
            }
            catch (err) {
                if (this.debugMode)
                    console.error("StatelogClient: failed to append to logFile:", err);
            }
        }
        if (!this.host)
            return;
        if (this.host.toLowerCase() === "stdout") {
            console.log(postBody);
            return;
        }
        // Remote sink is only attempted when an apiKey was provided. Without
        // it we silently skip the http POST so a configured logFile/stdout
        // sink keeps working without firing unauthenticated requests.
        if (!this.remoteEnabled)
            return;
        try {
            const fullUrl = new URL("/api/logs", this.host);
            const url = fullUrl.toString();
            // Bound each remote send by `requestTimeoutMs` so a slow or
            // unreachable statelog host cannot wedge process exit. The
            // request still completes asynchronously; on timeout it just
            // aborts with no retry.
            const request = fetch(url, {
                method: "POST",
                headers: {
                    "Content-Type": "application/json",
                    Authorization: `Bearer ${this.apiKey}`,
                },
                body: postBody,
                signal: AbortSignal.timeout(this.requestTimeoutMs),
            }).catch((err) => {
                if (this.debugMode)
                    console.error("Failed to send statelog:", err);
            });
            if (options?.noWait) {
                // Detach the in-flight request from the caller's await chain.
                // The `.catch` above guarantees no UnhandledPromiseRejection if
                // it later fails or aborts.
                return;
            }
            await request;
        }
        catch (err) {
            if (this.debugMode)
                console.error("Error sending log in statelog client:", err, {
                    host: this.host,
                });
        }
    }
}
export function getStatelogClient(config) {
    const statelogConfig = {
        host: config.host,
        traceId: config.traceId || nanoid(),
        apiKey: process.env.STATELOG_API_KEY || "",
        projectId: config.projectId,
        debugMode: config.debugMode || false,
        observability: config.observability,
        logFile: config.logFile,
    };
    const client = new StatelogClient(statelogConfig);
    return client;
}
