/*
 * Copyright (c) 2023, 2023, 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 com.oracle.svm.hosted.methodhandles;

import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

import com.oracle.graal.pointsto.BigBang;
import com.oracle.graal.pointsto.infrastructure.OriginalClassProvider;
import com.oracle.graal.pointsto.infrastructure.SubstitutionProcessor;
import com.oracle.graal.pointsto.meta.AnalysisType;
import com.oracle.svm.core.SubstrateUtil;
import com.oracle.svm.core.util.VMError;
import com.oracle.svm.util.ReflectionUtil;

import jdk.vm.ci.meta.ResolvedJavaType;

/**
 * A substitution processor that renames classes generated by {@code InvokerBytecodeGenerator},
 * which are assigned more or less arbitrary names by the host VM, to stable names that are based on
 * the {@code LambdaForm} which they were compiled from.
 */
public class MethodHandleInvokerRenamingSubstitutionProcessor extends SubstitutionProcessor {
    private static final Method CLASS_GET_CLASS_DATA_METHOD = ReflectionUtil.lookupMethod(Class.class, "getClassData");
    private static final Class<?> LAMBDA_FORM_CLASS = ReflectionUtil.lookupClass(false, "java.lang.invoke.LambdaForm");
    private static final Field LAMBDA_FORM_CUSTOMIZED_FIELD = ReflectionUtil.lookupField(LAMBDA_FORM_CLASS, "customized");
    private static final Class<?> DIRECT_METHOD_HANDLE_CLASS = ReflectionUtil.lookupClass(false, "java.lang.invoke.DirectMethodHandle");
    private static final Method DIRECT_METHOD_HANDLE_INTERNAL_MEMBER_NAME_METHOD = ReflectionUtil.lookupMethod(DIRECT_METHOD_HANDLE_CLASS, "internalMemberName");

    /*
     * We currently only replace the invokers of direct method handles which have a simpler
     * structure and appear to be reliably reused.
     */
    private static final String CLASS_NAME_SUBSTRING = "LambdaForm$DMH";
    private static final String STABLE_NAME_TEMPLATE = "Ljava/lang/invoke/LambdaForm$DMH.s";

    private final BigBang bb;

    private final ConcurrentMap<ResolvedJavaType, MethodHandleInvokerSubstitutionType> typeSubstitutions = new ConcurrentHashMap<>();
    private final Set<String> uniqueTypeNames = new HashSet<>();

    MethodHandleInvokerRenamingSubstitutionProcessor(BigBang bb) {
        this.bb = bb;
    }

    @Override
    public ResolvedJavaType lookup(ResolvedJavaType type) {
        if (!shouldReplace(type)) {
            return type;
        }
        return getSubstitution(type);
    }

    private static boolean shouldReplace(ResolvedJavaType type) {
        return !(type instanceof MethodHandleInvokerSubstitutionType) && type.isFinalFlagSet() && type.getName().contains(CLASS_NAME_SUBSTRING);
    }

    @Override
    public ResolvedJavaType resolve(ResolvedJavaType type) {
        if (type instanceof MethodHandleInvokerSubstitutionType) {
            return ((MethodHandleInvokerSubstitutionType) type).getOriginal();
        }
        return type;
    }

    private ResolvedJavaType getSubstitution(ResolvedJavaType type) {
        return typeSubstitutions.computeIfAbsent(type, original -> {
            Object lambdaForm;
            Object customizedMemberName = null;
            try {
                Class<?> clazz = OriginalClassProvider.getJavaClass(original);
                Object classData = CLASS_GET_CLASS_DATA_METHOD.invoke(clazz);
                if (LAMBDA_FORM_CLASS.isInstance(classData)) {
                    lambdaForm = classData;
                } else if (classData instanceof List<?> list && list.size() == 2) {
                    lambdaForm = list.get(0);
                    Object customizedHandle = list.get(1);
                    VMError.guarantee(LAMBDA_FORM_CLASS.isInstance(lambdaForm) && DIRECT_METHOD_HANDLE_CLASS.isInstance(customizedHandle) &&
                                    LAMBDA_FORM_CUSTOMIZED_FIELD.get(lambdaForm) == customizedHandle, "Expected classData to contain LambdaForm and its customization: %s", classData);
                    customizedMemberName = DIRECT_METHOD_HANDLE_INTERNAL_MEMBER_NAME_METHOD.invoke(customizedHandle);
                } else {
                    throw VMError.shouldNotReachHere("Unexpected classData: %s", classData);
                }
            } catch (ReflectiveOperationException e) {
                throw VMError.shouldNotReachHere(e);
            }
            /*
             * LambdaForm.hashCode() is not stable between image builds because it incorporates
             * identity hash codes of objects such as those of Class<?> that don't override
             * hashCode(). For that reason, we compute a hash code from LambdaForm.toString(). It
             * might also not be perfectly unique because the string contains unqualified class
             * names and can contain string representations of constraints that may be arbitrary
             * objects, but it should typically be distinct and stable.
             */
            int hash = lambdaForm.toString().hashCode();
            if (customizedMemberName != null) {
                /* MemberName.hashCode() also includes identity hash codes of Class<?> objects. */
                hash = hash * 31 + customizedMemberName.toString().hashCode();
            }
            return new MethodHandleInvokerSubstitutionType(original, findUniqueName(hash));
        });
    }

    private String findUniqueName(int hashCode) {
        CharSequence baseName = STABLE_NAME_TEMPLATE + Integer.toHexString(hashCode);
        String name = baseName + ";";
        synchronized (uniqueTypeNames) {
            int suffix = 1;
            while (uniqueTypeNames.contains(name)) {
                name = baseName + "_" + suffix + ";";
                suffix++;
            }
            uniqueTypeNames.add(name);
            return name;
        }
    }

    boolean checkAllTypeNames() {
        if (!SubstrateUtil.assertionsEnabled()) {
            throw new AssertionError("Expensive check: should only run with assertions enabled.");
        }

        List<AnalysisType> types = bb.getUniverse().getTypes();

        if (types.stream().anyMatch(aType -> shouldReplace(aType.getWrapped()))) {
            throw new AssertionError("All relevant types must have been substituted.");
        }

        Set<String> names = new HashSet<>();
        types.stream()
                        .map(AnalysisType::getName)
                        .filter(x -> x.contains(CLASS_NAME_SUBSTRING))
                        .forEach(name -> {
                            if (names.contains(name)) {
                                throw new AssertionError("Duplicate name: " + name);
                            }
                            names.add(name);
                        });
        return true;
    }
}
