/*

 * Copyright (c) 2017, 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.

 *

 * 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 p;

import java.io.File;

import java.io.IOException;

import java.io.OutputStream;

import java.nio.file.Files;

import java.nio.file.Path;

import java.nio.file.Paths;

import java.util.ArrayList;

import java.util.List;

import java.util.jar.Attributes;

import java.util.jar.JarEntry;

import java.util.jar.JarFile;

import java.util.jar.JarOutputStream;

import java.util.jar.Manifest;

import com.sun.tools.attach.VirtualMachine;

public class Main {

    public static void main(String[] args) throws Exception {

        System.out.println("#modules loaded: " + moduleInfoCont());

        String vmid = "" + ProcessHandle.current().pid();

        VirtualMachine vm = VirtualMachine.attach(vmid);

        for (String test : args) {

            switch (test) {

                case "jmx" :

                    startJMXAgent(vm);

                    break;

                case "javaagent" :

                    startJavaAgent(vm, createAgentJar());

                    break;

            }

            System.out.println("#modules loaded: " + moduleInfoCont());

        }

    }

    /**

     * Locates module-info.class resources to get a count of the module of system

     * modules.

     */

    static long moduleInfoCont() {

        ClassLoader scl = ClassLoader.getSystemClassLoader();

        return scl.resources("module-info.class").count();

    }

    /**

     * Starts a JMX agent and checks that java.management is loaded.

     */

    static void startJMXAgent(VirtualMachine vm) throws Exception {

        System.out.println("Start JMX agent");

        vm.startLocalManagementAgent();

        // types in java.management should be visible

        Class.forName("javax.management.MXBean");

    }

    /**

     * Loads a java agent into the VM and checks that java.instrument is loaded.

     */

    static void startJavaAgent(VirtualMachine vm, Path agent) throws Exception {

        System.out.println("Load java agent ...");

        vm.loadAgent(agent.toString());

        // the Agent class should be visible

        Class.forName("Agent");

        // types in java.instrument should be visible

        Class.forName("java.lang.instrument.Instrumentation");

    }

    /**

     * Creates a java agent, return the file path to the agent JAR file.

     */

    static Path createAgentJar() throws IOException {

        Manifest man = new Manifest();

        Attributes attrs = man.getMainAttributes();

        attrs.put(Attributes.Name.MANIFEST_VERSION, "1.0.0");

        attrs.put(new Attributes.Name("Agent-Class"), "Agent");

        Path agent = Paths.get("agent.jar");

        Path dir = Paths.get(System.getProperty("test.classes"));

        createJarFile(agent, man, dir, "Agent.class");

        return agent;

    }

    /**

     * Creates a JAR file.

     *

     * Equivalent to {@code jar cfm <jarfile> <manifest> -C <dir> file...}

     *

     * The input files are resolved against the given directory. Any input

     * files that are directories are processed recursively.

     */

    static void createJarFile(Path jarfile, Manifest man, Path dir, String... files)

        throws IOException

    {

        // create the target directory

        Path parent = jarfile.getParent();

        if (parent != null)

            Files.createDirectories(parent);

        List<Path> entries = new ArrayList<>();

        for (String file : files) {

            Files.find(dir.resolve(file), Integer.MAX_VALUE,

                    (p, attrs) -> attrs.isRegularFile())

                    .map(e -> dir.relativize(e))

                    .forEach(entries::add);

        }

        try (OutputStream out = Files.newOutputStream(jarfile);

             JarOutputStream jos = new JarOutputStream(out))

        {

            if (man != null) {

                JarEntry je = new JarEntry(JarFile.MANIFEST_NAME);

                jos.putNextEntry(je);

                man.write(jos);

                jos.closeEntry();

            }

            for (Path entry : entries) {

                String name = toJarEntryName(entry);

                jos.putNextEntry(new JarEntry(name));

                Files.copy(dir.resolve(entry), jos);

                jos.closeEntry();

            }

        }

    }

    /**

     * Map a file path to the equivalent name in a JAR file

     */

    static String toJarEntryName(Path file) {

        Path normalized = file.normalize();

        return normalized.subpath(0, normalized.getNameCount())

                .toString()

                .replace(File.separatorChar, '/');

    }

}