package generators;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import ast.*;
import designPatternExtensions.MediatorInsertion;
import models.dataConstraintModel.MapType;
import models.deltaAlgebra.*;

public class MediatorInsertionWeaver {

    static class MediatorContext {
        final Codebase codebase;
        final String srcName, dstName, mediatorName;
        String mediatorFieldName;
        final String dstFieldName;
        final models.algebra.Type dstType, mediatorType;
        final TypeDeclaration srcClass, dstClass, mediatorClass;

        boolean toMany;
        FieldDeclaration srcDstMapField;
        FieldDeclaration dstField;
        FieldDeclaration srcFieldToReplace;
        String oldFieldName;

        MediatorContext(Codebase codebase, ObjectNode srcNode, ObjectNode dstNode, ObjectNode mediatorNode) {
            this.codebase = codebase;
            this.srcName = toComponentName(srcNode.getName());
            this.dstName = toComponentName(dstNode.getName());
            this.mediatorName = toComponentName(mediatorNode.getName());
            this.mediatorFieldName = toVariableName(mediatorName);
            this.dstFieldName = toVariableName(dstName);
            this.dstType = codebase.getComponentType(dstName);
            this.mediatorType = codebase.getComponentType(mediatorName);
            this.mediatorClass = createClass(codebase, mediatorName);
            this.srcClass = createClass(codebase, srcName);
            this.dstClass = createClass(codebase, dstName);
        }
    }

    static Codebase weave(Codebase codebase, MediatorInsertion mediatorInsertion, ObjectNode srcNode, ObjectNode dstNode) {
        MediatorContext mediatorContext = new MediatorContext(codebase, srcNode, dstNode, mediatorInsertion.getMediator());

        determineMediatorField(mediatorContext);
        addMediatorConstructor(mediatorContext);
        addMediatorDelegationMethods(mediatorContext);
        findSrcFieldToReplace(mediatorContext);
        if (mediatorContext.srcFieldToReplace != null) {
            replaceSrcDependencyWithMediator(mediatorContext);
            moveSrcMethodsToMediator(mediatorContext);
            replaceRemainingFieldReferences(mediatorContext);
        }
        return codebase;
    }

    static void determineMediatorField(MediatorContext mediatorContext) {
        mediatorContext.srcDstMapField = null;
        for (FieldDeclaration field: mediatorContext.srcClass.getFields()) {
            if (field.getType() instanceof MapType) {
                MapType mapType = (MapType) field.getType();
                if (mapType.getValueType() != null && mapType.getValueType().getTypeName().equals(mediatorContext.dstType.getTypeName())) {
                    mediatorContext.srcDstMapField = field;
                    break;
                }
            }
        }
        mediatorContext.toMany = mediatorContext.srcDstMapField != null;

        if (mediatorContext.toMany) {
            MapType srcMapType = (MapType) mediatorContext.srcDstMapField.getType();
            FieldDeclaration existingMediatorMapField = null;
            for (FieldDeclaration field: mediatorContext.mediatorClass.getFields()) {
                if (field.getType() instanceof MapType) {
                    MapType mapType = (MapType) field.getType();
                    if (mapType.getValueType() != null && mapType.getValueType().getTypeName().equals(mediatorContext.dstType.getTypeName())
                            && mapType.getKeyType().getTypeName().equals(srcMapType.getKeyType().getTypeName())) {
                        existingMediatorMapField = field;
                        break;
                    }
                }
            }
            mediatorContext.dstField = (existingMediatorMapField != null)
                    ? existingMediatorMapField
                    : createField(mediatorContext.mediatorClass, mediatorContext.srcDstMapField.getName(), srcMapType);
        } else {
            mediatorContext.dstField = createField(mediatorContext.mediatorClass, mediatorContext.dstFieldName, mediatorContext.dstType);
        }
    }

    static void addMediatorConstructor(MediatorContext mediatorContext) {
        MethodDeclaration constructor = createConstructor(mediatorContext.mediatorClass);
        if (!hasParameterNamed(constructor, mediatorContext.dstField.getName())) {
            VariableDeclaration dstParam = new VariableDeclaration(mediatorContext.dstField.getType(), mediatorContext.dstField.getName());
            constructor.addParameter(dstParam);
            FieldAccess dstFieldAccess = new FieldAccess(new ThisExpression(), mediatorContext.dstField.getName());
            Assignment dstAssignment = new Assignment(dstFieldAccess, new Variable(mediatorContext.dstField.getName()));
            constructor.addUniqueStatement(new ExpressionStatement(dstAssignment));
        }
    }

    static void addMediatorDelegationMethods(MediatorContext mediatorContext) {
        if (mediatorContext.toMany) return;
        for (MethodDeclaration method: mediatorContext.dstClass.getMethods()) {
            if (method.isConstructor()) continue;
            MethodDeclaration delegateMethod = createMethod(mediatorContext.mediatorClass, method.getName());
            delegateMethod.setReturnType(method.getReturnType());
            List<Expression> args = new ArrayList<>();
            if (method.getParameters() != null) {
                for (VariableDeclaration param: method.getParameters()) {
                    delegateMethod.addParameter(param);
                    args.add(new Variable(param.getName()));
                }
            }
            MethodInvocation callDst = new MethodInvocation(new FieldAccess(mediatorContext.dstField.getName()), method.getName(), args);
            if (method.getReturnType() != null) {
                ReturnStatement returnStatement = new ReturnStatement();
                returnStatement.setExpression(callDst);
                delegateMethod.addUniqueStatement(returnStatement);
            } else {
                delegateMethod.addUniqueStatement(new ExpressionStatement(callDst));
            }
        }
    }

    static void findSrcFieldToReplace(MediatorContext mediatorContext) {
        if (mediatorContext.toMany) {
            mediatorContext.srcFieldToReplace = mediatorContext.srcDstMapField;
            return;
        }
        for (FieldDeclaration field: mediatorContext.srcClass.getFields()) {
            if (field.getType() != null && field.getType().getTypeName().equals(mediatorContext.dstType.getTypeName())) {
                mediatorContext.srcFieldToReplace = field;
                return;
            }
        }
        mediatorContext.srcFieldToReplace = null;
    }

    static void replaceSrcDependencyWithMediator(MediatorContext mediatorContext) {
        FieldDeclaration existingMediatorField = null;
        for (FieldDeclaration field: mediatorContext.srcClass.getFields()) {
            if (field != mediatorContext.srcFieldToReplace && field.getType() != null
                    && field.getType().getTypeName().equals(mediatorContext.mediatorType.getTypeName())) {
                existingMediatorField = field;
                break;
            }
        }
        mediatorContext.oldFieldName = mediatorContext.srcFieldToReplace.getName();
        if (existingMediatorField != null) {
            mediatorContext.srcClass.removeField(mediatorContext.srcFieldToReplace);
            mediatorContext.mediatorFieldName = existingMediatorField.getName();
        } else {
            mediatorContext.srcFieldToReplace.setType(mediatorContext.mediatorType);
            mediatorContext.srcFieldToReplace.setName(mediatorContext.mediatorFieldName);
        }

        MethodDeclaration mediatorSetter = createMethod(mediatorContext.srcClass, ASTGenerator.setterPrefix + mediatorContext.mediatorName);
        if (mediatorSetter.getParameters() == null || mediatorSetter.getParameters().isEmpty()) {
            VariableDeclaration mediatorParam = new VariableDeclaration(mediatorContext.mediatorType, mediatorContext.mediatorFieldName);
            mediatorSetter.addParameter(mediatorParam);
            FieldAccess mediatorFieldAccess = new FieldAccess(new ThisExpression(), mediatorContext.mediatorFieldName);
            Assignment mediatorAssignment = new Assignment(mediatorFieldAccess, new Variable(mediatorContext.mediatorFieldName));
            mediatorSetter.addUniqueStatement(new ExpressionStatement(mediatorAssignment));
        }

        for (MethodDeclaration method: mediatorContext.srcClass.getMethods()) {
            if (method.isConstructor() && method.getParameters() != null) {
                for (VariableDeclaration param: method.getParameters()) {
                    if (param.getType() != null && param.getType().getTypeName().equals(mediatorContext.dstType.getTypeName())) {
                        param.setType(mediatorContext.mediatorType);
                        param.setName(mediatorContext.mediatorFieldName);
                    }
                }
            }
        }
    }

    static void moveSrcMethodsToMediator(MediatorContext mediatorContext) {
        List<MethodDeclaration> methodsToMove = new ArrayList<>();
        for (MethodDeclaration method: mediatorContext.srcClass.getMethods()) {
            if (method.isConstructor()) continue;
            if (methodReferencesField(method, mediatorContext.oldFieldName)) {
                methodsToMove.add(method);
            }
        }
        for (MethodDeclaration method: methodsToMove) {
            moveSingleMethodToMediator(mediatorContext, method);
        }
    }

    private static void moveSingleMethodToMediator(MediatorContext mediatorContext, MethodDeclaration method) {
        mediatorContext.srcClass.removeMethod(method);
        List<FieldDeclaration> danglingFields = new ArrayList<>();
        if (method.getBody() != null) {
            for (Statement statement: method.getBody().getStatements2()) {
                collectDanglingValueFields(statement, mediatorContext.srcClass, mediatorContext.mediatorClass, mediatorContext.oldFieldName, danglingFields);
            }
        }
        Set<String> danglingFieldNames = new HashSet<>();
        for (FieldDeclaration extra: danglingFields) {
            danglingFieldNames.add(extra.getName());
            if (!hasParameterNamed(method, extra.getName())) {
                method.addParameter(new VariableDeclaration(extra.getType(), extra.getName()));
            }
        }
        if (!danglingFieldNames.isEmpty() && method.getBody() != null) {
            for (Statement statement: method.getBody().getStatements2()) {
                rewriteDanglingFieldsAsParams(statement, danglingFieldNames);
            }
        }

        MethodDeclaration existingOnMediator = null;
        for (MethodDeclaration existing: mediatorContext.mediatorClass.getMethods()) {
            if (existing.getName().equals(method.getName())) { existingOnMediator = existing; break; }
        }
        if (existingOnMediator == null) {
            mediatorContext.mediatorClass.addMethod(method);
        }

        createSrcStubForMovedMethod(mediatorContext, method, danglingFieldNames);
    }

    private static void createSrcStubForMovedMethod(MediatorContext mediatorContext, MethodDeclaration method, Set<String> danglingFieldNames) {
        MethodDeclaration stub = createMethod(mediatorContext.srcClass, method.getName());
        stub.setReturnType(method.getReturnType());
        List<Expression> args = new ArrayList<>();
        if (method.getParameters() != null) {
            for (VariableDeclaration param: method.getParameters()) {
                if (danglingFieldNames.contains(param.getName())) {
                    args.add(new FieldAccess(param.getName()));
                } else {
                    stub.addParameter(param);
                    args.add(new Variable(param.getName()));
                }
            }
        }
        MethodInvocation callMediator = new MethodInvocation(new FieldAccess(mediatorContext.mediatorFieldName), method.getName(), args);
        if (method.getReturnType() != null) {
            ReturnStatement returnStatement = new ReturnStatement();
            returnStatement.setExpression(callMediator);
            stub.addUniqueStatement(returnStatement);
        } else {
            stub.addUniqueStatement(new ExpressionStatement(callMediator));
        }
    }

    static void replaceRemainingFieldReferences(MediatorContext mediatorContext) {
        for (MethodDeclaration method: mediatorContext.srcClass.getMethods()) {
            if (method.getBody() == null) continue;
            for (Statement statement: method.getBody().getStatements2()) {
                replaceFieldNameInStatement(statement, mediatorContext.oldFieldName, mediatorContext.mediatorFieldName);
            }
        }
    }

    //==================================================================================
    // Field replacement, reference check, and reference cleanup
    //==================================================================================

    private static void replaceFieldNameInStatement(Statement statement, String oldName, String newName) {
        if (statement instanceof ExpressionStatement) {
            replaceFieldNameInExpression(((ExpressionStatement) statement).getExpression(), oldName, newName);
        }
    }

    private static void replaceFieldNameInExpression(Expression expr, String oldName, String newName) {
        if (expr instanceof MethodInvocation) {
            MethodInvocation mi = (MethodInvocation) expr;
            if (mi.getReceiver() instanceof FieldAccess) {
                FieldAccess fa = (FieldAccess) mi.getReceiver();
                if (fa.getFieldName().equals(oldName)) {
                    fa.setFieldName(newName);
                }
            }
            for (Expression arg: mi.getArguments()) {
                replaceFieldNameInExpression(arg, oldName, newName);
            }
        } else if (expr instanceof Assignment) {
            Assignment assignment = (Assignment) expr;
            if (assignment.getLeft() instanceof FieldAccess) {
                FieldAccess fa = (FieldAccess) assignment.getLeft();
                if (fa.getFieldName().equals(oldName)) {
                    fa.setFieldName(newName);
                }
            }
            if (assignment.getRight() instanceof Variable) {
                Variable var = (Variable) assignment.getRight();
                if (var.getName().equals(oldName)) {
                    assignment.setRight(new Variable(newName));
                }
            }
            replaceFieldNameInExpression(assignment.getLeft(), oldName, newName);
            replaceFieldNameInExpression(assignment.getRight(), oldName, newName);
        }
    }

    private static boolean methodReferencesField(MethodDeclaration method, String fieldName) {
        if (method.getBody() == null) return false;
        for (Statement statement: method.getBody().getStatements2()) {
            if (statementReferencesField(statement, fieldName)) return true;
        }
        return false;
    }

    private static boolean statementReferencesField(Statement statement, String fieldName) {
        if (statement instanceof ExpressionStatement) {
            return expressionReferencesField(((ExpressionStatement) statement).getExpression(), fieldName);
        } else if (statement instanceof ReturnStatement) {
            Expression expr = ((ReturnStatement) statement).getExpression();
            return expr != null && expressionReferencesField(expr, fieldName);
        } else if (statement instanceof EnhancedForStatement) {
            EnhancedForStatement forStatement = (EnhancedForStatement) statement;
            if (expressionReferencesField(forStatement.getExpression(), fieldName)) return true;
            Statement body = forStatement.getBody();
            if (body == null) return false;
            if (body instanceof Block) {
                for (Statement inner: ((Block) body).getStatements2()) {
                    if (statementReferencesField(inner, fieldName)) return true;
                }
                return false;
            }
            return statementReferencesField(body, fieldName);
        }
        return false;
    }

    private static boolean expressionReferencesField(Expression expr, String fieldName) {
        if (expr == null) return false;
        if (expr instanceof FieldAccess) {
            FieldAccess fa = (FieldAccess) expr;
            if (fieldName.equals(fa.getFieldName())) return true;
        }
        if (expr instanceof MethodInvocation) {
            MethodInvocation mi = (MethodInvocation) expr;
            if (expressionReferencesField(mi.getReceiver(), fieldName)) return true;
            if (mi.getArguments() != null) {
                for (Expression arg: mi.getArguments()) {
                    if (expressionReferencesField(arg, fieldName)) return true;
                }
            }
        } else if (expr instanceof Assignment) {
            Assignment assignment = (Assignment) expr;
            if (expressionReferencesField(assignment.getLeft(), fieldName)) return true;
            if (expressionReferencesField(assignment.getRight(), fieldName)) return true;
        }
        return false;
    }

    private static void collectDanglingValueFields(Statement statement, TypeDeclaration srcClass,
                                                   TypeDeclaration mediatorClass, String movedFieldName, List<FieldDeclaration> out) {
        if (statement instanceof ExpressionStatement) {
            walkExpressionForDanglingFields(((ExpressionStatement) statement).getExpression(), false, srcClass, mediatorClass, movedFieldName, out);
        } else if (statement instanceof ReturnStatement) {
            walkExpressionForDanglingFields(((ReturnStatement) statement).getExpression(), true, srcClass, mediatorClass, movedFieldName, out);
        } else if (statement instanceof EnhancedForStatement) {
            EnhancedForStatement forStatement = (EnhancedForStatement) statement;
            walkExpressionForDanglingFields(forStatement.getExpression(), false, srcClass, mediatorClass, movedFieldName, out);
            Statement body = forStatement.getBody();
            if (body instanceof Block) {
                for (Statement inner: ((Block) body).getStatements2()) {
                    collectDanglingValueFields(inner, srcClass, mediatorClass, movedFieldName, out);
                }
            } else if (body != null) {
                collectDanglingValueFields(body, srcClass, mediatorClass, movedFieldName, out);
            }
        }
    }

    private static void walkExpressionForDanglingFields(Expression expr, boolean asValue, TypeDeclaration srcClass,
                                                        TypeDeclaration mediatorClass, String movedFieldName, List<FieldDeclaration> out) {
        if (expr == null) return;
        if (expr instanceof FieldAccess) {
            if (!asValue) return;
            String name = ((FieldAccess) expr).getFieldName();
            if (name.equals(movedFieldName)) return;
            for (FieldDeclaration f: mediatorClass.getFields()) if (f.getName().equals(name)) return;
            for (FieldDeclaration f: out) if (f.getName().equals(name)) return;
            for (FieldDeclaration f: srcClass.getFields()) {
                if (f.getName().equals(name)) { out.add(f); return; }
            }
        } else if (expr instanceof MethodInvocation) {
            MethodInvocation mi = (MethodInvocation) expr;
            walkExpressionForDanglingFields(mi.getReceiver(), false, srcClass, mediatorClass, movedFieldName, out);
            if (mi.getArguments() != null) {
                for (Expression arg: mi.getArguments()) {
                    walkExpressionForDanglingFields(arg, true, srcClass, mediatorClass, movedFieldName, out);
                }
            }
        } else if (expr instanceof Assignment) {
            Assignment assignment = (Assignment) expr;
            walkExpressionForDanglingFields(assignment.getLeft(), false, srcClass, mediatorClass, movedFieldName, out);
            walkExpressionForDanglingFields(assignment.getRight(), true, srcClass, mediatorClass, movedFieldName, out);
        }
    }

    private static void rewriteDanglingFieldsAsParams(Statement statement, Set<String> fieldNames) {
        if (statement instanceof ExpressionStatement) {
            rewriteExpressionValues(((ExpressionStatement) statement).getExpression(), false, fieldNames);
        } else if (statement instanceof ReturnStatement) {
            ReturnStatement rs = (ReturnStatement) statement;
            Expression expr = rs.getExpression();
            if (expr instanceof FieldAccess && fieldNames.contains(((FieldAccess) expr).getFieldName())) {
                rs.setExpression(new Variable(((FieldAccess) expr).getFieldName()));
            } else {
                rewriteExpressionValues(expr, true, fieldNames);
            }
        } else if (statement instanceof EnhancedForStatement) {
            EnhancedForStatement forStatement = (EnhancedForStatement) statement;
            Statement body = forStatement.getBody();
            if (body instanceof Block) {
                for (Statement inner: ((Block) body).getStatements2()) {
                    rewriteDanglingFieldsAsParams(inner, fieldNames);
                }
            } else if (body != null) {
                rewriteDanglingFieldsAsParams(body, fieldNames);
            }
        }
    }

    private static void rewriteExpressionValues(Expression expr, boolean asValue, Set<String> fieldNames) {
        if (expr instanceof MethodInvocation) {
            MethodInvocation mi = (MethodInvocation) expr;
            rewriteExpressionValues(mi.getReceiver(), false, fieldNames);
            if (mi.getArguments() != null) {
                List<Expression> args = mi.getArguments();
                for (int i = 0; i < args.size(); i++) {
                    Expression arg = args.get(i);
                    if (arg instanceof FieldAccess && fieldNames.contains(((FieldAccess) arg).getFieldName())) {
                        args.set(i, new Variable(((FieldAccess) arg).getFieldName()));
                    } else {
                        rewriteExpressionValues(arg, true, fieldNames);
                    }
                }
            }
        } else if (expr instanceof Assignment) {
            Assignment assignment = (Assignment) expr;
            rewriteExpressionValues(assignment.getLeft(), false, fieldNames);
            Expression right = assignment.getRight();
            if (right instanceof FieldAccess && fieldNames.contains(((FieldAccess) right).getFieldName())) {
                assignment.setRight(new Variable(((FieldAccess) right).getFieldName()));
            } else {
                rewriteExpressionValues(right, true, fieldNames);
            }
        }
    }

    //==================================================================================
    // Same logic as ASTGenerator create methods
    //==================================================================================

    private static String toComponentName(String name) {
        return name.substring(0, 1).toUpperCase() + name.substring(1);
    }

    private static String toVariableName(String name) {
        return name.substring(0, 1).toLowerCase() + name.substring(1);
    }

    private static TypeDeclaration createClass(Codebase codebase, String name) {
        if (codebase.getCompilationUnit(name) != null) return codebase.getCompilationUnit(name).types().getFirst();
        TypeDeclaration type = new TypeDeclaration(name);
        CompilationUnit compilationUnit = new CompilationUnit(type);
        codebase.addCompilationUnit(name, compilationUnit);
        return type;
    }

    private static MethodDeclaration createMethod(TypeDeclaration type, String name) {
        for (MethodDeclaration method: type.getMethods()) {
            if (method.getName().equals(name)) return method;
        }
        MethodDeclaration method = new MethodDeclaration(name);
        type.addMethod(method);
        return method;
    }

    private static MethodDeclaration createConstructor(TypeDeclaration type) {
        for (MethodDeclaration method: type.getMethods()) {
            if (method.getName().equals(type.getTypeName())) return method;
        }
        MethodDeclaration method = new MethodDeclaration(type.getTypeName(), true);
        type.addMethod(method);
        return method;
    }

    private static FieldDeclaration createField(TypeDeclaration type, String fieldName, models.algebra.Type fieldType) {
        for (FieldDeclaration field: type.getFields()) {
            if (field.getName().equals(fieldName)) return field;
        }
        FieldDeclaration field = new FieldDeclaration(fieldType, fieldName);
        type.addField(field);
        return field;
    }

    private static boolean hasParameterNamed(MethodDeclaration method, String name) {
        if (method.getParameters() == null) return false;
        for (VariableDeclaration param: method.getParameters()) {
            if (param.getName().equals(name)) return true;
        }
        return false;
    }
}