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

package io.opentelemetry.instrumentation.awssdk.v2_2;

import static io.opentelemetry.instrumentation.awssdk.v2_2.TracingExecutionInterceptor.SDK_HTTP_REQUEST_ATTRIBUTE;
import static io.opentelemetry.instrumentation.awssdk.v2_2.TracingExecutionInterceptor.SDK_REQUEST_ATTRIBUTE;

import io.opentelemetry.context.Scope;
import io.opentelemetry.context.propagation.TextMapPropagator;
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
import io.opentelemetry.instrumentation.api.internal.InstrumenterUtil;
import io.opentelemetry.instrumentation.api.internal.Timer;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import javax.annotation.Nullable;
import software.amazon.awssdk.core.SdkRequest;
import software.amazon.awssdk.core.SdkResponse;
import software.amazon.awssdk.core.interceptor.Context;
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
import software.amazon.awssdk.core.interceptor.SdkExecutionAttribute;
import software.amazon.awssdk.services.sqs.SqsAsyncClient;
import software.amazon.awssdk.services.sqs.SqsClient;
import software.amazon.awssdk.services.sqs.model.Message;
import software.amazon.awssdk.services.sqs.model.MessageAttributeValue;
import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest;
import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse;
import software.amazon.awssdk.services.sqs.model.SendMessageBatchRequest;
import software.amazon.awssdk.services.sqs.model.SendMessageBatchRequestEntry;
import software.amazon.awssdk.services.sqs.model.SendMessageRequest;
import software.amazon.awssdk.services.sqs.model.SendMessageResponse;

// this class is only used from SqsAccess from method with @NoMuzzle annotation
final class SqsImpl {
  static {
    // Force loading of SqsClient; this ensures that an exception is thrown at this point when the
    // SQS library is not present, which will cause SqsAccess to have enabled=false in library mode.
    @SuppressWarnings("unused")
    String ensureLoadedDummy = SqsClient.class.getName();
  }

  private SqsImpl() {}

  static boolean afterReceiveMessageExecution(
      Context.AfterExecution context,
      ExecutionAttributes executionAttributes,
      TracingExecutionInterceptor config,
      Timer timer) {

    SdkResponse rawResponse = context.response();
    if (!(rawResponse instanceof ReceiveMessageResponse)) {
      return false;
    }

    ReceiveMessageResponse response = (ReceiveMessageResponse) rawResponse;
    if (response.messages().isEmpty()) {
      return false;
    }

    io.opentelemetry.context.Context parentContext =
        TracingExecutionInterceptor.getParentContext(executionAttributes);
    Instrumenter<SqsReceiveRequest, Response> consumerReceiveInstrumenter =
        config.getConsumerReceiveInstrumenter();
    io.opentelemetry.context.Context receiveContext = null;
    SqsReceiveRequest receiveRequest =
        SqsReceiveRequest.create(executionAttributes, SqsMessageImpl.wrap(response.messages()));
    if (timer != null && consumerReceiveInstrumenter.shouldStart(parentContext, receiveRequest)) {
      receiveContext =
          InstrumenterUtil.startAndEnd(
              consumerReceiveInstrumenter,
              parentContext,
              receiveRequest,
              new Response(context.httpResponse(), response),
              null,
              timer.startTime(),
              timer.now());
    }
    // copy ExecutionAttributes as these will get cleared before the process spans are created
    ExecutionAttributes copy = new ExecutionAttributes();
    copy.putAttribute(
        SDK_HTTP_REQUEST_ATTRIBUTE, executionAttributes.getAttribute(SDK_HTTP_REQUEST_ATTRIBUTE));
    copy.putAttribute(
        SDK_REQUEST_ATTRIBUTE, executionAttributes.getAttribute(SDK_REQUEST_ATTRIBUTE));
    copy.putAttribute(
        SdkExecutionAttribute.SERVICE_NAME,
        executionAttributes.getAttribute(SdkExecutionAttribute.SERVICE_NAME));
    copy.putAttribute(
        SdkExecutionAttribute.OPERATION_NAME,
        executionAttributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME));

    TracingList tracingList =
        TracingList.wrap(
            response.messages(),
            config.getConsumerProcessInstrumenter(),
            copy,
            new Response(context.httpResponse(), response),
            config,
            receiveContext);

    // store tracing list in context so that our proxied SqsClient/SqsAsyncClient could pick it up
    SqsTracingContext.set(parentContext, tracingList);

    return true;
  }

  private static final Field messagesField = getMessagesField();

  private static Field getMessagesField() {
    try {
      Field field = ReceiveMessageResponse.class.getDeclaredField("messages");
      field.setAccessible(true);
      return field;
    } catch (Exception e) {
      return null;
    }
  }

  public static void setMessages(
      ReceiveMessageResponse receiveMessageResponse, List<Message> messages) {
    try {
      messagesField.set(receiveMessageResponse, messages);
    } catch (IllegalAccessException ignored) {
      // should not happen, we call setAccessible on the field
    }
  }

  @Nullable
  static SdkRequest modifyRequest(
      SdkRequest request,
      io.opentelemetry.context.Context otelContext,
      boolean useXrayPropagator,
      TextMapPropagator messagingPropagator) {
    if (request instanceof ReceiveMessageRequest) {
      return modifyReceiveMessageRequest(
          (ReceiveMessageRequest) request, useXrayPropagator, messagingPropagator);
    } else if (messagingPropagator != null) {
      if (request instanceof SendMessageRequest) {
        return injectIntoSendMessageRequest(
            (SendMessageRequest) request, otelContext, messagingPropagator);
      } else if (request instanceof SendMessageBatchRequest) {
        return injectIntoSendMessageBatchRequest(
            (SendMessageBatchRequest) request, otelContext, messagingPropagator);
      }
    }
    return null;
  }

  private static SdkRequest injectIntoSendMessageBatchRequest(
      SendMessageBatchRequest request,
      io.opentelemetry.context.Context otelContext,
      TextMapPropagator messagingPropagator) {
    ArrayList<SendMessageBatchRequestEntry> entries = new ArrayList<>(request.entries());
    for (int i = 0; i < entries.size(); ++i) {
      SendMessageBatchRequestEntry entry = entries.get(i);
      Map<String, MessageAttributeValue> messageAttributes =
          new HashMap<>(entry.messageAttributes());

      // TODO: Per https://github.com/open-telemetry/oteps/pull/220, each message should get
      //  a separate context. We don't support this yet, also because it would be inconsistent
      //  with the header-based X-Ray propagation
      //  (probably could override it here by setting the X-Ray message system attribute)
      if (injectIntoMessageAttributes(messageAttributes, otelContext, messagingPropagator)) {
        entries.set(i, entry.toBuilder().messageAttributes(messageAttributes).build());
      }
    }
    return request.toBuilder().entries(entries).build();
  }

  private static SdkRequest injectIntoSendMessageRequest(
      SendMessageRequest request,
      io.opentelemetry.context.Context otelContext,
      TextMapPropagator messagingPropagator) {
    Map<String, MessageAttributeValue> messageAttributes =
        new HashMap<>(request.messageAttributes());
    if (!injectIntoMessageAttributes(messageAttributes, otelContext, messagingPropagator)) {
      return request;
    }
    return request.toBuilder().messageAttributes(messageAttributes).build();
  }

  private static boolean injectIntoMessageAttributes(
      Map<String, MessageAttributeValue> messageAttributes,
      io.opentelemetry.context.Context otelContext,
      TextMapPropagator messagingPropagator) {
    messagingPropagator.inject(
        otelContext,
        messageAttributes,
        (carrier, k, v) -> {
          carrier.put(k, MessageAttributeValue.builder().stringValue(v).dataType("String").build());
        });

    // Return whether the injection resulted in an attribute count that is still supported.
    // See
    // https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/sqs-message-metadata.html#sqs-message-attributes
    return messageAttributes.size() <= 10;
  }

  private static SdkRequest modifyReceiveMessageRequest(
      ReceiveMessageRequest request,
      boolean useXrayPropagator,
      TextMapPropagator messagingPropagator) {
    boolean hasXrayAttribute = true;
    List<String> existingAttributeNames = null;
    if (useXrayPropagator) {
      existingAttributeNames = request.attributeNamesAsStrings();
      hasXrayAttribute =
          existingAttributeNames.contains(SqsParentContext.AWS_TRACE_SYSTEM_ATTRIBUTE);
    }

    boolean hasMessageAttribute = true;
    List<String> existingMessageAttributeNames = null;
    if (messagingPropagator != null) {
      existingMessageAttributeNames = request.messageAttributeNames();
      hasMessageAttribute = existingMessageAttributeNames.containsAll(messagingPropagator.fields());
    }

    if (hasMessageAttribute && hasXrayAttribute) {
      return request;
    }

    ReceiveMessageRequest.Builder builder = request.toBuilder();
    if (!hasXrayAttribute) {
      List<String> attributeNames = new ArrayList<>(existingAttributeNames);
      attributeNames.add(SqsParentContext.AWS_TRACE_SYSTEM_ATTRIBUTE);
      builder.attributeNamesWithStrings(attributeNames);
    }
    if (messagingPropagator != null) {
      List<String> messageAttributeNames = new ArrayList<>(existingMessageAttributeNames);
      for (String field : messagingPropagator.fields()) {
        if (!existingMessageAttributeNames.contains(field)) {
          messageAttributeNames.add(field);
        }
      }
      builder.messageAttributeNames(messageAttributeNames);
    }
    return builder.build();
  }

  static boolean isSqsProducerRequest(SdkRequest request) {
    return request instanceof SendMessageRequest || request instanceof SendMessageBatchRequest;
  }

  static String getQueueUrl(SdkRequest request) {
    if (request instanceof SendMessageRequest) {
      return ((SendMessageRequest) request).queueUrl();
    } else if (request instanceof SendMessageBatchRequest) {
      return ((SendMessageBatchRequest) request).queueUrl();
    } else if (request instanceof ReceiveMessageRequest) {
      return ((ReceiveMessageRequest) request).queueUrl();
    }
    return null;
  }

  static String getMessageAttribute(SdkRequest request, String name) {
    if (request instanceof SendMessageRequest) {
      MessageAttributeValue value = ((SendMessageRequest) request).messageAttributes().get(name);
      return value != null ? value.stringValue() : null;
    }
    return null;
  }

  static String getMessageId(SdkResponse response) {
    if (response instanceof SendMessageResponse) {
      return ((SendMessageResponse) response).messageId();
    }
    return null;
  }

  static SqsClient wrap(SqsClient sqsClient) {
    // proxy SqsClient so we could replace the messages list in ReceiveMessageResponse returned from
    // receiveMessage call
    return (SqsClient)
        Proxy.newProxyInstance(
            sqsClient.getClass().getClassLoader(),
            new Class<?>[] {SqsClient.class},
            (proxy, method, args) -> {
              if ("receiveMessage".equals(method.getName())) {
                SqsTracingContext sqsTracingContext = new SqsTracingContext();
                try (Scope ignored =
                    io.opentelemetry.context.Context.current()
                        .with(sqsTracingContext)
                        .makeCurrent()) {
                  Object result = invokeProxyMethod(method, sqsClient, args);
                  TracingList tracingList = sqsTracingContext.get();
                  if (tracingList != null) {
                    ReceiveMessageResponse response = (ReceiveMessageResponse) result;
                    SqsImpl.setMessages(response, tracingList);
                    return response;
                  }
                  return result;
                }
              } else {
                return invokeProxyMethod(method, sqsClient, args);
              }
            });
  }

  @SuppressWarnings("unchecked")
  static SqsAsyncClient wrap(SqsAsyncClient sqsClient) {
    // proxy SqsAsyncClient so we could replace the messages list in ReceiveMessageResponse returned
    // from receiveMessage call
    return (SqsAsyncClient)
        Proxy.newProxyInstance(
            sqsClient.getClass().getClassLoader(),
            new Class<?>[] {SqsAsyncClient.class},
            (proxy, method, args) -> {
              if ("receiveMessage".equals(method.getName())) {
                SqsTracingContext sqsTracingContext = new SqsTracingContext();
                try (Scope ignored =
                    io.opentelemetry.context.Context.current()
                        .with(sqsTracingContext)
                        .makeCurrent()) {
                  Object result = invokeProxyMethod(method, sqsClient, args);
                  CompletableFuture<ReceiveMessageResponse> originalFuture =
                      (CompletableFuture<ReceiveMessageResponse>) result;
                  CompletableFuture<ReceiveMessageResponse> resultFuture =
                      new CompletableFuture<>();
                  originalFuture.whenComplete(
                      (response, throwable) -> {
                        if (throwable != null) {
                          resultFuture.completeExceptionally(throwable);
                        } else {
                          TracingList tracingList = sqsTracingContext.get();
                          if (tracingList != null) {
                            SqsImpl.setMessages(response, tracingList);
                          }
                          resultFuture.complete(response);
                        }
                      });

                  return resultFuture;
                } catch (InvocationTargetException exception) {
                  throw exception.getCause();
                }
              } else {
                return invokeProxyMethod(method, sqsClient, args);
              }
            });
  }

  private static Object invokeProxyMethod(Method method, Object target, Object[] args)
      throws Throwable {
    try {
      return method.invoke(target, args);
    } catch (InvocationTargetException exception) {
      throw exception.getCause();
    }
  }
}
