bridge58

package comet.agent;

import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.matcher.ElementMatchers;
import net.bytebuddy.agent.ByteBuddyAgent;
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;

public class ProfilerAgent {
    // CRITICAL: Must be public and static to avoid synthetic accessors (access$000)
    public static final ConcurrentHashMap<String, Stats> metrics = new ConcurrentHashMap<>();
    public static final Set<String> targetClassSet = new CopyOnWriteArraySet<>();
    public static final SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");

    // Add this for dynamic attachment
    private static volatile boolean isActive = false;

    public static void agentmain(String agentArgs, Instrumentation inst) {
        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];
        }

        loadClasses(configPath);
        startReporter(interval, outputDir);

        // Check if this is a dynamic attachment
        boolean isDynamic = isDynamicAttachment();
        System.out.println("PROFILER: Installing agent (dynamic=" + isDynamic + ")");

        // For dynamic attachment, we need to ensure classes are retransformed
        AgentBuilder.RedefinitionStrategy redefStrategy = isDynamic ? 
            AgentBuilder.RedefinitionStrategy.RETRANSFORMATION :
            AgentBuilder.RedefinitionStrategy.DISABLED;

        new AgentBuilder.Default()
            .with(redefStrategy)
            .with(AgentBuilder.TypeStrategy.Default.REDEFINE) // Explicit type strategy
            .with(AgentBuilder.InitializationStrategy.NoOp.INSTANCE) // Better for dynamic attachment
            .type(typeDescription -> targetClassSet.contains(typeDescription.getActualName()))
            .transform((builder, typeDescription, classLoader, module, protectionDomain) -> 
                builder.visit(Advice.to(ProfilerAdvice.class)
                    .on(ElementMatchers.any()
                        .and(ElementMatchers.not(ElementMatchers.isAbstract()))
                        .and(ElementMatchers.not(ElementMatchers.isNative()))
                        .and(ElementMatchers.not(ElementMatchers.isConstructor()))))
            )
            .installOn(inst);

        // For dynamic attachment, try to retransform already loaded classes
        if (isDynamic) {
            retransformLoadedClasses(inst);
        }

        isActive = true;
    }

    private static boolean isDynamicAttachment() {
        // Check if we're in a dynamic attachment scenario
        try {
            // If ByteBuddyAgent is available and we're not in premain, it's dynamic
            return !isPremain();
        } catch (Throwable t) {
            return false;
        }
    }

    private static boolean isPremain() {
        try {
            StackTraceElement[] stackTrace = Thread.currentThread().getStackTrace();
            for (StackTraceElement element : stackTrace) {
                if ("premain".equals(element.getMethodName())) {
                    return true;
                }
            }
        } catch (Exception e) {
            // Ignore
        }
        return false;
    }

    private static void retransformLoadedClasses(Instrumentation inst) {
        try {
            // Get all loaded classes
            Class<?>[] loadedClasses = inst.getAllLoadedClasses();
            List<Class<?>> classesToRetransform = new ArrayList<>();

            for (Class<?> clazz : loadedClasses) {
                if (targetClassSet.contains(clazz.getName())) {
                    classesToRetransform.add(clazz);
                }
            }

            if (!classesToRetransform.isEmpty()) {
                System.out.println("PROFILER: Retransforming " + classesToRetransform.size() + " already loaded classes");
                inst.retransformClasses(classesToRetransform.toArray(new Class[0]));
            }
        } catch (Exception e) {
            System.err.println("PROFILER: Error during retransformation: " + e.getMessage());
            e.printStackTrace();
        }
    }

    public static void loadClasses(String path) {
        try {
            // Try multiple locations for the config file
            Path configPath = findConfigFile(path);
            if (configPath == null) {
                System.err.println("PROFILER: Could not find config file: " + path);
                return;
            }

            List<String> lines = Files.readAllLines(configPath);
            for (String line : lines) {
                line = line.trim();
                if (!line.isEmpty() && !line.startsWith("#")) {
                    targetClassSet.add(line);
                    System.out.println("PROFILER: Target class added: " + line);
                }
            }
            System.out.println("PROFILER: Loaded " + targetClassSet.size() + " target classes");
        } catch (Exception e) {
            System.err.println("PROFILER: Could not read config: " + path);
            e.printStackTrace();
        }
    }

    private static Path findConfigFile(String path) {
        try {
            Path directPath = Paths.get(path);
            if (Files.exists(directPath)) {
                return directPath;
            }

            // Try in the current working directory
            Path cwdPath = Paths.get(System.getProperty("user.dir"), path);
            if (Files.exists(cwdPath)) {
                return cwdPath;
            }

            // Try in the agent jar directory
            String agentLocation = ProfilerAgent.class.getProtectionDomain()
                .getCodeSource().getLocation().getPath();
            Path agentDirPath = Paths.get(new File(agentLocation).getParent(), path);
            if (Files.exists(agentDirPath)) {
                return agentDirPath;
            }
        } catch (Exception e) {
            // Ignore
        }
        return null;
    }

    public static class ProfilerAdvice {
        @Advice.OnMethodEnter
        static long enter() {
            if (!ProfilerAgent.isActive) return 0L;
            return System.nanoTime();
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class)
        static void exit(@Advice.Enter long start, 
                         @Advice.Origin("#t") String className,
                         @Advice.Origin("#m") String methodName) {
            if (start == 0L || !ProfilerAgent.isActive) return;

            long duration = System.nanoTime() - start;

            // Create fully qualified method name
            String fullMethodName = className + "." + methodName;

            // Use Fully Qualified Name to avoid any ambiguity
            Stats s = metrics.get(fullMethodName);
            if (s == null) {
                metrics.putIfAbsent(fullMethodName, new Stats());
                s = metrics.get(fullMethodName);
            }
            if (s != null) {
                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) {
        // Ensure output directory exists
        File dir = new File(outputDir);
        if (!dir.exists()) {
            dir.mkdirs();
        }

        File csvFile = new File(outputDir, "profiler_combined_report.csv");
        Executors.newSingleThreadScheduledExecutor().scheduleAtFixedRate(() -> {
            try {
                String timestamp = sdf.format(new Date());
                boolean writeHeader = !csvFile.exists() || csvFile.length() == 0;

                try (PrintWriter pw = new PrintWriter(new BufferedWriter(new FileWriter(csvFile, true)))) {
                    if (writeHeader) {
                        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) / 1_000_000.0;
                            double totalMs = t / 1_000_000.0;
                            pw.printf("%s,%s,%d,%.4f,%.2f%n", timestamp, method, c, avgMs, totalMs);
                        }
                    });

                    if (tally[0] > 0) {
                        System.out.println("PROFILER: Flush @ " + timestamp + ". Captured: " + tally[0] + " methods");
                    }

                    pw.flush(); // Ensure data is written
                }
            } catch (Exception e) {
                System.err.println("PROFILER: Error in reporter: " + e.getMessage());
                e.printStackTrace();
            }
        }, seconds, seconds, TimeUnit.SECONDS);
    }

    // Add a shutdown hook to flush remaining stats
    static {
        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            isActive = false;
            // Final flush could be added here
        }));
    }
}