bridge46

package comet.agent;

import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.matcher.ElementMatchers;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.utility.JavaModule;

import java.lang.instrument.Instrumentation;
import java.io.*;
import java.nio.file.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.LongAdder;
import java.text.SimpleDateFormat;
import java.util.jar.JarFile;

public class ProfilerAgent {
    private static final String MAP_KEY = "comet.metrics.global";
    private static final SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");

    @SuppressWarnings("unchecked")
    public static ConcurrentHashMap<String, Stats> getGlobalMap() {
        if (System.getProperties().get(MAP_KEY) == null) {
            System.getProperties().put(MAP_KEY, new ConcurrentHashMap<String, Stats>());
        }
        return (ConcurrentHashMap<String, Stats>) System.getProperties().get(MAP_KEY);
    }

    public static void agentmain(String agentArgs, Instrumentation inst) {
        System.out.println("PROFILER: [MODE] Dynamic Attach");
        setupBootstrap(inst);
        dumpLoadedClasses(inst); // Diagnostic Line
        premain(agentArgs, inst);
    }

    public static void premain(String agentArgs, Instrumentation inst) {
        if (!System.getProperties().containsKey("profiler.started")) {
            System.out.println("PROFILER: [MODE] Startup Agent");
            System.getProperties().put("profiler.started", "true");
            setupBootstrap(inst);
            dumpLoadedClasses(inst); // Diagnostic Line
        }

        String folderPath = "build/profiler-results";
        new File(folderPath).mkdirs();

        List<String> targets = loadClasses("profiler_targets.txt");
        startReporter(30, folderPath);

        new AgentBuilder.Default()
            .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
            .with(AgentBuilder.TypeStrategy.Default.REDEFINE) 
            .with(new AgentBuilder.CircularityLock() {
                private final ThreadLocal<Boolean> lock = new ThreadLocal<Boolean>() {
                    @Override protected Boolean initialValue() { return false; }
                };
                @Override public boolean acquire() { if (lock.get()) return false; lock.set(true); return true; }
                @Override public void release() { lock.set(false); }
            })
            .ignore(ElementMatchers.none())
            .type(builder -> {
                String name = builder.getName();
                if (name.startsWith("comet.agent.") || name.startsWith("net.bytebuddy.")) return false;
                return targets.stream().anyMatch(t -> name.trim().equals(t.trim()) || 
                       (t.endsWith("*") && name.startsWith(t.substring(0, t.length()-1))));
            })
            .transform((builder, typeDescription, classLoader, module) ->
                builder.method(ElementMatchers.any()
                        .and(ElementMatchers.not(ElementMatchers.isAbstract()))
                        .and(ElementMatchers.not(ElementMatchers.isNative())))
                        .intercept(Advice.to(ProfilerAdvice.class))
            )
            .with(new AgentBuilder.Listener.Adapter() {
                @Override
                public void onTransformation(TypeDescription td, ClassLoader cl, JavaModule m, boolean loaded, DynamicType dt) {
                    System.out.println("PROFILER: [SUCCESS] Instrumented " + td.getName());
                    System.out.flush();
                }
            })
            .installOn(inst);
    }

    private static void dumpLoadedClasses(Instrumentation inst) {
        System.out.println("--- START LOADED CLASS DUMP ---");
        Class<?>[] allLoaded = inst.getAllLoadedClasses();
        for (Class<?> clazz : allLoaded) {
            System.out.println("LOADED: " + clazz.getName());
        }
        System.out.println("--- END LOADED CLASS DUMP ---");
        System.out.flush();
    }

    private static void setupBootstrap(Instrumentation inst) {
        try {
            File agentJar = new File(ProfilerAgent.class.getProtectionDomain().getCodeSource().getLocation().toURI());
            if (agentJar.exists()) {
                inst.appendToBootstrapClassLoaderSearch(new JarFile(agentJar));
            }
        } catch (Exception ignored) {}
    }

    private static List<String> loadClasses(String path) {
        try { return Files.readAllLines(Paths.get(path)); }
        catch (Exception e) { return Collections.emptyList(); }
    }

    public static class ProfilerAdvice {
        @Advice.OnMethodEnter
        static long enter() { return System.nanoTime(); }

        @Advice.OnMethodExit(onThrowable = Throwable.class)
        static void exit(@Advice.Enter long start, @Advice.Origin("#t.#m") String methodName) {
            if (start == 0L) return;
            long duration = System.nanoTime() - start;
            ConcurrentHashMap<String, Stats> metrics = comet.agent.ProfilerAgent.getGlobalMap();
            Stats s = metrics.get(methodName);
            if (s == null) {
                metrics.putIfAbsent(methodName, new comet.agent.ProfilerAgent.Stats());
                s = metrics.get(methodName);
            }
            s.record(duration);
        }
    }

    public static class Stats implements Serializable {
        public final LongAdder count = new LongAdder();
        public final LongAdder totalTime = new LongAdder();
        public void record(long nanos) {
            count.increment();
            totalTime.add(nanos);
        }
    }

    private static void startReporter(int seconds, String outputDir) {
        Executors.newSingleThreadScheduledExecutor().scheduleAtFixedRate(() -> {
            ConcurrentHashMap<String, Stats> metrics = getGlobalMap();
            File csvFile = new File(outputDir, "profiler-report.csv");
            try (PrintWriter pw = new PrintWriter(new BufferedWriter(new FileWriter(csvFile, true)))) {
                String ts = sdf.format(new Date());
                final int[] active = {0};
                metrics.forEach((method, stats) -> {
                    long c = stats.count.sumThenReset();
                    long t = stats.totalTime.sumThenReset();
                    if (c > 0) {
                        active[0]++;
                        pw.printf("%s,%s,%d,%.4f,%.2f%n", ts, method, c, (t/(double)c)/1000000.0, t/1000000.0);
                    }
                });
                pw.flush();
                System.out.println("PROFILER: [" + ts + "] Active hooks: " + active[0]);
                System.out.flush();
            } catch (Exception e) { e.printStackTrace(); }
        }, seconds, seconds, TimeUnit.SECONDS);
    }
}