"use strict";
/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation and GitHub. All rights reserved.
 *--------------------------------------------------------------------------------------------*/
Object.defineProperty(exports, "__esModule", { value: true });
exports.PromptRenderer = exports.MetadataMap = void 0;
const JSONT = require("./jsonTypes");
const materialized_1 = require("./materialized");
const mode_1 = require("./output/mode");
const promptElements_1 = require("./promptElements");
const results_1 = require("./results");
var MetadataMap;
(function (MetadataMap) {
    MetadataMap.empty = {
        get: () => undefined,
        getAll: () => [],
    };
    MetadataMap.from = (metadata) => {
        return {
            get: ctor => metadata.find(m => m instanceof ctor),
            getAll: ctor => metadata.filter(m => m instanceof ctor),
        };
    };
})(MetadataMap || (exports.MetadataMap = MetadataMap = {}));
/**
 * A prompt renderer is responsible for rendering a {@link PromptElementCtor prompt element} to {@link ChatMessagePromptElement chat messages}.
 *
 * Note: You must create a fresh prompt renderer instance for each prompt element you want to render.
 */
class PromptRenderer {
    _endpoint;
    _ctor;
    _props;
    _tokenizer;
    _usedContext = [];
    _ignoredFiles = [];
    _growables = [];
    _root = new PromptTreeElement(null, 0);
    _tokenLimits = [];
    /** Epoch used to tracing the order in which elements render. */
    tracer = undefined;
    /**
     * @param _endpoint The chat endpoint that the rendered prompt will be sent to.
     * @param _ctor The prompt element constructor to render.
     * @param _props The props to pass to the prompt element.
     */
    constructor(_endpoint, _ctor, _props, _tokenizer) {
        this._endpoint = _endpoint;
        this._ctor = _ctor;
        this._props = _props;
        this._tokenizer = _tokenizer;
    }
    getIgnoredFiles() {
        return Array.from(new Set(this._ignoredFiles));
    }
    getUsedContext() {
        return this._usedContext;
    }
    createElement(element) {
        return new element.ctor(element.props);
    }
    async _processPromptPieces(sizing, pieces, progress, token) {
        // Collect all prompt elements in the next flex group to render, grouping
        // by the flex order in which they're rendered.
        const promptElements = new Map();
        for (const [i, element] of pieces.entries()) {
            // Set any jsx children as the props.children
            if (Array.isArray(element.children)) {
                element.props = element.props ?? {};
                element.props.children = element.children; // todo@joyceerhl clean up any
            }
            // Instantiate the prompt part
            if (!element.ctor) {
                const loc = atPath(element.path);
                throw new Error(`Invalid ChatMessage child! Child must be a TSX component that extends PromptElement at ${loc}`);
            }
            const promptElement = this.createElement(element);
            let tokenLimit;
            if (promptElement instanceof promptElements_1.TokenLimit) {
                tokenLimit = element.props.max;
                this._tokenLimits.push({ limit: tokenLimit, id: element.node.id });
            }
            element.node.setObj(promptElement);
            // Prepare rendering
            const flexGroupValue = element.props.flexGrow ?? Infinity;
            let flexGroup = promptElements.get(flexGroupValue);
            if (!flexGroup) {
                flexGroup = [];
                promptElements.set(flexGroupValue, flexGroup);
            }
            flexGroup.push({ element, promptElementInstance: promptElement, tokenLimit });
        }
        if (promptElements.size === 0) {
            return;
        }
        const flexGroups = [...promptElements.entries()]
            .sort(([a], [b]) => b - a)
            .map(([_, group]) => group);
        const setReserved = (groupIndex) => {
            let reservedTokens = 0;
            for (let i = groupIndex + 1; i < flexGroups.length; i++) {
                for (const { element } of flexGroups[i]) {
                    if (!element.props.flexReserve) {
                        continue;
                    }
                    const reserve = typeof element.props.flexReserve === 'string'
                        ? // Typings ensure the string is `/${number}`
                            Math.floor(sizing.remainingTokenBudget / Number(element.props.flexReserve.slice(1)))
                        : element.props.flexReserve;
                    reservedTokens += reserve;
                }
            }
            sizing.consume(reservedTokens);
            return reservedTokens;
        };
        // Prepare all currently known prompt elements in parallel
        for (const [groupIndex, promptElements] of flexGroups.entries()) {
            // Temporarily consume any reserved budget for later elements so that the sizing is calculated correctly here.
            const reservedTokens = setReserved(groupIndex);
            // Calculate the flex basis for dividing the budget amongst siblings in this group.
            let flexBasisSum = 0;
            for (const { element } of promptElements) {
                flexBasisSum += element.props.flexBasis ?? 1;
            }
            let constantTokenLimits = 0;
            //.For elements that limit their token usage and would use less than we
            // otherwise would assign to them, 'cap' their usage at the limit and
            // remove their share directly from the budget in distribution.
            const useConstantLimitsForIndex = promptElements.map(e => {
                if (e.tokenLimit === undefined) {
                    return false;
                }
                const flexBasis = e.element.props.flexBasis ?? 1;
                const proportion = flexBasis / flexBasisSum;
                const proportionateUsage = Math.floor(sizing.remainingTokenBudget * proportion);
                if (proportionateUsage < e.tokenLimit) {
                    return false;
                }
                flexBasisSum -= flexBasis;
                constantTokenLimits += e.tokenLimit;
                return true;
            });
            // Finally calculate the final sizing for each element in this group.
            const elementSizings = promptElements.map((e, i) => {
                const proportion = (e.element.props.flexBasis ?? 1) / flexBasisSum;
                return {
                    tokenBudget: useConstantLimitsForIndex[i]
                        ? e.tokenLimit
                        : Math.floor((sizing.remainingTokenBudget - constantTokenLimits) * proportion),
                    endpoint: sizing.endpoint,
                    countTokens: (text, cancellation) => this._tokenizer.tokenLength(typeof text === 'string'
                        ? { type: mode_1.Raw.ChatCompletionContentPartKind.Text, text }
                        : text, cancellation),
                };
            });
            // Free the previously-reserved budget now that we calculated sizing
            sizing.consume(-reservedTokens);
            this.tracer?.addRenderEpoch?.({
                inNode: promptElements[0].element.node.parent?.id,
                flexValue: promptElements[0].element.props.flexGrow ?? 0,
                tokenBudget: sizing.remainingTokenBudget,
                reservedTokens,
                elements: promptElements.map((e, i) => ({
                    id: e.element.node.id,
                    tokenBudget: elementSizings[i].tokenBudget,
                })),
            });
            await Promise.all(promptElements.map(async ({ element, promptElementInstance }, i) => {
                const state = await annotateError(element, () => promptElementInstance.prepare?.(elementSizings[i], progress, token));
                element.node.setState(state);
            }));
            const templates = await Promise.all(promptElements.map(async ({ element, promptElementInstance }, i) => {
                const elementSizing = elementSizings[i];
                return await annotateError(element, () => promptElementInstance.render(element.node.getState(), elementSizing, progress, token));
            }));
            // Render
            for (const [i, { element, promptElementInstance }] of promptElements.entries()) {
                const elementSizing = elementSizings[i];
                const template = templates[i];
                if (!template) {
                    // it doesn't want to render anything
                    continue;
                }
                const childConsumption = await this._processPromptRenderPiece(new PromptSizingContext(elementSizing.tokenBudget, this._endpoint), element, promptElementInstance, template, progress, token);
                // Append growables here so that when we go back and expand them we do so in render order.
                if (promptElementInstance instanceof promptElements_1.Expandable) {
                    this._growables.push({ initialConsume: childConsumption, elem: element.node });
                }
                // Tally up the child consumption into the parent context for any subsequent flex group
                sizing.consume(childConsumption);
            }
        }
    }
    async _processPromptRenderPiece(elementSizing, element, promptElementInstance, template, progress, token) {
        const pieces = flattenAndReduce(template);
        // Compute token budget for the pieces that this child wants to render
        const childSizing = new PromptSizingContext(elementSizing.tokenBudget, this._endpoint);
        const { tokensConsumed } = await computeTokensConsumedByLiterals(this._tokenizer, element, promptElementInstance, pieces);
        childSizing.consume(tokensConsumed);
        await this._handlePromptChildren(element, pieces, childSizing, progress, token);
        // Tally up the child consumption into the parent context for any subsequent flex group
        return childSizing.consumed;
    }
    /**
     * Renders the prompt element and its children to a JSON-serializable state.
     * @returns A promise that resolves to an object containing the rendered chat messages and the total token count.
     * The total token count is guaranteed to be less than or equal to the token budget.
     */
    async renderElementJSON(token) {
        await this._processPromptPieces(new PromptSizingContext(this._endpoint.modelMaxPromptTokens, this._endpoint), [
            {
                node: this._root,
                ctor: this._ctor,
                props: this._props,
                children: [],
                path: [this._ctor],
            },
        ], undefined, token);
        // todo@connor4312: should ignored files, used context, etc. be passed here?
        return {
            node: this._root.toJSON(),
        };
    }
    /**
     * Renders the prompt element and its children.
     * @returns A promise that resolves to an object containing the rendered chat messages and the total token count.
     * The total token count is guaranteed to be less than or equal to the token budget.
     */
    async render(progress, token) {
        const result = await this.renderRaw(progress, token);
        return { ...result, messages: (0, mode_1.toMode)(this._tokenizer.mode, result.messages) };
    }
    /**
     * Renders the prompt element and its children. Similar to {@link render}, but
     * returns the original message representation.
     */
    async renderRaw(progress, token) {
        // Convert root prompt element to prompt pieces
        await this._processPromptPieces(new PromptSizingContext(this._endpoint.modelMaxPromptTokens, this._endpoint), [
            {
                node: this._root,
                ctor: this._ctor,
                props: this._props,
                children: [],
                path: [this._ctor],
            },
        ], progress, token);
        const { container, allMetadata, removed } = await this._getFinalElementTree(this._endpoint.modelMaxPromptTokens, token);
        this.tracer?.didMaterializeTree?.({
            budget: this._endpoint.modelMaxPromptTokens,
            renderedTree: { container, removed, budget: this._endpoint.modelMaxPromptTokens },
            tokenizer: this._tokenizer,
            renderTree: budget => this._getFinalElementTree(budget, undefined).then(r => ({ ...r, budget })),
        });
        // Then finalize the chat messages
        const messageResult = [...container.toChatMessages()];
        const tokenCount = await container.tokenCount(this._tokenizer);
        const remainingMetadata = [...container.allMetadata()];
        // Remove undefined and duplicate references
        const referenceNames = new Set();
        const references = remainingMetadata
            .map(m => {
            if (!(m instanceof ReferenceMetadata)) {
                return;
            }
            const ref = m.reference;
            const isVariableName = 'variableName' in ref.anchor;
            if (isVariableName && !referenceNames.has(ref.anchor.variableName)) {
                referenceNames.add(ref.anchor.variableName);
                return ref;
            }
            else if (!isVariableName) {
                return ref;
            }
        })
            .filter(isDefined);
        // Collect the references for chat message chunks that did not survive prioritization
        const omittedReferences = allMetadata
            .map(m => {
            if (!(m instanceof ReferenceMetadata) || remainingMetadata.includes(m)) {
                return;
            }
            const ref = m.reference;
            const isVariableName = 'variableName' in ref.anchor;
            if (isVariableName && !referenceNames.has(ref.anchor.variableName)) {
                referenceNames.add(ref.anchor.variableName);
                return ref;
            }
            else if (!isVariableName) {
                return ref;
            }
        })
            .filter(isDefined);
        return {
            metadata: MetadataMap.from(remainingMetadata),
            messages: messageResult,
            hasIgnoredFiles: this._ignoredFiles.length > 0,
            tokenCount,
            references,
            omittedReferences,
        };
    }
    /**
     * Note: this may be called multiple times from the tracer as users play
     * around with budgets. It should be side-effect-free.
     */
    async _getFinalElementTree(tokenBudget, token) {
        const root = this._root.materialize();
        const originalMessages = [...root.toChatMessages()];
        const allMetadata = [...root.allMetadata()];
        const limits = [{ limit: tokenBudget, id: this._root.id }, ...this._tokenLimits];
        let removed = 0;
        for (let i = limits.length - 1; i >= 0; i--) {
            const limit = limits[i];
            if (limit.limit > tokenBudget) {
                continue;
            }
            const container = root.findById(limit.id);
            if (!container) {
                continue;
            }
            const initialTokenCount = await container.tokenCount(this._tokenizer);
            if (initialTokenCount < limit.limit) {
                const didChange = await this._grow(container, initialTokenCount, limit.limit, token);
                // if nothing grew, we already counted tokens so we can safely return
                if (!didChange) {
                    continue;
                }
            }
            // Trim the elements to fit within the token budget. The "upper bound" count
            // is a cachable count derived from the individual token counts of each component.
            // The actual token count is <= the upper bound count due to BPE merging of tokens
            // at component boundaries.
            //
            // To avoid excess tokenization, we first calculate the precise token
            // usage of the message, and then remove components, subtracting their
            // "upper bound" usage from the count until it's <= the budget. We then
            // repeat this and refine as necessary, though most of the time we only
            // need a single iteration of this.<sup>[citation needed]</sup>
            try {
                let tokenCount = await container.tokenCount(this._tokenizer);
                while (tokenCount > limit.limit) {
                    const overhead = await container.baseMessageTokenCount(this._tokenizer);
                    do {
                        for (const node of container.removeLowestPriorityChild()) {
                            removed++;
                            const rmCount = node.upperBoundTokenCount(this._tokenizer);
                            // buffer an extra 25% to roughly account for any potential undercount
                            tokenCount -= (typeof rmCount === 'number' ? rmCount : await rmCount) * 1.25;
                        }
                    } while (tokenCount - overhead > limit.limit);
                    tokenCount = await container.tokenCount(this._tokenizer);
                }
            }
            catch (e) {
                if (e instanceof materialized_1.BudgetExceededError) {
                    e.metadata = MetadataMap.from([...root.allMetadata()]);
                    e.messages = originalMessages;
                }
                throw e;
            }
        }
        return { container: root, allMetadata, removed };
    }
    /** Grows all Expandable elements, returns if any changes were made. */
    async _grow(tree, tokensUsed, tokenBudget, token) {
        if (!this._growables.length) {
            return false;
        }
        for (const growable of this._growables) {
            if (!tree.findById(growable.elem.id)) {
                continue; // not in this subtree
            }
            const obj = growable.elem.getObj();
            if (!(obj instanceof promptElements_1.Expandable)) {
                throw new Error('unreachable: expected growable');
            }
            const tempRoot = new PromptTreeElement(null, 0, growable.elem.id);
            // Sizing for the grow is the remaining excess plus the initial consumption,
            // since the element consuming the initial amount of tokens will be replaced
            const sizing = new PromptSizingContext(tokenBudget - tokensUsed + growable.initialConsume, this._endpoint);
            const newConsumed = await this._processPromptRenderPiece(sizing, { node: tempRoot, ctor: this._ctor, props: {}, children: [], path: [this._ctor] }, obj, await obj.render(undefined, {
                tokenBudget: sizing.tokenBudget,
                endpoint: this._endpoint,
                countTokens: (text, cancellation) => this._tokenizer.tokenLength(typeof text === 'string'
                    ? { type: mode_1.Raw.ChatCompletionContentPartKind.Text, text }
                    : text, cancellation),
            }), undefined, token);
            const newContainer = tempRoot.materialize();
            const oldContainer = tree.replaceNode(growable.elem.id, newContainer);
            if (!oldContainer) {
                throw new Error('unreachable: could not find old element to replace');
            }
            tokensUsed -= growable.initialConsume;
            tokensUsed += newConsumed;
            if (tokensUsed >= tokenBudget) {
                break;
            }
        }
        return true;
    }
    _handlePromptChildren(element, pieces, sizing, progress, token) {
        if (element.ctor === promptElements_1.TextChunk) {
            this._handleExtrinsicTextChunkChildren(element.node, element.node, element.props, pieces);
            return;
        }
        let todo = [];
        for (const piece of pieces) {
            if (piece.kind === 'literal') {
                element.node.appendStringChild(piece.value, element.props.priority ?? Number.MAX_SAFE_INTEGER);
                continue;
            }
            if (piece.kind === 'intrinsic') {
                // intrinsic element
                this._handleIntrinsic(element.node, piece.name, {
                    priority: element.props.priority ?? Number.MAX_SAFE_INTEGER,
                    ...piece.props,
                }, flattenAndReduceArr(piece.children));
                continue;
            }
            const childNode = element.node.createChild();
            todo.push({
                node: childNode,
                ctor: piece.ctor,
                props: piece.props,
                children: piece.children,
                path: [...element.path, piece.ctor],
            });
        }
        return this._processPromptPieces(sizing, todo, progress, token);
    }
    _handleIntrinsic(node, name, props, children, sortIndex) {
        switch (name) {
            case 'meta':
                return this._handleIntrinsicMeta(node, props, children);
            case 'br':
                return this._handleIntrinsicLineBreak(node, props, children, props.priority, sortIndex);
            case 'usedContext':
                return this._handleIntrinsicUsedContext(node, props, children);
            case 'references':
                return this._handleIntrinsicReferences(node, props, children);
            case 'ignoredFiles':
                return this._handleIntrinsicIgnoredFiles(node, props, children);
            case 'elementJSON':
                return this._handleIntrinsicElementJSON(node, props.data);
            case 'cacheBreakpoint':
                return this._handleIntrinsicCacheBreakpoint(node, props, children, sortIndex);
            case 'opaque':
                return this._handleIntrinsicOpaque(node, props, sortIndex);
        }
        throw new Error(`Unknown intrinsic element ${name}!`);
    }
    _handleIntrinsicCacheBreakpoint(node, props, children, sortIndex) {
        if (children.length > 0) {
            throw new Error(`<cacheBreakpoint /> must not have children!`);
        }
        node.addCacheBreakpoint(props, sortIndex);
    }
    _handleIntrinsicMeta(node, props, children) {
        if (children.length > 0) {
            throw new Error(`<meta /> must not have children!`);
        }
        if (props.local) {
            node.addMetadata(props.value);
        }
        else {
            this._root.addMetadata(props.value);
        }
    }
    _handleIntrinsicLineBreak(node, props, children, inheritedPriority, sortIndex) {
        if (children.length > 0) {
            throw new Error(`<br /> must not have children!`);
        }
        node.appendLineBreak(inheritedPriority ?? Number.MAX_SAFE_INTEGER, sortIndex);
    }
    _handleIntrinsicOpaque(node, props, sortIndex) {
        node.appendOpaque(props.value, props.tokenUsage, props.priority, sortIndex);
    }
    _handleIntrinsicElementJSON(node, data) {
        const appended = node.appendPieceJSON(data.node);
        if (this.tracer?.includeInEpoch) {
            for (const child of appended.elements()) {
                // tokenBudget is just 0 because we don't know the renderer state on the tool side.
                this.tracer.includeInEpoch({ id: child.id, tokenBudget: 0 });
            }
        }
    }
    _handleIntrinsicUsedContext(node, props, children) {
        if (children.length > 0) {
            throw new Error(`<usedContext /> must not have children!`);
        }
        this._usedContext.push(...props.value);
    }
    _handleIntrinsicReferences(node, props, children) {
        if (children.length > 0) {
            throw new Error(`<reference /> must not have children!`);
        }
        for (const ref of props.value) {
            node.addMetadata(new ReferenceMetadata(ref));
        }
    }
    _handleIntrinsicIgnoredFiles(node, props, children) {
        if (children.length > 0) {
            throw new Error(`<ignoredFiles /> must not have children!`);
        }
        this._ignoredFiles.push(...props.value);
    }
    /**
     * @param node Parent of the <TextChunk />
     * @param textChunkNode The <TextChunk /> node. All children are in-order
     * appended to the parent using the same sort index to ensure order is preserved.
     * @param props Props of the <TextChunk />
     * @param children Rendered children of the <TextChunk />
     */
    _handleExtrinsicTextChunkChildren(node, textChunkNode, props, children) {
        const content = [];
        const metadata = [];
        for (const child of children) {
            if (child.kind === 'extrinsic') {
                throw new Error('TextChunk cannot have extrinsic children!');
            }
            if (child.kind === 'literal') {
                content.push(child.value);
            }
            if (child.kind === 'intrinsic') {
                if (child.name === 'br') {
                    // Preserve newlines
                    content.push('\n');
                }
                else if (child.name === 'references') {
                    // For TextChunks, references must be propagated through the PromptText element that is appended to the node
                    for (const reference of child.props.value) {
                        metadata.push(new ReferenceMetadata(reference));
                    }
                }
                else {
                    this._handleIntrinsic(node, child.name, child.props, flattenAndReduceArr(child.children), textChunkNode.childIndex);
                }
            }
        }
        node.appendStringChild(content.join(''), props?.priority ?? Number.MAX_SAFE_INTEGER, metadata, textChunkNode.childIndex, true);
    }
}
exports.PromptRenderer = PromptRenderer;
async function computeTokensConsumedByLiterals(tokenizer, element, instance, pieces) {
    let tokensConsumed = 0;
    if ((0, promptElements_1.isChatMessagePromptElement)(instance)) {
        const raw = {
            role: element.props.role,
            content: [],
            ...(element.props.name ? { name: element.props.name } : undefined),
            ...(element.props.toolCalls ? { toolCalls: element.props.toolCalls } : undefined),
            ...(element.props.toolCallId ? { toolCallId: element.props.toolCallId } : undefined),
        };
        tokensConsumed += await tokenizer.countMessageTokens((0, mode_1.toMode)(tokenizer.mode, raw));
    }
    for (const piece of pieces) {
        if (piece.kind === 'literal') {
            tokensConsumed += await tokenizer.tokenLength({
                type: mode_1.Raw.ChatCompletionContentPartKind.Text,
                text: piece.value,
            });
        }
    }
    return { tokensConsumed };
}
// Flatten nested fragments and normalize children
function flattenAndReduce(c, into = []) {
    if (typeof c === 'undefined' || typeof c === 'boolean') {
        // booleans are ignored to allow for the pattern: { cond && <Element ... /> }
        return [];
    }
    else if (typeof c === 'string' || typeof c === 'number') {
        into.push(new LiteralPromptPiece(String(c)));
    }
    else if (isFragmentCtor(c)) {
        flattenAndReduceArr(c.children, into);
    }
    else if (isIterable(c)) {
        flattenAndReduceArr(c, into);
    }
    else if (typeof c.ctor === 'string') {
        // intrinsic element
        into.push(new IntrinsicPromptPiece(c.ctor, c.props, c.children));
    }
    else {
        // extrinsic element
        into.push(new ExtrinsicPromptPiece(c.ctor, c.props, c.children));
    }
    return into;
}
function flattenAndReduceArr(arr, into = []) {
    for (const entry of arr) {
        flattenAndReduce(entry, into);
    }
    return into;
}
class IntrinsicPromptPiece {
    name;
    props;
    children;
    kind = 'intrinsic';
    constructor(name, props, children) {
        this.name = name;
        this.props = props;
        this.children = children;
    }
}
class ExtrinsicPromptPiece {
    ctor;
    props;
    children;
    kind = 'extrinsic';
    constructor(ctor, props, children) {
        this.ctor = ctor;
        this.props = props;
        this.children = children;
    }
}
class LiteralPromptPiece {
    value;
    priority;
    kind = 'literal';
    constructor(value, priority) {
        this.value = value;
        this.priority = priority;
    }
}
class PromptOpaque {
    parent;
    childIndex;
    value;
    tokenUsage;
    priority;
    static fromJSON(parent, index, json) {
        return new PromptOpaque(parent, index, json.value, json.tokenUsage, json.priority);
    }
    kind = 2 /* PromptNodeType.Text */;
    constructor(parent, childIndex, value, tokenUsage, priority) {
        this.parent = parent;
        this.childIndex = childIndex;
        this.value = value;
        this.tokenUsage = tokenUsage;
        this.priority = priority;
    }
    materialize(parent) {
        return new materialized_1.MaterializedChatMessageOpaque(parent, {
            type: mode_1.Raw.ChatCompletionContentPartKind.Opaque,
            value: this.value,
            tokenUsage: this.tokenUsage,
        }, this.priority);
    }
    toJSON() {
        return {
            type: 3 /* JSONT.PromptNodeType.Opaque */,
            value: this.value,
            tokenUsage: this.tokenUsage,
            priority: this.priority,
        };
    }
}
/**
 * A shared instance given to each PromptTreeElement that contains information
 * about the parent sizing and its children.
 */
class PromptSizingContext {
    tokenBudget;
    endpoint;
    _consumed = 0;
    constructor(tokenBudget, endpoint) {
        this.tokenBudget = tokenBudget;
        this.endpoint = endpoint;
    }
    get consumed() {
        return this._consumed > this.tokenBudget ? this.tokenBudget : this._consumed;
    }
    get remainingTokenBudget() {
        return Math.max(0, this.tokenBudget - this._consumed);
    }
    /** Marks part of the budget as having been consumed by a render() call. */
    consume(budget) {
        this._consumed += budget;
    }
}
class PromptTreeElement {
    parent;
    childIndex;
    id;
    static _nextId = 0;
    static fromJSON(index, json, keepWithMap) {
        const element = new PromptTreeElement(null, index);
        element._metadata =
            json.references?.map(r => new ReferenceMetadata(results_1.PromptReference.fromJSON(r))) ?? [];
        element._children = json.children
            .map((childJson, i) => {
            switch (childJson.type) {
                case 1 /* JSONT.PromptNodeType.Piece */:
                    return PromptTreeElement.fromJSON(i, childJson, keepWithMap);
                case 2 /* JSONT.PromptNodeType.Text */:
                    return PromptText.fromJSON(element, i, childJson);
                case 3 /* JSONT.PromptNodeType.Opaque */:
                    return PromptOpaque.fromJSON(element, i, childJson);
                default:
                    softAssertNever(childJson);
            }
        })
            .filter(isDefined);
        switch (json.ctor) {
            case 1 /* JSONT.PieceCtorKind.BaseChatMessage */:
                element._objFlags = json.flags ?? 0;
                element._obj = new promptElements_1.BaseChatMessage(json.props);
                break;
            case 2 /* JSONT.PieceCtorKind.Other */: {
                if (json.keepWithId !== undefined) {
                    let kw = keepWithMap.get(json.keepWithId);
                    if (!kw) {
                        kw = (0, promptElements_1.useKeepWith)();
                        keepWithMap.set(json.keepWithId, kw);
                    }
                    element._obj = new kw(json.props || {});
                }
                else {
                    element._obj = new promptElements_1.LogicalWrapper(json.props || {});
                }
                element._objFlags = json.flags ?? 0;
                break;
            }
            case 3 /* JSONT.PieceCtorKind.ImageChatMessage */:
                element._obj = new promptElements_1.Image(json.props);
                break;
            case 4 /* JSONT.PieceCtorKind.DocumentChatMessage */:
                element._obj = new promptElements_1.Document(json.props);
                break;
            default:
                softAssertNever(json);
        }
        return element;
    }
    kind = 1 /* PromptNodeType.Piece */;
    _obj = null;
    _state = undefined;
    _children = [];
    _metadata = [];
    _objFlags = 0;
    constructor(parent = null, childIndex, id = PromptTreeElement._nextId++) {
        this.parent = parent;
        this.childIndex = childIndex;
        this.id = id;
    }
    setObj(obj) {
        this._obj = obj;
        // todo@connor4312: clean this up so we don't actually hold _obj but instead
        // just hold metadata that can be more cleanly rehydrated
        if (this._obj instanceof promptElements_1.LegacyPrioritization)
            this._objFlags |= 1 /* ContainerFlags.IsLegacyPrioritization */;
        if (this._obj instanceof promptElements_1.Chunk)
            this._objFlags |= 2 /* ContainerFlags.IsChunk */;
        if (this._obj instanceof promptElements_1.IfEmpty)
            this._objFlags |= 8 /* ContainerFlags.EmptyAlternate */;
        if (this._obj.props.passPriority)
            this._objFlags |= 4 /* ContainerFlags.PassPriority */;
    }
    /** @deprecated remove when Expandable is gone */
    getObj() {
        return this._obj;
    }
    setState(state) {
        this._state = state;
    }
    getState() {
        return this._state;
    }
    createChild() {
        const child = new PromptTreeElement(this, this._children.length);
        this._children.push(child);
        return child;
    }
    appendPieceJSON(data) {
        const child = PromptTreeElement.fromJSON(this._children.length, data, new Map());
        this._children.push(child);
        return child;
    }
    appendStringChild(text, priority, metadata, sortIndex = this._children.length, lineBreakBefore = false) {
        this._children.push(new PromptText(this, sortIndex, text, priority, metadata, lineBreakBefore));
    }
    appendLineBreak(priority, sortIndex = this._children.length) {
        this._children.push(new PromptText(this, sortIndex, '\n', priority));
    }
    appendOpaque(value, tokenUsage, priority, sortIndex = this._children.length) {
        this._children.push(new PromptOpaque(this, sortIndex, value, tokenUsage, priority));
    }
    toJSON() {
        const json = {
            type: 1 /* JSONT.PromptNodeType.Piece */,
            ctor: 2 /* JSONT.PieceCtorKind.Other */,
            ctorName: this._obj?.constructor.name,
            children: this._children
                .slice()
                .sort((a, b) => a.childIndex - b.childIndex)
                .map(c => c.toJSON())
                .filter(isDefined),
            props: {},
            references: this._metadata
                .filter(m => m instanceof ReferenceMetadata)
                .map(r => r.reference.toJSON()),
        };
        if (this._obj) {
            json.props = pickProps(this._obj.props, JSONT.jsonRetainedProps);
        }
        if (this._obj instanceof promptElements_1.BaseChatMessage) {
            json.ctor = 1 /* JSONT.PieceCtorKind.BaseChatMessage */;
            Object.assign(json.props, pickProps(this._obj.props, ['role', 'name', 'toolCalls', 'toolCallId']));
        }
        else if (this._obj instanceof promptElements_1.Image) {
            return {
                ...json,
                ctor: 3 /* JSONT.PieceCtorKind.ImageChatMessage */,
                props: {
                    ...json.props,
                    ...pickProps(this._obj.props, ['src', 'detail', 'mimeType']),
                },
            };
        }
        else if (this._obj instanceof promptElements_1.Document) {
            return {
                ...json,
                ctor: 4 /* JSONT.PieceCtorKind.DocumentChatMessage */,
                props: {
                    ...json.props,
                    ...pickProps(this._obj.props, ['data', 'mediaType']),
                },
            };
        }
        else if (this._obj instanceof promptElements_1.AbstractKeepWith) {
            json.keepWithId = this._obj.id;
        }
        if (this._objFlags !== 0) {
            json.flags = this._objFlags;
        }
        return json;
    }
    materialize(parent) {
        this._children.sort((a, b) => a.childIndex - b.childIndex);
        if (this._obj instanceof promptElements_1.Image) {
            // #region materialize baseimage
            return new materialized_1.MaterializedChatMessageImage(parent, this.id, this._obj.props.src, this._obj.props.priority ?? Number.MAX_SAFE_INTEGER, this._metadata, 0 /* LineBreakBefore.None */, this._obj.props.detail ?? undefined, this._obj.props.mimeType ?? undefined);
        }
        if (this._obj instanceof promptElements_1.Document) {
            return new materialized_1.MaterializedChatMessageDocument(parent, this.id, this._obj.props.data, this._obj.props.mediaType, this._obj.props.priority ?? Number.MAX_SAFE_INTEGER, this._metadata, 0 /* LineBreakBefore.None */);
        }
        if (this._obj instanceof promptElements_1.BaseChatMessage) {
            if (this._obj.props.role === undefined || typeof this._obj.props.role !== 'number') {
                throw new Error(`Invalid ChatMessage!`);
            }
            return new materialized_1.MaterializedChatMessage(parent, this.id, this._obj.props.role, this._obj.props.name, this._obj instanceof promptElements_1.AssistantMessage ? this._obj.props.toolCalls : undefined, this._obj instanceof promptElements_1.ToolMessage ? this._obj.props.toolCallId : undefined, this._obj.props.priority ?? Number.MAX_SAFE_INTEGER, this._metadata, parent => this._children.map(child => child.materialize(parent)));
        }
        else {
            const container = new materialized_1.GenericMaterializedContainer(parent, this.id, this._obj?.constructor.name, this._obj?.props.priority ?? (this._obj?.props.passPriority ? 0 : Number.MAX_SAFE_INTEGER), parent => this._children.map(child => child.materialize(parent)), this._metadata, this._objFlags);
            if (this._obj instanceof promptElements_1.AbstractKeepWith) {
                container.keepWithId = this._obj.id;
            }
            return container;
        }
    }
    addMetadata(metadata) {
        this._metadata.push(metadata);
    }
    addCacheBreakpoint(breakpoint, sortIndex = this._children.length) {
        if (!(this._obj instanceof promptElements_1.BaseChatMessage)) {
            throw new Error('Cache breakpoints may only be direct children of chat messages');
        }
        this._children.push(new PromptCacheBreakpoint({ type: mode_1.Raw.ChatCompletionContentPartKind.CacheBreakpoint, cacheType: breakpoint.type }, sortIndex));
    }
    *elements() {
        yield this;
        for (const child of this._children) {
            if (child instanceof PromptTreeElement) {
                yield* child.elements();
            }
        }
    }
}
class PromptCacheBreakpoint {
    part;
    childIndex;
    constructor(part, childIndex) {
        this.part = part;
        this.childIndex = childIndex;
    }
    toJSON() {
        return undefined;
    }
    materialize(parent) {
        return new materialized_1.MaterializedChatMessageBreakpoint(parent, this.part);
    }
}
class PromptText {
    parent;
    childIndex;
    text;
    priority;
    metadata;
    lineBreakBefore;
    static fromJSON(parent, index, json) {
        return new PromptText(parent, index, json.text, json.priority, json.references?.map(r => new ReferenceMetadata(results_1.PromptReference.fromJSON(r))), json.lineBreakBefore);
    }
    kind = 2 /* PromptNodeType.Text */;
    constructor(parent, childIndex, text, priority, metadata, lineBreakBefore = false) {
        this.parent = parent;
        this.childIndex = childIndex;
        this.text = text;
        this.priority = priority;
        this.metadata = metadata;
        this.lineBreakBefore = lineBreakBefore;
    }
    materialize(parent) {
        const lineBreak = this.lineBreakBefore
            ? 1 /* LineBreakBefore.Always */
            : this.childIndex === 0
                ? 2 /* LineBreakBefore.IfNotTextSibling */
                : 0 /* LineBreakBefore.None */;
        return new materialized_1.MaterializedChatMessageTextChunk(parent, this.text, this.priority ?? Number.MAX_SAFE_INTEGER, this.metadata || [], lineBreak);
    }
    toJSON() {
        return {
            type: 2 /* JSONT.PromptNodeType.Text */,
            priority: this.priority,
            text: this.text,
            references: this.metadata
                ?.filter(m => m instanceof ReferenceMetadata)
                .map(r => r.reference.toJSON()),
            lineBreakBefore: this.lineBreakBefore,
        };
    }
}
function isFragmentCtor(template) {
    return (typeof template.ctor === 'function' && template.ctor.isFragment) ?? false;
}
function softAssertNever(x) {
    // note: does not actually throw, because we want to handle any unknown cases
    // gracefully for forwards-compatibility
}
function isDefined(x) {
    return x !== undefined;
}
class InternalMetadata extends results_1.PromptMetadata {
}
class ReferenceMetadata extends InternalMetadata {
    reference;
    constructor(reference) {
        super();
        this.reference = reference;
    }
}
function iterableToArray(t) {
    if (isIterable(t)) {
        return Array.from(t);
    }
    return t;
}
function isIterable(t) {
    return !!t && typeof t[Symbol.iterator] === 'function';
}
function pickProps(obj, keys) {
    const result = {};
    for (const key of keys) {
        if (obj.hasOwnProperty(key)) {
            result[key] = obj[key];
        }
    }
    return result;
}
function atPath(path) {
    return path
        .map(p => (typeof p === 'string' ? p : p ? p.name || '<anonymous>' : String(p)))
        .join(' > ');
}
const annotatedErrors = new WeakSet();
async function annotateError(q, fn) {
    try {
        return await fn();
    }
    catch (e) {
        // Add a path to errors, except cancellation errors which are generally expected
        if (e instanceof Error &&
            !annotatedErrors.has(e) &&
            e.constructor.name !== 'CancellationError') {
            annotatedErrors.add(e);
            e.message += ` (at tsx element ${atPath(q.path)})`;
        }
        throw e;
    }
}
