content: /*
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package cc.shanruifeng.functions.udfs.presto.json;
  
  import com.facebook.presto.spi.PrestoException;
  import com.fasterxml.jackson.core.JsonFactory;
  import com.fasterxml.jackson.core.JsonGenerator;
  import com.fasterxml.jackson.core.JsonParseException;
  import com.fasterxml.jackson.core.JsonParser;
  import com.fasterxml.jackson.core.JsonToken;
  import com.fasterxml.jackson.core.io.SerializedString;
  import com.google.common.base.Throwables;
  import com.google.common.collect.ImmutableList;
  import io.airlift.slice.DynamicSliceOutput;
  import io.airlift.slice.Slice;
  import java.io.IOException;
  
  import static cc.shanruifeng.functions.udfs.utils.JsonUtil.createJsonGenerator;
  import static cc.shanruifeng.functions.udfs.utils.JsonUtil.createJsonParser;
  import static com.facebook.presto.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT;
  import static com.fasterxml.jackson.core.JsonFactory.Feature.CANONICALIZE_FIELD_NAMES;
  import static com.fasterxml.jackson.core.JsonToken.END_ARRAY;
  import static com.fasterxml.jackson.core.JsonToken.END_OBJECT;
  import static com.fasterxml.jackson.core.JsonToken.FIELD_NAME;
  import static com.fasterxml.jackson.core.JsonToken.START_ARRAY;
  import static com.fasterxml.jackson.core.JsonToken.START_OBJECT;
  import static com.fasterxml.jackson.core.JsonToken.VALUE_NULL;
  import static io.airlift.slice.Slices.utf8Slice;
  import static java.util.Objects.requireNonNull;
  
  /**
   * Extracts values from JSON
   * <p>
   * Supports the following JSON path primitives:
   * <pre>
   *    $ : Root object
   *    . or [] : Child operator
   *   [] : Subscript operator for array
   * </pre>
   * <p>
   * Supported JSON Path Examples:
   * <pre>
   *    { "store": {
   *        "book": [
   *          { "category": "reference",
   *            "author": "Nigel Rees",
   *            "title": "Sayings of the Century",
   *            "price": 8.95,
   *            "contributors": [["Adam", "Levine"], ["Bob", "Strong"]]
   *          },
   *          { "category": "fiction",
   *            "author": "Evelyn Waugh",
   *            "title": "Sword of Honour",
   *            "price": 12.99,
   *            "isbn": "0-553-21311-3",
   *            "last_owner": null
   *          }
   *        ],
   *        "bicycle": {
   *          "color": "red",
   *          "price": 19.95
   *        }
   *      }
   *    }
   * </pre>
   * <p>
   * With only scalar values using dot-notation of path:
   * <pre>
   *    $.store.book[0].author => Nigel Rees
   *    $.store.bicycle.price => 19.95
   *    $.store.book[0].isbn => NULL (Doesn't exist becomes java null)
   *    $.store.book[1].last_owner => NULL (json null becomes java null)
   *    $.store.book[0].contributors[0][1] => Levine
   * </pre>
   * <p>
   * With json values using dot-notation of path:
   * <pre>
   *    $.store.book[0].author => "Nigel Rees"
   *    $.store.bicycle.price => 19.95
   *    $.store.book[0].isbn => NULL (Doesn't exist becomes java null)
   *    $.store.book[1].last_owner => null (json null becomes the string "null")
   *    $.store.book[0].contributors[0] => ["Adam", "Levine"]
   *    $.store.bicycle => {"color": "red", "price": 19.95}
   * </pre>
   * With only scalar values using bracket-notation of path:
   * <pre>
   *    $["store"]["book"][0]["author"] => Nigel Rees
   *    $["store"]["bicycle"]["price"] => 19.95
   *    $["store"]["book"][0]["isbn"] => NULL (Doesn't exist becomes java null)
   *    $["store"]["book"][1]["last_owner"] => NULL (json null becomes java null)
   *    $["store"]["book"][0]["contributors"][0][1] => Levine
   * </pre>
   * <p>
   * With json values using bracket-notation of path:
   * <pre>
   *    $["store"]["book"][0]["author"] => "Nigel Rees"
   *    $["store"]["bicycle"]["price"] => 19.95
   *    $["store"]["book"][0]["isbn"] => NULL (Doesn't exist becomes java null)
   *    $["store"]["book"][1]["last_owner"] => null (json null becomes the string "null")
   *    $["store"]["book"][0]["contributors"][0] => ["Adam", "Levine"]
   *    $["store"]["bicycle"] => {"color": "red", "price": 19.95}
   * </pre>
   */
  public final class JsonExtract {
      private static final int ESTIMATED_JSON_OUTPUT_SIZE = 512;
  
      private static final JsonFactory JSON_FACTORY = new JsonFactory()
              .disable(CANONICALIZE_FIELD_NAMES);
  
      private JsonExtract() {
      }
  
      public static <T> T extract(Slice jsonInput, JsonExtractor<T> jsonExtractor) {
          requireNonNull(jsonInput, "jsonInput is null");
          try {
              try (JsonParser jsonParser = createJsonParser(JSON_FACTORY, jsonInput)) {
                  // Initialize by advancing to first token and make sure it exists
                  if (jsonParser.nextToken() == null) {
                      return null;
                  }
  
                  return jsonExtractor.extract(jsonParser);
              }
          } catch (JsonParseException e) {
              // Return null if we failed to parse something
              return null;
          } catch (IOException e) {
              throw Throwables.propagate(e);
          }
      }
  
      public static <T> JsonExtractor<T> generateExtractor(String path, JsonExtractor<T> rootExtractor) {
          return generateExtractor(path, rootExtractor, false);
      }
  
      public static <T> JsonExtractor<T> generateExtractor(String path, JsonExtractor<T> rootExtractor, boolean exceptionOnOutOfBounds) {
          ImmutableList<String> tokens = ImmutableList.copyOf(new JsonPathTokenizer(path));
  
          JsonExtractor<T> jsonExtractor = rootExtractor;
          for (String token : tokens.reverse()) {
              jsonExtractor = new ObjectFieldJsonExtractor<>(token, jsonExtractor, exceptionOnOutOfBounds);
          }
          return jsonExtractor;
      }
  
      private static int tryParseInt(String fieldName, int defaultValue) {
          int index = defaultValue;
          try {
              index = Integer.parseInt(fieldName);
          } catch (NumberFormatException ignored) {
          }
          return index;
      }
  
      public interface JsonExtractor<T> {
          /**
           * Executes the extraction on the existing content of the JsonParser and outputs the match.
           * <p>
           * Notes:
           * <ul>
           * <li>JsonParser must be on the FIRST token of the value to be processed when extract is called</li>
           * <li>INVARIANT: when extract() returns, the current token of the parser will be the LAST token of the value</li>
           * </ul>
           *
           * @return the value, or null if not applicable
           */
          T extract(JsonParser jsonParser)
                  throws IOException;
      }
  
      public static class ObjectFieldJsonExtractor<T>
              implements JsonExtractor<T> {
          private final SerializedString fieldName;
          private final JsonExtractor<? extends T> delegate;
          private final int index;
          private final boolean exceptionOnOutOfBounds;
  
          public ObjectFieldJsonExtractor(String fieldName, JsonExtractor<? extends T> delegate) {
              this(fieldName, delegate, false);
          }
  
          public ObjectFieldJsonExtractor(String fieldName, JsonExtractor<? extends T> delegate, boolean exceptionOnOutOfBounds) {
              this.fieldName = new SerializedString(requireNonNull(fieldName, "fieldName is null"));
              this.delegate = requireNonNull(delegate, "delegate is null");
              this.exceptionOnOutOfBounds = exceptionOnOutOfBounds;
              this.index = tryParseInt(fieldName, -1);
          }
  
          @Override
          public T extract(JsonParser jsonParser)
                  throws IOException {
              if (jsonParser.getCurrentToken() == START_OBJECT) {
                  return processJsonObject(jsonParser);
              }
  
              if (jsonParser.getCurrentToken() == START_ARRAY) {
                  return processJsonArray(jsonParser);
              }
  
              throw new JsonParseException("Expected a JSON object or array", jsonParser.getCurrentLocation());
          }
  
          public T processJsonObject(JsonParser jsonParser)
                  throws IOException {
              while (!jsonParser.nextFieldName(fieldName)) {
                  if (!jsonParser.hasCurrentToken()) {
                      throw new JsonParseException("Unexpected end of object", jsonParser.getCurrentLocation());
                  }
                  if (jsonParser.getCurrentToken() == END_OBJECT) {
                      // Unable to find matching field
                      return null;
                  }
                  jsonParser.skipChildren(); // Skip nested structure if currently at the start of one
              }
  
              jsonParser.nextToken(); // Shift to first token of the value
  
              return delegate.extract(jsonParser);
          }
  
          public T processJsonArray(JsonParser jsonParser)
                  throws IOException {
              int currentIndex = 0;
              while (true) {
                  JsonToken token = jsonParser.nextToken();
                  if (token == null) {
                      throw new JsonParseException("Unexpected end of array", jsonParser.getCurrentLocation());
                  }
                  if (token == END_ARRAY) {
                      // Index out of bounds
                      if (exceptionOnOutOfBounds) {
                          throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "Index out of bounds");
                      }
                      return null;
                  }
                  if (currentIndex == index) {
                      break;
                  }
                  currentIndex++;
                  jsonParser.skipChildren(); // Skip nested structure if currently at the start of one
              }
  
              return delegate.extract(jsonParser);
          }
      }
  
      public static class ScalarValueJsonExtractor
              implements JsonExtractor<Slice> {
          @Override
          public Slice extract(JsonParser jsonParser)
                  throws IOException {
              JsonToken token = jsonParser.getCurrentToken();
              if (token == null) {
                  throw new JsonParseException("Unexpected end of value", jsonParser.getCurrentLocation());
              }
              if (!token.isScalarValue() || token == VALUE_NULL) {
                  return null;
              }
              return utf8Slice(jsonParser.getText());
          }
      }
  
      public static class JsonValueJsonExtractor
              implements JsonExtractor<Slice> {
          @Override
          public Slice extract(JsonParser jsonParser)
                  throws IOException {
              if (!jsonParser.hasCurrentToken()) {
                  throw new JsonParseException("Unexpected end of value", jsonParser.getCurrentLocation());
              }
  
              DynamicSliceOutput dynamicSliceOutput = new DynamicSliceOutput(ESTIMATED_JSON_OUTPUT_SIZE);
              try (JsonGenerator jsonGenerator = createJsonGenerator(JSON_FACTORY, dynamicSliceOutput)) {
                  jsonGenerator.copyCurrentStructure(jsonParser);
              }
              return dynamicSliceOutput.slice();
          }
      }
  
      public static class JsonSizeExtractor
              implements JsonExtractor<Long> {
          @Override
          public Long extract(JsonParser jsonParser)
                  throws IOException {
              if (!jsonParser.hasCurrentToken()) {
                  throw new JsonParseException("Unexpected end of value", jsonParser.getCurrentLocation());
              }
  
              if (jsonParser.getCurrentToken() == START_ARRAY) {
                  long length = 0;
                  while (true) {
                      JsonToken token = jsonParser.nextToken();
                      if (token == null) {
                          return null;
                      }
                      if (token == END_ARRAY) {
                          return length;
                      }
                      jsonParser.skipChildren();
  
                      length++;
                  }
              }
  
              if (jsonParser.getCurrentToken() == START_OBJECT) {
                  long length = 0;
                  while (true) {
                      JsonToken token = jsonParser.nextToken();
                      if (token == null) {
                          return null;
                      }
                      if (token == END_OBJECT) {
                          return length;
                      }
  
                      if (token == FIELD_NAME) {
                          length++;
                      } else {
                          jsonParser.skipChildren();
                      }
                  }
              }
  
              return 0L;
          }
      }
  }
  