import com.fasterxml.jackson.core.util.DefaultIndenter;
import com.fasterxml.jackson.core.util.DefaultPrettyPrinter;
import com.fasterxml.jackson.databind.MapperFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import org.junit.jupiter.api.extension.BeforeEachCallback;
import org.junit.jupiter.api.extension.ExtensionContext;

import java.lang.reflect.Method;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;

import static org.assertj.core.api.Assertions.assertThat;

public final class SnapshotAssertion {

	private static Method method;
	private static int index;
	private static int calls;

	private final ObjectMapper objectMapper;
	private final boolean forceUpdate;

	public SnapshotAssertion(ObjectMapper objectMapper) {
		// 美化输出的 JSON
		var indenter = new DefaultIndenter().withIndent("\t");
		var printer = new DefaultPrettyPrinter()
				.withArrayIndenter(indenter)
				.withObjectIndenter(indenter);

		this.objectMapper = objectMapper.copy()
				.enable(SerializationFeature.INDENT_OUTPUT)
				.setDefaultPrettyPrinter(printer)
				// 对属性排序，确保结果稳定
				.enable(SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS)
				.enable(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY);

		this.forceUpdate = System.getProperty("updateSnapshot") != null;
	}

	// 断言对象与快照里保存的相等
	public void assertMatch(Object object) {
		createOrLoad(object).assertEquals();
	}

	// 加载快照，如果快照不存在则创建
	@SneakyThrows
	private Matching createOrLoad(Object object) {
		var actual = objectMapper.writeValueAsString(object);
		var expect = actual;
		var path = getSnapshotPath();

		if (!Files.exists(path) || forceUpdate) {
			Files.createDirectories(path.getParent());
			Files.writeString(path, actual);
		} else {
			// 重新格式化，避免手动美化快照文件造成的影响。
			var tree = objectMapper.readTree(Files.newInputStream(path));
			expect = objectMapper.writeValueAsString(tree);
		}

		return new Matching(expect, actual);
	}

	// 获取当前断言对应的快照文件路径
	private Path getSnapshotPath() {
		if (method == null) {
			throw new Error("必须把 ContextHolder 注册到 JUnit 扩展");
		}
		var template = "src/test/resources/snapshots/%s/%s-%d-%d.json";
		var c = method.getDeclaringClass().getSimpleName();
		var m = method.getName();
		return Paths.get(String.format(template, c, m, index, calls++));
	}

	// 用于获取当前测试方法的 JUnit5 扩展
	public static final class ContextHolder implements BeforeEachCallback {

		@Override
		public void beforeEach(ExtensionContext context) {
			var latestMethod = method;
			method = context.getRequiredTestMethod();
			calls = 0;
			index = method == latestMethod ? index + 1 : 0;
		}
	}

	@RequiredArgsConstructor
	private static final class Matching {

		private final String expect;
		private final String actual;

		public void assertEquals() {
			assertThat(actual).isEqualTo(expect);
		}
	}
}

// 需要添加一个扩展
@ExtendWith(SnapshotAssertion.ContextHolder.class)
final class SnapshotTestExample {
	
	private final ObjectMapper objectMapper = new ObjectMapper()
			.setDefaultPropertyInclusion(JsonInclude.Include.NON_NULL); // 可以自定义序列化器

	private final SnapshotAssertion snapshot = new SnapshotAssertion(objectMapper);

	@Test
	void test() {
		// 测试运行的结果，实际中可以是复杂的对象
		var result = Map.of("field1", 0, "field2", "text");

		snapshot.assertMatch(result);
	}
}

{
	"field1" : 0,
	"field2" : "text"
}