import { wrapManagedAgentSessionEvents } from "./utils/anthropic_managed.js";
import { isTraceableFunction, traceable, } from "../traceable.js";
import { convertAnthropicUsageToInputTokenDetails } from "../utils/usage.js";
const TRACED_INVOCATION_KEYS = ["top_k", "top_p", "stream", "thinking"];
/**
 * Create usage metadata from Anthropic's token usage format.
 */
export function createUsageMetadata(anthropicUsage) {
    if (!anthropicUsage) {
        return undefined;
    }
    const inputTokens = typeof anthropicUsage.input_tokens === "number"
        ? anthropicUsage.input_tokens
        : 0;
    const outputTokens = typeof anthropicUsage.output_tokens === "number"
        ? anthropicUsage.output_tokens
        : 0;
    const inputTokenDetails = convertAnthropicUsageToInputTokenDetails(
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    anthropicUsage);
    // Anthropic cache tokens are ADDITIVE (not subsets of input_tokens like OpenAI).
    // Sum them into input_tokens so the backend cost calculation is correct.
    const cacheTokenSum = Object.values(inputTokenDetails).reduce((sum, v) => sum + (v ?? 0), 0);
    const adjustedInputTokens = inputTokens + cacheTokenSum;
    const adjustedTotalTokens = adjustedInputTokens + outputTokens;
    return {
        input_tokens: adjustedInputTokens,
        output_tokens: outputTokens,
        total_tokens: adjustedTotalTokens,
        ...(Object.keys(inputTokenDetails).length > 0 && {
            input_token_details: inputTokenDetails,
        }),
    };
}
/**
 * Process Anthropic message outputs
 */
function processMessageOutput(outputs) {
    const message = outputs;
    const result = { ...message };
    delete result.type;
    if (message.usage) {
        result.usage_metadata = createUsageMetadata(message.usage);
        delete result.usage;
    }
    return result;
}
/**
 * Accumulate a single content block delta into the content array.
 */
function accumulateContentBlockDelta(content, event) {
    const block = content[event.index];
    if (!block)
        return;
    if (block.type === "text" && event.delta.type === "text_delta") {
        block.text += event.delta.text;
    }
    else if (block.type === "tool_use" &&
        event.delta.type === "input_json_delta") {
        // Accumulate JSON input for tool use
        const toolBlock = block;
        toolBlock._partial_json =
            (toolBlock._partial_json ?? "") + event.delta.partial_json;
    }
}
/**
 * Aggregate streaming chunks into a complete message response
 */
const messageAggregator = (chunks) => {
    if (!chunks || chunks.length === 0) {
        return {
            role: "assistant",
            content: [],
        };
    }
    let message = {
        role: "assistant",
        content: [],
        model: "",
        stop_reason: null,
        stop_sequence: null,
    };
    // Track usage
    let usage = {
        input_tokens: 0,
        output_tokens: 0,
    };
    for (const chunk of chunks) {
        switch (chunk.type) {
            case "message_start":
                // Initialize message
                message = {
                    id: chunk.message.id,
                    role: chunk.message.role,
                    content: [],
                    model: chunk.message.model,
                    stop_reason: chunk.message.stop_reason,
                    stop_sequence: chunk.message.stop_sequence,
                };
                // Capture initial usage
                if (chunk.message.usage) {
                    usage = chunk.message.usage;
                }
                break;
            case "content_block_start":
                // Add new content block
                if (message.content) {
                    message.content[chunk.index] =
                        chunk.content_block;
                }
                break;
            case "content_block_delta":
                // Accumulate delta
                if (message.content) {
                    accumulateContentBlockDelta(message.content, chunk);
                }
                break;
            case "content_block_stop":
                // Finalize content block
                if (message.content) {
                    const block = message.content[chunk.index];
                    if (block?.type === "tool_use") {
                        const toolBlock = block;
                        if (toolBlock._partial_json) {
                            try {
                                toolBlock.input = JSON.parse(toolBlock._partial_json);
                            }
                            catch {
                                // Keep partial JSON as-is if parsing fails
                                toolBlock.input = toolBlock._partial_json;
                            }
                            delete toolBlock._partial_json;
                        }
                    }
                }
                break;
            case "message_delta":
                // Update message metadata
                message.stop_reason = chunk.delta.stop_reason;
                message.stop_sequence = chunk.delta.stop_sequence ?? null;
                if (chunk.usage) {
                    // Override only non-null keys
                    for (const [key, value] of Object.entries(chunk.usage)) {
                        if (value != null) {
                            usage[key] = value;
                        }
                    }
                }
                break;
            case "message_stop":
                // Message complete
                break;
        }
    }
    // Build final output
    const result = {
        ...message,
    };
    delete result.type;
    // Add usage metadata
    result.usage_metadata = createUsageMetadata(usage);
    return result;
};
/**
 * Wraps an Anthropic client's completion methods, enabling automatic LangSmith
 * tracing. Method signatures are unchanged, with the exception that you can pass
 * an additional and optional "langsmithExtra" field within the second parameter.
 *
 * @param anthropic An Anthropic client instance.
 * @param options LangSmith options.
 * @returns The wrapped client.
 *
 * @example
 * ```ts
 * import Anthropic from "@anthropic-ai/sdk";
 * import { wrapAnthropic } from "langsmith/wrappers/anthropic";
 *
 * const anthropic = wrapAnthropic(new Anthropic());
 *
 * // Non-streaming
 * const message = await anthropic.messages.create({
 *   model: "claude-sonnet-4-20250514",
 *   max_tokens: 1024,
 *   messages: [{ role: "user", content: "Hello!" }],
 * });
 *
 * // Streaming
 * const messageStream = anthropic.messages.stream({
 *   model: "claude-sonnet-4-20250514",
 *   max_tokens: 1024,
 *   messages: [{ role: "user", content: "Hello!" }],
 * });
 * const finalMessage = await messageStream.finalMessage();
 * ```
 */
export const wrapAnthropic = (anthropic, options) => {
    if (isTraceableFunction(anthropic.messages.create) ||
        isTraceableFunction(anthropic.messages.stream) ||
        (anthropic.beta?.sessions?.events?.stream &&
            isTraceableFunction(anthropic.beta.sessions.events.stream)) ||
        (anthropic.beta?.sessions?.events?.send &&
            isTraceableFunction(anthropic.beta.sessions.events.send))) {
        throw new Error("This instance of Anthropic client has been already wrapped once.");
    }
    const tracedAnthropicClient = { ...anthropic };
    // Extract ls_invocation_params from metadata
    const prepopulatedInvocationParams = typeof options?.metadata?.ls_invocation_params === "object"
        ? options.metadata.ls_invocation_params
        : {};
    // Remove ls_invocation_params from metadata to avoid duplication
    const { ls_invocation_params, ...restMetadata } = options?.metadata ?? {};
    const cleanedOptions = {
        ...options,
        metadata: restMetadata,
    };
    /**
     * Transform system parameter into visible message for playground editability.
     * This provides parity with the Python SDK behavior and enables system prompts
     * to be viewed and edited in the LangSmith playground.
     */
    function processSystemMessage(params) {
        if (!params.system) {
            return params;
        }
        const processed = { ...params };
        // Handle both string and ContentBlock[] formats
        const systemContent = Array.isArray(params.system)
            ? params.system
                .map((block) => typeof block === "string" ? block : block.text)
                .join("\n")
            : params.system;
        // Transform into first message
        processed.messages = [
            { role: "system", content: systemContent },
            // eslint-disable-next-line @typescript-eslint/no-explicit-any
            ...(params.messages || []),
        ];
        delete processed.system;
        return processed;
    }
    // Common configuration for messages.create
    const messagesCreateConfig = {
        name: "ChatAnthropic",
        run_type: "llm",
        aggregator: messageAggregator,
        argsConfigPath: [1, "langsmithExtra"],
        processInputs: processSystemMessage,
        getInvocationParams: (payload) => {
            if (typeof payload !== "object" || payload == null)
                return undefined;
            const params = payload;
            const ls_stop = (typeof params.stop_sequences === "string"
                ? [params.stop_sequences]
                : params.stop_sequences) ?? undefined;
            const ls_invocation_params = {};
            for (const [key, value] of Object.entries(params)) {
                if (TRACED_INVOCATION_KEYS.includes(key)) {
                    ls_invocation_params[key] = value;
                }
            }
            return {
                ls_provider: "anthropic",
                ls_model_type: "chat",
                ls_model_name: params.model,
                ls_max_tokens: params.max_tokens ?? undefined,
                ls_temperature: params.temperature ?? undefined,
                ls_stop,
                ls_invocation_params: {
                    ...prepopulatedInvocationParams,
                    ...ls_invocation_params,
                },
            };
        },
        processOutputs: processMessageOutput,
        ...cleanedOptions,
    };
    // Create a new messages object preserving the prototype
    tracedAnthropicClient.messages = Object.create(Object.getPrototypeOf(anthropic.messages));
    // Copy all own properties
    Object.assign(tracedAnthropicClient.messages, anthropic.messages);
    // Wrap messages.create
    tracedAnthropicClient.messages.create = traceable(anthropic.messages.create.bind(anthropic.messages), messagesCreateConfig);
    // Shared function to wrap stream methods
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    const wrapStreamMethod = (originalStreamFn) => {
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        return function (...args) {
            const stream = originalStreamFn(...args);
            if ("finalMessage" in stream &&
                typeof stream.finalMessage === "function") {
                const originalFinalMessage = stream.finalMessage.bind(stream);
                // eslint-disable-next-line @typescript-eslint/no-explicit-any
                stream.finalMessage = async (...args) => {
                    if ("done" in stream && typeof stream.done === "function") {
                        await stream.done();
                    }
                    for await (const _ of stream) {
                        // Finish consuming the stream if it has not already been consumed
                        // It should be relatively uncommon to consume an iterator after calling
                        // .finalMessage()
                    }
                    return originalFinalMessage(...args);
                };
            }
            return stream;
        };
    };
    // Wrap messages.stream
    tracedAnthropicClient.messages.stream = traceable(wrapStreamMethod(anthropic.messages.stream.bind(anthropic.messages)), {
        name: "ChatAnthropic",
        run_type: "llm",
        aggregator: messageAggregator,
        argsConfigPath: [1, "langsmithExtra"],
        processInputs: processSystemMessage,
        getInvocationParams: messagesCreateConfig.getInvocationParams,
        processOutputs: processMessageOutput,
        ...cleanedOptions,
    });
    if (anthropic.beta) {
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        const tracedBeta = { ...anthropic.beta };
        // Wrap beta.messages if it exists
        if (anthropic.beta.messages &&
            typeof anthropic.beta.messages.create === "function") {
            tracedBeta.messages = Object.create(Object.getPrototypeOf(anthropic.beta.messages));
            Object.assign(tracedBeta.messages, anthropic.beta.messages);
            // Wrap beta.messages.create
            tracedBeta.messages.create = traceable(anthropic.beta.messages.create.bind(anthropic.beta.messages), messagesCreateConfig);
            // Wrap beta.messages.parse if it exists
            if (typeof anthropic.beta.messages.parse === "function") {
                tracedBeta.messages.parse = traceable(anthropic.beta.messages.parse.bind(anthropic.beta.messages), messagesCreateConfig);
            }
            // Wrap beta.messages.stream if it exists
            if (typeof anthropic.beta.messages.stream === "function") {
                tracedBeta.messages.stream = traceable(wrapStreamMethod(anthropic.beta.messages.stream.bind(anthropic.beta.messages)), {
                    name: "ChatAnthropic",
                    run_type: "llm",
                    aggregator: messageAggregator,
                    argsConfigPath: [1, "langsmithExtra"],
                    processInputs: processSystemMessage,
                    getInvocationParams: messagesCreateConfig.getInvocationParams,
                    processOutputs: processMessageOutput,
                    ...cleanedOptions,
                });
            }
        }
        // Wrap Claude Managed Agents session events if they exist
        wrapManagedAgentSessionEvents({
            originalBeta: anthropic.beta,
            tracedBeta,
            cleanedOptions,
            prepopulatedInvocationParams,
        });
        tracedAnthropicClient.beta = tracedBeta;
    }
    return tracedAnthropicClient;
};
