package com.strobel.decompiler.languages.java.ast.transforms;

import com.strobel.assembler.Collection;
import com.strobel.assembler.metadata.*;
import com.strobel.core.VerifyArgument;
import com.strobel.decompiler.DecompilerContext;
import com.strobel.decompiler.ast.AstCode;
import com.strobel.decompiler.ast.Variable;
import com.strobel.decompiler.languages.TextLocation;
import com.strobel.decompiler.languages.java.ast.*;

import java.util.Collections;
import java.util.List;

public class InvokeDynamicRewriter extends AbstractHelperClassTransform {
    private final static String T_DESC_THROWABLE = "java/lang/Throwable";
    private final static String T_DESC_THROWABLE_WRAPPER = "java/lang/reflect/UndeclaredThrowableException";
    private final static String T_DESC_METHOD_HANDLE = "java/lang/invoke/MethodHandle";
    private final static String T_DESC_METHOD_HANDLES = "java/lang/invoke/MethodHandles";
    private final static String T_DESC_LOOKUP = "java/lang/invoke/MethodHandles$Lookup";
    private final static String T_DESC_CALL_SITE = "java/lang/invoke/CallSite";
    private final static String T_SIGNATURE_LOOKUP = "()L" + T_DESC_LOOKUP + ';';
    private final static String M_DESC_INVOKE_EXACT = "([Ljava/lang/Object;)Ljava/lang/Object;";
    private final static String M_DESC_DYNAMIC_INVOKER = "()L" + T_DESC_METHOD_HANDLE + ";";

    public InvokeDynamicRewriter(final DecompilerContext context) {
        super(context);
    }

    @Override
    public Void visitInvocationExpression(final InvocationExpression node, final Void data) {
        super.visitInvocationExpression(node, data);

        final Expression target = node.getTarget();

        if (target instanceof InlinedBytecodeExpression) {
            final InlinedBytecodeExpression bytecode = (InlinedBytecodeExpression) target;
            final TypeDeclaration currentType = this.currentType;
            final TypeDefinition currentTypeDefinition = currentType != null ? currentType.getUserData(Keys.TYPE_DEFINITION) : null;
            final DynamicCallSite callSite = node.getUserData(Keys.DYNAMIC_CALL_SITE);

            if (currentTypeDefinition == null || callSite == null || bytecode.getCode() != AstCode.InvokeDynamic || resolver == null) {
                return null;
            }

            final IndyHelperBuilder hb = new IndyHelperBuilder(currentType, currentTypeDefinition, callSite);

            if (hb.build()) {
                final AstNodeCollection<Expression> originalArguments = node.getArguments();
                final InvocationExpression replacement = makeType(hb.definition).invoke(hb.definition.mdInvoke);
                final AstNodeCollection<Expression> newArguments = replacement.getArguments();

//                originalArguments.firstOrNullObject().remove();

                for (final Expression argument : originalArguments) {
                    argument.remove();
                    newArguments.add(argument);
                }

                node.replaceWith(replacement);

                replacement.getParent().insertChildBefore(replacement, new Comment(" invokedynamic(!) ", CommentType.MultiLine), Roles.COMMENT);
                currentType.getMembers().insertAfter(currentType.getMembers().lastOrNullObject(), hb.declaration);

                final AstNode commentAnchor = hb.declaration.getFirstChild();

                final List<Object> bootstrapArguments = callSite.getBootstrapArguments();

                if (bootstrapArguments.size() >= 3 && bootstrapArguments.get(1) instanceof MethodHandle) {
                    final MethodReference method = ((MethodHandle) bootstrapArguments.get(1)).getMethod();
                    final MethodDefinition methodDefinition = method.resolve();

                    if (methodDefinition != null) {
                        context.getForcedVisibleMembers().add(methodDefinition);
                    }
                }

                hb.declaration.insertChildBefore(commentAnchor,
                                                 new Comment(" This helper class was generated by Procyon to approximate the behavior of an"),
                                                 Roles.COMMENT);

                hb.declaration.insertChildBefore(commentAnchor,
                                                 new Comment(" 'invokedynamic' instruction that it doesn't know how to interpret."),
                                                 Roles.COMMENT);
            }
        }

        return null;
    }

    @SuppressWarnings("DuplicatedCode")
    protected final class IndyHelperBuilder {
        final static String T_DESC_METHOD_HANDLE = "java/lang/invoke/MethodHandle";
        final static String F_DESC_ENSURE_HANDLE = "L" + T_DESC_METHOD_HANDLE + ";";
        final static String M_DESC_ENSURE_HANDLE = "()" + F_DESC_ENSURE_HANDLE;

        final TypeDeclaration parentDeclaration;
        final TypeReference parentType;
        final DynamicCallSite callSite;
        final TypeReference callSiteType;
        final TypeReference methodHandleType;
        final TypeReference methodHandlesType;
        final TypeReference lookupType;
        final MethodReference handleMethod;
        final MethodReference ensureHandleMethod;
        final HelperTypeDefinition definition;
        final Variable lookupVariable;
        final int uniqueTypeId;

        TypeDeclaration declaration;
        MethodDeclaration handleDeclaration;
        MethodDeclaration invokeDeclaration;
        MethodDeclaration ensureHandleDeclaration;
        InvocationExpression bootstrapCall;

        IndyHelperBuilder(
            final TypeDeclaration parentDeclaration,
            final TypeReference parentType,
            final DynamicCallSite callSite) {

            this.parentDeclaration = VerifyArgument.notNull(parentDeclaration, "parentDeclaration");
            this.parentType = VerifyArgument.notNull(parentType, "parentType");
            this.callSite = VerifyArgument.notNull(callSite, "callSite");
            this.uniqueTypeId = nextUniqueId();

            final TypeReference helperType = parser.parseTypeDescriptor("ProcyonInvokeDynamicHelper_" + uniqueTypeId);

            this.callSiteType = parser.parseTypeDescriptor(T_DESC_CALL_SITE);
            this.methodHandleType = parser.parseTypeDescriptor(T_DESC_METHOD_HANDLE);
            this.methodHandlesType = parser.parseTypeDescriptor(T_DESC_METHOD_HANDLES);
            this.definition = new HelperTypeDefinition(helperType, parentType);
            this.handleMethod = parser.parseMethod(definition, "handle", M_DESC_ENSURE_HANDLE);
            this.ensureHandleMethod = parser.parseMethod(definition, "ensureHandle", M_DESC_ENSURE_HANDLE);
            this.lookupType = parser.parseTypeDescriptor(T_DESC_LOOKUP);

            final Variable vLookup = new Variable();

            vLookup.setGenerated(true);
            vLookup.setName("lookup");
            vLookup.setType(lookupType);

            this.lookupVariable = vLookup;
        }

        boolean build() {
            this.bootstrapCall = makeBootstrapCall(callSite, lookupVariable);

            final TypeDeclaration declaration = new TypeDeclaration();

            this.declaration = declaration;

            final AstNodeCollection<JavaModifierToken> modifiers = this.declaration.getModifiers();

            for (final Flags.Flag modifier : Flags.asFlagSet((Flags.AccessFlags | Flags.MemberStaticClassFlags) &
                                                             (Flags.STATIC | Flags.FINAL | Flags.PRIVATE))) {
                modifiers.add(new JavaModifierToken(TextLocation.EMPTY, modifier));
            }

            declaration.setClassType(ClassType.CLASS);
            declaration.setName(definition.getSimpleName());
            declaration.putUserData(Keys.TYPE_REFERENCE, definition.selfReference);
            declaration.putUserData(Keys.TYPE_DEFINITION, definition);

            final AstNodeCollection<EntityDeclaration> members = declaration.getMembers();

            members.add(buildLookupField());
            members.add(buildHandleField());
            members.add(buildFenceField());
            members.add(buildHandleMethod());
            members.add(buildEnsureHandleMethod());
            members.add(buildInvokeMethod());

            return true;
        }

        FieldDeclaration buildHandleField() {
            return declareField(definition.fdHandle, Expression.NULL, 0);
        }

        FieldDeclaration buildFenceField() {
            return declareField(definition.fdFence, Expression.NULL, Flags.VOLATILE);
        }

        FieldDeclaration buildLookupField() {
            final MethodReference lookup = parser.parseMethod(methodHandlesType, "lookup", T_SIGNATURE_LOOKUP);
            return declareField(definition.fdLookup, makeType(methodHandlesType).invoke(lookup), Flags.FINAL);
        }

        VariableDeclarationStatement makeHandleVariableDeclaration() {
            final Variable v = new Variable();
            final VariableDeclarationStatement vd = new VariableDeclarationStatement(makeType(methodHandleType), "handle", makeReference(definition.fdHandle));

            v.setGenerated(false);
            v.setName("handle");
            v.setType(methodHandleType);
            vd.putUserData(Keys.VARIABLE, v);

            return vd;
        }

        MethodDeclaration buildHandleMethod() {
            final MethodDeclaration declaration = newMethod(definition.mdHandle);
            final VariableDeclarationStatement vd = makeHandleVariableDeclaration();

            vd.setModifiers(Collections.singletonList(Flags.Flag.FINAL));

            declaration.setBody(
                new BlockStatement(
                    vd,
                    new IfElseStatement(
                        Expression.MYSTERY_OFFSET,
                        new BinaryOperatorExpression(varReference(vd), BinaryOperatorType.INEQUALITY, new NullReferenceExpression()),
                        new ReturnStatement(varReference(vd))
                    ),
                    new ReturnStatement(makeReference(ensureHandleMethod).invoke())
                )
            );

            this.handleDeclaration = declaration;

            return declaration;
        }

        MethodDeclaration buildInvokeMethod() {
            final MethodDeclaration declaration = newMethod(definition.mdInvoke);

            final TryCatchStatement tryCatch = new TryCatchStatement(Expression.MYSTERY_OFFSET);
            final CatchClause cc = new CatchClause();
            final AstType throwableType = makeType(T_DESC_THROWABLE);
            final Variable t = new Variable();

            t.setType(throwableType.toTypeReference());
            t.setName("t");
            t.setGenerated(true);

            cc.putUserData(Keys.VARIABLE, t);
            cc.setVariableName("t");
            cc.getExceptionTypes().add(throwableType);
            cc.setBody(new BlockStatement(new ThrowStatement(makeType(T_DESC_THROWABLE_WRAPPER).makeNew(varReference(t)))));

            final MethodReference invokeExact = parser.parseMethod(parser.parseTypeDescriptor(T_DESC_METHOD_HANDLE), "invokeExact", M_DESC_INVOKE_EXACT);
            final InvocationExpression invoke = makeReference(handleMethod).invoke().invoke(invokeExact);
            final AstNodeCollection<Expression> arguments = invoke.getArguments();

            for (final ParameterDeclaration p : declaration.getParameters()) {
                arguments.add(varReference(p));
            }

            //invocation.cast(callSiteType).invoke(dynamicInvoker).invoke(invokeExact, ArrayUtilities.remove(nodeArguments, 0))
            tryCatch.setTryBlock(new BlockStatement(new ReturnStatement(invoke)));
            tryCatch.getCatchClauses().add(cc);

            declaration.setBody(new BlockStatement(tryCatch));

            this.invokeDeclaration = declaration;

            return declaration;
        }

        MethodDeclaration buildEnsureHandleMethod() {
            final MethodDeclaration declaration = newMethod(definition.mdEnsureHandle);
            final VariableDeclarationStatement vd = makeHandleVariableDeclaration();
            final MethodReference dynamicInvoker = parser.parseMethod(callSiteType, "dynamicInvoker", M_DESC_DYNAMIC_INVOKER);
            final VariableDeclarationStatement vdLookup = new VariableDeclarationStatement(makeType(lookupType), "lookup", makeReference(definition.fdLookup));

            final TryCatchStatement tryCatch = new TryCatchStatement(Expression.MYSTERY_OFFSET);
            final CatchClause cc = new CatchClause();
            final AstType throwableType = makeType(T_DESC_THROWABLE);
            final Variable t = new Variable();

            t.setType(throwableType.toTypeReference());
            t.setName("t");
            t.setGenerated(true);

            cc.putUserData(Keys.VARIABLE, t);
            cc.setVariableName("t");
            cc.getExceptionTypes().add(throwableType);
            cc.setBody(new BlockStatement(new ThrowStatement(makeType(T_DESC_THROWABLE_WRAPPER).makeNew(varReference(t)))));

            tryCatch.setTryBlock(
                new BlockStatement(
                    new ExpressionStatement(
                        new AssignmentExpression(
                            varReference(vd),
                            bootstrapCall.cast(makeType(callSiteType)).invoke(dynamicInvoker)
                        )
                    )
                )
            );

            tryCatch.getCatchClauses().add(cc);

            vdLookup.putUserData(Keys.VARIABLE, lookupVariable);

            declaration.setBody(
                new BlockStatement(
                    new ExpressionStatement(
                        new AssignmentExpression(
                            makeReference(definition.fdFence),
                            new PrimitiveExpression(Expression.MYSTERY_OFFSET, 0)
                        )
                    ),
                    vd,
                    new IfElseStatement(
                        Expression.MYSTERY_OFFSET,
                        new BinaryOperatorExpression(
                            varReference(vd),
                            BinaryOperatorType.EQUALITY,
                            new NullReferenceExpression(Expression.MYSTERY_OFFSET)
                        ),
                        new BlockStatement(
                            vdLookup,
                            tryCatch,
                            new ExpressionStatement(
                                new AssignmentExpression(
                                    makeReference(definition.fdFence),
                                    new PrimitiveExpression(Expression.MYSTERY_OFFSET, 1)
                                )
                            ),
                            new ExpressionStatement(
                                new AssignmentExpression(
                                    makeReference(definition.fdHandle),
                                    varReference(vd)
                                )
                            ),
                            new ExpressionStatement(
                                new AssignmentExpression(
                                    makeReference(definition.fdFence),
                                    new PrimitiveExpression(Expression.MYSTERY_OFFSET, 0)
                                )
                            )
                        )
                    ),
                    new ReturnStatement(varReference(vd))
                )
            );

            this.ensureHandleDeclaration = declaration;

            return declaration;
        }

        // <editor-fold defaultstate="collapsed" desc="HelperTypeDefinition Class">

        private final class HelperTypeDefinition extends TypeDefinition {
            final TypeReference selfReference;
            final FieldDefinition fdLookup;
            final FieldDefinition fdHandle;
            final FieldDefinition fdFence;
            final MethodDefinition mdHandle;
            final MethodDefinition mdEnsureHandle;
            final MethodDefinition mdInvoke;

            HelperTypeDefinition(final TypeReference selfReference, final TypeReference parentType) {
                super(resolver(parentType));

                this.selfReference = selfReference;

                setPackageName(selfReference.getPackageName());
                setSimpleName(selfReference.getSimpleName());
                setBaseType(BuiltinTypes.Object);
                setFlags((Flags.AccessFlags | Flags.MemberClassFlags) & (Flags.PRIVATE | Flags.FINAL | Flags.STATIC));
                setDeclaringType(parentType);

                final TypeDefinition resolvedParent = parentType.resolve();

                if (resolvedParent != null) {
                    setResolver(resolvedParent.getResolver());
                    setCompilerVersion(resolvedParent.getCompilerMajorVersion(), resolvedParent.getCompilerMinorVersion());
                }
                else {
                    setResolver(resolver());
                    setCompilerVersion(CompilerTarget.JDK1_7.majorVersion, CompilerTarget.JDK1_7.minorVersion);
                }

                final TypeDefinition self = this;

                fdHandle = new FieldDefinition(methodHandleType) {{
                    this.setName("handle");
                    this.setDeclaringType(self);
                    this.setFlags((Flags.AccessFlags | Flags.VarFlags) & (Flags.PRIVATE | Flags.STATIC));
                }};

                fdFence = new FieldDefinition(BuiltinTypes.Integer) {{
                    this.setName("fence");
                    this.setDeclaringType(self);
                    this.setFlags((Flags.AccessFlags | Flags.VarFlags) & (Flags.PRIVATE | Flags.STATIC));
                }};

                fdLookup = new FieldDefinition(resolver.lookupType(T_DESC_LOOKUP)) {{
                    this.setName("LOOKUP");
                    this.setDeclaringType(self);
                    this.setFlags((Flags.AccessFlags | Flags.VarFlags) & (Flags.PRIVATE | Flags.STATIC | Flags.FINAL));
                }};

                mdHandle = new MethodDefinition() {{
                    this.setName("handle");
                    this.setDeclaringType(self);
                    this.setFlags((Flags.AccessFlags | Flags.MethodFlags) & (Flags.PRIVATE | Flags.STATIC));
                    this.setReturnType(methodHandleType);
                }};

                mdEnsureHandle = new MethodDefinition() {{
                    this.setName("ensureHandle");
                    this.setDeclaringType(self);
                    this.setFlags((Flags.AccessFlags | Flags.MethodFlags) & (Flags.PRIVATE | Flags.STATIC));
                    this.setReturnType(methodHandleType);
                }};

                mdInvoke = new MethodDefinition() {{
                    final IMethodSignature methodType = callSite.getMethodType();

                    this.setName("invoke");
                    this.setDeclaringType(self);
                    this.setFlags((Flags.AccessFlags | Flags.MethodFlags) & (Flags.PRIVATE | Flags.STATIC));
                    this.setReturnType(methodType.getReturnType());

                    final ParameterDefinitionCollection ps = this.getParametersInternal();
                    final List<ParameterDefinition> tps = methodType.getParameters();

                    for (int i = 0, n = tps.size(); i < n; i++) {
                        final ParameterDefinition template = tps.get(i);
                        final ParameterDefinition pd = new ParameterDefinition(template.getSlot(), "p" + i, template.getParameterType());

                        ps.add(pd);
                    }
                }};

                final Collection<FieldDefinition> fields = getDeclaredFieldsInternal();

                fields.add(fdFence);
                fields.add(fdHandle);
                fields.add(fdLookup);
            }
        }

        // </editor-fold>
    }
}
