/*
 *  Licensed to the Apache Software Foundation (ASF) under one
 *  or more contributor license agreements.  See the NOTICE file
 *  distributed with this work for additional information
 *  regarding copyright ownership.  The ASF licenses this file
 *  to you under the Apache License, Version 2.0 (the
 *  "License"); you may not use this file except in compliance
 *  with the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing,
 *  software distributed under the License is distributed on an
 *  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 *  KIND, either express or implied.  See the License for the
 *  specific language governing permissions and limitations
 *  under the License.
 */
package org.codehaus.groovy.classgen.asm.sc;

import org.codehaus.groovy.ast.ClassHelper;
import org.codehaus.groovy.ast.ClassNode;
import org.codehaus.groovy.ast.MethodNode;
import org.codehaus.groovy.ast.Parameter;
import org.codehaus.groovy.ast.PropertyNode;
import org.codehaus.groovy.ast.expr.AttributeExpression;
import org.codehaus.groovy.ast.expr.BinaryExpression;
import org.codehaus.groovy.ast.expr.ConstructorCallExpression;
import org.codehaus.groovy.ast.expr.Expression;
import org.codehaus.groovy.ast.expr.MethodCallExpression;
import org.codehaus.groovy.ast.expr.PropertyExpression;
import org.codehaus.groovy.ast.expr.VariableExpression;
import org.codehaus.groovy.ast.stmt.EmptyStatement;
import org.codehaus.groovy.ast.stmt.ForStatement;
import org.codehaus.groovy.ast.tools.WideningCategories;
import org.codehaus.groovy.classgen.asm.BinaryExpressionMultiTypeDispatcher;
import org.codehaus.groovy.classgen.asm.CompileStack;
import org.codehaus.groovy.classgen.asm.OperandStack;
import org.codehaus.groovy.classgen.asm.VariableSlotLoader;
import org.codehaus.groovy.classgen.asm.WriterController;
import org.codehaus.groovy.syntax.Token;
import org.codehaus.groovy.syntax.TokenUtil;
import org.objectweb.asm.Label;
import org.objectweb.asm.MethodVisitor;

import java.util.concurrent.atomic.AtomicInteger;

import static org.apache.groovy.ast.tools.ExpressionUtils.isThisExpression;
import static org.codehaus.groovy.ast.ClassHelper.isNumberType;
import static org.codehaus.groovy.ast.ClassHelper.isPrimitiveChar;
import static org.codehaus.groovy.ast.ClassHelper.isPrimitiveDouble;
import static org.codehaus.groovy.ast.ClassHelper.isPrimitiveFloat;
import static org.codehaus.groovy.ast.ClassHelper.isPrimitiveLong;
import static org.codehaus.groovy.ast.ClassHelper.isPrimitiveType;
import static org.codehaus.groovy.ast.tools.GeneralUtils.args;
import static org.codehaus.groovy.ast.tools.GeneralUtils.binX;
import static org.codehaus.groovy.ast.tools.GeneralUtils.callX;
import static org.codehaus.groovy.ast.tools.GeneralUtils.ctorX;
import static org.codehaus.groovy.ast.tools.GeneralUtils.declX;
import static org.codehaus.groovy.ast.tools.GeneralUtils.getSetterName;
import static org.codehaus.groovy.ast.tools.GeneralUtils.isOrImplements;
import static org.codehaus.groovy.ast.tools.GeneralUtils.stmt;
import static org.codehaus.groovy.ast.tools.GeneralUtils.varX;
import static org.codehaus.groovy.classgen.AsmClassGenerator.ELIDE_EXPRESSION_VALUE;
import static org.codehaus.groovy.transform.sc.StaticCompilationVisitor.ARRAYLIST_ADD_METHOD;
import static org.codehaus.groovy.transform.sc.StaticCompilationVisitor.ARRAYLIST_CLASSNODE;
import static org.codehaus.groovy.transform.sc.StaticCompilationVisitor.ARRAYLIST_CONSTRUCTOR;
import static org.codehaus.groovy.transform.stc.StaticTypeCheckingSupport.isAssignment;
import static org.codehaus.groovy.transform.stc.StaticTypeCheckingVisitor.inferLoopElementType;
import static org.codehaus.groovy.transform.stc.StaticTypesMarker.COMPOUND_ASSIGN_TARGET;
import static org.codehaus.groovy.transform.stc.StaticTypesMarker.DIRECT_METHOD_CALL_TARGET;
import static org.codehaus.groovy.transform.stc.StaticTypesMarker.INFERRED_TYPE;
import static org.objectweb.asm.Opcodes.ACC_PUBLIC;
import static org.objectweb.asm.Opcodes.DADD;
import static org.objectweb.asm.Opcodes.DCONST_1;
import static org.objectweb.asm.Opcodes.DSUB;
import static org.objectweb.asm.Opcodes.FADD;
import static org.objectweb.asm.Opcodes.FCONST_1;
import static org.objectweb.asm.Opcodes.FSUB;
import static org.objectweb.asm.Opcodes.IADD;
import static org.objectweb.asm.Opcodes.ICONST_1;
import static org.objectweb.asm.Opcodes.IFNULL;
import static org.objectweb.asm.Opcodes.ISUB;
import static org.objectweb.asm.Opcodes.LADD;
import static org.objectweb.asm.Opcodes.LCONST_1;
import static org.objectweb.asm.Opcodes.LSUB;

/**
 * A specialized version of the multi type binary expression dispatcher which is aware of static compilation.
 * It is able to generate optimized bytecode for some operations using JVM instructions when available.
 */
public class StaticTypesBinaryExpressionMultiTypeDispatcher extends BinaryExpressionMultiTypeDispatcher {

    private final AtomicInteger labelCounter = new AtomicInteger();

    /**
     * Creates a binary-expression dispatcher for statically compiled bytecode generation.
     */
    public StaticTypesBinaryExpressionMultiTypeDispatcher(final WriterController wc) {
        super(wc);
    }

    /** {@inheritDoc} */
    @Override
    protected void writePostOrPrefixMethod(final int op, final String method, final Expression expression, final Expression orig) {
        MethodNode mn = orig.getNodeMetaData(DIRECT_METHOD_CALL_TARGET);
        if (mn != null) {
            controller.getOperandStack().pop();
            MethodCallExpression call = callX(expression, method);
            call.setMethodTarget(mn);
            call.visit(controller.getAcg());
            return;
        }

        ClassNode top = controller.getOperandStack().getTopOperand();
        if (isPrimitiveType(top) && (isNumberType(top) || isPrimitiveChar(top))) {
            MethodVisitor mv = controller.getMethodVisitor();
            visitInsnByType(top, mv, ICONST_1, LCONST_1, FCONST_1, DCONST_1);
            if ("next".equals(method)) {
                visitInsnByType(top, mv, IADD, LADD, FADD, DADD);
            } else {
                visitInsnByType(top, mv, ISUB, LSUB, FSUB, DSUB);
            }
            return;
        }

        super.writePostOrPrefixMethod(op, method, expression, orig);
    }

    private static void visitInsnByType(final ClassNode top, final MethodVisitor mv, final int iInsn, final int lInsn, final int fInsn, final int dInsn) {
        if (WideningCategories.isIntCategory(top) || isPrimitiveChar(top)) {
            mv.visitInsn(iInsn);
        } else if (isPrimitiveLong(top)) {
            mv.visitInsn(lInsn);
        } else if (isPrimitiveFloat(top)) {
            mv.visitInsn(fInsn);
        } else if (isPrimitiveDouble(top)) {
            mv.visitInsn(dInsn);
        }
    }

    /** {@inheritDoc} */
    @Override
    protected void evaluateBinaryExpressionWithAssignment(final String method, final BinaryExpression expression) {
        if (tryStaticCompoundAssignPaths(method, expression)) return;
        super.evaluateBinaryExpressionWithAssignment(method, expression);
    }

    /** {@inheritDoc} */
    @Override
    protected void evaluateCompoundAssign(final String assignName, final String baseName, final BinaryExpression expression) {
        if (tryStaticCompoundAssignPaths(baseName, expression)) return;
        super.evaluateCompoundAssign(assignName, baseName, expression);
    }

    /**
     * GEP-15 + legacy setter fast-path. Returns true when codegen has been emitted
     * (no further dispatch required).
     */
    private boolean tryStaticCompoundAssignPaths(final String baseName, final BinaryExpression expression) {
        MethodNode assignTarget = expression.getNodeMetaData(COMPOUND_ASSIGN_TARGET);
        if (assignTarget != null) {
            // GEP-15: receiver.<assignMethod>(arg); receiver remains the expression value.
            // The setter (for property LHS) is intentionally skipped.
            emitCompoundAssignCall(assignTarget, expression);
            return true;
        }
        Expression leftExpression = expression.getLeftExpression();
        if (leftExpression instanceof PropertyExpression pexp
                && !(leftExpression instanceof AttributeExpression)) {

            BinaryExpression expressionWithoutAssignment = binX(
                    leftExpression,
                    Token.newSymbol(
                            TokenUtil.removeAssignment(expression.getOperation().getType()),
                            expression.getOperation().getStartLine(),
                            expression.getOperation().getStartColumn()
                    ),
                    expression.getRightExpression()
            );
            expressionWithoutAssignment.copyNodeMetaData(expression);
            expressionWithoutAssignment.setSafe(expression.isSafe());
            expressionWithoutAssignment.setSourcePosition(expression);

            if (makeSetProperty(
                    pexp.getObjectExpression(),
                    pexp.getProperty(),
                    expressionWithoutAssignment,
                    pexp.isSafe(),
                    pexp.isSpreadSafe(),
                    pexp.isImplicitThis(),
                    true)) { // TODO: GROOVY-11843
                return true;
            }
        }
        return false;
    }

    private void emitCompoundAssignCall(final MethodNode target, final BinaryExpression expression) {
        OperandStack operandStack = controller.getOperandStack();
        Expression leftExpression = expression.getLeftExpression();
        Expression rightExpression = expression.getRightExpression();

        leftExpression.visit(controller.getAcg());
        ClassNode receiverType = operandStack.getTopOperand();
        int slot = controller.getCompileStack().defineTemporaryVariable("$gep15recv", receiverType, true);

        VariableSlotLoader callReceiver = new VariableSlotLoader(receiverType, slot, operandStack);
        MethodCallExpression call = callX(callReceiver, target.getName(), rightExpression);
        call.setMethodTarget(target);
        call.setImplicitThis(false);
        call.setSourcePosition(expression);
        call.visit(controller.getAcg());
        operandStack.pop(); // discard the *Assign return value

        new VariableSlotLoader(receiverType, slot, operandStack).visit(controller.getAcg());
        controller.getCompileStack().removeVar(slot);
    }

    /** {@inheritDoc} */
    @Override
    public void evaluateEqual(final BinaryExpression expression, final boolean defineVariable) {
        Expression leftExpression = expression.getLeftExpression();
        if (leftExpression instanceof PropertyExpression pexp
                && !(leftExpression instanceof AttributeExpression)) {
            if (!defineVariable && makeSetProperty(
                    pexp.getObjectExpression(),
                    pexp.getProperty(),
                    expression.getRightExpression(),
                    pexp.isSafe(),
                    pexp.isSpreadSafe(),
                    pexp.isImplicitThis(),
                    !Boolean.TRUE.equals(expression.getNodeMetaData(ELIDE_EXPRESSION_VALUE)))) { // GROOVY-11843
                return;
            }
            // GROOVY-5620: spread-safe operator on LHS is not supported
            if (pexp.isSpreadSafe() && isAssignment(expression.getOperation().getType())) {
                // rewrite it so that it can be statically compiled
                transformSpreadOnLHS(expression);
                return;
            }
        }
        super.evaluateEqual(expression, defineVariable);
    }

    private void transformSpreadOnLHS(final BinaryExpression expression) {
        PropertyExpression spreadExpression = (PropertyExpression) expression.getLeftExpression();
        Expression receiver = spreadExpression.getObjectExpression();

        int counter = labelCounter.incrementAndGet();
        CompileStack compileStack = controller.getCompileStack();
        OperandStack operandStack = controller.getOperandStack();

        // create an empty arraylist
        VariableExpression result = varX(this.getClass().getSimpleName() + "$spreadresult" + counter, ARRAYLIST_CLASSNODE);
        ConstructorCallExpression newArrayList = ctorX(ARRAYLIST_CLASSNODE);
        newArrayList.setNodeMetaData(DIRECT_METHOD_CALL_TARGET, ARRAYLIST_CONSTRUCTOR);
        Expression decl = declX(result, newArrayList);
        decl.visit(controller.getAcg());
        // if (receiver != null)
        receiver.visit(controller.getAcg());
        Label ifnull = compileStack.createLocalLabel("ifnull_" + counter);
        MethodVisitor mv = controller.getMethodVisitor();
        mv.visitJumpInsn(IFNULL, ifnull);
        operandStack.remove(1); // receiver consumed by if()
        Label nonull = compileStack.createLocalLabel("nonull_" + counter);
        mv.visitLabel(nonull);
        ClassNode componentType = inferLoopElementType(controller.getTypeChooser().resolveType(receiver, controller.getClassNode()));
        Parameter iterator = new Parameter(componentType, "for$it$" + counter);
        VariableExpression iteratorAsVar = varX(iterator);
        PropertyExpression pexp = spreadExpression instanceof AttributeExpression
            ? new AttributeExpression(iteratorAsVar, spreadExpression.getProperty(), true)
            : new PropertyExpression(iteratorAsVar, spreadExpression.getProperty(), true);
        pexp.setImplicitThis(spreadExpression.isImplicitThis());
        pexp.setSourcePosition(spreadExpression);
        BinaryExpression assignment = binX(pexp, expression.getOperation(), expression.getRightExpression());
        MethodCallExpression add = callX(result, "add", assignment);
        add.setMethodTarget(ARRAYLIST_ADD_METHOD);
        // for (e in receiver) { result.add(e?.method(arguments) }
        ForStatement stmt = new ForStatement(
                iterator,
                receiver,
                stmt(add)
        );
        stmt.visit(controller.getAcg());
        // else { empty list }
        mv.visitLabel(ifnull);
        // end of if/else
        // return result list
        result.visit(controller.getAcg());
    }

    private boolean makeSetProperty(final Expression receiver, final Expression message, final Expression arguments, final boolean safe, final boolean spreadSafe, final boolean implicitThis, final boolean returnValue) {
        var receiverType = controller.getTypeChooser().resolveType(receiver, controller.getClassNode());
        var thisReceiver = isThisExpression(receiver);
        var propertyName = message.getText();

        String setterName = getSetterName(propertyName);
        MethodNode setterMethod = receiverType.getSetterMethod(setterName, false);
        if (setterMethod != null) {
            if ((thisReceiver && setterMethod.getDeclaringClass().equals(controller.getClassNode()))
                || (!setterMethod.isPublic() && isOrImplements(receiverType, ClassHelper.MAP_TYPE))) {
                // this.x = ... should not use same-class setter
                // that.x = ... should not use non-public setter for map
                setterMethod = null;
            } else { // GROOVY-11119
                java.util.List<MethodNode> setters = receiverType.getMethods(setterName);
                setters.removeIf(s -> s.isAbstract() || s.getParameters().length != 1);
                if (setters.size() > 1) setterMethod = null;
            }
        } else {
            PropertyNode propertyNode = receiverType.getProperty(propertyName);
            if (propertyNode != null && !propertyNode.isFinal()) {
                setterMethod = new MethodNode(
                        setterName,
                        ACC_PUBLIC,
                        ClassHelper.VOID_TYPE,
                        new Parameter[]{new Parameter(propertyNode.getOriginType(), "value")},
                        ClassNode.EMPTY_ARRAY,
                        EmptyStatement.INSTANCE
                );
                setterMethod.setDeclaringClass(receiverType);
                setterMethod.setSynthetic(true);
            }
        }
        if (setterMethod != null) {
            Expression call = StaticPropertyAccessHelper.transformToSetterCall(
                    receiver,
                    setterMethod,
                    arguments,
                    implicitThis,
                    safe,
                    spreadSafe,
                    returnValue,
                    message
            );
            call.visit(controller.getAcg());
            return true;
        }

        return false;
    }

    /** {@inheritDoc} */
    @Override
    protected void assignToArray(final Expression enclosing, final Expression receiver, final Expression subscript, final Expression rhsValueLoader, final boolean safe) {
        ClassNode receiverType = controller.getTypeChooser().resolveType(receiver, controller.getClassNode());

        if (!safe && receiverType.isArray() && binExpWriter[getOperandType(receiverType.getComponentType())].arraySet(true)) {
            super.assignToArray(enclosing, receiver, subscript, rhsValueLoader, false);
        } else { // handle safe subscript and other cases by calling the "putAt" method
            if (rhsValueLoader instanceof VariableSlotLoader && enclosing instanceof BinaryExpression) { // GROOVY-6061
                rhsValueLoader.putNodeMetaData(INFERRED_TYPE, controller.getTypeChooser().resolveType(enclosing, controller.getClassNode()));
            }

            // replace assignment to a subscript operator with a method call
            // e.g. x[5] = 10 --> x.putAt(5, 10)
            MethodCallExpression call = callX(receiver, "putAt", args(subscript, rhsValueLoader));
            call.setSafe(safe);
            call.setSourcePosition(enclosing);

            call.visit(controller.getAcg());
            controller.getOperandStack().pop(); // method return value

            if (!Boolean.TRUE.equals(enclosing.getNodeMetaData(ELIDE_EXPRESSION_VALUE)))
                rhsValueLoader.visit(controller.getAcg()); // assignment expression value
        }
    }
}
