/*
 * Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.  Oracle designates this
 * particular file as subject to the "Classpath" exception as provided
 * by Oracle in the LICENSE file that accompanied this code.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */
package org.openjdk.jmh.profile;

import joptsimple.OptionException;
import joptsimple.OptionParser;
import joptsimple.OptionSpec;
import org.openjdk.jmh.infra.BenchmarkParams;
import org.openjdk.jmh.results.BenchmarkResult;
import org.openjdk.jmh.results.BenchmarkResultMetaData;
import org.openjdk.jmh.results.Result;
import org.openjdk.jmh.util.*;

import java.io.File;
import java.io.IOException;
import java.io.PrintWriter;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * macOS permasm profiler based on xctrace utility shipped with Xcode Instruments.
 * Depending on the chosen tracing configuration, so-called template, xctrace may sample program counter
 * of a profiled process using either timer interrupts, or PMU interrupts.
 * <p>
 * The profiling process consists of several steps:
 * 1) launching a program that needs to be profiled using `xctrace record` command; in case of success,
 * the output of this step is a "trace-file", which is in fact a directory containing multiple files
 * representing the recorded trace, the trace may contain multiple resulting tables, depending on the template;
 * 2) inspecting a recorded trace to check if it contains a table supported by the profiler; this information
 * could be obtained from the trace's table of contents (`xctrace export --toc`);
 * 3) extracting the table with profiling results from the trace file using xctrace export and parsing it
 * to extract individual samples.
 * <p>
 * `xctrace export` command extracts data only in XML format, thus both the table of contents and the table
 * with profiling results need to be parsed as XML documents.
 * <p>
 * This profiler currently supports only three table types:
 * 1) time-profile tables generated by the Time Profiler instrument performing sampling by timer interrupts;
 * 2) cpu-profile tables generated by the CPU Profiler instrument performing sampling by the cycles PMU interrupts;
 * 3) counters-profile tables generated by the CPU Counters instrument performing sampling either by
 * timer interrupts, or interrupts generated by a PMU counter, depending on particular configuration.
 * <p>
 * Tracing configuration, or template, mentioned earlier is a file that could be configured and saved
 * using Instruments application. Unfortunately, there is no other way to set it up.
 * <p>
 * There are several builtin templates, and some of them, namely "Time Profiler" and "CPU Profiler", could be
 * used without any additional setup.
 * <p>
 */
public class XCTraceAsmProfiler extends AbstractPerfAsmProfiler {
    private final boolean shouldFixStartTime;
    private final Path temporaryDirectory;
    private final String template;
    private final String xctracePath;
    private OptionSpec<String> templateOpt;
    private OptionSpec<Boolean> correctOpt;
    private XCTraceTableHandler.ProfilingTableType resultsTable;

    private long recordStartMs;
    private long forkStartTimeMs;

    public XCTraceAsmProfiler(String initLine) throws ProfilerException {
        super(initLine, "sampled_pc");
        xctracePath = XCTraceSupport.getXCTracePath();
        try {
            template = set.valueOf(templateOpt);
            shouldFixStartTime = set.valueOf(correctOpt);
            temporaryDirectory = XCTraceSupport.createTemporaryDirectoryName();
        } catch (OptionException e) {
            throw new ProfilerException(e.getMessage());
        }
    }

    @Override
    protected void addMyOptions(OptionParser parser) {
        templateOpt = parser.accepts("template",
                        "Path to or name of an Instruments template. " +
                                "Use `xctrace list templates` to view available templates.")
                .withRequiredArg()
                .describedAs("string")
                .ofType(String.class)
                .defaultsTo("Time Profiler");

        correctOpt = parser.accepts("fixStartTime", "Fix the start time by the time it took to launch.")
                .withRequiredArg()
                .describedAs("bool")
                .ofType(Boolean.class)
                .defaultsTo(true);
    }

    @Override
    protected void parseEvents() {
        Path profile = XCTraceSupport.findTraceFile(temporaryDirectory);
        XCTraceSupport.exportTableOfContents(xctracePath, profile.toAbsolutePath().toString(),
                perfParsedData.getAbsolutePath());
        XCTraceTableOfContentsHandler handler = new XCTraceTableOfContentsHandler();
        handler.parse(perfParsedData.file());
        recordStartMs = handler.getRecordStartMs();
        List<XCTraceTableHandler.XCTraceTableDesc> tables = handler.getSupportedTables();
        if (tables.isEmpty()) {
            throw new IllegalStateException("Profiling results does not contain tables supported by this profiler.");
        }
        if (tables.size() != 1) {
            String supportedSchemas = Stream.of(XCTraceTableHandler.ProfilingTableType.values())
                    .map(t -> t.tableName)
                    .sorted()
                    .collect(Collectors.joining(", "));
            throw new IllegalStateException("There are multiple supported tables in output, " +
                    "please make sure a template you chose produces only a single table with one of the following " +
                    "schemas: " + supportedSchemas);
        }
        resultsTable = tables.get(0).getTableType();
        XCTraceSupport.exportTable(xctracePath, profile.toAbsolutePath().toString(), perfParsedData.getAbsolutePath(),
                resultsTable);
    }

    @Override
    public void beforeTrial(BenchmarkParams params) {
        super.beforeTrial(params);
        temporaryDirectory.toFile().mkdirs();
    }

    @Override
    public Collection<? extends Result> afterTrial(BenchmarkResult br, long pid, File stdOut, File stdErr) {
        try {
            BenchmarkResultMetaData md = br.getMetadata();
            if (md != null) {
                forkStartTimeMs = md.getStartTime();
            }
            return super.afterTrial(br, pid, stdOut, stdErr);
        } finally {
            XCTraceSupport.removeDirectory(temporaryDirectory);
        }
    }

    @Override
    protected PerfEvents readEvents(double skipMs, double lenMs) {
        Deduplicator<MethodDesc> dedup = new Deduplicator<>();
        Map<MethodDesc, AddressInterval> methods = new HashMap<>();
        Multiset<Long> events = new TreeMultiset<>();

        /*
         * BenchmarkResultMetaData captures the time when a fork was launched.
         * xctrace saves timestamps relative to the start of the tracing process.
         * It may take a considerable time to start xctrace (up to several seconds in some cases), so we
         * can't directly use measurementDelayMs.
         * However, xctrace trace's table of contents contains the timestamp corresponding to
         * the actual traced process start time. It could be used to correct measurementDelayMs.
         *
         *               /<-- delta -->/
         *              /             /
         *  time -------|-------------|------------------------------|----------->
         *              |             |                              |
         *       *fork launched*      |                     *measurements started*
         *        getStartTime()      |                      getMeasurementTime()
         *                      xctrace started java
         *                    tocHandler.getRecordStartMs()
         */
        long timeCorrectionMs = 0;
        if (shouldFixStartTime) {
            timeCorrectionMs = recordStartMs - forkStartTimeMs;
        }

        double correctedSkipNs = (skipMs - timeCorrectionMs) * 1e6;
        double endTimeNs = correctedSkipNs + lenMs * 1e6;
        XCTraceTableProfileHandler handler = new XCTraceTableProfileHandler(resultsTable, sample -> {
            double sampleTimeNs = sample.getTimeFromStartNs();
            if (sampleTimeNs < correctedSkipNs || sampleTimeNs >= endTimeNs) {
                return;
            }
            if (sample.getAddress() == 0L) {
                return;
            }
            events.add(sample.getAddress());

            if (sample.getBinary() != null) {
                String name = sample.getBinary();
                if (name.isEmpty()) {
                    throw new IllegalStateException("Empty binary name");
                }
                MethodDesc method = dedup.dedup(MethodDesc.nativeMethod(sample.getSymbol(), name));

                methods.compute(method, (key, value) -> {
                    if (value == null) {
                        return new AddressInterval(sample.getAddress());
                    }
                    value.add(sample.getAddress());
                    return value;
                });
            }
        });
        handler.parse(perfParsedData.file());

        IntervalMap<MethodDesc> methodMap = new IntervalMap<>();
        methods.forEach((method, addresses) -> methodMap.add(method, addresses.getMin(), addresses.getMax()));

        Map<String, Multiset<Long>> allEvents = new TreeMap<>();
        assert requestedEventNames.size() == 1;
        allEvents.put(requestedEventNames.get(0), events);
        return new PerfEvents(requestedEventNames, allEvents, methodMap);
    }

    @Override
    protected String perfBinaryExtension() {
        return ".trace";
    }

    @Override
    public Collection<String> addJVMInvokeOptions(BenchmarkParams params) {
        return XCTraceSupport.recordCommandPrefix(xctracePath, temporaryDirectory.toAbsolutePath().toString(),
                template);
    }

    @Override
    public String getDescription() {
        return "macOS xctrace (Instruments) + PrintAssembly profiler";
    }

    @Override
    protected void savePerfBinaryOutput(BenchmarkParams runParams, PrintWriter logger) {
        String target = (savePerfBinFile == null) ?
                savePerfBinTo + "/" + runParams.id() + perfBinaryExtension() :
                savePerfBinFile;
        try {
            XCTraceSupport.copyDirectory(XCTraceSupport.findTraceFile(temporaryDirectory), Paths.get(target));
            logger.println("xctrace binary output saved to " + target);
        } catch (IOException e) {
            logger.println("Unable to save xctrace binary output to " + target);
        }
    }

    private static class AddressInterval {
        private long min;
        private long max;

        public AddressInterval(long address) {
            min = address;
            max = address;
        }

        void add(long address) {
            min = Math.min(min, address);
            max = Math.max(max, address);
        }

        public long getMin() {
            return min;
        }

        public long getMax() {
            return max;
        }
    }
}
