import { __awaiter, __rest } from "tslib";
import { TOOL_FOLD_SYMBOL } from '../utils/tool-fold';
/**
 * Walk a state object and snapshot any machine-shaped values found on it.
 * Recognises foundation-state-machine instances by structural shape rather
 * than `instanceof`, so the helper stays free of a runtime dep on
 * foundation-state-machine.
 *
 * Strips the `errors`/`error` framework noise from each machine's context —
 * these are foundation-state-machine internals (an ErrorMap instance and its
 * last-error pointer) that bloat the debug payload without aiding debugging.
 */
function defaultStatefulDebugSnapshot(state) {
    var _a;
    if (!state || typeof state !== 'object')
        return state;
    const result = {};
    for (const [key, val] of Object.entries(state)) {
        if (val && typeof val === 'object' && 'state' in val && 'context' in val && 'complete' in val) {
            const m = val;
            const ctx = (_a = m.context) !== null && _a !== void 0 ? _a : {};
            const { errors: _e, error: _err } = ctx, userContext = __rest(ctx, ["errors", "error"]);
            result[key] = {
                state: m.state,
                context: userContext,
                complete: m.complete,
                output: m.output,
            };
        }
    }
    return Object.keys(result).length
        ? result
        : '<no auto-snapshot — state has no machine-shaped properties; provide getDebugSnapshot>';
}
/**
 * Build a restore-ready snapshot of agent state. Each property exposing a
 * `getPersistedSnapshot()` method (a foundation-state-machine `AbstractMachine`)
 * is replaced by its persisted snapshot — the lossless, JSON-safe image
 * {@link restoreMachine} feeds back via `start({ state })`. Every other property
 * is kept as-is, so plain serializable state (e.g. a planner's `{ plan, cursor }`)
 * round-trips directly with no per-agent `serialize`.
 *
 * Distinct from `defaultStatefulDebugSnapshot`, which abbreviates for the debug
 * log and is not guaranteed lossless.
 */
function defaultStatefulSerialize(state) {
    if (!state || typeof state !== 'object')
        return state;
    const result = {};
    for (const [key, val] of Object.entries(state)) {
        result[key] =
            val &&
                typeof val === 'object' &&
                typeof val.getPersistedSnapshot === 'function'
                ? val.getPersistedSnapshot()
                : val;
    }
    return result;
}
/**
 * Restore a foundation-state-machine instance from a snapshot produced by the
 * default serializer (or `machine.getPersistedSnapshot()`). Call inside a
 * stateful agent's `init` to resume the machine in its exact prior place:
 *
 * ```ts
 * init: ({ restored }) => {
 *   const machine = new MyMachine();
 *   restoreMachine(machine, restored); // resume, or no-op on a fresh start
 *   return { machine };
 * }
 * ```
 *
 * No-ops when there is nothing to restore (a fresh session), leaving the
 * freshly-constructed machine in its initial state.
 *
 * @param machine - the machine to restore (already constructed/started).
 * @param restored - the agent's serialized image (`ctx.restored` in `init`).
 * @param key - the property the machine occupies in the snapshot; defaults to
 *   `'machine'` to match the conventional `{ machine }` state bag.
 * @beta
 */
export function restoreMachine(machine, restored, key = 'machine') {
    const snapshot = restored === null || restored === void 0 ? void 0 : restored[key];
    if (snapshot === undefined || snapshot === null) {
        return;
    }
    machine.stop();
    machine.start({ state: snapshot });
}
/**
 * Build an `AgentConfig` whose `systemPrompt`, `toolDefinitions`, and tool
 * handlers all close over a long-lived state object created on activation.
 *
 * The framework wires the lifecycle: `init` on `onActivate`, `dispose` on
 * `onDeactivate`. State is held inside the helper's closure — never exposed on
 * the resulting `AgentConfig` — so the redux serializer doesn't see it.
 *
 * @example
 * ```ts
 * const guidedBooking = defineStatefulAgent<{ machine: GuidedBookingMachine }>({
 *   name: 'Guided Booking',
 *   description: 'Books a trade via a guided wizard.',
 *   excludeFromClassifier: true,
 *   manualSelection: { enabled: true, hint: 'Step-by-step trade booking' },
 *   init: () => ({ machine: new GuidedBookingMachine() }),
 *   dispose: ({ state }) => state.machine.stop(),
 *   systemPrompt: ({ state }) => composeFromMachine(state.machine),
 *   toolDefinitions: ({ state }) => toolsForState(state.machine.state),
 *   toolHandlers: ({ machine }) => ({ ... }),
 * });
 * ```
 *
 * @beta
 */
export function defineStatefulAgent(opts) {
    let state;
    const assertNoFolds = (handlers) => {
        for (const [name, handler] of Object.entries(handlers)) {
            if (handler[TOOL_FOLD_SYMBOL]) {
                throw new Error(`Stateful agent "${opts.name}" tool "${name}" carries fold metadata. ` +
                    `Folds and state-machine-driven tool filtering both try to control the ` +
                    `LLM's tool view — pick one. Remove the fold or migrate this agent to ` +
                    `a non-stateful AgentConfig.`);
            }
        }
    };
    // Wrap the optional dynamic-tools factory so it threads `state` in and asserts
    // we never accidentally resolved a fold-tagged tool. Static arrays are passed
    // through unchanged.
    const wrappedTools = (() => {
        const td = opts.toolDefinitions;
        if (typeof td === 'function') {
            return (ctx) => __awaiter(this, void 0, void 0, function* () {
                if (!state) {
                    throw new Error(`Stateful agent "${opts.name}" tools called before init`);
                }
                return td(Object.assign(Object.assign({}, ctx), { state }));
            });
        }
        return td;
    })();
    // Wrap the handler factory into the ctx-shaped form the driver expects, and
    // validate folds on each resolve. The driver re-resolves this per tool-loop
    // iteration, so the handler set the driver dispatches against always
    // matches the LLM-visible tools for the current state — handlers don't need
    // defensive `if (!machine.matches(...))` guards.
    const wrappedHandlers = opts.toolHandlers
        ? () => __awaiter(this, void 0, void 0, function* () {
            if (!state) {
                throw new Error(`Stateful agent "${opts.name}" handlers called before init`);
            }
            const handlers = yield opts.toolHandlers(state);
            assertNoFolds(handlers);
            return handlers;
        })
        : undefined;
    const wrappedDisplayName = opts.displayName
        ? (ctx) => __awaiter(this, void 0, void 0, function* () {
            if (!state) {
                throw new Error(`Stateful agent "${opts.name}" displayName called before init`);
            }
            return opts.displayName(Object.assign(Object.assign({}, ctx), { state }));
        })
        : undefined;
    // Function-form `provider` needs `state` injected the same way the other
    // per-turn resolvers do; the static-string form is passed through unchanged
    // and validated up-front by OrchestratingDriver.
    const wrappedProvider = typeof opts.provider === 'function'
        ? (ctx) => __awaiter(this, void 0, void 0, function* () {
            if (!state) {
                throw new Error(`Stateful agent "${opts.name}" provider called before init`);
            }
            return opts.provider(Object.assign(Object.assign({}, ctx), { state }));
        })
        : opts.provider;
    // `temperature` and `toolChoice` thread `state` in exactly like `provider`;
    // static values pass through unchanged.
    const wrappedTemperature = typeof opts.temperature === 'function'
        ? (ctx) => __awaiter(this, void 0, void 0, function* () {
            if (!state) {
                throw new Error(`Stateful agent "${opts.name}" temperature called before init`);
            }
            return opts.temperature(Object.assign(Object.assign({}, ctx), { state }));
        })
        : opts.temperature;
    const wrappedToolChoice = typeof opts.toolChoice === 'function'
        ? (ctx) => __awaiter(this, void 0, void 0, function* () {
            if (!state) {
                throw new Error(`Stateful agent "${opts.name}" toolChoice called before init`);
            }
            return opts.toolChoice(Object.assign(Object.assign({}, ctx), { state }));
        })
        : opts.toolChoice;
    const wrappedCachePolicy = typeof opts.cachePolicy === 'function'
        ? (ctx) => __awaiter(this, void 0, void 0, function* () {
            if (!state) {
                throw new Error(`Stateful agent "${opts.name}" cachePolicy called before init`);
            }
            return opts.cachePolicy(Object.assign(Object.assign({}, ctx), { state }));
        })
        : opts.cachePolicy;
    const wrappedTailContext = typeof opts.tailContext === 'function'
        ? (ctx) => __awaiter(this, void 0, void 0, function* () {
            if (!state) {
                throw new Error(`Stateful agent "${opts.name}" tailContext called before init`);
            }
            return opts.tailContext(Object.assign(Object.assign({}, ctx), { state }));
        })
        : opts.tailContext;
    // `onUnresolvedTool` threads `state` in like the resolvers above, but unlike
    // them returns `undefined` before init rather than throwing: it fires on the
    // tool-dispatch path, so a pre-init call must degrade to the framework
    // default instead of breaking dispatch.
    const wrappedOnUnresolvedTool = typeof opts.onUnresolvedTool === 'function'
        ? (ctx) => __awaiter(this, void 0, void 0, function* () {
            if (!state) {
                return undefined;
            }
            return opts.onUnresolvedTool(Object.assign(Object.assign({}, ctx), { state }));
        })
        : opts.onUnresolvedTool;
    // `resumable` threads the LIVE state bag in (like the resolvers above) rather
    // than the serialized image the raw `AgentConfig` form receives — so a machine
    // agent gates on `state.machine.state` directly. Consulted only while active
    // (state is set); a pre-init call degrades to `true` (resumable, fail-safe).
    const wrappedResumable = opts.resumable
        ? (ctx) => __awaiter(this, void 0, void 0, function* () {
            if (state === undefined)
                return true;
            return opts.resumable({
                agentName: ctx.agentName,
                sessionKey: ctx.sessionKey,
                messages: ctx.messages,
                isFlowOwner: ctx.isFlowOwner,
                state,
            });
        })
        : undefined;
    const base = {
        name: opts.name,
        displayName: wrappedDisplayName,
        primerHistory: opts.primerHistory,
        subAgents: opts.subAgents,
        manualSelection: opts.manualSelection,
        chatInputDuringExecution: opts.chatInputDuringExecution,
        toolDefinitions: wrappedTools,
        toolHandlers: wrappedHandlers,
        provider: wrappedProvider,
        temperature: wrappedTemperature,
        toolChoice: wrappedToolChoice,
        cachePolicy: wrappedCachePolicy,
        tailContext: wrappedTailContext,
        onUnresolvedTool: wrappedOnUnresolvedTool,
        resumable: wrappedResumable,
        notResumableMessage: opts.notResumableMessage,
        onActivate: (ctx) => __awaiter(this, void 0, void 0, function* () {
            state = yield opts.init(ctx);
            // Sample once with the init state to fail loud on fold misuse at the
            // agent-switch boundary instead of waiting for the first tool call.
            // The driver re-resolves the factory each iteration, so misuse on
            // handlers gated behind a non-init state surfaces when that state is
            // reached — same loud failure mode, narrower coverage.
            if (opts.toolHandlers) {
                const sampled = yield opts.toolHandlers(state);
                assertNoFolds(sampled);
            }
        }),
        onDeactivate: (ctx) => __awaiter(this, void 0, void 0, function* () {
            if (state !== undefined && opts.dispose) {
                yield opts.dispose(Object.assign(Object.assign({}, ctx), { state }));
            }
            // Orchestrator serializes onActivate/onDeactivate; concurrent calls
            // cannot interleave the read-then-write of `state`.
            // eslint-disable-next-line require-atomic-updates
            state = undefined;
        }),
        systemPrompt: opts.systemPrompt
            ? (ctx) => __awaiter(this, void 0, void 0, function* () {
                if (!state) {
                    throw new Error(`Stateful agent "${opts.name}" systemPrompt called before init`);
                }
                return opts.systemPrompt(Object.assign(Object.assign({}, ctx), { state }));
            })
            : undefined,
        // Called per LLM call by ChatDriver, and once at debug-log export time.
        // Returns `undefined` before activation (state hasn't been built yet) and
        // after deactivation (state was disposed) — the snapshotter is best-effort
        // and the driver tolerates undefined.
        getDebugSnapshot: () => {
            if (!state)
                return undefined;
            return opts.getDebugSnapshot
                ? opts.getDebugSnapshot(state)
                : defaultStatefulDebugSnapshot(state);
        },
        // Lossless restore image for session persistence (see `serialize` on the
        // config). Returns undefined before activation / after disposal.
        serialize: () => {
            if (state === undefined)
                return undefined;
            return opts.serialize ? opts.serialize(state) : defaultStatefulSerialize(state);
        },
    };
    return Object.assign(Object.assign({}, base), { description: opts.description, excludeFromClassifier: opts.excludeFromClassifier });
}
