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

package io.opentelemetry.javaagent.instrumentation.netty.v4_1;

import io.netty.resolver.AddressResolver;
import io.netty.resolver.AddressResolverGroup;
import io.netty.util.concurrent.EventExecutor;
import io.netty.util.concurrent.Future;
import io.netty.util.concurrent.Promise;
import io.opentelemetry.context.Context;
import io.opentelemetry.instrumentation.netty.common.internal.NettyConnectionRequest;
import io.opentelemetry.instrumentation.netty.v4.common.internal.client.NettyConnectionInstrumenter;
import java.net.SocketAddress;
import java.util.List;
import java.util.function.Supplier;

public final class InstrumentedAddressResolverGroup<T extends SocketAddress>
    extends AddressResolverGroup<T> {

  public static <T extends SocketAddress> AddressResolverGroup<T> wrap(
      NettyConnectionInstrumenter instrumenter, AddressResolverGroup<T> delegate) {
    if (delegate == null || delegate instanceof InstrumentedAddressResolverGroup) {
      return delegate;
    }
    return new InstrumentedAddressResolverGroup<>(instrumenter, delegate);
  }

  private final NettyConnectionInstrumenter instrumenter;
  private final AddressResolverGroup<T> delegate;

  private InstrumentedAddressResolverGroup(
      NettyConnectionInstrumenter instrumenter, AddressResolverGroup<T> delegate) {
    this.instrumenter = instrumenter;
    this.delegate = delegate;
  }

  @Override
  public AddressResolver<T> getResolver(EventExecutor executor) {
    return new InstrumentedResolver<>(instrumenter, delegate.getResolver(executor));
  }

  @Override
  public void close() {
    delegate.close();
  }

  @Override
  protected AddressResolver<T> newResolver(EventExecutor eventExecutor) {
    // this method is called from the super class's implementation of `getResolver` which is
    // overridden by this class
    throw new UnsupportedOperationException("This method should never be called");
  }

  private static final class InstrumentedResolver<T extends SocketAddress>
      implements AddressResolver<T> {

    private final NettyConnectionInstrumenter instrumenter;
    private final AddressResolver<T> delegate;

    private InstrumentedResolver(
        NettyConnectionInstrumenter instrumenter, AddressResolver<T> delegate) {
      this.instrumenter = instrumenter;
      this.delegate = delegate;
    }

    @Override
    public boolean isSupported(SocketAddress socketAddress) {
      return delegate.isSupported(socketAddress);
    }

    @Override
    public boolean isResolved(SocketAddress socketAddress) {
      return delegate.isResolved(socketAddress);
    }

    @Override
    public Future<T> resolve(SocketAddress socketAddress) {
      return instrumentResolve(socketAddress, () -> delegate.resolve(socketAddress));
    }

    @Override
    public Future<T> resolve(SocketAddress socketAddress, Promise<T> promise) {
      return instrumentResolve(socketAddress, () -> delegate.resolve(socketAddress, promise));
    }

    @Override
    public Future<List<T>> resolveAll(SocketAddress socketAddress) {
      return instrumentResolve(socketAddress, () -> delegate.resolveAll(socketAddress));
    }

    @Override
    public Future<List<T>> resolveAll(SocketAddress socketAddress, Promise<List<T>> promise) {
      return instrumentResolve(socketAddress, () -> delegate.resolveAll(socketAddress, promise));
    }

    private <U> Future<U> instrumentResolve(
        SocketAddress socketAddress, Supplier<Future<U>> resolveFunc) {
      Context parentContext = Context.current();
      NettyConnectionRequest request = NettyConnectionRequest.resolve(socketAddress);
      if (!instrumenter.shouldStart(parentContext, request)) {
        return resolveFunc.get();
      }
      Context context = instrumenter.start(parentContext, request);
      try {
        Future<U> future = resolveFunc.get();
        return future.addListener(f -> instrumenter.end(context, request, null, f.cause()));
      } catch (Throwable t) {
        instrumenter.end(context, request, null, t);
        throw t;
      }
    }

    @Override
    public void close() {
      delegate.close();
    }
  }
}
