import type { RunTreeConfig } from "../index.js";
import { traceable } from "../traceable.js";
import { _wrapClient } from "./generic.js";

/**
 * Wrap a Vercel AI SDK model, enabling automatic LangSmith tracing.
 * After wrapping a model, you can use it with the Vercel AI SDK Core
 * methods as normal.
 *
 * @example
 * ```ts
 * import { anthropic } from "@ai-sdk/anthropic";
 * import { streamText } from "ai";
 * import { wrapAISDKModel } from "langsmith/wrappers/vercel";
 *
 * const anthropicModel = anthropic("claude-3-haiku-20240307");
 *
 * const modelWithTracing = wrapAISDKModel(anthropicModel);
 *
 * const { textStream } = await streamText({
 *   model: modelWithTracing,
 *   prompt: "Write a vegetarian lasagna recipe for 4 people.",
 * });
 *
 * for await (const chunk of textStream) {
 *   console.log(chunk);
 * }
 * ```
 * @param model An AI SDK model instance.
 * @param options LangSmith options.
 * @returns
 */
export const wrapAISDKModel = <T extends object>(
  model: T,
  options?: Partial<RunTreeConfig>
): T => {
  if (
    !("doStream" in model) ||
    typeof model.doStream !== "function" ||
    !("doGenerate" in model) ||
    typeof model.doGenerate !== "function"
  ) {
    throw new Error(
      `Received invalid input. This version of wrapAISDKModel only supports Vercel LanguageModelV1 instances.`
    );
  }
  const runName = options?.name ?? model.constructor?.name;
  return new Proxy(model, {
    get(target, propKey, receiver) {
      const originalValue = target[propKey as keyof T];
      if (typeof originalValue === "function") {
        let __finalTracedIteratorKey;
        let aggregator;
        if (propKey === "doStream") {
          __finalTracedIteratorKey = "stream";
          aggregator = (chunks: any[]) => {
            return chunks.reduce(
              (aggregated, chunk) => {
                if (chunk.type === "text-delta") {
                  return {
                    ...aggregated,
                    text: aggregated.text + chunk.textDelta,
                  };
                } else if (chunk.type === "tool-call") {
                  return {
                    ...aggregated,
                    ...chunk,
                  };
                } else if (chunk.type === "finish") {
                  return {
                    ...aggregated,
                    usage: chunk.usage,
                    finishReason: chunk.finishReason,
                  };
                } else {
                  return aggregated;
                }
              },
              {
                text: "",
              }
            );
          };
        }
        return traceable(originalValue.bind(target), {
          run_type: "llm",
          name: runName,
          ...options,
          __finalTracedIteratorKey,
          aggregator,
        });
      } else if (
        originalValue != null &&
        !Array.isArray(originalValue) &&
        // eslint-disable-next-line no-instanceof/no-instanceof
        !(originalValue instanceof Date) &&
        typeof originalValue === "object"
      ) {
        return _wrapClient(
          originalValue,
          [runName, propKey.toString()].join("."),
          options
        );
      } else {
        return Reflect.get(target, propKey, receiver);
      }
    },
  });
};
