bridge36

package comet.agent;

import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.matcher.ElementMatchers;
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 {
    public static final ConcurrentHashMap<String, Stats> metrics = new ConcurrentHashMap<>();
    private static final SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");

    public static void agentmain(String agentArgs, Instrumentation inst) {
        System.out.println("PROFILER: Dynamic attach successful.");
        System.out.flush();

        // 1. Add Agent JAR to Bootstrap Search Path
        try {
            File agentJar = new File(ProfilerAgent.class.getProtectionDomain()
                    .getCodeSource().getLocation().toURI());
            if (agentJar.exists() && agentJar.getName().endsWith(".jar")) {
                inst.appendToBootstrapClassLoaderSearch(new JarFile(agentJar));
                System.out.println("PROFILER: Injected JAR into Bootstrap path.");
            }
        } catch (Exception e) {
            System.err.println("PROFILER ERROR: Bootstrap injection failed: " + e.getMessage());
        }

        // 2. Dump currently loaded classes to file for verification
        dumpLoadedClasses(inst, "loaded_classes_dump.txt");

        premain(agentArgs, inst);
    }

    public static void premain(String agentArgs, Instrumentation inst) {
        String configPath = "profiler_targets.txt";
        int interval = 60;
        String outputDir = ".";

        if (agentArgs != null && !agentArgs.isEmpty()) {
            String[] parts = agentArgs.split(";");
            if (parts.length > 0) configPath = parts[0];
            if (parts.length > 1) interval = Integer.parseInt(parts[1]);
            if (parts.length > 2) outputDir = parts[2];
        }

        List<String> targetClasses = loadClasses(configPath);
        startReporter(interval, outputDir);

        new AgentBuilder.Default()
            .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
            .with(AgentBuilder.CircularityLock.Default.INSTANCE)
            // 3. Disable the default ignore filter to see Bootstrap classes
            .ignore(ElementMatchers.none())
            .type(builder -> {
                String name = builder.getName();
                // 4. Safe-Mode Check: Skip self and ByteBuddy internals
                if (name.startsWith("comet.agent.") || name.startsWith("net.bytebuddy.")) {
                    return false;
                }
                return targetClasses.stream().anyMatch(t -> name.equals(t.trim()));
            })
            .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: Transformed " + td.getName() + " [" + (cl == null ? "Bootstrap" : cl) + "]");
                    System.out.flush();
                }
            })
            .installOn(inst);
    }

    private static void dumpLoadedClasses(Instrumentation inst, String fileName) {
        try (PrintWriter writer = new PrintWriter(new BufferedWriter(new FileWriter(fileName)))) {
            Class<?>[] allLoadedClasses = inst.getAllLoadedClasses();
            writer.println("Total Classes: " + allLoadedClasses.length);
            for (Class<?> clazz : allLoadedClasses) {
                writer.println(clazz.getName());
            }
            writer.flush();
            System.out.println("PROFILER: Dumped loaded classes to " + fileName);
            System.out.flush();
        } catch (IOException e) {
            System.err.println("PROFILER ERROR: Class dump failed: " + e.getMessage());
        }
    }

    private static List<String> loadClasses(String path) {
        try {
            return Files.readAllLines(Paths.get(path));
        } catch (Exception e) {
            System.err.println("PROFILER ERROR: Could not read config: " + path);
            return Collections.emptyList();
        }
    }

    public static class ProfilerAdvice {
        // 5. Thread-safe set for Heartbeat tracking
        private static final Set<String> seenMethods = Collections.newSetFromMap(new ConcurrentHashMap<String, Boolean>());

        @Advice.OnMethodEnter
        static long enter(@Advice.Origin("#t.#m") String methodName) {
            if (seenMethods.add(methodName)) {
                System.out.println(">>> HEARTBEAT: First call detected in " + methodName);
                System.out.flush();
            }
            return System.nanoTime();
        }

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

    public static class Stats {
        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, final String outputDir) {
        File csvFile = new File(outputDir, "profiler_combined_report.csv");
        Executors.newSingleThreadScheduledExecutor().scheduleAtFixedRate(() -> {
            String timestamp = sdf.format(new Date());
            try (PrintWriter pw = new PrintWriter(new BufferedWriter(new FileWriter(csvFile, true)))) {
                if (!csvFile.exists() || csvFile.length() == 0) {
                    pw.println("Timestamp,Method,Calls,Avg_MS,Total_MS");
                }

                final int[] tally = {0};
                metrics.forEach((method, stats) -> {
                    long c = stats.count.sumThenReset();
                    long t = stats.totalTime.sumThenReset();
                    if (c > 0) {
                        tally[0]++;
                        double avgMs = (t / (double) c) / 1000000.0;
                        double totalMs = t / 1000000.0;
                        pw.printf("%s,%s,%d,%.4f,%.2f%n", timestamp, method, c, avgMs, totalMs);
                    }
                });
                // 6. Force Flush for Console and File
                pw.flush();
                System.out.println("PROFILER: Report flushed at " + timestamp + ". Methods recorded this interval: " + tally[0]);
                System.out.flush();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }, seconds, seconds, TimeUnit.SECONDS);
    }
}