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

package io.opentelemetry.javaagent.instrumentation.reactornetty.v0_9;

import io.netty.channel.Channel;
import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import io.opentelemetry.instrumentation.netty.v4_1.internal.AttributeKeys;
import java.util.function.BiConsumer;
import javax.annotation.Nullable;
import reactor.netty.Connection;
import reactor.netty.http.client.HttpClientRequest;
import reactor.netty.http.client.HttpClientResponse;

public final class DecoratorFunctions {

  // ignore already decorated functions
  public static boolean shouldDecorate(Class<?> callbackClass) {
    return callbackClass != OnRequestDecorator.class
        && callbackClass != OnResponseDecorator.class
        && callbackClass != OnRequestErrorDecorator.class
        && callbackClass != OnResponseErrorDecorator.class;
  }

  private abstract static class OnMessageDecorator<M> implements BiConsumer<M, Connection> {
    private final BiConsumer<? super M, ? super Connection> delegate;
    private final boolean forceParentContext;

    protected OnMessageDecorator(
        BiConsumer<? super M, ? super Connection> delegate, boolean forceParentContext) {
      this.delegate = delegate;
      this.forceParentContext = forceParentContext;
    }

    @Override
    public final void accept(M message, Connection connection) {
      Channel channel = connection.channel();
      // don't try to get the client span from the netty channel when forceParentSpan is true
      // this way the parent context will always be propagated
      if (forceParentContext) {
        channel = null;
      }
      Context context = getChannelContext(currentContext(message), channel);
      if (context == null) {
        delegate.accept(message, connection);
      } else {
        try (Scope ignored = context.makeCurrent()) {
          delegate.accept(message, connection);
        }
      }
    }

    abstract reactor.util.context.Context currentContext(M message);
  }

  public static final class OnRequestDecorator extends OnMessageDecorator<HttpClientRequest> {
    public OnRequestDecorator(BiConsumer<? super HttpClientRequest, ? super Connection> delegate) {
      super(delegate, /* forceParentContext= */ false);
    }

    @Override
    reactor.util.context.Context currentContext(HttpClientRequest message) {
      return message.currentContext();
    }
  }

  public static final class OnResponseDecorator extends OnMessageDecorator<HttpClientResponse> {
    public OnResponseDecorator(
        BiConsumer<? super HttpClientResponse, ? super Connection> delegate,
        boolean forceParentContext) {
      super(delegate, forceParentContext);
    }

    @Override
    reactor.util.context.Context currentContext(HttpClientResponse message) {
      return message.currentContext();
    }
  }

  private abstract static class OnMessageErrorDecorator<M> implements BiConsumer<M, Throwable> {
    private final BiConsumer<? super M, ? super Throwable> delegate;

    protected OnMessageErrorDecorator(BiConsumer<? super M, ? super Throwable> delegate) {
      this.delegate = delegate;
    }

    @Override
    public final void accept(M message, Throwable throwable) {
      Context context = getChannelContext(currentContext(message), null);
      if (context == null) {
        delegate.accept(message, throwable);
      } else {
        try (Scope ignored = context.makeCurrent()) {
          delegate.accept(message, throwable);
        }
      }
    }

    abstract reactor.util.context.Context currentContext(M message);
  }

  public static final class OnRequestErrorDecorator
      extends OnMessageErrorDecorator<HttpClientRequest> {
    public OnRequestErrorDecorator(
        BiConsumer<? super HttpClientRequest, ? super Throwable> delegate) {
      super(delegate);
    }

    @Override
    reactor.util.context.Context currentContext(HttpClientRequest message) {
      return message.currentContext();
    }
  }

  public static final class OnResponseErrorDecorator
      extends OnMessageErrorDecorator<HttpClientResponse> {
    public OnResponseErrorDecorator(
        BiConsumer<? super HttpClientResponse, ? super Throwable> delegate) {
      super(delegate);
    }

    @Override
    reactor.util.context.Context currentContext(HttpClientResponse message) {
      return message.currentContext();
    }
  }

  @Nullable
  private static Context getChannelContext(
      reactor.util.context.Context reactorContext, @Nullable Channel channel) {
    // try to get the client span context from the channel if it's available
    if (channel != null) {
      Context context = channel.attr(AttributeKeys.CLIENT_CONTEXT).get();
      if (context != null) {
        return context;
      }
    }
    // otherwise use the parent span context
    return reactorContext.getOrDefault(MapConnect.CONTEXT_ATTRIBUTE, null);
  }

  private DecoratorFunctions() {}
}
