/*
 * Copyright (c) 2016, 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 jdk.tools.jlink.internal.plugins;

import java.io.ByteArrayInputStream;
import java.lang.reflect.Method;
import java.util.Arrays;
import java.util.Collections;
import java.util.EnumSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import jdk.internal.org.objectweb.asm.ClassWriter;
import jdk.internal.org.objectweb.asm.Label;
import jdk.internal.org.objectweb.asm.MethodVisitor;
import jdk.internal.org.objectweb.asm.Opcodes;
import jdk.internal.org.objectweb.asm.Type;
import static jdk.internal.org.objectweb.asm.Opcodes.*;
import jdk.tools.jlink.plugin.PluginException;
import jdk.tools.jlink.plugin.Pool;
import jdk.tools.jlink.plugin.TransformerPlugin;

/**
 * Plugin to generate BoundMethodHandle classes.
 */
public final class GenerateBMHClassesPlugin implements TransformerPlugin {

    private static final String NAME = "generate-bmh";

    private static final String DESCRIPTION = PluginsResourceBundle.getDescription(NAME);

    private static final String BMH = "java/lang/invoke/BoundMethodHandle";

    private static final Method FACTORY_METHOD;

    List<String> speciesTypes;

    public GenerateBMHClassesPlugin() {
    }

    @Override
    public Set<PluginType> getType() {
        return Collections.singleton(CATEGORY.TRANSFORMER);
    }

    @Override
    public String getName() {
        return NAME;
    }

    @Override
    public String getDescription() {
        return DESCRIPTION;
    }

    @Override
    public Set<STATE> getState() {
        return EnumSet.of(STATE.AUTO_ENABLED, STATE.FUNCTIONAL);
    }

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

    @Override
    public String getArgumentsDescription() {
       return PluginsResourceBundle.getArgument(NAME);
    }

    @Override
    public void configure(Map<String, String> config) {
        String args = config.get(NAME);
        if (args != null && !args.isEmpty()) {
            speciesTypes = Arrays.stream(args.split(","))
                .map(String::trim)
                .filter(s -> !s.isEmpty())
                .collect(Collectors.toList());
        } else {
            // Default set of Species forms to generate
            //speciesTypes = Collections.emptyList();
            speciesTypes = List.of("LL", "L3", "L4", "L5", "L6", "L7", "L7I", "L7II",
                    "L8", "L9", "L10", "L11", "L11I", "L11II", "L11IIL",
                    "LI", "D", "L3I", "LIL", "LLI", "LLIL", "LILL", "I", "LLILL"
            );
        }
    }

    @Override
    public void visit(Pool in, Pool out) {
        for (Pool.ModuleData data : in.getContent()) {
            if (("/java.base/" + BMH + "$Factory$SpeciesLookup.class").equals(data.getPath())) {
                speciesTypes.forEach(types -> generateConcreteClass(types, data, out));
                generateSpeciesLookupClass(data, out);
            } else {
                out.add(data);
            }
        }
    }

    private void generateSpeciesLookupClass(Pool.ModuleData data, Pool out) {
        System.out.println("Generating SpeciesLookup");
        final String CLASSNAME = "java/lang/invoke/BoundMethodHandle$Factory$SpeciesLookup";
        final String GENERATED_SPECIES = "SPECIES";
        final String STRING_BUILDER = "java/lang/StringBuilder";

        ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_MAXS + ClassWriter.COMPUTE_FRAMES);
        cw.visit(Opcodes.V1_8, ACC_STATIC + ACC_SUPER, CLASSNAME,
                "Ljava/lang/Object;Ljava/util/function/Function<Ljava/lang/String;"
                        + "Ljava/lang/Class<+Ljava/lang/invoke/BoundMethodHandle;>;>;",
                "java/lang/Object",
                new String[] { "java/util/function/Function" });

        cw.visitInnerClass("java/lang/invoke/BoundMethodHandle$Factory",
                "java/lang/invoke/BoundMethodHandle",
                "Factory",
                ACC_STATIC);

        cw.visitInnerClass("java/lang/invoke/BoundMethodHandle$Factory$SpeciesLookup",
                "java/lang/invoke/BoundMethodHandle$Factory",
                "SpeciesLookup",
                ACC_STATIC);

        // static HashSet<String> GENERATED_SPECIES
        cw.visitField(ACC_STATIC, GENERATED_SPECIES,
                    "Ljava/util/HashSet;", "Ljava/util/HashSet<Ljava/lang/String;>;", null)
                    .visitEnd();

        // static {
        MethodVisitor mv = cw.visitMethod(ACC_STATIC, "<clinit>", "()V",
                    null, null);
        mv.visitCode();
        mv.visitTypeInsn(NEW, "java/util/HashSet");
        mv.visitInsn(DUP);
        mv.visitMethodInsn(INVOKESPECIAL, "java/util/HashSet",
                "<init>", "()V", false);
        mv.visitVarInsn(ASTORE, 0);

        for (String types : speciesTypes) {
            mv.visitVarInsn(ALOAD, 0);
            mv.visitLdcInsn(expandSignature(types));
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/util/HashSet",
                    "add", "(Ljava/lang/Object;)Z", false);
            mv.visitInsn(POP);
        }

        mv.visitVarInsn(ALOAD, 0);
        mv.visitFieldInsn(PUTSTATIC, CLASSNAME, GENERATED_SPECIES, "Ljava/util/HashSet;");
        mv.visitInsn(RETURN);
        mv.visitMaxs(0, 0);
        mv.visitEnd();

        // Empty constructor
        mv = cw.visitMethod(0, "<init>", "()V", null, null);
        mv.visitCode();
        mv.visitVarInsn(ALOAD, 0);
        mv.visitMethodInsn(INVOKESPECIAL, "java/lang/Object", "<init>", "()V", false);
        mv.visitInsn(RETURN);
        mv.visitMaxs(0, 0);
        mv.visitEnd();

        /* apply method, equivalent of:
         *
         *   @Override
         *   @SuppressWarnings("unchecked")
         *   public Class<? extends BoundMethodHandle> apply(String types) {
         *       if (GENERATED_SPECIES.contains(types)) {
         *           String cl = "java.lang.invoke.BoundMethodHandle$Species_"
         *                   + LambdaForm.shortenSignature(types);
         *           try {
         *               return (Class<? extends BoundMethodHandle>)
         *                       Class.forName(cl);
         *           } catch (ClassNotFoundException cnf) {
         *               throw new InternalError(cnf);
         *           }
         *       }
         *       return generateConcreteBMHClass(types);
         *   }
         */
        mv = cw.visitMethod(ACC_PUBLIC, "apply",
                "(Ljava/lang/String;)Ljava/lang/Class;",
                "(Ljava/lang/String;)Ljava/lang/Class<+Ljava/lang/invoke/BoundMethodHandle;>;",
                null);
        mv.visitCode();
        Label start = new Label();
        Label end = new Label();
        Label handler = new Label();
        mv.visitTryCatchBlock(start, end, handler, "java/lang/ClassNotFoundException");
        mv.visitFieldInsn(GETSTATIC, CLASSNAME, GENERATED_SPECIES, "Ljava/util/HashSet;");
        mv.visitVarInsn(ALOAD, 1);
        mv.visitMethodInsn(INVOKEVIRTUAL, "java/util/HashSet", "contains",
                "(Ljava/lang/Object;)Z", false);
        Label generateClass = new Label();
        mv.visitJumpInsn(IFEQ, generateClass);
        mv.visitTypeInsn(NEW, STRING_BUILDER);
        mv.visitInsn(DUP);
        mv.visitMethodInsn(INVOKESPECIAL, STRING_BUILDER, "<init>", "()V", false);
        mv.visitLdcInsn("java.lang.invoke.BoundMethodHandle$Species_");
        mv.visitMethodInsn(INVOKEVIRTUAL, STRING_BUILDER, "append",
                "(Ljava/lang/String;)Ljava/lang/StringBuilder;", false);
        mv.visitVarInsn(ALOAD, 1);

        mv.visitMethodInsn(INVOKESTATIC, "java/lang/invoke/LambdaForm", "shortenSignature",
                "(Ljava/lang/String;)Ljava/lang/String;", false);
        mv.visitMethodInsn(INVOKEVIRTUAL, STRING_BUILDER, "append",
                "(Ljava/lang/String;)Ljava/lang/StringBuilder;", false);
        mv.visitMethodInsn(INVOKEVIRTUAL, STRING_BUILDER, "toString",
                "()Ljava/lang/String;", false);
        mv.visitVarInsn(ASTORE, 2);
        mv.visitLabel(start);
        mv.visitVarInsn(ALOAD, 2);
        mv.visitMethodInsn(INVOKESTATIC, "java/lang/Class", "forName",
                "(Ljava/lang/String;)Ljava/lang/Class;", false);
        mv.visitLabel(end);
        mv.visitInsn(ARETURN);
        mv.visitLabel(handler);
        mv.visitFrame(Opcodes.F_FULL, 3,
                new Object[] {"java/lang/invoke/BoundMethodHandle$Factory$SpeciesLookup",
                    "java/lang/String", "java/lang/String"}, 1,
                    new Object[] {"java/lang/ClassNotFoundException"});
        mv.visitVarInsn(ASTORE, 3);
        mv.visitTypeInsn(NEW, "java/lang/InternalError");
        mv.visitInsn(DUP);
        mv.visitVarInsn(ALOAD, 3);
        mv.visitMethodInsn(INVOKESPECIAL, "java/lang/InternalError", "<init>",
                "(Ljava/lang/Throwable;)V", false);
        mv.visitInsn(ATHROW);
        mv.visitLabel(generateClass);
        mv.visitFrame(Opcodes.F_CHOP, 1, null, 0, null);

        mv.visitVarInsn(ALOAD, 1);
        mv.visitMethodInsn(INVOKESTATIC, "java/lang/invoke/BoundMethodHandle$Factory",
                "generateConcreteBMHClass", "(Ljava/lang/String;)Ljava/lang/Class;", false);

        mv.visitInsn(ARETURN);
        mv.visitMaxs(0, 0);
        mv.visitEnd();

        // Synthetic bridge method to make apply non-abstract
        mv = cw.visitMethod(ACC_PUBLIC + ACC_BRIDGE + ACC_SYNTHETIC,
                "apply", "(Ljava/lang/Object;)Ljava/lang/Object;", null, null);
        mv.visitCode();
        mv.visitVarInsn(ALOAD, 0);
        mv.visitVarInsn(ALOAD, 1);
        mv.visitTypeInsn(CHECKCAST, "java/lang/String");
        mv.visitMethodInsn(INVOKEVIRTUAL,
                "java/lang/invoke/BoundMethodHandle$Factory$SpeciesLookup",
                "apply", "(Ljava/lang/String;)Ljava/lang/Class;", false);
        mv.visitInsn(ARETURN);
        mv.visitMaxs(0, 0);
        mv.visitEnd();
        cw.visitEnd();

        System.out.println("Generated SpeciesLookup");
        byte[] bytes = cw.toByteArray();
        Pool.ModuleData ndata = new Pool.ModuleData(data.getModule(),
                    "/java.base/" + BMH + "$Factory$SpeciesLookup.class",
                    Pool.ModuleDataType.CLASS_OR_RESOURCE,
                    new ByteArrayInputStream(bytes), bytes.length);
        out.add(ndata);
    }

    private void generateConcreteClass(String shortTypes, Pool.ModuleData data, Pool out) {
        try {
            String types = expandSignature(shortTypes);

            // Generate class
            byte[] bytes = (byte[])FACTORY_METHOD.invoke(null,
                    BMH + "$Species_" + shortTypes, // class name
                    "Species_" + shortTypes, // source name
                    types);

            // Add class to pool
            Pool.ModuleData ndata = new Pool.ModuleData(data.getModule(),
                    "/java.base/" + BMH + "$Species_" + shortTypes + ".class",
                    Pool.ModuleDataType.CLASS_OR_RESOURCE,
                    new ByteArrayInputStream(bytes), bytes.length);
            out.add(ndata);
        } catch (Exception ex) {
            throw new PluginException(ex);
        }
    }

    static {
        try {
            Class<?> BMHFactory = Class.forName("java.lang.invoke.BoundMethodHandle$Factory");
            Method genClassMethod = BMHFactory.getDeclaredMethod("generateConcreteBMHClassBytes",
                    String.class, String.class, String.class);
            genClassMethod.setAccessible(true);
            FACTORY_METHOD = genClassMethod;
        } catch (Exception e) {
            throw new PluginException(e);
        }
    }

    // Convert LL -> LL, L3 -> LLL
    private static String expandSignature(String signature) {
        StringBuilder sb = new StringBuilder();
        char last = 'X';
        int count = 0;
        for (int i = 0; i < signature.length(); i++) {
            char c = signature.charAt(i);
            if (c >= '0' && c <= '9') {
                count *= 10;
                count += (c - '0');
            } else {
                for (int j = 1; j < count; j++) {
                    sb.append(last);
                }
                sb.append(c);
                last = c;
                count = 0;
            }
        }
        for (int j = 1; j < count; j++) {
            sb.append(last);
        }
        return sb.toString();
    }

}