bridge-6

package comet.agent;

import net.bytebuddy.agent.ByteBuddyAgent;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.implementation.MethodDelegation;
import net.bytebuddy.implementation.bind.annotation.*;
import net.bytebuddy.matcher.ElementMatchers;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Method;
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 {

    public static final ConcurrentHashMap<String, Stats> metrics = new ConcurrentHashMap<>();
    private static final SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");

    /**
     * Call this method as the very first line of your main(String[] args) method.
     */
    public static void initProfilerAgent(String configPath, int intervalSeconds, String outputDir) {
        try {
            System.out.println("[PROFILER-AGENT] Dynamically attaching ByteBuddy agent...");

            // 1. Programmatically acquire the Instrumentation instance via the Attach API
            Instrumentation inst = ByteBuddyAgent.install();

            // 2. Fallbacks for configuration if parameters are omitted
            String finalConfigPath = (configPath != null && !configPath.isEmpty()) ? configPath : "profiler_targets.txt";
            String finalOutputDir = (outputDir != null && !outputDir.isEmpty()) ? outputDir : ".";
            int finalInterval = (intervalSeconds > 0) ? intervalSeconds : 60;

            // 3. Load targeted classes and spin up reporting thread
            List<String> targetClasses = loadClasses(finalConfigPath);
            if (targetClasses.isEmpty()) {
                System.out.println("[PROFILER-AGENT] WARNING: No target classes loaded from " + finalConfigPath);
                return;
            }

            startReporter(finalInterval, finalOutputDir);

            // 4. Initialize AgentBuilder with RedefinitionStrategy for runtime attachment
            AgentBuilder agentBuilder = new AgentBuilder.Default()
                .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION);

            for (String className : targetClasses) {
                agentBuilder = agentBuilder
                    .type(ElementMatchers.named(className.trim()))
                    .transform((builder, td, cl, m, pd) -> builder
                        .method(ElementMatchers.any()
                            .and(ElementMatchers.not(ElementMatchers.isAbstract()))
                            .and(ElementMatchers.not(ElementMatchers.isNative())))
                        .intercept(MethodDelegation.to(Interceptor.class))
                    );
            }

            agentBuilder.installOn(inst);
            System.out.println("[PROFILER-AGENT] Agent successfully attached and instrumentation applied.");

        } catch (Exception e) {
            System.err.println("[PROFILER-AGENT] Critical error trying to initialize programmatic agent injection:");
            e.printStackTrace();
        }
    }

    private static List<String> loadClasses(String path) {
        try {
            return Files.readAllLines(Paths.get(path));
        } catch (Exception e) {
            System.err.println("[PROFILER-AGENT] Failed to read target classes file: " + path);
            return Collections.emptyList();
        }
    }

    public static class Interceptor {
        @RuntimeType
        public static Object intercept(@Origin Method method, @SuperCall Callable<?> callable) throws Exception {
            long start = System.nanoTime();
            try {
                return callable.call();
            } finally {
                long duration = System.nanoTime() - start;
                String key = method.getDeclaringClass().getSimpleName() + "." + method.getName();
                metrics.computeIfAbsent(key, k -> new Stats()).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(() -> {
            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");
                }

                String timestamp = sdf.format(new Date());

                metrics.forEach((k, v) -> {
                    long c = v.count.sumThenReset();
                    long t = v.totalTime.sumThenReset();
                    if (c > 0) {
                        pw.printf("%s,%s,%d,%.4f,%.2f\n", 
                            timestamp, 
                            k, 
                            c, 
                            (t / (double) c) / 1e6, 
                            t / 1e6
                        );
                    }
                });
                pw.flush();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }, seconds, seconds, TimeUnit.SECONDS);
    }
}