/*
 * Copyright The OpenTelemetry Authors
 * SPDX-License-Identifier: Apache-2.0
 */

package io.opentelemetry.javaagent.instrumentation.lettuce.v5_0;

import static io.opentelemetry.javaagent.instrumentation.lettuce.v5_0.LettuceSingletons.instrumenter;

import io.lettuce.core.protocol.RedisCommand;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.bootstrap.internal.InstrumentationConfig;
import java.util.concurrent.CancellationException;
import java.util.function.BiFunction;

/**
 * Callback class to close the command span on an error or a success in the RedisFuture returned by
 * the lettuce async API.
 *
 * @param <T> the normal completion result
 * @param <U> the error
 * @param <R> the return type, should be null since nothing else should happen from tracing
 *     standpoint after the span is closed
 */
public class EndCommandAsyncBiFunction<T, U extends Throwable, R>
    implements BiFunction<T, Throwable, R> {

  private static final boolean CAPTURE_EXPERIMENTAL_SPAN_ATTRIBUTES =
      InstrumentationConfig.get()
          .getBoolean("otel.instrumentation.lettuce.experimental-span-attributes", false);

  private final Context context;
  private final RedisCommand<?, ?, ?> command;

  public EndCommandAsyncBiFunction(Context context, RedisCommand<?, ?, ?> command) {
    this.context = context;
    this.command = command;
  }

  @Override
  public R apply(T t, Throwable throwable) {
    end(context, command, throwable);
    return null;
  }

  public static void end(Context context, RedisCommand<?, ?, ?> command, Throwable throwable) {
    if (throwable instanceof CancellationException) {
      if (CAPTURE_EXPERIMENTAL_SPAN_ATTRIBUTES) {
        Span.fromContext(context).setAttribute("lettuce.command.cancelled", true);
      }
      // and don't report this as an error
      throwable = null;
    }
    instrumenter().end(context, command, null, throwable);
  }
}
