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

package io.opentelemetry.instrumentation.api.annotation.support.async;

import io.opentelemetry.context.Context;
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
import javax.annotation.Nullable;

/**
 * A wrapper over {@link Instrumenter} that is able to defer {@link Instrumenter#end(Context,
 * Object, Object, Throwable)} until asynchronous computation finishes.
 */
public final class AsyncOperationEndSupport<REQUEST, RESPONSE> {

  /**
   * Returns a new {@link AsyncOperationEndSupport} that wraps over passed {@code syncInstrumenter},
   * configured for usage with asynchronous computations that are instances of {@code asyncType}. If
   * the result of the async computation ends up being an instance of {@code responseType} it will
   * be passed as the response to the {@code syncInstrumenter} call; otherwise {@code null} value
   * will be used as the response.
   */
  public static <REQUEST, RESPONSE> AsyncOperationEndSupport<REQUEST, RESPONSE> create(
      Instrumenter<REQUEST, RESPONSE> syncInstrumenter,
      Class<RESPONSE> responseType,
      Class<?> asyncType) {
    return new AsyncOperationEndSupport<>(
        syncInstrumenter,
        responseType,
        asyncType,
        AsyncOperationEndStrategies.instance().resolveStrategy(asyncType));
  }

  private final Instrumenter<REQUEST, RESPONSE> instrumenter;
  private final Class<RESPONSE> responseType;
  private final Class<?> asyncType;
  @Nullable private final AsyncOperationEndStrategy asyncOperationEndStrategy;

  private AsyncOperationEndSupport(
      Instrumenter<REQUEST, RESPONSE> instrumenter,
      Class<RESPONSE> responseType,
      Class<?> asyncType,
      @Nullable AsyncOperationEndStrategy asyncOperationEndStrategy) {
    this.instrumenter = instrumenter;
    this.responseType = responseType;
    this.asyncType = asyncType;
    this.asyncOperationEndStrategy = asyncOperationEndStrategy;
  }

  /**
   * Attempts to compose over passed {@code asyncValue} and delay the {@link
   * Instrumenter#end(Context, Object, Object, Throwable)} call until the async operation completes.
   *
   * <p>This method will end the operation immediately if {@code throwable} is passed, if there is
   * no {@link AsyncOperationEndStrategy} for the {@code asyncType} used, or if there is a type
   * mismatch between passed {@code asyncValue} and the {@code asyncType} that was used to create
   * this object.
   *
   * <p>If the passed {@code asyncValue} is recognized as an asynchronous computation, the operation
   * won't be {@link Instrumenter#end(Context, Object, Object, Throwable) ended} until {@code
   * asyncValue} completes.
   */
  @SuppressWarnings("unchecked")
  @Nullable
  public <ASYNC> ASYNC asyncEnd(
      Context context, REQUEST request, @Nullable ASYNC asyncValue, @Nullable Throwable throwable) {
    // we can end early if an exception was thrown
    if (throwable != null) {
      instrumenter.end(context, request, null, throwable);
      return asyncValue;
    }

    // use the configured strategy to compose over the asyncValue
    if (asyncOperationEndStrategy != null && asyncType.isInstance(asyncValue)) {
      return (ASYNC)
          asyncOperationEndStrategy.end(instrumenter, context, request, asyncValue, responseType);
    }

    // fall back to sync end() if asyncValue type doesn't match
    instrumenter.end(context, request, tryToGetResponse(responseType, asyncValue), null);
    return asyncValue;
  }

  @Nullable
  public static <RESPONSE> RESPONSE tryToGetResponse(
      Class<RESPONSE> responseType, @Nullable Object asyncValue) {
    if (responseType.isInstance(asyncValue)) {
      return responseType.cast(asyncValue);
    }
    return null;
  }
}
