package profilers;

import org.openjdk.jmh.infra.BenchmarkParams;
import org.openjdk.jmh.profile.ExternalProfiler;
import org.openjdk.jmh.results.*;
import org.openjdk.jmh.util.FileUtils;

import java.io.File;
import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.lang.management.ManagementFactory;
import java.util.*;

// Effectively equivalent to passing jvm args to append for each benchmark. e.g.,
//
//@Fork(jvmArgsAppend =
//{"-XX:+UnlockCommercialFeatures",
//  "-XX:+FlightRecorder",
//  "-XX:StartFlightRecording=duration=60s,filename=./profiling-data.jfr,name=profile,settings=profile",
//  "-XX:FlightRecorderOptions=settings=./openjdk/jdk1.8.0/jre/lib/jfr/profile.jfc,samplethreads=true"
//})

/**
 * The flight recording profiler enables flight recording for benchmarks and starts recording right away.
 */
public class FlightRecordingProfiler implements ExternalProfiler {

    private String startFlightRecordingOptions = "duration=60s,name=profile,settings=profile,";
    private String flightRecorderOptions = "samplethreads=true,";

    /**
     * Directory to contain all generated reports.
     */
    private static final String SAVE_FLIGHT_OUTPUT_TO = System.getProperty("jmh.jfr.saveTo", ".");

    /**
     * Temporary location to record data
     */
    private final String jfrData;

    /**
     * Holds whether recording is supported (checking the existence of the needed unlocking flag)
     */
    private static final boolean IS_SUPPORTED;

    static {
        IS_SUPPORTED = ManagementFactory.getRuntimeMXBean().getInputArguments().contains("-XX:+UnlockCommercialFeatures");
    }

    public FlightRecordingProfiler() throws IOException {
        jfrData = FileUtils.tempFile("jfrData").getAbsolutePath();
    }

    @Override
    public Collection<String> addJVMInvokeOptions(BenchmarkParams params) {
        return Collections.emptyList();
    }

    @Override
    public Collection<String> addJVMOptions(BenchmarkParams params) {

        startFlightRecordingOptions += "filename=" + jfrData;
        flightRecorderOptions       += "settings=" + params.getJvm().replace("bin/java", "lib/jfr/profile.jfc");

        return Arrays.asList(
                "-XX:+FlightRecorder",
                "-XX:StartFlightRecording="  + startFlightRecordingOptions,
                "-XX:FlightRecorderOptions=" + flightRecorderOptions);
    }

    @Override
    public void beforeTrial(BenchmarkParams benchmarkParams) {

    }

    static int currentId;

    @Override
    public Collection<? extends Result> afterTrial(BenchmarkResult benchmarkResult, long l, File stdOut, File stdErr) {
        String target = SAVE_FLIGHT_OUTPUT_TO + "/" + benchmarkResult.getParams().id() + "-" + currentId++ + ".jfr";

        StringWriter sw = new StringWriter();
        PrintWriter pw = new PrintWriter(sw);

        try {
            FileUtils.copy(jfrData, target);
            pw.println("Flight Recording output saved to " + target);
        } catch (IOException e) {
            pw.println("Unable to save flight output to " + target);
        }

        pw.flush();
        pw.close();

        NoResult r = new NoResult(sw.toString());

        return Collections.singleton(r);
    }

    @Override
    public boolean allowPrintOut() {
        return true;
    }

    @Override
    public boolean allowPrintErr() {
        return false;
    }

    @Override
    public boolean checkSupport(List<String> msgs) {
        msgs.add("Commercial features of the JVM need to be enabled for this profiler.");
        return IS_SUPPORTED;
    }

    @Override
    public String label() {
        return "jfr";
    }

    @Override
    public String getDescription() {
        return "Java Flight Recording profiler runs for every benchmark.";
    }

    private class NoResult extends Result<NoResult> {
        private final String output;

        public NoResult(String output) {
            super(ResultRole.SECONDARY, "JFR", of(Double.NaN), "N/A", AggregationPolicy.SUM);

            this.output = output;
        }

        @Override
        protected Aggregator<NoResult> getThreadAggregator() {
            return new NoResultAggregator();
        }

        @Override
        protected Aggregator<NoResult> getIterationAggregator() {
            return new NoResultAggregator();
        }

        @Override
        public String extendedInfo() {
            return "JFR Messages:\n--------------------------------------------\n" + output;
        }

        private class NoResultAggregator implements Aggregator<NoResult> {
            @Override
            public NoResult aggregate(Collection<NoResult> results) {
                String output = "";
                for (NoResult r : results) {
                    output += r.output;
                }
                return new NoResult(output);
            }
        }
    }
}

@@ -335,16 +335,92 @@
 actory;%0A
+import java.net.URI;%0Aimport java.nio.file.Path;%0Aimport java.nio.file.Paths;%0A
 import j
@@ -2155,53 +2155,35 @@
     
-flightRecorderOptions       += %22settings=%22 + 
+String jfcPath = Paths.get(
 para
@@ -2197,30 +2197,25 @@
 vm()
+)
 .re
-place(%22bin/java%22, %22
+solve(%22../../
 lib/
@@ -2231,16 +2231,118 @@
 le.jfc%22)
+.normalize().toAbsolutePath().toString();%0A        flightRecorderOptions       += %22settings=%22 + jfcPath
 ;%0A%0A     
@@ -2829,16 +2829,26 @@
 arget = 
+Paths.get(
 SAVE_FLI
@@ -2864,17 +2864,18 @@
 T_TO
- + %22/%22 + 
+).resolve(
 benc
@@ -2898,18 +2898,28 @@
 arams().
-id
+getBenchmark
 () + %22-%22
@@ -2941,16 +2941,45 @@
 + %22.jfr%22
+).toAbsolutePath().toString()
 ;%0A%0A     
@@ -3225,24 +3225,57 @@
 eption e) %7B%0A
+            e.printStackTrace();%0A
             
@@ -3323,32 +3323,112 @@
 to %22 + target);%0A
+            pw.println(%22Did you miss the system property: -Djmh.jfr.saveTo ?%22);%0A
         %7D%0A%0A     

Use Path api instead of string concatenate