var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
    return new (P || (P = Promise))(function (resolve, reject) {
        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
        step((generator = generator.apply(thisArg, _arguments || [])).next());
    });
};
var __rest = (this && this.__rest) || function (s, e) {
    var t = {};
    for (var p in s) if (Object.prototype.hasOwnProperty.call(s, p) && e.indexOf(p) < 0)
        t[p] = s[p];
    if (s != null && typeof Object.getOwnPropertySymbols === "function")
        for (var i = 0, p = Object.getOwnPropertySymbols(s); i < p.length; i++) {
            if (e.indexOf(p[i]) < 0 && Object.prototype.propertyIsEnumerable.call(s, p[i]))
                t[p[i]] = s[p[i]];
        }
    return t;
};
import { createMiddleware } from "langchain";
import YAML from "yaml";
function createPendingSequenceDebug(sequenceId) {
    return {
        sequenceId,
        startedAt: new Date().toISOString(),
        prompt: "",
        promptTokens: 0,
        reasoning: "",
        reasoningTokens: 0,
        text: "",
        textTokens: 0,
        cachedTokens: 0,
        responseId: null,
        toolCalls: [],
        pendingToolCalls: 0,
        resultType: null,
    };
}
function hasSequenceDebugContent(sequence) {
    return Boolean(sequence &&
        (sequence.prompt ||
            sequence.reasoning ||
            sequence.text ||
            sequence.toolCalls.length > 0));
}
function finalizeSequenceDebug(sequence) {
    var _a;
    return {
        sequenceId: sequence.sequenceId,
        startedAt: sequence.startedAt,
        prompt: sequence.prompt,
        promptTokens: sequence.promptTokens,
        reasoning: sequence.reasoning,
        reasoningTokens: sequence.reasoningTokens,
        text: sequence.text,
        textTokens: sequence.textTokens,
        cachedTokens: sequence.cachedTokens,
        responseId: sequence.responseId,
        toolCalls: sequence.toolCalls.map((_a) => {
            var { completed: _completed } = _a, toolCall = __rest(_a, ["completed"]);
            return toolCall;
        }),
        endedAt: new Date().toISOString(),
        resultType: (_a = sequence.resultType) !== null && _a !== void 0 ? _a : "final_text",
    };
}
function getDebugModelName(model) {
    return typeof model.getName === "function" ? model.getName() : undefined;
}
function stringifyPromptForDebug(params) {
    var _a, _b, _c;
    const { model, systemMessage, messages, tools, toolChoice, modelSettings } = params;
    const invocationParams = typeof model.invocationParams === "function"
        ? model.invocationParams(Object.assign(Object.assign(Object.assign({}, (modelSettings !== null && modelSettings !== void 0 ? modelSettings : {})), (tools.length > 0 ? { tools } : {})), (toolChoice !== undefined ? { tool_choice: toolChoice } : {})))
        : null;
    return YAML.stringify(Object.assign(Object.assign(Object.assign(Object.assign({ model: {
            name: (_a = getDebugModelName(model)) !== null && _a !== void 0 ? _a : null,
            provider: (_c = (_b = model._defaultConfig) === null || _b === void 0 ? void 0 : _b.modelProvider) !== null && _c !== void 0 ? _c : null,
            configuredModel: typeof model.model === "string" ? model.model : null,
        }, systemMessage,
        messages }, (tools.length > 0 ? { tools } : {})), (toolChoice !== undefined ? { toolChoice } : {})), (modelSettings ? { modelSettings } : {})), (invocationParams ? { invocationParams } : {})));
}
function getMessageBlocks(message) {
    if (Array.isArray(message.contentBlocks)) {
        return message.contentBlocks;
    }
    if (Array.isArray(message.content)) {
        return message.content;
    }
    return [];
}
function hasToolCallSignal(message) {
    var _a;
    return Boolean((Array.isArray(message.tool_calls) && message.tool_calls.length > 0) ||
        (Array.isArray(message.tool_call_chunks) && message.tool_call_chunks.length > 0) ||
        (Array.isArray((_a = message.additional_kwargs) === null || _a === void 0 ? void 0 : _a.tool_calls) &&
            message.additional_kwargs.tool_calls.length > 0));
}
function hasTokenCounter(model) {
    return (typeof model === "object" &&
        model !== null &&
        "getNumTokens" in model &&
        typeof model.getNumTokens === "function");
}
function countTokens(model, content) {
    return __awaiter(this, void 0, void 0, function* () {
        if (!content) {
            return 0;
        }
        if (!hasTokenCounter(model)) {
            return 0;
        }
        return yield model.getNumTokens(content);
    });
}
function extractSequenceResponseDebug(message) {
    var _a, _b, _c, _d, _e, _f, _g;
    const blocks = getMessageBlocks(message);
    const reasoning = blocks
        .filter((block) => (block === null || block === void 0 ? void 0 : block.type) === "reasoning")
        .map((block) => { var _a; return String((_a = block.reasoning) !== null && _a !== void 0 ? _a : ""); })
        .join("");
    const textFromBlocks = blocks
        .filter((block) => (block === null || block === void 0 ? void 0 : block.type) === "text")
        .map((block) => { var _a; return String((_a = block.text) !== null && _a !== void 0 ? _a : ""); })
        .join("");
    return {
        promptTokens: (_b = (_a = message.usage_metadata) === null || _a === void 0 ? void 0 : _a.input_tokens) !== null && _b !== void 0 ? _b : 0,
        reasoning,
        reasoningTokens: 0,
        text: textFromBlocks || (typeof message.content === "string" ? message.content : ""),
        textTokens: 0,
        cachedTokens: (_e = (_d = (_c = message.usage_metadata) === null || _c === void 0 ? void 0 : _c.input_token_details) === null || _d === void 0 ? void 0 : _d.cache_read) !== null && _e !== void 0 ? _e : 0,
        responseId: (_g = (_f = message.response_metadata) === null || _f === void 0 ? void 0 : _f.id) !== null && _g !== void 0 ? _g : null,
        resultType: hasToolCallSignal(message) ? "tool_calls" : "final_text",
    };
}
export function createSequenceDebugCollector() {
    let nextSequenceId = 1;
    let currentSequenceDebug = null;
    const history = [];
    const ensureSequenceDebug = () => {
        if (!currentSequenceDebug) {
            currentSequenceDebug = createPendingSequenceDebug(nextSequenceId++);
        }
        return currentSequenceDebug;
    };
    const flush = () => {
        if (!hasSequenceDebugContent(currentSequenceDebug)) {
            currentSequenceDebug = null;
            return;
        }
        const finalizedSequenceDebug = finalizeSequenceDebug(currentSequenceDebug);
        history.push(finalizedSequenceDebug);
        currentSequenceDebug = null;
    };
    return {
        handleModelCallStart(prompt) {
            if ((currentSequenceDebug === null || currentSequenceDebug === void 0 ? void 0 : currentSequenceDebug.resultType) &&
                currentSequenceDebug.pendingToolCalls === 0) {
                flush();
            }
            const sequenceDebug = ensureSequenceDebug();
            sequenceDebug.prompt = prompt;
        },
        handleModelCallComplete(params) {
            const sequenceDebug = ensureSequenceDebug();
            sequenceDebug.promptTokens = params.promptTokens;
            sequenceDebug.reasoning = params.reasoning;
            sequenceDebug.reasoningTokens = params.reasoningTokens;
            sequenceDebug.text = params.text;
            sequenceDebug.textTokens = params.textTokens;
            sequenceDebug.cachedTokens = params.cachedTokens;
            sequenceDebug.responseId = params.responseId;
            sequenceDebug.resultType = params.resultType;
            if (sequenceDebug.resultType === "final_text" &&
                sequenceDebug.pendingToolCalls === 0) {
                flush();
            }
        },
        handleToolCallEvent(event) {
            const sequenceDebug = ensureSequenceDebug();
            if (event.phase === "start") {
                sequenceDebug.toolCalls.push({
                    toolCallId: event.toolCallId,
                    toolName: event.toolName,
                    input: event.input,
                    output: null,
                    error: null,
                    completed: false,
                });
                sequenceDebug.pendingToolCalls += 1;
                return;
            }
            const pendingToolCall = sequenceDebug.toolCalls.find((toolCall) => toolCall.toolCallId === event.toolCallId && !toolCall.completed);
            if (pendingToolCall) {
                pendingToolCall.output = event.output;
                pendingToolCall.error = event.error;
                pendingToolCall.completed = true;
            }
            else {
                sequenceDebug.toolCalls.push({
                    toolCallId: event.toolCallId,
                    toolName: event.toolName,
                    input: "",
                    output: event.output,
                    error: event.error,
                    completed: true,
                });
            }
            sequenceDebug.pendingToolCalls = Math.max(0, sequenceDebug.pendingToolCalls - 1);
            if (sequenceDebug.resultType === "tool_calls" &&
                sequenceDebug.pendingToolCalls === 0) {
                flush();
            }
        },
        flush,
        getHistory() {
            return history;
        },
    };
}
export function createSequenceDebugMiddleware(sink) {
    return createMiddleware({
        name: "SequenceDebugMiddleware",
        wrapModelCall(request, handler) {
            return __awaiter(this, void 0, void 0, function* () {
                const prompt = stringifyPromptForDebug({
                    model: request.model,
                    systemMessage: request.systemMessage,
                    messages: request.messages,
                    tools: request.tools,
                    toolChoice: request.toolChoice,
                    modelSettings: request.modelSettings,
                });
                sink.handleModelCallStart(prompt);
                const response = yield handler(request);
                const debug = extractSequenceResponseDebug(response);
                const [promptTokens, reasoningTokens, textTokens] = yield Promise.all([
                    debug.promptTokens || countTokens(request.model, prompt),
                    countTokens(request.model, debug.reasoning),
                    countTokens(request.model, debug.text),
                ]);
                sink.handleModelCallComplete(Object.assign(Object.assign({}, debug), { promptTokens,
                    reasoningTokens,
                    textTokens }));
                return response;
            });
        },
    });
}
