content: package us.abstracta.jmeter.javadsl.graphql;
  
  import com.fasterxml.jackson.core.JsonProcessingException;
  import com.fasterxml.jackson.databind.JsonNode;
  import com.fasterxml.jackson.databind.ObjectMapper;
  import com.fasterxml.jackson.databind.util.RawValue;
  import java.lang.reflect.Method;
  import java.util.Iterator;
  import java.util.LinkedHashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Spliterator;
  import java.util.Spliterators;
  import java.util.stream.Stream;
  import java.util.stream.StreamSupport;
  import org.apache.http.entity.ContentType;
  import org.apache.jmeter.config.Arguments;
  import org.apache.jmeter.protocol.http.config.GraphQLRequestParams;
  import org.apache.jmeter.protocol.http.config.gui.GraphQLUrlConfigGui;
  import org.apache.jmeter.protocol.http.control.HeaderManager;
  import org.apache.jmeter.protocol.http.control.gui.GraphQLHTTPSamplerGui;
  import org.apache.jmeter.protocol.http.sampler.HTTPSamplerBase;
  import org.apache.jmeter.protocol.http.sampler.HTTPSamplerProxy;
  import org.apache.jmeter.protocol.http.util.GraphQLRequestParamUtils;
  import org.apache.jmeter.protocol.http.util.HTTPArgument;
  import org.apache.jmeter.protocol.http.util.HTTPConstants;
  import us.abstracta.jmeter.javadsl.codegeneration.MethodCall;
  import us.abstracta.jmeter.javadsl.codegeneration.MethodCallContext;
  import us.abstracta.jmeter.javadsl.codegeneration.MethodParam;
  import us.abstracta.jmeter.javadsl.codegeneration.TestElementParamBuilder;
  import us.abstracta.jmeter.javadsl.codegeneration.params.FixedParam;
  import us.abstracta.jmeter.javadsl.codegeneration.params.StringParam;
  import us.abstracta.jmeter.javadsl.http.DslBaseHttpSampler;
  
  /**
   * Allows configuring a GraphQL HTTP Sampler which significantly eases creating GraphQL requests
   * compared to httpSampler.
   * <p>
   * This sampler is based on HTTP Sampler, so all configurations affecting http sampler (headers,
   * cookies, cache, system properties, etc.) affect this sampler.
   *
   * @since 0.52
   */
  public class DslGraphqlSampler extends DslBaseHttpSampler<DslGraphqlSampler> {
  
    private static final String DEFAULT_NAME = "GraphQL HTTP Request";
    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
  
    protected String query;
    protected String operationName;
    protected final Map<String, Object> variables = new LinkedHashMap<>();
    protected String variablesJson;
    protected String method = HTTPConstants.POST;
  
    public DslGraphqlSampler(String name, String url, String query) {
      super(name != null ? name : DEFAULT_NAME, url, GraphQLHTTPSamplerGui.class);
      this.query = query;
      headers.contentType(ContentType.APPLICATION_JSON);
    }
  
    /**
     * Builds a GraphQL HTTP Request sampler to invoke GraphQL operations.
     *
     * @param url specifies URL where the GraphQL service is attending requests.
     * @return the sampler for further configuration or usage.
     * @see DslGraphqlSampler
     */
    public static DslGraphqlSampler graphqlSampler(String url, String query) {
      return graphqlSampler(null, url, query);
    }
  
    /**
     * Same as {@link #graphqlSampler(String, String)} but allowing to set a name to the GraphQL HTTP
     * Request sampler.
     * <p>
     * Setting a proper name allows to easily identify the requests generated by this sampler and
     * check its particular statistics.
     *
     * @see #graphqlSampler(String, String)
     */
    public static DslGraphqlSampler graphqlSampler(String name, String url, String query) {
      return new DslGraphqlSampler(name, url, query);
    }
  
    /**
     * Specifies the GraphQL operation name to send in the request.
     * <p>
     * <a href="https://graphql.org/learn/queries/#operation-name">Here</a> you can find additional
     * information about operation name purposes.
     *
     * @param operationName is the name to send in the GraphQL request.
     * @return the sampler for further configuration or usage.
     */
    public DslGraphqlSampler operationName(String operationName) {
      this.operationName = operationName;
      return this;
    }
  
    /**
     * Specifies a GraphQL variable to send in the request.
     * <p>
     * <a href="https://graphql.org/learn/queries/#variablese">Here</a> you can find additional
     * information about variables purposes.
     *
     * @param name  specifies the name of the variable.
     * @param value specifies the value associated to the variable. This value is serialized into
     *              json, so for non-primitive values (int, long, boolean, float, double, string,
     *              etc.) make sure they can be properly serialized by jackson library or use
     *              {@link #rawVariable(String, String)} or {@link #variablesJson(String)} method
     *              instead.
     * @return the sampler for further configuration or usage.
     * @throws IllegalArgumentException when provided value object can't be serialized by jackson.
     */
    public DslGraphqlSampler variable(String name, Object value) {
      checkValueSerialization(value);
      this.variables.put(name, value);
      return this;
    }
  
    private void checkValueSerialization(Object value) {
      try {
        OBJECT_MAPPER.writeValueAsString(value);
      } catch (JsonProcessingException e) {
        throw new IllegalArgumentException("Could not serialize a given GraphQL variable", e);
      }
    }
  
    /**
     * Allows specifying variables values which should not be escaped in final JSON string.
     * <p>
     * This method is handy when you want to use JMeter expressions for ints, longs, booleans, floats
     * and doubles (which should not include quotes in final json) and when you want to specify
     * complex json strings like arrays &amp; objects.
     *
     * @param name  specifies the name of the variable.
     * @param value the raw JSON string part to be set as variable value.
     * @return the sampler for further configuration or usage.
     * @since 0.54
     */
    public DslGraphqlSampler rawVariable(String name, String value) {
      this.variables.put(name, new RawValue(value));
      return this;
    }
  
    /**
     * Allows specifying variables in a string in JSON format as JMeter sampler does.
     * <p>
     * In general prefer using {@link #variable(String, Object)} instead, which is simpler, and avoids
     * potential format issues. Use this method in corner cases when {@link #variable(String, Object)}
     * is a limiting factor.
     *
     * @param variables specifies all variables to be sent in GraphQL request in JSON Object format.
     * @return the sampler for further configuration or usage.
     * @see #variable(String, Object)
     */
    public DslGraphqlSampler variablesJson(String variables) {
      this.variablesJson = variables;
      return this;
    }
  
    /**
     * Allows specifying to use HTTP GET method instead of POST method.
     * <p>
     * This method is helpful in some particular cases where you might want to send request with GET
     * method instead of POST method to properly emulate a client using such HTTP method.
     *
     * @return the sampler for further configuration or usage.
     * @since 1.7
     */
    public DslGraphqlSampler httpGet() {
      method = HTTPConstants.GET;
      return this;
    }
  
    @Override
    public HTTPSamplerProxy configureHttpTestElement(HTTPSamplerProxy elem) {
      elem.setMethod(method);
      elem.setProperty(GraphQLUrlConfigGui.OPERATION_NAME, operationName);
      elem.setProperty(GraphQLUrlConfigGui.QUERY, query);
      elem.setProperty(HTTPSamplerBase.POST_BODY_RAW, !HTTPConstants.GET.equals(method));
      String varsJson = buildVariablesJson();
      elem.setProperty(GraphQLUrlConfigGui.VARIABLES, varsJson);
      elem.setArguments(HTTPConstants.GET.equals(method) ? buildHttpGetArguments(varsJson)
          : buildHttpPostArguments(varsJson));
      return elem;
    }
  
    private String buildVariablesJson() {
      if (variablesJson != null) {
        return variablesJson;
      } else if (variables.isEmpty()) {
        return "";
      } else {
        try {
          return OBJECT_MAPPER.writeValueAsString(variables);
        } catch (JsonProcessingException e) {
          // this should never happen due to checks in variable setting method
          throw new RuntimeException(e);
        }
      }
    }
  
    private Arguments buildHttpGetArguments(String jsonVariables) {
      Arguments args = new Arguments();
      if (operationName != null && !operationName.trim().isEmpty()) {
        args.addArgument(new HTTPArgument("operationName", operationName.trim()));
      }
      args.addArgument(
          new HTTPArgument("query", GraphQLRequestParamUtils.queryToGetParamValue(query)));
      if (variablesJson != null && !variablesJson.trim().isEmpty()) {
        args.addArgument(new HTTPArgument("variables", jsonVariables));
      }
      return args;
    }
  
    private Arguments buildHttpPostArguments(String jsonVariables) {
      Arguments args = new Arguments();
      HTTPArgument arg = new HTTPArgument("", buildGraphqlBody(jsonVariables));
      arg.setAlwaysEncoded(false);
      args.addArgument(arg);
      return args;
    }
  
    private String buildGraphqlBody(String jsonVariables) {
      GraphQLRequestParams params = new GraphQLRequestParams();
      params.setOperationName(operationName);
      params.setQuery(query);
      params.setVariables(jsonVariables);
      return GraphQLRequestParamUtils.toPostBodyString(params);
    }
  
    public static class CodeBuilder extends BaseHttpSamplerCodeBuilder {
  
      public CodeBuilder(List<Method> builderMethods) {
        super(DEFAULT_NAME, GraphQLHTTPSamplerGui.class, builderMethods);
      }
  
      @Override
      protected MethodCall buildBaseHttpMethodCall(MethodParam name, MethodParam url,
          TestElementParamBuilder paramBuilder) {
        return buildMethodCall(name, url, paramBuilder.stringParam(GraphQLUrlConfigGui.QUERY));
      }
  
      @Override
      protected void chainRequestCalls(MethodCall ret, HTTPSamplerProxy testElem,
          MethodCallContext buildContext) {
        TestElementParamBuilder paramBuilder = new TestElementParamBuilder(testElem);
        ret.chain("operationName",
            paramBuilder.stringParam(GraphQLUrlConfigGui.OPERATION_NAME));
        MethodParam vars = paramBuilder.stringParam(GraphQLUrlConfigGui.VARIABLES);
        if (!(vars instanceof StringParam)) {
          chainRawVariables(ret, vars);
        } else if (!vars.isDefault()) {
          try {
            JsonNode varsNode = OBJECT_MAPPER.readValue(((StringParam) vars).getValue(),
                JsonNode.class);
            if (!varsNode.isObject()) {
              chainRawVariables(ret, vars);
            } else {
              iterator2Stream(varsNode.fields())
                  .forEach(var -> ret.chain(var.getValue().isValueNode() ? "variable" : "rawVariable",
                      new StringParam(var.getKey()), new JsonValueParam(var.getValue())));
            }
          } catch (JsonProcessingException e) {
            chainRawVariables(ret, vars);
          }
        }
        setHeaders(ret, buildContext);
      }
  
      private void setHeaders(MethodCall ret, MethodCallContext buildContext) {
        MethodCallContext headersContext = buildContext.removeChild(HeaderManager.class::isInstance);
        String contentType = removeContentTypeHeader(headersContext);
        if (!ContentType.APPLICATION_JSON.toString().equals(contentType)) {
          chainContentType(ret, contentType);
        }
        chainHeaders(ret, headersContext);
      }
  
      private void chainRawVariables(MethodCall ret, MethodParam vars) {
        ret.chain("variablesJson", vars);
      }
  
      private <T> Stream<T> iterator2Stream(Iterator<T> iterator) {
        return StreamSupport.stream(
            Spliterators.spliteratorUnknownSize(iterator, Spliterator.ORDERED), false);
      }
  
      @Override
      protected void chainAdditionalOptions(MethodCall ret, TestElementParamBuilder paramBuilder) {
      }
  
      private static class JsonValueParam extends FixedParam<Object> {
  
        private JsonValueParam(JsonNode value) {
          super(Object.class, value, null);
        }
  
        @Override
        public String buildCode(String indent) {
          JsonNode node = (JsonNode) value;
          if (node.isValueNode()) {
            return node.isTextual() ? string2Code(node.textValue()) : node.asText();
          } else {
            try {
              return string2Code(OBJECT_MAPPER.writeValueAsString(value));
            } catch (JsonProcessingException e) {
              // this should never happen since value was parsed from json.
              throw new RuntimeException(e);
            }
          }
        }
  
        private String string2Code(String value) {
          return buildStringLiteral(value, "");
        }
  
      }
  
    }
  
  }
  