import type { Attributes, Span, Tracer, TracerProvider } from "@opentelemetry/api";
import { context, diag, SpanKind, SpanStatusCode, trace } from "@opentelemetry/api";
import { isTracingSuppressed } from "@opentelemetry/core";
import type {
  InstrumentationConfig,
  InstrumentationModuleDefinition,
} from "@opentelemetry/instrumentation";
import {
  InstrumentationBase,
  InstrumentationNodeModuleDefinition,
  safeExecuteInTheMiddle,
} from "@opentelemetry/instrumentation";
import type openai from "openai";
import { APIPromise } from "openai";
import type { Completion, CreateEmbeddingResponse, EmbeddingCreateParams } from "openai/resources";
import type {
  ChatCompletion,
  ChatCompletionChunk,
  ChatCompletionCreateParamsBase,
  ChatCompletionMessage,
  ChatCompletionMessageParam,
} from "openai/resources/chat/completions";
import type { CompletionCreateParamsBase } from "openai/resources/completions";
import type {
  Response as ResponseType,
  ResponseCreateParamsBase,
  ResponseStreamEvent,
} from "openai/resources/responses/responses";
import type { Stream } from "openai/streaming";

import type { TraceConfigOptions } from "@arizeai/openinference-core";
import { OITracer, safelyJSONStringify } from "@arizeai/openinference-core";
import {
  LLMProvider,
  LLMSystem,
  MimeType,
  OpenInferenceSpanKind,
  SemanticConventions,
} from "@arizeai/openinference-semantic-conventions";

import {
  consumeResponseStreamEvents,
  getResponsesInputMessagesAttributes,
  getResponsesOutputMessagesAttributes,
  getResponsesUsageAttributes,
} from "./responsesAttributes";
import { assertUnreachable, isString } from "./typeUtils";
import { VERSION } from "./version";

const MODULE_NAME = "openai";

const INSTRUMENTATION_NAME = "@arizeai/openinference-instrumentation-openai";

/**
 * Maps hostname suffixes to their corresponding LLM provider value.
 */
export const HOST_SUFFIX_TO_PROVIDER: Record<string, LLMProvider> = {
  "api.openai.com": LLMProvider.OPENAI,
  "openai.azure.com": LLMProvider.AZURE,
  "services.ai.azure.com": LLMProvider.AZURE,
  "cognitiveservices.azure.com": LLMProvider.AZURE,
  "api.anthropic.com": LLMProvider.ANTHROPIC,
  "api.cohere.com": LLMProvider.COHERE,
  "api.cohere.ai": LLMProvider.COHERE,
  "api.mistral.ai": LLMProvider.MISTRALAI,
  "generativelanguage.googleapis.com": LLMProvider.GOOGLE,
  "aiplatform.googleapis.com": LLMProvider.GOOGLE,
  "amazonaws.com": LLMProvider.AWS,
  "api.x.ai": LLMProvider.XAI,
  "api.deepseek.com": LLMProvider.DEEPSEEK,
  "api.groq.com": LLMProvider.GROQ,
  "api.fireworks.ai": LLMProvider.FIREWORKS,
  "api.moonshot.cn": LLMProvider.MOONSHOT,
  "api.cerebras.ai": LLMProvider.CEREBRAS,
  "api.perplexity.ai": LLMProvider.PERPLEXITY,
  "api.together.ai": LLMProvider.TOGETHER,
  "api.together.xyz": LLMProvider.TOGETHER,
  "ollama.com": LLMProvider.OLLAMA,
};

/**
 * Return the LLM provider name for the given API hostname.
 */
export function getProviderFromHost(host: string): LLMProvider | undefined {
  const normalised = host.toLowerCase().trim();
  for (const [suffix, provider] of Object.entries(HOST_SUFFIX_TO_PROVIDER)) {
    // Anchor at a label boundary so e.g. "smollama.com" does not match the
    // "ollama.com" suffix.
    if (normalised === suffix || normalised.endsWith("." + suffix)) {
      return provider;
    }
  }
  return undefined;
}

/**
 * Flag to check if the openai module has been patched
 * Note: This is a fallback in case the module is made immutable (e.x. Deno, webpack, etc.)
 */
let _isOpenInferencePatched = false;

/**
 * function to check if instrumentation is enabled / disabled
 */
export function isPatched() {
  return _isOpenInferencePatched;
}

/**
 * Resolves the execution context for the current span
 * If tracing is suppressed, the span is dropped and the current context is returned
 * @param span
 */
function getExecContext(span: Span) {
  const activeContext = context.active();
  const suppressTracing = isTracingSuppressed(activeContext);
  const execContext = suppressTracing ? trace.setSpan(context.active(), span) : activeContext;
  // Drop the span from the context
  if (suppressTracing) {
    trace.deleteSpan(activeContext);
  }
  return execContext;
}

/**
 * Gets the appropriate LLM provider based on the OpenAI client instance
 * Follows the same logic as the Python implementation by checking the baseURL host
 * @param clientInstance The OpenAI client instance
 * @returns LLM provider based on the API hostname
 */
function getLLMProvider(clientInstance: unknown): LLMProvider | undefined {
  try {
    // The clientInstance might be a sub-object (like Completions) that has a _client property
    // pointing to the actual OpenAI/AzureOpenAI client
    let host: string | undefined;
    let baseURL: unknown;

    // First try to get baseURL from the instance itself
    if (clientInstance != null && typeof clientInstance === "object") {
      baseURL = Reflect.get(clientInstance, "baseURL");
      // If not found (or empty), try the _client property (this is where Azure OpenAI stores it)
      if (!baseURL) {
        const nestedClient = Reflect.get(clientInstance, "_client");
        if (nestedClient != null && typeof nestedClient === "object") {
          baseURL = Reflect.get(nestedClient, "baseURL");
        }
      }
    }

    if (typeof baseURL === "string") {
      // Extract host from URL string
      try {
        const url = new URL(baseURL);
        host = url.hostname;
      } catch {
        // If URL parsing fails, fallback to string matching
        host = baseURL;
      }
    } else if (baseURL && typeof baseURL === "object" && "host" in baseURL) {
      // Direct host property
      host = typeof baseURL.host === "string" ? baseURL.host : undefined;
    }

    if (host && typeof host === "string") {
      return getProviderFromHost(host);
    }
  } catch (error) {
    diag.debug("Failed to determine LLM provider from instance", error);
  }
}

/**
 * An auto instrumentation class for OpenAI that creates {@link https://github.com/Arize-ai/openinference/blob/main/spec/semantic_conventions.md|OpenInference} Compliant spans for the OpenAI API
 * @param instrumentationConfig The config for the instrumentation @see {@link InstrumentationConfig}
 * @param traceConfig The OpenInference trace configuration. Can be used to mask or redact sensitive information on spans. @see {@link TraceConfigOptions}
 */
export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
  private oiTracer: OITracer;
  private tracerProvider?: TracerProvider;
  private traceConfig?: TraceConfigOptions;
  constructor({
    instrumentationConfig,
    traceConfig,
    tracerProvider,
  }: {
    /**
     * The config for the instrumentation
     * @see {@link InstrumentationConfig}
     */
    instrumentationConfig?: InstrumentationConfig;
    /**
     * The OpenInference trace configuration. Can be used to mask or redact sensitive information on spans.
     * @see {@link TraceConfigOptions}
     */
    traceConfig?: TraceConfigOptions;
    /**
     * An optional custom trace provider to be used for tracing. If not provided, a tracer will be created using the global tracer provider.
     * This is useful if you want to use a non-global tracer provider.
     *
     * @see {@link TracerProvider}
     */
    tracerProvider?: TracerProvider;
  } = {}) {
    super(INSTRUMENTATION_NAME, VERSION, Object.assign({}, instrumentationConfig));
    this.tracerProvider = tracerProvider;
    this.traceConfig = traceConfig;
    this.oiTracer = new OITracer({
      tracer: this.tracerProvider?.getTracer(INSTRUMENTATION_NAME, VERSION) ?? this.tracer,
      traceConfig,
    });
  }

  protected init(): InstrumentationModuleDefinition<typeof openai> {
    const module = new InstrumentationNodeModuleDefinition<typeof openai>(
      "openai",
      // 5.x is best effort
      ["^6.0.0", "^5.0.0"],
      this.patch.bind(this),
      this.unpatch.bind(this),
    );
    return module;
  }

  /**
   * Manually instruments the OpenAI module. This is needed when the module is not loaded via require (commonjs)
   * @param {openai} module
   */
  manuallyInstrument(module: typeof openai) {
    diag.debug(`Manually instrumenting ${MODULE_NAME}`);
    this.patch(module);
  }

  get tracer(): Tracer {
    if (this.tracerProvider) {
      return this.tracerProvider.getTracer(this.instrumentationName, this.instrumentationVersion);
    }
    return super.tracer;
  }

  setTracerProvider(tracerProvider: TracerProvider): void {
    super.setTracerProvider(tracerProvider);
    this.tracerProvider = tracerProvider;
    this.oiTracer = new OITracer({
      tracer: this.tracer,
      traceConfig: this.traceConfig,
    });
  }

  /**
   * Patches the OpenAI module
   */
  private patch(
    module: typeof openai & { openInferencePatched?: boolean },
    moduleVersion?: string,
  ) {
    diag.debug(`Applying patch for ${MODULE_NAME}@${moduleVersion}`);
    if (module?.openInferencePatched || _isOpenInferencePatched) {
      return module;
    }
    // eslint-disable-next-line @typescript-eslint/no-this-alias
    const instrumentation: OpenAIInstrumentation = this;

    // Patch create chat completions
    type ChatCompletionCreateType = typeof module.OpenAI.Chat.Completions.prototype.create;

    this._wrap(
      module.OpenAI.Chat.Completions.prototype,
      "create",
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      (original: ChatCompletionCreateType): any => {
        return function patchedCreate(
          this: unknown,
          ...args: Parameters<ChatCompletionCreateType>
        ) {
          const body = args[0];
          const { messages: _messages, ...invocationParameters } = body;
          const span = instrumentation.oiTracer.startSpan(`OpenAI Chat Completions`, {
            kind: SpanKind.INTERNAL,
            attributes: {
              [SemanticConventions.OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.LLM,
              [SemanticConventions.LLM_MODEL_NAME]: body.model,
              [SemanticConventions.INPUT_VALUE]: JSON.stringify(body),
              [SemanticConventions.INPUT_MIME_TYPE]: MimeType.JSON,
              [SemanticConventions.LLM_INVOCATION_PARAMETERS]: JSON.stringify(invocationParameters),
              [SemanticConventions.LLM_SYSTEM]: LLMSystem.OPENAI,
              [SemanticConventions.LLM_PROVIDER]: getLLMProvider(this),
              ...getLLMInputMessagesAttributes(body),
              ...getLLMToolsJSONSchema(body),
            },
          });
          const execContext = getExecContext(span);
          const execPromise = safeExecuteInTheMiddle(
            () => {
              return context.with(trace.setSpan(execContext, span), () => {
                return original.apply(this, args);
              });
            },
            (error) => {
              // Push the error to the span
              if (error) {
                span.recordException(error);
                span.setStatus({
                  code: SpanStatusCode.ERROR,
                  message: error.message,
                });
                span.end();
              }
            },
          );

          const wrappedPromiseThen = (
            result:
              | Stream<openai.Chat.Completions.ChatCompletionChunk>
              | openai.Chat.Completions.ChatCompletion,
          ) => {
            if (isChatCompletionResponse(result)) {
              // Record the results
              span.setAttributes({
                [SemanticConventions.OUTPUT_VALUE]: JSON.stringify(result),
                [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.JSON,
                // Override the model from the value sent by the server
                [SemanticConventions.LLM_MODEL_NAME]: result.model,
                ...getChatCompletionLLMOutputMessagesAttributes(result),
                ...getChatCompletionFinishReasonAttributes(result),
                ...getUsageAttributes(result),
              });
              span.setStatus({ code: SpanStatusCode.OK });
              span.end();
            } else {
              // This is a streaming response
              // handle the chunks and add them to the span
              // First split the stream via tee
              const [leftStream, rightStream] = result.tee();
              void consumeChatCompletionStreamChunks(rightStream, span);
              result = leftStream;
            }

            return result;
          };
          const wrappedPromise = invokeMaybeAPIPromise(execPromise, wrappedPromiseThen);
          return context.bind(execContext, wrappedPromise);
        };
      },
    );

    // Patch create completions
    type CompletionsCreateType = typeof module.OpenAI.Completions.prototype.create;

    this._wrap(
      module.OpenAI.Completions.prototype,
      "create",
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      (original: CompletionsCreateType): any => {
        return function patchedCreate(this: unknown, ...args: Parameters<CompletionsCreateType>) {
          const body = args[0];
          const { prompt: _prompt, ...invocationParameters } = body;
          const span = instrumentation.oiTracer.startSpan(`OpenAI Completions`, {
            kind: SpanKind.INTERNAL,
            attributes: {
              [SemanticConventions.OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.LLM,
              [SemanticConventions.LLM_MODEL_NAME]: body.model,
              [SemanticConventions.LLM_INVOCATION_PARAMETERS]: JSON.stringify(invocationParameters),
              [SemanticConventions.LLM_SYSTEM]: LLMSystem.OPENAI,
              [SemanticConventions.LLM_PROVIDER]: getLLMProvider(this),
              ...getCompletionInputValueAndMimeType(body),
            },
          });
          const execContext = getExecContext(span);

          const execPromise = safeExecuteInTheMiddle(
            () => {
              return context.with(trace.setSpan(execContext, span), () => {
                return original.apply(this, args);
              });
            },
            (error) => {
              // Push the error to the span
              if (error) {
                span.recordException(error);
                span.setStatus({
                  code: SpanStatusCode.ERROR,
                  message: error.message,
                });
                span.end();
              }
            },
          );
          const wrappedPromiseThen = (result: Completion | Stream<Completion>) => {
            if (isCompletionResponse(result)) {
              // Record the results
              span.setAttributes({
                [SemanticConventions.OUTPUT_VALUE]: JSON.stringify(result),
                [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.JSON,
                // Override the model from the value sent by the server
                [SemanticConventions.LLM_MODEL_NAME]: result.model,
                ...getCompletionOutputValueAndMimeType(result),
                ...getUsageAttributes(result),
              });
              span.setStatus({ code: SpanStatusCode.OK });
              span.end();
            }
            return result;
          };
          const wrappedPromise = invokeMaybeAPIPromise(execPromise, wrappedPromiseThen);
          return context.bind(execContext, wrappedPromise);
        };
      },
    );

    // Patch embeddings
    type EmbeddingsCreateType = typeof module.OpenAI.Embeddings.prototype.create;
    this._wrap(
      module.OpenAI.Embeddings.prototype,
      "create",
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      (original: EmbeddingsCreateType): any => {
        return function patchedEmbeddingCreate(
          this: unknown,
          ...args: Parameters<EmbeddingsCreateType>
        ) {
          const body = args[0];
          const { input } = body;
          const isStringInput = typeof input === "string";
          const span = instrumentation.oiTracer.startSpan(`OpenAI Embeddings`, {
            kind: SpanKind.INTERNAL,
            attributes: {
              [SemanticConventions.OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.EMBEDDING,
              [SemanticConventions.EMBEDDING_MODEL_NAME]: body.model,
              [SemanticConventions.INPUT_VALUE]: isStringInput ? input : JSON.stringify(input),
              [SemanticConventions.INPUT_MIME_TYPE]: isStringInput ? MimeType.TEXT : MimeType.JSON,
              [SemanticConventions.LLM_SYSTEM]: LLMSystem.OPENAI,
              [SemanticConventions.LLM_PROVIDER]: getLLMProvider(this),
              ...getEmbeddingTextAttributes(body),
            },
          });
          const execContext = getExecContext(span);
          const execPromise = safeExecuteInTheMiddle(
            () => {
              return context.with(trace.setSpan(execContext, span), () => {
                return original.apply(this, args);
              });
            },
            (error) => {
              // Push the error to the span
              if (error) {
                span.recordException(error);
                span.setStatus({
                  code: SpanStatusCode.ERROR,
                  message: error.message,
                });
                span.end();
              }
            },
          );
          const wrappedPromiseThen = (result: CreateEmbeddingResponse) => {
            if (result) {
              // Record the results
              span.setAttributes({
                // Do not record the output data as it can be large
                ...getEmbeddingEmbeddingsAttributes(result),
              });
            }
            span.setStatus({ code: SpanStatusCode.OK });
            span.end();
            return result;
          };
          const wrappedPromise = invokeMaybeAPIPromise(execPromise, wrappedPromiseThen);
          return context.bind(execContext, wrappedPromise);
        };
      },
    );

    // Patch responses (if the patched module contains the Responses interface)
    if (module.OpenAI.Responses) {
      type ResponsesCreateType = typeof module.OpenAI.Responses.prototype.create;

      this._wrap(
        module.OpenAI.Responses.prototype,
        "create",
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        (original: ResponsesCreateType): any => {
          return function patchedCreate(this: unknown, ...args: Parameters<ResponsesCreateType>) {
            const body = args[0];
            const { input: _messages, ...invocationParameters } = body;
            const span = instrumentation.oiTracer.startSpan(`OpenAI Responses`, {
              kind: SpanKind.INTERNAL,
              attributes: {
                [SemanticConventions.OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.LLM,
                [SemanticConventions.LLM_MODEL_NAME]: body.model,
                [SemanticConventions.INPUT_VALUE]: JSON.stringify(body),
                [SemanticConventions.INPUT_MIME_TYPE]: MimeType.JSON,
                [SemanticConventions.LLM_INVOCATION_PARAMETERS]:
                  JSON.stringify(invocationParameters),
                [SemanticConventions.LLM_SYSTEM]: LLMSystem.OPENAI,
                [SemanticConventions.LLM_PROVIDER]: getLLMProvider(this),
                ...getResponsesInputMessagesAttributes(body),
                ...getLLMToolsJSONSchema(body),
              },
            });
            const execContext = getExecContext(span);
            const execPromise = safeExecuteInTheMiddle(
              () => {
                return context.with(trace.setSpan(execContext, span), () => {
                  return original.apply(this, args);
                });
              },
              (error) => {
                // Push the error to the span
                if (error) {
                  span.recordException(error);
                  span.setStatus({
                    code: SpanStatusCode.ERROR,
                    message: error.message,
                  });
                  span.end();
                }
              },
            );
            const wrappedPromiseThen = (result: Stream<ResponseStreamEvent> | ResponseType) => {
              const recordSpan = (result?: ResponseType) => {
                if (!result) {
                  span.setStatus({ code: SpanStatusCode.ERROR });
                  span.end();
                  return;
                }
                span.setAttributes({
                  [SemanticConventions.OUTPUT_VALUE]: JSON.stringify(result),
                  [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.JSON,
                  // Override the model from the value sent by the server
                  [SemanticConventions.LLM_MODEL_NAME]: result.model,
                  ...getResponsesOutputMessagesAttributes(result),
                  ...getResponsesUsageAttributes(result),
                });
                span.setStatus({ code: SpanStatusCode.OK });
                span.end();
              };
              if (isResponseCreateResponse(result)) {
                // Record the results, as we have the final result
                recordSpan(result);
              } else {
                // This is a streaming response
                // First split the stream via tee
                const [leftStream, rightStream] = result.tee();
                // take the right stream, consuming it and then recording the final chunk
                // into the span
                void consumeResponseStreamEvents(rightStream).then(recordSpan);
                // give the left stream back to the caller
                result = leftStream;
              }

              return result;
            };
            const wrappedPromise = invokeMaybeAPIPromise(execPromise, wrappedPromiseThen);
            return context.bind(execContext, wrappedPromise);
          };
        },
      );
    }

    _isOpenInferencePatched = true;
    try {
      // This can fail if the module is made immutable via the runtime or bundler
      module.openInferencePatched = true;
    } catch (e) {
      diag.debug(`Failed to set ${MODULE_NAME} patched flag on the module`, e);
    }

    return module;
  }
  /**
   * Un-patches the OpenAI module's chat completions API
   */
  private unpatch(
    moduleExports: typeof openai & { openInferencePatched?: boolean },
    moduleVersion?: string,
  ) {
    diag.debug(`Removing patch for ${MODULE_NAME}@${moduleVersion}`);
    this._unwrap(moduleExports.OpenAI.Chat.Completions.prototype, "create");
    this._unwrap(moduleExports.OpenAI.Completions.prototype, "create");
    this._unwrap(moduleExports.OpenAI.Embeddings.prototype, "create");

    _isOpenInferencePatched = false;
    try {
      // This can fail if the module is made immutable via the runtime or bundler
      moduleExports.openInferencePatched = false;
    } catch (e) {
      diag.warn(`Failed to unset ${MODULE_NAME} patched flag on the module`, e);
    }
  }
}

function isResponseCreateResponse(
  response: Stream<ResponseStreamEvent> | ResponseType,
): response is ResponseType {
  return "object" in response && response.object === "response";
}

/**
 * type-guard that checks if the response is a chat completion response
 */
function isChatCompletionResponse(
  response: Stream<ChatCompletionChunk> | ChatCompletion,
): response is ChatCompletion {
  return "choices" in response;
}

/**
 * type-guard that checks if the response is a completion response
 */
function isCompletionResponse(response: Stream<Completion> | Completion): response is Completion {
  return "choices" in response;
}

/**
 * type-guard that checks if completion prompt attribute is an array of strings
 */
function isPromptStringArray(
  prompt: CompletionCreateParamsBase["prompt"],
): prompt is Array<string> {
  return Array.isArray(prompt) && prompt.every((item) => typeof item === "string");
}

/**
 * Converts the body of a chat completions request to LLM input messages
 */
function getLLMInputMessagesAttributes(body: ChatCompletionCreateParamsBase): Attributes {
  return body.messages.reduce((acc, message, index) => {
    const messageAttributes = getChatCompletionInputMessageAttributes(message);
    const indexPrefix = `${SemanticConventions.LLM_INPUT_MESSAGES}.${index}.`;
    // Flatten the attributes on the index prefix
    for (const [key, value] of Object.entries(messageAttributes)) {
      acc[`${indexPrefix}${key}`] = value;
    }
    return acc;
  }, {} as Attributes);
}

/**
 * Converts each tool definition into a json schema
 */
function getLLMToolsJSONSchema(
  body: ChatCompletionCreateParamsBase | ResponseCreateParamsBase,
): Attributes {
  if (!body.tools) {
    // If tools is undefined, return an empty object
    return {};
  }
  return body.tools.reduce((acc: Attributes, tool, index) => {
    const toolJsonSchema = safelyJSONStringify(tool);
    const key = `${SemanticConventions.LLM_TOOLS}.${index}.${SemanticConventions.TOOL_JSON_SCHEMA}`;
    if (toolJsonSchema) {
      acc[key] = toolJsonSchema;
    }
    return acc;
  }, {});
}

function getChatCompletionInputMessageAttributes(message: ChatCompletionMessageParam): Attributes {
  const role = message.role;
  const attributes: Attributes = {
    [SemanticConventions.MESSAGE_ROLE]: role,
  };
  // Add the content only if it is a string
  if (typeof message.content === "string") {
    attributes[SemanticConventions.MESSAGE_CONTENT] = message.content;
  } else if (Array.isArray(message.content)) {
    message.content.forEach((part, index) => {
      const contentsIndexPrefix = `${SemanticConventions.MESSAGE_CONTENTS}.${index}.`;
      if (part.type === "text") {
        attributes[`${contentsIndexPrefix}${SemanticConventions.MESSAGE_CONTENT_TYPE}`] = "text";
        attributes[`${contentsIndexPrefix}${SemanticConventions.MESSAGE_CONTENT_TEXT}`] = part.text;
      } else if (part.type === "image_url") {
        attributes[`${contentsIndexPrefix}${SemanticConventions.MESSAGE_CONTENT_TYPE}`] = "image";
        attributes[
          `${contentsIndexPrefix}${SemanticConventions.MESSAGE_CONTENT_IMAGE}.${SemanticConventions.IMAGE_URL}`
        ] = part.image_url.url;
      }
    });
  }
  switch (role) {
    case "user":
      // There's nothing to add for the user
      break;
    case "assistant":
      if (message.tool_calls) {
        message.tool_calls.forEach((toolCall, index) => {
          const toolCallIndexPrefix = `${SemanticConventions.MESSAGE_TOOL_CALLS}.${index}.`;

          // Add the tool call id if it exists
          if (toolCall.id) {
            attributes[`${toolCallIndexPrefix}${SemanticConventions.TOOL_CALL_ID}`] = toolCall.id;
          }
          // Make sure the tool call has a function
          if (toolCall.type === "function") {
            attributes[`${toolCallIndexPrefix}${SemanticConventions.TOOL_CALL_FUNCTION_NAME}`] =
              toolCall.function.name;
            attributes[
              `${toolCallIndexPrefix}${SemanticConventions.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}`
            ] = toolCall.function.arguments;
          } else {
            // TODO: add exhaustive checks
            diag.warn(`Unsupported tool type: ${toolCall.type}`);
          }
        });
      }
      break;
    case "function":
      attributes[SemanticConventions.MESSAGE_FUNCTION_CALL_NAME] = message.name;
      break;
    case "tool":
      if (message.tool_call_id) {
        attributes[`${SemanticConventions.MESSAGE_TOOL_CALL_ID}`] = message.tool_call_id;
      }
      break;
    case "system":
      // There's nothing to add for the system. Content is captured above
      break;
    case "developer":
      // There's nothing to add for the developer. Content is captured above
      break;
    default:
      assertUnreachable(role);
      break;
  }
  return attributes;
}

/**
 * Converts the body of a completions request to input attributes
 */
function getCompletionInputValueAndMimeType(body: CompletionCreateParamsBase): Attributes {
  if (typeof body.prompt === "string") {
    return {
      [SemanticConventions.INPUT_VALUE]: body.prompt,
      [SemanticConventions.INPUT_MIME_TYPE]: MimeType.TEXT,
    };
  } else if (isPromptStringArray(body.prompt)) {
    const prompt = body.prompt[0]; // Only single prompts are currently supported
    if (prompt === undefined) {
      return {};
    }
    return {
      [SemanticConventions.INPUT_VALUE]: prompt,
      [SemanticConventions.INPUT_MIME_TYPE]: MimeType.TEXT,
    };
  }
  // Other cases in which the prompt is a token or array of tokens are currently unsupported
  return {};
}

/**
 * Get usage attributes
 */
function getUsageAttributes(completion: ChatCompletion | Completion): Attributes {
  if (completion.usage) {
    const usageAttributes: Attributes = {
      [SemanticConventions.LLM_TOKEN_COUNT_COMPLETION]: completion.usage.completion_tokens,
      [SemanticConventions.LLM_TOKEN_COUNT_PROMPT]: completion.usage.prompt_tokens,
      [SemanticConventions.LLM_TOKEN_COUNT_TOTAL]: completion.usage.total_tokens,
      [SemanticConventions.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ]:
        completion.usage.prompt_tokens_details?.cached_tokens,
      [SemanticConventions.LLM_TOKEN_COUNT_PROMPT_DETAILS_AUDIO]:
        completion.usage.prompt_tokens_details?.audio_tokens,
      [SemanticConventions.LLM_TOKEN_COUNT_COMPLETION_DETAILS_AUDIO]:
        completion.usage.completion_tokens_details?.audio_tokens,
      [SemanticConventions.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING]:
        completion.usage.completion_tokens_details?.reasoning_tokens,
    };
    return usageAttributes;
  }
  return {};
}

/**
 * Converts the chat completion result to LLM output attributes
 */
function getChatCompletionLLMOutputMessagesAttributes(chatCompletion: ChatCompletion): Attributes {
  // Right now support just the first choice
  const choice = chatCompletion.choices[0];
  if (!choice) {
    return {};
  }
  return [choice.message].reduce((acc, message, index) => {
    const indexPrefix = `${SemanticConventions.LLM_OUTPUT_MESSAGES}.${index}.`;
    const messageAttributes = getChatCompletionOutputMessageAttributes(message);
    // Flatten the attributes on the index prefix
    for (const [key, value] of Object.entries(messageAttributes)) {
      acc[`${indexPrefix}${key}`] = value;
    }
    return acc;
  }, {} as Attributes);
}

/**
 * Extracts the llm.finish_reason attribute from a chat completion response.
 */
function getChatCompletionFinishReasonAttributes(chatCompletion: ChatCompletion): Attributes {
  const choice = chatCompletion.choices[0];
  if (!choice?.finish_reason) {
    return {};
  }
  return {
    [SemanticConventions.LLM_FINISH_REASON]: choice.finish_reason,
  };
}

function getChatCompletionOutputMessageAttributes(message: ChatCompletionMessage): Attributes {
  const role = message.role;
  const attributes: Attributes = {
    [SemanticConventions.MESSAGE_ROLE]: role,
  };
  if (message.content) {
    attributes[SemanticConventions.MESSAGE_CONTENT] = message.content;
  }
  if (message.tool_calls) {
    message.tool_calls.forEach((toolCall, index) => {
      const toolCallIndexPrefix = `${SemanticConventions.MESSAGE_TOOL_CALLS}.${index}.`;
      // Add the tool call id if it exists
      if (toolCall.id) {
        attributes[`${toolCallIndexPrefix}${SemanticConventions.TOOL_CALL_ID}`] = toolCall.id;
      }
      if (toolCall.type === "function") {
        attributes[toolCallIndexPrefix + SemanticConventions.TOOL_CALL_FUNCTION_NAME] =
          toolCall.function.name;
        attributes[toolCallIndexPrefix + SemanticConventions.TOOL_CALL_FUNCTION_ARGUMENTS_JSON] =
          toolCall.function.arguments;
      } else {
        // TODO: switch to exhaustive checks
        diag.warn(`Unsupported tool type: ${toolCall.type}`);
      }
    });
  }
  if (message.function_call) {
    attributes[SemanticConventions.MESSAGE_FUNCTION_CALL_NAME] = message.function_call.name;
    attributes[SemanticConventions.MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON] =
      message.function_call.arguments;
  }
  return attributes;
}

/**
 * Converts the completion result to output attributes
 */
function getCompletionOutputValueAndMimeType(completion: Completion): Attributes {
  // Right now support just the first choice
  const choice = completion.choices[0];
  if (!choice) {
    return {};
  }
  return {
    [SemanticConventions.OUTPUT_VALUE]: String(choice.text),
    [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.TEXT,
  };
}

/**
 * Converts the embedding result payload to embedding attributes
 */
function getEmbeddingTextAttributes(request: EmbeddingCreateParams): Attributes {
  if (typeof request.input === "string") {
    return {
      [`${SemanticConventions.EMBEDDING_EMBEDDINGS}.0.${SemanticConventions.EMBEDDING_TEXT}`]:
        request.input,
    };
  } else if (
    Array.isArray(request.input) &&
    request.input.length > 0 &&
    typeof request.input[0] === "string"
  ) {
    return request.input.reduce((acc, input, index) => {
      const indexPrefix = `${SemanticConventions.EMBEDDING_EMBEDDINGS}.${index}.`;
      acc[`${indexPrefix}${SemanticConventions.EMBEDDING_TEXT}`] = input;
      return acc;
    }, {} as Attributes);
  }
  // Ignore other cases where input is a number or an array of numbers
  return {};
}

/**
 * Converts the embedding result payload to embedding attributes
 */
function getEmbeddingEmbeddingsAttributes(response: CreateEmbeddingResponse): Attributes {
  return response.data.reduce((acc, embedding, index) => {
    const indexPrefix = `${SemanticConventions.EMBEDDING_EMBEDDINGS}.${index}.`;
    acc[`${indexPrefix}${SemanticConventions.EMBEDDING_VECTOR}`] = embedding.embedding;
    return acc;
  }, {} as Attributes);
}

/**
 * Consumes the stream chunks and adds them to the span
 */
async function consumeChatCompletionStreamChunks(stream: Stream<ChatCompletionChunk>, span: Span) {
  let streamResponse = "";
  // Tool and function call attributes can also arrive in the stream
  // NB: the tools and function calls arrive in partial diffs
  // So the final tool and function calls need to be aggregated
  // across chunks
  const toolAndFunctionCallAttributes: Attributes = {};
  let finishReason: string | undefined;
  // The first message is for the assistant response so we start at 1
  for await (const chunk of stream) {
    if (chunk.choices.length <= 0) {
      continue;
    }
    const choice = chunk.choices[0];
    if (choice.delta.content) {
      streamResponse += choice.delta.content;
    }
    if (choice.finish_reason) {
      finishReason = choice.finish_reason;
    }
    // Accumulate the tool and function call attributes
    const toolAndFunctionCallAttributesDiff =
      getToolAndFunctionCallAttributesFromStreamChunk(chunk);
    for (const [key, value] of Object.entries(toolAndFunctionCallAttributesDiff)) {
      if (isString(toolAndFunctionCallAttributes[key]) && isString(value)) {
        toolAndFunctionCallAttributes[key] += value;
      } else if (isString(value)) {
        toolAndFunctionCallAttributes[key] = value;
      }
    }
  }
  const messageIndexPrefix = `${SemanticConventions.LLM_OUTPUT_MESSAGES}.0.`;

  // Append the attributes to the span as a message
  const attributes: Attributes = {
    [SemanticConventions.OUTPUT_VALUE]: streamResponse,
    [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.TEXT,
    [`${messageIndexPrefix}${SemanticConventions.MESSAGE_CONTENT}`]: streamResponse,
    [`${messageIndexPrefix}${SemanticConventions.MESSAGE_ROLE}`]: "assistant",
  };
  if (finishReason) {
    attributes[SemanticConventions.LLM_FINISH_REASON] = finishReason;
  }
  // Add the tool and function call attributes
  for (const [key, value] of Object.entries(toolAndFunctionCallAttributes)) {
    attributes[`${messageIndexPrefix}${key}`] = value;
  }
  span.setAttributes(attributes);
  span.end();
}

/**
 * Extracts the semantic attributes from the stream chunk for tool_calls and function_calls
 */
function getToolAndFunctionCallAttributesFromStreamChunk(chunk: ChatCompletionChunk): Attributes {
  if (chunk.choices.length <= 0) {
    return {};
  }
  const choice = chunk.choices[0];
  const attributes: Attributes = {};
  if (choice.delta.tool_calls) {
    choice.delta.tool_calls.forEach((toolCall, index) => {
      const toolCallIndexPrefix = `${SemanticConventions.MESSAGE_TOOL_CALLS}.${index}.`;
      // Add the tool call id if it exists
      if (toolCall.id) {
        attributes[`${toolCallIndexPrefix}${SemanticConventions.TOOL_CALL_ID}`] = toolCall.id;
      }
      // Double check that the tool call has a function
      // NB: OpenAI only supports tool calls with functions right now but this may change
      if (toolCall.function) {
        attributes[toolCallIndexPrefix + SemanticConventions.TOOL_CALL_FUNCTION_NAME] =
          toolCall.function.name;
        attributes[toolCallIndexPrefix + SemanticConventions.TOOL_CALL_FUNCTION_ARGUMENTS_JSON] =
          toolCall.function.arguments;
      }
    });
  }
  if (choice.delta.function_call) {
    attributes[SemanticConventions.MESSAGE_FUNCTION_CALL_NAME] = choice.delta.function_call.name;
    attributes[SemanticConventions.MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON] =
      choice.delta.function_call.arguments;
  }
  return attributes;
}

/**
 * Type-guard that checks if the promise is an APIPromise.
 *
 * APIPromise is a class from openai that wraps promises with special behavior.
 *
 * @param promise - The promise to check
 * @returns True if the promise is an APIPromise, false otherwise
 */
function isAPIPromise<T>(promise: unknown): promise is APIPromise<T> {
  return promise instanceof APIPromise;
}

/**
 * Invokes the thennable of a promise or an APIPromise.
 *
 * This is necessary to safely invoke promises returned by openai sdk methods.
 *
 * Without this wrapper, instrumentation will "consume" returned APIPromise instances,
 * Making them unusable by other openai sdk methods, such as when completions.parse internally
 * calls completions.create, expecting an APIPromise which we would otherwise consume by calling `then` on it.
 *
 * @param promise - The promise to invoke the thennable of
 * @param then - The thennable to invoke
 * @returns The promise with the thennable invoked
 */
function invokeMaybeAPIPromise<T>(promise: APIPromise<T>, then: (value: T) => T): APIPromise<T>;
function invokeMaybeAPIPromise<T>(promise: Promise<T>, then: (value: T) => T): Promise<T>;
function invokeMaybeAPIPromise<T>(promise: T, then: (value: T) => T): T;
function invokeMaybeAPIPromise(promise: unknown, then: (value: unknown) => unknown): unknown {
  if (isAPIPromise<unknown>(promise)) {
    return promise._thenUnwrap(then);
  } else if (promise instanceof Promise) {
    return promise.then(then);
  } else {
    // eslint-disable-next-line no-console
    console.warn("Promise is not an APIPromise or a regular promise, cannot instrument.");
    return promise;
  }
}