package com.github.sulir.runtimesearch.agent;

import com.github.sulir.runtimesearch.agent.transformer.ClassTransformer;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.regex.Pattern;

public class SearchAgent {
    private static final List<String> defaultExclude = Arrays.asList(
            "com.sun.", "java.", "javax.", "jdk.", "sun.",
            "com.intellij.rt.", "org.jetbrains.capture.", "org.groovy.debug.", "groovyResetJarjarAsm.",
            "org.objectweb.asm.", "com.github.sulir.runtimesearch."
    );

    public static void premain(String agentArgs, Instrumentation inst) {
        Check.initialize();
        Server.getInstance().start();

        boolean filterSupplied = agentArgs != null && !agentArgs.isEmpty();
        Pattern include = Pattern.compile(filterSupplied ? agentArgs : ".*");
        List<String> exclude = filterSupplied ? Collections.emptyList() : defaultExclude;

        inst.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                    ProtectionDomain protectionDomain, byte[] classfileBuffer) {
                try {
                    if (className == null)
                        return null;
                    String name = className.replace('/', '.');

                    if (include.matcher(name).matches() && exclude.stream().noneMatch(name::startsWith))
                        return new ClassTransformer(className, classfileBuffer).transform();
                } catch (Throwable e) {
                    e.printStackTrace();
                }

                return null;
            }
        });
    }
}

@@ -781,62 +781,117 @@
     
-Check.initialize();%0A        Server.getInstance().start
+SearchAgent agent = new SearchAgent();%0A        agent.setupCommunication();%0A        agent.fixJbossClassLoading
 ();%0A
@@ -1135,16 +1135,739 @@
 xclude;%0A
+        agent.setupInstrumentation(inst, include, exclude);%0A    %7D%0A%0A    private void setupCommunication() %7B%0A        Check.initialize();%0A        Server.getInstance().start();%0A    %7D%0A%0A    private void fixJbossClassLoading() %7B%0A        String SYSTEM_PACKAGES = %22jboss.modules.system.pkgs%22;%0A        String currentPackages = System.getProperty(SYSTEM_PACKAGES);%0A        boolean isSet = currentPackages != null && !currentPackages.isEmpty();%0A        String agentPackage = Check.class.getPackage().getName();%0A        System.setProperty(SYSTEM_PACKAGES, isSet ? agentPackage : currentPackages + %22,%22 + agentPackage);%0A    %7D%0A%0A    private void setupInstrumentation(Instrumentation instrumentation, Pattern include, List%3CString%3E exclude) %7B
 %0A       
@@ -1871,16 +1871,27 @@
     inst
+rumentation
 .addTran
@@ -2053,16 +2053,32 @@
 efined,%0A
+                
         

Fix ClassNotFoundException in JBoss module classloader