import { isRunTree, RunTree } from "../../run_trees.js";
import { getCurrentRunTree, withRunTree } from "../../singletons/traceable.js";
import { isTracingEnabled } from "../../env.js";
import { convertMessageToTracedFormat } from "./utils.js";
import { setUsageMetadataOnRunTree } from "./utils.js";
import { isPrimitive, isRecord } from "../../utils/types.js";
function _formatMessages(messages) {
    if (!Array.isArray(messages))
        return messages;
    return messages.map((msg) => convertMessageToTracedFormat(msg));
}
// oxlint-disable-next-line typescript/no-explicit-any
function _formatToolCalls(toolCalls) {
    return toolCalls.map((tc) => ({
        id: tc.toolCallId,
        type: "function",
        function: {
            name: tc.toolName,
            arguments: isPrimitive(tc.input)
                ? String(tc.input)
                : JSON.stringify(tc.input),
        },
    }));
}
function _formatStepOutput(
// oxlint-disable-next-line typescript/no-explicit-any
event, traceRawHttp) {
    // Build an assistant-style message from the step result
    const output = { role: "assistant" };
    // Text content
    if (event.content != null) {
        output.content = event.content;
    }
    else if (event.text != null) {
        output.content = event.text;
    }
    // Tool calls
    if (Array.isArray(event.toolCalls) && event.toolCalls.length > 0) {
        output.tool_calls = _formatToolCalls(event.toolCalls);
    }
    if (event.finishReason != null) {
        output.finish_reason = event.finishReason;
    }
    if (traceRawHttp) {
        if (event.request != null)
            output.request = event.request;
        if (event.response != null)
            output.response = event.response;
    }
    return convertMessageToTracedFormat(output);
}
function _getLsAgentType(parentRunTree) {
    if (isRunTree(parentRunTree) && parentRunTree.run_type === "tool") {
        return "subagent";
    }
    return "root";
}
/**
 * Creates a LangSmith `Telemetry` for the Vercel AI SDK.
 *
 * This adapter implements the Vercel AI SDK's `Telemetry` interface
 * and maps lifecycle events to LangSmith traces. It creates a root span for
 * the entire generation, child LLM spans for each step, and tool spans for
 * tool calls.
 *
 * ```ts
 * import { generateText, registerTelemetry } from "ai";
 * import { LangSmithTelemetry } from "langsmith/experimental/vercel";
 *
 * registerTelemetry(LangSmithTelemetry());
 *
 * const result = await generateText({
 *   model: openai("gpt-4o"),
 *   prompt: "Hello!",
 * });
 * ```
 *
 * @experimental Only available in Vercel AI SDK 7.
 */
export function LangSmithTelemetry(config) {
    const { name: customName, runType = "chain", metadata: customMetadata, tags: customTags, client, projectName, processInputs, processOutputs, processChildLLMRunInputs, processChildLLMRunOutputs, traceResponseMetadata, traceRawHttp, tracingEnabled, extra: customExtra, } = config ?? {};
    function getOpenStepOrRoot(state) {
        let openStep;
        state.stepRunTrees.forEach((stepRt) => {
            if (stepRt.end_time == null) {
                openStep = stepRt;
            }
        });
        return openStep ?? state.rootRunTree;
    }
    async function finalizeOpenToolRuns(state, opts) {
        const entries = Array.from(state.toolRunTrees.entries());
        for (let i = 0; i < entries.length; i++) {
            const [, toolRt] = entries[i];
            if (toolRt.end_time == null) {
                if (opts?.error != null) {
                    await toolRt.end(undefined, opts.error);
                }
                else {
                    await toolRt.end(opts?.note != null ? { note: opts.note } : undefined);
                }
                await toolRt.patchRun({ excludeInputs: true });
            }
        }
        state.toolRunTrees.clear();
    }
    /** Per-generation state keyed by AI SDK `callId` (stable across nested calls). */
    const invocationsByCallId = new Map();
    const onStart = async (event) => {
        if (!isTracingEnabled(tracingEnabled))
            return;
        if (!("callId" in event) || typeof event.callId !== "string")
            return;
        // If called within an existing traceable context, nest under it
        const parentRunTree = getCurrentRunTree(true);
        let inputs = {};
        if (event.recordInputs !== false) {
            if ("messages" in event && event.messages != null) {
                inputs.messages = _formatMessages(event.messages);
            }
            if ("prompt" in event && event.prompt != null) {
                inputs.prompt = event.prompt;
            }
            if ("instructions" in event && event.instructions != null) {
                inputs.instructions = event.instructions;
            }
            if ("system" in event && event.system != null) {
                inputs.system = event.system;
            }
            if ("tools" in event && event.tools != null) {
                inputs.tools = Object.keys(event.tools);
            }
            if ("runtimeContext" in event && event.runtimeContext != null) {
                inputs.runtimeContext = event.runtimeContext;
            }
            if ("toolsContext" in event && event.toolsContext != null) {
                inputs.toolsContext = event.toolsContext;
            }
            // Apply user-provided input processing
            if (processInputs) {
                try {
                    inputs = processInputs(inputs);
                }
                catch (e) {
                    console.error("Error in processInputs, using raw inputs:", e);
                }
            }
        }
        const runTreeConfig = {
            name: customName ?? event.functionId ?? event.provider,
            run_type: runType,
            inputs,
            tracingEnabled: true,
            extra: {
                ...customExtra,
                metadata: {
                    ...customMetadata,
                    ai_sdk_method: event.operationId,
                    ls_agent_type: _getLsAgentType(parentRunTree),
                    ls_model_name: event.modelId,
                    ls_provider: event.provider,
                    ls_integration: "vercel-ai-sdk-telemetry",
                },
            },
            tags: customTags,
            ...(client ? { client } : {}),
            ...(projectName ? { project_name: projectName } : {}),
        };
        let rootRunTree;
        if (isRunTree(parentRunTree)) {
            rootRunTree = parentRunTree.createChild(runTreeConfig);
        }
        else {
            rootRunTree = new RunTree(runTreeConfig);
        }
        await rootRunTree.postRun();
        invocationsByCallId.set(event.callId, {
            rootRunTree,
            stepRunTrees: new Map(),
            toolRunTrees: new Map(),
        });
    };
    const onStepStart = async (event) => {
        const state = invocationsByCallId.get(event.callId);
        if (!state)
            return;
        const stepNumber = event.stepNumber ?? 0;
        let inputs = {};
        if (event.recordInputs !== false) {
            if ("messages" in event && event.messages != null) {
                inputs.messages = _formatMessages(event.messages);
            }
            if ("runtimeContext" in event && event.runtimeContext != null) {
                inputs.runtimeContext = event.runtimeContext;
            }
            if ("toolsContext" in event && event.toolsContext != null) {
                inputs.toolsContext = event.toolsContext;
            }
            if (processChildLLMRunInputs) {
                try {
                    inputs = processChildLLMRunInputs(inputs);
                }
                catch (e) {
                    console.error("Error in processChildLLMRunInputs, using raw inputs:", e);
                }
            }
        }
        const stepRunTree = state.rootRunTree.createChild({
            name: event.provider,
            run_type: "llm",
            inputs,
            extra: { metadata: { step_number: stepNumber } },
        });
        state.stepRunTrees.set(stepNumber, stepRunTree);
        await stepRunTree.postRun();
    };
    const onLanguageModelCallStart = async (event) => {
        const state = invocationsByCallId.get(event.callId);
        if (!state)
            return;
        const stepRunTree = getOpenStepOrRoot(state);
        if (stepRunTree.run_type !== "llm")
            return;
        const prevParams = isRecord(stepRunTree.extra?.invocation_params)
            ? stepRunTree.extra.invocation_params
            : {};
        const nextParams = { ...event };
        // Remove properties that are already in the step run tree
        delete nextParams.messages;
        delete nextParams.provider;
        delete nextParams.modelId;
        // Remove telemetry options (except functionId)
        delete nextParams.recordInputs;
        delete nextParams.recordOutputs;
        delete nextParams.includeToolsContext;
        // Massage tools for LangSmith to render schema nicely
        nextParams.tools = nextParams.tools?.map((tool) => {
            const newTool = { ...tool };
            if ("inputSchema" in newTool) {
                newTool.input_schema = newTool.inputSchema;
                delete newTool.inputSchema;
            }
            return newTool;
        });
        stepRunTree.extra = {
            ...stepRunTree.extra,
            invocation_params: { ...prevParams, ...nextParams },
        };
    };
    const onToolExecutionStart = async (event) => {
        const state = invocationsByCallId.get(event.callId);
        if (!state)
            return;
        const parentRunTree = getOpenStepOrRoot(state);
        let inputs = {};
        if (event.recordInputs !== false) {
            if (isRecord(event.toolCall.input)) {
                inputs = { ...event.toolCall.input };
            }
            else if (typeof event.toolCall.input !== "undefined") {
                inputs = { input: event.toolCall.input };
            }
            if ("toolContext" in event && event.toolContext != null) {
                inputs.toolContext = event.toolContext;
            }
        }
        else {
            inputs = {};
        }
        const toolRunTree = parentRunTree.createChild({
            name: event.toolCall.toolName,
            run_type: "tool",
            inputs,
            extra: {
                metadata: {
                    tool_call_id: event.toolCall.toolCallId,
                    ai_sdk_call_id: event.callId,
                },
            },
        });
        await toolRunTree.postRun();
        state.toolRunTrees.set(event.toolCall.toolCallId, toolRunTree);
    };
    const onToolExecutionEnd = async (event) => {
        const state = invocationsByCallId.get(event.callId);
        if (!state)
            return;
        const toolRunTree = state.toolRunTrees.get(event.toolCall.toolCallId);
        if (!toolRunTree)
            return;
        state.toolRunTrees.delete(event.toolCall.toolCallId);
        let outputs;
        let error;
        if (event.recordOutputs !== false) {
            if (event.toolOutput.type === "tool-result") {
                outputs = { output: event.toolOutput.output };
            }
            else if (event.toolOutput.type === "tool-error") {
                const err = event.toolOutput.error;
                error = err instanceof Error ? err.message : String(err);
            }
        }
        else {
            outputs = {};
        }
        await toolRunTree.end(outputs, error, Math.floor(toolRunTree.start_time + event.toolExecutionMs));
        await toolRunTree.patchRun({ excludeInputs: true });
    };
    const onStepFinish = async (event) => {
        const state = invocationsByCallId.get(event.callId);
        if (!state)
            return;
        const stepNumber = event.stepNumber ?? 0;
        const stepRunTree = state.stepRunTrees.get(stepNumber);
        if (!stepRunTree)
            return;
        let outputs = {};
        if (event.recordOutputs !== false) {
            outputs = _formatStepOutput(event, traceRawHttp);
            if (processChildLLMRunOutputs) {
                try {
                    outputs = processChildLLMRunOutputs(outputs);
                }
                catch (e) {
                    console.error("Error in processChildLLMRunOutputs, using raw outputs:", e);
                }
            }
        }
        // Set usage metadata
        // @ts-expect-error SharedV4ProviderMetadata is not assignable to SharedV2ProviderMetadata
        setUsageMetadataOnRunTree(event, stepRunTree);
        await stepRunTree.end(outputs);
        await stepRunTree.patchRun({ excludeInputs: true });
        state.stepRunTrees.delete(stepNumber);
    };
    const onEnd = async (event) => {
        if (!("callId" in event) || typeof event.callId !== "string")
            return;
        const state = invocationsByCallId.get(event.callId);
        if (!state)
            return;
        const { rootRunTree } = state;
        await finalizeOpenToolRuns(state, { note: "closed on finish" });
        // Ensure any remaining step runs are closed
        const remainingSteps = Array.from(state.stepRunTrees.entries());
        for (let i = 0; i < remainingSteps.length; i++) {
            const [stepNumber, stepRt] = remainingSteps[i];
            if (stepRt.end_time == null) {
                await stepRt.end({ note: "closed on finish" });
                await stepRt.patchRun({ excludeInputs: true });
            }
            state.stepRunTrees.delete(stepNumber);
        }
        let outputs = {};
        if (event.recordOutputs !== false) {
            // Final result output
            if ("text" in event && event.text != null) {
                outputs.content = event.text;
            }
            else if ("content" in event && event.content != null) {
                outputs.content = event.content;
            }
            if (outputs.content != null) {
                outputs.role = "assistant";
            }
            if ("object" in event && event.object != null) {
                outputs.object = event.object;
            }
            if ("toolCalls" in event &&
                Array.isArray(event.toolCalls) &&
                event.toolCalls.length > 0) {
                outputs.tool_calls = _formatToolCalls(event.toolCalls);
            }
            if ("finishReason" in event && event.finishReason != null) {
                outputs.finish_reason = event.finishReason;
            }
            if (traceResponseMetadata &&
                "steps" in event &&
                Array.isArray(event.steps)) {
                outputs.steps = event.steps.map((step, idx) => ({
                    step_number: idx,
                    ..._formatStepOutput(step, traceRawHttp),
                }));
            }
            if (processOutputs) {
                try {
                    outputs = processOutputs(outputs);
                }
                catch (e) {
                    console.error("Error in processOutputs, using raw outputs:", e);
                }
            }
        }
        // Set aggregated usage on root
        if ("totalUsage" in event && event.totalUsage != null) {
            setUsageMetadataOnRunTree(
            // @ts-expect-error SharedV4ProviderMetadata is not assignable to SharedV2ProviderMetadata
            { usage: event.totalUsage, providerMetadata: event.providerMetadata }, rootRunTree);
        }
        else if ("usage" in event && event.usage != null) {
            // @ts-expect-error SharedV4ProviderMetadata is not assignable to SharedV2ProviderMetadata
            setUsageMetadataOnRunTree(event, rootRunTree);
        }
        await rootRunTree.end(outputs);
        await rootRunTree.patchRun({ excludeInputs: true });
        invocationsByCallId.delete(event.callId);
    };
    const onError = async (payload) => {
        const callId = typeof payload === "object" &&
            payload !== null &&
            "callId" in payload &&
            typeof payload.callId === "string"
            ? payload.callId
            : undefined;
        const error = typeof payload === "object" && payload !== null && "error" in payload
            ? payload.error
            : payload;
        if (callId === undefined)
            return;
        const state = invocationsByCallId.get(callId);
        if (!state)
            return;
        const { rootRunTree } = state;
        const errorMsg = error instanceof Error ? error.message : String(error);
        await finalizeOpenToolRuns(state, { error: errorMsg });
        // Close any open step runs with error
        const errorSteps = Array.from(state.stepRunTrees.entries());
        for (let i = 0; i < errorSteps.length; i++) {
            const [stepNumber, stepRt] = errorSteps[i];
            if (stepRt.end_time == null) {
                await stepRt.end(undefined, errorMsg);
                await stepRt.patchRun({ excludeInputs: true });
            }
            state.stepRunTrees.delete(stepNumber);
        }
        await rootRunTree.end(undefined, errorMsg);
        await rootRunTree.patchRun({ excludeInputs: true });
        invocationsByCallId.delete(callId);
    };
    const executeTool = async (params) => {
        const state = invocationsByCallId.get(params.callId);
        const toolRunTree = state?.toolRunTrees.get(params.toolCallId);
        if (toolRunTree != null) {
            return withRunTree(toolRunTree, () => params.execute());
        }
        return params.execute();
    };
    return {
        onStart,
        onStepStart,
        onLanguageModelCallStart,
        onToolExecutionStart,
        onToolExecutionEnd,
        onStepFinish,
        onEnd,
        onError,
        executeTool,
    };
}
