package datadog.trace.instrumentation.servlet2;

import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE;
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator._500;
import static datadog.trace.instrumentation.servlet2.HttpServletRequestExtractAdapter.GETTER;
import static datadog.trace.instrumentation.servlet2.Servlet2Decorator.DECORATE;
import static datadog.trace.instrumentation.servlet2.Servlet2Decorator.SERVLET_REQUEST;

import datadog.trace.api.Config;
import datadog.trace.api.CorrelationIdentifier;
import datadog.trace.api.DDTags;
import datadog.trace.api.GlobalTracer;
import datadog.trace.bootstrap.InstrumentationContext;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.api.Tags;
import java.security.Principal;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.implementation.bytecode.assign.Assigner;

public class Servlet2Advice {

  @Advice.OnMethodEnter(suppress = Throwable.class)
  public static AgentScope onEnter(
      @Advice.This final Object servlet,
      @Advice.Argument(0) final ServletRequest request,
      @Advice.Argument(value = 1, typing = Assigner.Typing.DYNAMIC)
          final ServletResponse response) {

    final boolean hasServletTrace = request.getAttribute(DD_SPAN_ATTRIBUTE) instanceof AgentSpan;
    final boolean invalidRequest = !(request instanceof HttpServletRequest);
    if (invalidRequest || hasServletTrace) {
      // Tracing might already be applied by the FilterChain or a parent request (forward/include).
      return null;
    }

    final HttpServletRequest httpServletRequest = (HttpServletRequest) request;

    if (response instanceof HttpServletResponse) {
      // Default value for checking for uncaught error later
      InstrumentationContext.get(ServletResponse.class, Integer.class).put(response, 200);
    }

    final AgentSpan.Context extractedContext = propagate().extract(httpServletRequest, GETTER);

    final AgentSpan span = startSpan(SERVLET_REQUEST, extractedContext).setMeasured(true);

    DECORATE.afterStart(span);
    DECORATE.onConnection(span, httpServletRequest);
    DECORATE.onRequest(span, httpServletRequest);

    final AgentScope scope = activateSpan(span);
    scope.setAsyncPropagation(true);

    httpServletRequest.setAttribute(DD_SPAN_ATTRIBUTE, span);
    httpServletRequest.setAttribute(
        CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
    httpServletRequest.setAttribute(
        CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());

    return scope;
  }

  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
  public static void stopSpan(
      @Advice.Argument(0) final ServletRequest request,
      @Advice.Argument(1) final ServletResponse response,
      @Advice.Enter final AgentScope scope,
      @Advice.Thrown final Throwable throwable) {
    // Set user.principal regardless of who created this span.
    final Object spanAttr = request.getAttribute(DD_SPAN_ATTRIBUTE);
    if (Config.get().isServletPrincipalEnabled()
        && spanAttr instanceof AgentSpan
        && request instanceof HttpServletRequest) {
      final Principal principal = ((HttpServletRequest) request).getUserPrincipal();
      if (principal != null) {
        ((AgentSpan) spanAttr).setTag(DDTags.USER_NAME, principal.getName());
      }
    }

    if (scope == null) {
      return;
    }
    final AgentSpan span = scope.span();

    if (response instanceof HttpServletResponse) {
      DECORATE.onResponse(
          span, InstrumentationContext.get(ServletResponse.class, Integer.class).get(response));
    } else {
      DECORATE.onResponse(span, null);
    }

    if (throwable != null) {
      if (response instanceof HttpServletResponse
          && InstrumentationContext.get(ServletResponse.class, Integer.class).get(response)
              == HttpServletResponse.SC_OK) {
        // exception was thrown but status code wasn't set
        span.setTag(Tags.HTTP_STATUS, _500);
      }
      DECORATE.onError(span, throwable);
    }
    DECORATE.beforeFinish(span);

    scope.setAsyncPropagation(false);
    scope.close();
    span.finish();
  }
}
