import java.io.IOException;
import java.io.InputStream;
import java.lang.classfile.*;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicBoolean;

public class Example {

    public static void main(String[] args) {

        byte[] classBytes;
        String classPath = Example.class.getName().replace('.', '/') + ".class";
        
        try(InputStream stream = Objects.requireNonNull(Example.class.getClassLoader().getResourceAsStream(classPath))) {
            classBytes = stream.readAllBytes();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        var classFile = ClassFile.of();
        var classModel = classFile.parse(classBytes);

        AtomicBoolean triggered = new AtomicBoolean(false);

        classFile.transformClass(classModel,
                ClassTransform.transformingMethodBodies(
                        method -> method.methodName().equalsString("foo"),
                        CodeTransform.endHandler(_ -> triggered.set(true))
                )
        );

        if (!triggered.getAndSet(false)) {
            throw new RuntimeException("Expected triggered first");
        }

        classFile.transformClass(classModel,
                ClassTransform.transformingMethodBodies(
                        method -> method.methodName().equalsString("foo"),
                        CodeTransform.endHandler(_ -> triggered.set(true))
                ).andThen(ClassTransform.ofStateful(() -> ClassTransform.transformingMethodBodies(
                        method -> method.methodName().equalsString("bar"),
                        CodeTransform.endHandler(_ -> triggered.set(true))
                )))
        );

        if (!triggered.getAndSet(false)) {
            throw new RuntimeException("Expected triggered second");
        }

        classFile.transformClass(classModel,
                ClassTransform.transformingMethodBodies(
                        method -> method.methodName().equalsString("foo"),
                        CodeTransform.endHandler(_ -> triggered.set(true))
                ).andThen(ClassTransform.transformingMethodBodies(
                        method -> method.methodName().equalsString("bar"),
                        CodeTransform.endHandler(_ -> triggered.set(true))
                ))
        );

        if (!triggered.getAndSet(false)) {
            throw new RuntimeException("Expected triggered third");
        }
    }

    public void foo() { }
    public void bar() { }
}