package com.hubspot.jackson.test;

import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.CaseFormat;
import com.google.protobuf.Descriptors.Descriptor;
import com.google.protobuf.Descriptors.FieldDescriptor;
import com.google.protobuf.ExtensionRegistry;
import com.google.protobuf.ExtensionRegistry.ExtensionInfo;
import com.hubspot.jackson.datatype.protobuf.ExtensionRegistryWrapper;
import com.hubspot.jackson.test.util.TestExtensionRegistry;
import com.hubspot.jackson.test.util.TestProtobuf.AllFields;
import org.junit.BeforeClass;
import org.junit.Test;

import java.util.EnumSet;
import java.util.HashSet;
import java.util.Set;

import static com.hubspot.jackson.test.util.ObjectMapperHelper.camelCase;
import static org.assertj.core.api.Assertions.assertThat;

public class JsonInclusionTest {
  private static final ExtensionRegistry EXTENSION_REGISTRY = TestExtensionRegistry.getInstance();

  private static Set<String> allFields;
  private static Set<String> allExtensionFields;
  private static Set<String> arrayFields;
  private static Set<String> arrayExtensionFields;

  @BeforeClass
  public static void setup() {
    allFields = new HashSet<>();
    arrayFields = new HashSet<>();

    Descriptor descriptor = AllFields.getDescriptor();
    for (FieldDescriptor field : descriptor.getFields()) {
      allFields.add(translate(field.getName()));
      if (field.isRepeated()) {
        arrayFields.add(translate(field.getName()));
      }
    }

    allExtensionFields = new HashSet<>();
    arrayExtensionFields = new HashSet<>();

    ExtensionRegistryWrapper extensionRegistry = ExtensionRegistryWrapper.wrap(EXTENSION_REGISTRY);
    for (ExtensionInfo extensionInfo : extensionRegistry.findExtensionsByDescriptor(descriptor)) {
      allExtensionFields.add(translate(extensionInfo.descriptor.getName()));
      if (extensionInfo.descriptor.isRepeated()) {
        arrayExtensionFields.add(translate(extensionInfo.descriptor.getName()));
      }
    }
  }

  @Test
  public void itWritesMissingFieldsAsNullWhenSerializationIncludeIsAlways() {
    AllFields message = AllFields.getDefaultInstance();

    JsonNode node = mapper(Include.ALWAYS).valueToTree(message);

    for (String field : allFields) {
      assertThat(node.has(field)).isTrue();
      if (arrayFields.contains(field)) {
        assertThat(node.get(field).isArray());
      } else {
        assertThat(node.get(field).isNull());
      }
    }

    for (String field : allExtensionFields) {
      assertThat(node.has(field)).isFalse();
    }
  }

  @Test
  public void itOnlyWritesArrayFieldsWhenSerializationIncludeIsNotAlways() {
    AllFields message = AllFields.getDefaultInstance();

    for (Include inclusion : EnumSet.complementOf(EnumSet.of(Include.ALWAYS))) {
      JsonNode node = mapper(inclusion).valueToTree(message);

      for (String field : allFields) {
        if (arrayFields.contains(field)) {
          assertThat(node.has(field)).isTrue();
          assertThat(node.get(field).isArray());
        } else {
          assertThat(node.has(field)).isFalse();
        }
      }

      for (String field : allExtensionFields) {
        assertThat(node.has(field)).isFalse();
      }
    }
  }

  @Test
  public void itWritesMissingExtensionFieldsAsNullWhenSerializationIncludeIsAlways() {
    AllFields message = AllFields.getDefaultInstance();

    JsonNode node = mapper(Include.ALWAYS, EXTENSION_REGISTRY).valueToTree(message);

    for (String field : allExtensionFields) {
      assertThat(node.has(field)).isTrue();
      if (arrayExtensionFields.contains(field)) {
        assertThat(node.get(field).isArray());
      } else {
        assertThat(node.get(field).isNull());
      }
    }
  }

  @Test
  public void itOnlyWritesArrayExtensionFieldsWhenSerializationIncludeIsNotAlways() {
    AllFields message = AllFields.getDefaultInstance();

    for (Include inclusion : EnumSet.complementOf(EnumSet.of(Include.ALWAYS))) {
      JsonNode node = mapper(inclusion, EXTENSION_REGISTRY).valueToTree(message);

      for (String field : allExtensionFields) {
        if (arrayExtensionFields.contains(field)) {
          assertThat(node.has(field)).isTrue();
          assertThat(node.get(field).isArray());
        } else {
          assertThat(node.has(field)).isFalse();
        }
      }
    }
  }

  private static String translate(String fieldName) {
    return CaseFormat.LOWER_UNDERSCORE.to(CaseFormat.LOWER_CAMEL, fieldName);
  }

  private static ObjectMapper mapper(Include inclusion) {
    return camelCase().copy().setSerializationInclusion(inclusion);
  }

  private static ObjectMapper mapper(Include inclusion, ExtensionRegistry extensionRegistry) {
    return camelCase(extensionRegistry).copy().setSerializationInclusion(inclusion);
  }
}

@@ -24,24 +24,286 @@
 kson.test;%0A%0A
+import static com.hubspot.jackson.test.util.ObjectMapperHelper.camelCase;%0Aimport static org.assertj.core.api.Assertions.assertThat;%0A%0Aimport java.util.EnumSet;%0Aimport java.util.HashSet;%0Aimport java.util.Set;%0A%0Aimport org.junit.BeforeClass;%0Aimport org.junit.Test;%0A%0A
 import com.f
@@ -493,16 +493,53 @@
 Format;%0A
+import com.google.common.base.Enums;%0A
 import c
@@ -939,302 +939,150 @@
 ds;%0A
-import org.junit.BeforeClass;%0Aimport org.junit.Test;%0A%0Aimport java.util.EnumSet;%0Aimport java.util.HashSet;%0Aimport java.util.Set;%0A%0Aimport static com.hubspot.jackson.test.util.ObjectMapperHelper.camelCase;%0Aimport static org.assertj.core.api.Assertions.assertThat;%0A%0Apublic class JsonInclusionTest %7B
+%0Apublic class JsonInclusionTest %7B%0A  private static final EnumSet%3CInclude%3E EXCLUDED_VALUES = presentValues(%22ALWAYS%22, %22USE_DEFAULTS%22, %22CUSTOM%22);
 %0A  p
@@ -3006,33 +3006,22 @@
 Of(E
-numSet.of(Include.ALWAYS)
+XCLUDED_VALUES
 )) %7B
@@ -4173,33 +4173,22 @@
 Of(E
-numSet.of(Include.ALWAYS)
+XCLUDED_VALUES
 )) %7B
@@ -5014,11 +5014,268 @@
 on);%0A  %7D
+%0A%0A  private static EnumSet%3CInclude%3E presentValues(String... values) %7B%0A    EnumSet%3CInclude%3E set = EnumSet.noneOf(Include.class);%0A%0A    for (String value : values) %7B%0A      set.addAll(Enums.getIfPresent(Include.class, value).asSet());%0A    %7D%0A%0A    return set;%0A  %7D
 %0A%7D%0A

Make this work against old and new Jackson