package de.dhbwstuttgart.target.generate;

import de.dhbwstuttgart.bytecode.FunNGenerator;
import de.dhbwstuttgart.core.JavaTXCompiler;
import de.dhbwstuttgart.environment.ByteArrayClassLoader;
import de.dhbwstuttgart.environment.IByteArrayClassLoader;
import de.dhbwstuttgart.parser.NullToken;
import de.dhbwstuttgart.syntaxtree.*;
import de.dhbwstuttgart.syntaxtree.Record;
import de.dhbwstuttgart.syntaxtree.factory.ASTFactory;
import de.dhbwstuttgart.syntaxtree.statement.*;
import de.dhbwstuttgart.syntaxtree.type.*;
import de.dhbwstuttgart.target.tree.*;
import de.dhbwstuttgart.target.tree.expression.*;
import de.dhbwstuttgart.target.tree.type.*;
import de.dhbwstuttgart.typeinference.result.*;

import java.util.*;
import java.util.stream.Collectors;

/**
 * @author dholle
 */
public class ASTToTargetAST {

    public static RefType OBJECT = ASTFactory.createObjectType(); // TODO It would be better if I could call this directly but the hashcode seems to change

    protected List<Generics> all;
    protected Generics generics;
    final Map<ClassOrInterface, Set<GenericTypeVar>> userDefinedGenerics = new HashMap<>();

    public final JavaTXCompiler compiler;

    protected ClassOrInterface currentClass; // TODO This is only needed because of SuperCall, maybe there's

    public List<GenericsResult> txGenerics() {
        return all.stream().map(generics -> new GenericsResult(generics.txGenerics)).toList();
    }

    public List<GenericsResult> javaGenerics() {
        return all.stream().map(generics -> new GenericsResult(generics.javaGenerics)).toList();
    }

    public TargetExpression convert(Pattern pattern) {
        var converter = new StatementToTargetExpression(this);
        pattern.accept(converter);
        return converter.result;
    }

    record Generics(JavaGenerics javaGenerics, TxGenerics txGenerics) {
    }


    protected IByteArrayClassLoader classLoader;
    protected SourceFile sourceFile;

    public ASTToTargetAST(List<ResultSet> resultSets) {
        this(null, resultSets);
    }
    public ASTToTargetAST(JavaTXCompiler compiler, List<ResultSet> resultSets) {
        this(compiler, resultSets, null, new ByteArrayClassLoader());
    }

    public ASTToTargetAST(JavaTXCompiler compiler, List<ResultSet> resultSets, SourceFile sourceFile, IByteArrayClassLoader classLoader) {
        this.compiler = compiler;
        this.classLoader = classLoader;
        this.sourceFile = sourceFile;

        all = new ArrayList<>();
        for (var set : resultSets) {
            all.add(new Generics(new JavaGenerics(this, set), new TxGenerics(this, set)));
        }
        this.generics = all.get(0);
    }

    Optional<Method> findMethod(ClassOrInterface owner, String name, List<TargetType> argumentList) {
        Optional<Method> method = Optional.empty();
        while (method.isEmpty() && !owner.getClassName().toString().equals("java.lang.Object")) {
            method = owner.getMethods().stream().filter(m -> m.name.equals(name) && parameterEquals(m.getParameterList(), argumentList)).findFirst();
            owner = compiler.getClass(owner.getSuperClass().getName());
        }
        return method;
    }

    boolean parameterEquals(ParameterList parameterList, List<TargetType> arguments) {
        var pars = parameterList.getFormalparalist();
        if (pars.size() != arguments.size())
            return false;

        for (var i = 0; i < pars.size(); i++) {
            var type1 = convert(pars.get(i).getType(), generics.javaGenerics);
            var type2 = arguments.get(i);
            if (type1 instanceof TargetGenericType)
                return true;
            if (TargetType.toPrimitive(type2).equals(type1))
                return true;
            if (!type1.equals(type2))
                return false;
        }

        return true;
    }

    Set<TargetGeneric> convert(Set<GenerateGenerics.Pair> result, GenerateGenerics generics) {
        return result.stream().map(p -> {
            if (p instanceof GenerateGenerics.PairLT pair) {
                return new TargetGeneric(pair.left.resolve().getName(), convert(pair.right.resolve(), generics));
            } else if (p instanceof GenerateGenerics.PairEQ pair) {
                return new TargetGeneric(pair.left.resolve().getName(), convert(pair.right, generics));
            } else {
                throw new IllegalArgumentException();
            }
        }).collect(Collectors.toSet());
    }

    public List<TargetGeneric> convert(GenericTypeVar typeVar, GenerateGenerics generics) {
        var ret = new ArrayList<TargetGeneric>();
        for (var bound : typeVar.getBounds()) {
            ret.add(new TargetGeneric(typeVar.getName(), generics.getTargetType(bound)));
        }
        return ret;
    }

    public TargetStructure convert(ClassOrInterface input) {
        currentClass = input;
        Set<TargetGeneric> javaGenerics = new HashSet<>();
        Set<TargetGeneric> txGenerics = new HashSet<>();

        var genericsIter = input.getGenerics().iterator();
        if (genericsIter.hasNext()) {
            // Add empty set of generics to cache so that it doesn't try to calculate it later
            var userDefinedGenerics = new HashSet<GenericTypeVar>();
            this.userDefinedGenerics.put(input, userDefinedGenerics);
            while (genericsIter.hasNext()) {
                var next = genericsIter.next();
                userDefinedGenerics.add(next);
                // TODO Support multiple bounds
                javaGenerics.add(new TargetGeneric(next.getName(), convert(next.getBounds().get(0))));
            }
        } else {
            this.userDefinedGenerics.put(input, new HashSet<>());
            // Generate generics only if there are no user defined ones
            javaGenerics = convert(generics.javaGenerics.generics(input), generics.javaGenerics);
            txGenerics = convert(generics.txGenerics.generics(input), generics.txGenerics);
        }

        TargetBlock fieldInitializer = null;
        if (input.getfieldInitializations().isPresent())
            fieldInitializer = convert(input.getfieldInitializations().get().block);
        TargetBlock finalFieldInitializer = fieldInitializer;

        var superInterfaces = input.getSuperInterfaces().stream().map(clazz -> convert(clazz, generics.javaGenerics)).toList();
        var constructors = input.getConstructors().stream().map(constructor -> this.convert(constructor, finalFieldInitializer)).flatMap(List::stream).toList();
        var fields = input.getFieldDecl().stream().map(this::convert).toList();
        var methods = groupOverloads(input.getMethods()).stream().map(this::convert).flatMap(List::stream).toList();

        TargetMethod staticConstructor = null;
        if (input.getStaticInitializer().isPresent())
            staticConstructor = this.convert(input.getStaticInitializer().get()).get(0);

        if (input instanceof Record)
            return new TargetRecord(input.getModifiers(), input.getClassName(), javaGenerics, txGenerics, superInterfaces, constructors, staticConstructor, fields, methods);
        else if (input.isInterface())
            return new TargetInterface(input.getModifiers(), input.getClassName(), javaGenerics, txGenerics, methods, superInterfaces, staticConstructor);
        else return new TargetClass(input.getModifiers(), input.getClassName(), convert(input.getSuperClass(), generics.javaGenerics), javaGenerics, txGenerics, superInterfaces, constructors, staticConstructor, fields, methods);
    }

    private List<MethodParameter> convert(ParameterList input, GenerateGenerics generics) {
        return input.getFormalparalist().stream().map(param ->
            new MethodParameter((TargetPattern) convert(param))
        ).toList();
    }

    private boolean hasGeneric(Set<TargetGeneric> generics, GenericRefType type) {
        return generics.stream().anyMatch(g -> g.name().equals(type.getParsedName()));
    }

    private Set<TargetGeneric> collectMethodGenerics(GenerateGenerics generateGenerics, Set<GenerateGenerics.Pair> generics, Method input) {
        var convertedGenerics = new HashSet<>(convert(generics, generateGenerics));
        outer: for (GenericTypeVar typeVar : input.getGenerics()) {
            for (var classGeneric : currentClass.getGenerics()) {
                if (classGeneric.equals(typeVar)) {
                    continue outer;
                }
            }
            convertedGenerics.addAll(convert(typeVar, generateGenerics));
        }
        /*
         * var returnType = sigma.getType(input.getReturnType(), equality); if ((returnType instanceof GenericRefType refType) && !hasGeneric(convertedGenerics, refType)) { convertedGenerics.add(new TargetGeneric(refType.getParsedName(), convert(OBJECT))); } for (var param : input.getParameterList()) { var type = sigma.getType(param.getType(), equality); if (type instanceof GenericRefType refType && !hasGeneric(convertedGenerics, refType)) { convertedGenerics.add(new
         * TargetGeneric(refType.getParsedName(), convert(OBJECT))); } }
         */

        return convertedGenerics;
    }

    private List<TargetConstructor> convert(Constructor input, TargetBlock fieldInitializer) {
        generics = all.get(0);
        List<TargetConstructor> result = new ArrayList<>();
        Set<List<MethodParameter>> parameterSet = new HashSet<>();

        for (var s : all) {
            generics = s;
            var javaGenerics = this.generics.javaGenerics.generics(currentClass, input);
            var txGenerics = this.generics.txGenerics.generics(currentClass, input);
            List<MethodParameter> params = convert(input.getParameterList(), this.generics.javaGenerics);
            if (parameterSet.stream().noneMatch(p -> p.equals(params))) {
                List<MethodParameter> txParams = convert(input.getParameterList(), this.generics.txGenerics);
                var javaMethodGenerics = collectMethodGenerics(generics.javaGenerics(), javaGenerics, input);
                var txMethodGenerics = collectMethodGenerics(generics.txGenerics(), txGenerics, input);

                result.add(new TargetConstructor(input.modifier, javaMethodGenerics, txMethodGenerics, params, txParams, convert(input.block), fieldInitializer));
                parameterSet.add(params);
            }
        }

        return result;
    }

    /**
     * This only considers type patterns, all other methods aren't grouped together
     * @param a
     * @param b
     * @return
     */
    private boolean signatureEquals(Method a, Method b) {
        if (!a.name.equals(b.name)) return false;
        var para = a.getParameterList().getFormalparalist();
        var parb = b.getParameterList().getFormalparalist();
        if (para.size() != parb.size()) return false;

        for (var i = 0; i < para.size(); i++) {
            var pa = para.get(i);
            var pb = parb.get(i);

            if (pa instanceof RecordPattern rpa) {
                if (pb instanceof RecordPattern rpb) {
                    if (rpa.getType().equals(rpb.getType())) continue;
                }
                return false;
            } else if (pa.getType().equals(pb.getType())) {
                continue;
            }
            return false;
        }

        return true;
    }

    // TODO Nested patterns
    private List<List<Method>> groupOverloads(List<Method> input) {
        var done = new HashSet<Method>();
        var res = new ArrayList<List<Method>>();
        for (var method : input) {
            if (done.contains(method)) continue;
            var overloads = new ArrayList<Method>();
            overloads.add(method);
            done.add(method);
            for (var method2 : input) {
                if (!done.contains(method2) && signatureEquals(method, method2)) {
                    done.add(method2);
                    overloads.add(method2);
                }
            }
            res.add(overloads);
        }
        return res;
    }

    private String encodeName(String name, ParameterList params) {
        var res = new StringBuilder();
        res.append(name);
        res.append('$');
        for (var param : params.getFormalparalist()) {
            if (param instanceof RecordPattern rp) {
                res.append(FunNGenerator.encodeType(convert(param.getType())));
                for (var pattern : rp.getSubPattern()) {
                    res.append(FunNGenerator.encodeType(convert(pattern.getType())));
                }
            }
        }
        return res.toString();
    }

    private List<TargetMethod> convert(List<Method> overloadedMethods) {
        if (overloadedMethods.size() == 1) {
            return convert(overloadedMethods.get(0));
        }
        var res = new ArrayList<Method>();
        for (var method : overloadedMethods) {
            var newMethod = new Method(
                    method.modifier,
                    encodeName(method.name, method.getParameterList()),
                    method.getReturnType(),
                    method.getParameterList(),
                    method.block,
                    method.getGenerics(),
                    method.getOffset()
            );
            res.add(newMethod);
        }

        var template = overloadedMethods.get(0);

        var pParams = new ArrayList<Pattern>();
        var i = 0;
        for (var par : template.getParameterList()) {
            pParams.add(switch (par) {
                case RecordPattern rp -> new RecordPattern(rp.getSubPattern(), "par" + i, rp.getType(), new NullToken());
                default -> par;
            });
            i++;
        }
        var params = new ParameterList(pParams, new NullToken());

        var statements = new ArrayList<Statement>();
        statements.add(new Return(makeRecordSwitch(template.getReturnType(), params, res), new NullToken()));
        var block = new Block(statements, new NullToken());
        var entryPoint = new Method(template.modifier, template.name, template.getReturnType(), params, block, template.getGenerics(), new NullToken());

        res.add(entryPoint); // TODO
        return res.stream().map(this::convert).flatMap(List::stream).toList();
    }

    private Expression makeRecordSwitch(RefTypeOrTPHOrWildcardOrGeneric returnType, ParameterList params, List<Method> overloadedMethods) {
        var param = params.getFormalparalist().get(0);
        assert param instanceof RecordPattern; // TODO

        var cases = new ArrayList<SwitchBlock>();
        for (var method : overloadedMethods) {
            var statements = new ArrayList<Statement>();
            /*statements.add(new MethodCall(
                method.getReturnType(), new ExpressionReceiver(new This(new NullToken())), method.name,
                params,
            ));*/

            var block = new Block(statements, new NullToken());
            var labels = new ArrayList<SwitchLabel>();
            cases.add(new SwitchBlock(labels, block, true, new NullToken()));
        }
        var swtch = new Switch(new LocalVar("par0", param.getType(), new NullToken()), cases, returnType, false, new NullToken());

        return swtch;
    }

    private List<TargetMethod> convert(Method method) {
        generics = all.get(0);
        List<TargetMethod> result = new ArrayList<>();
        Set<List<MethodParameter>> parameterSet = new HashSet<>();

        for (var s : all) {
            generics = s;
            var javaGenerics = this.generics.javaGenerics.generics(currentClass, method);
            var txGenerics = this.generics.txGenerics.generics(currentClass, method);
            List<MethodParameter> params = convert(method.getParameterList(), this.generics.javaGenerics);
            if (parameterSet.stream().noneMatch(p -> p.equals(params))) {
                List<MethodParameter> txParams = convert(method.getParameterList(), this.generics.txGenerics);

                var javaMethodGenerics = collectMethodGenerics(generics.javaGenerics(), javaGenerics, method);
                var txMethodGenerics = collectMethodGenerics(generics.txGenerics(), txGenerics, method);

                var javaSignature = new TargetMethod.Signature(javaMethodGenerics, params, convert(method.getReturnType(), this.generics.javaGenerics));
                var txSignature = new TargetMethod.Signature(txMethodGenerics, txParams, convert(method.getReturnType(), this.generics.txGenerics));
                result.add(new TargetMethod(method.modifier, method.name, convert(method.block), javaSignature, txSignature));
                parameterSet.add(params);
            }
        }
        return result;
    }

    protected TargetSwitch.Case convert(SwitchBlock block) {
        return new TargetSwitch.Case(block.getLabels().stream().map(this::convert).toList(), convert((Block) block), block.isExpression);
    }

    protected TargetBlock convert(Block block) {
        if (block == null) return null;
        return new TargetBlock(block.statements.stream().map(this::convert).toList());
    }

    protected TargetExpression convert(Expression expr) {
        var converter = new StatementToTargetExpression(this);
        expr.accept(converter);
        return converter.result;
    }

    private TargetField convert(Field input) {
        return new TargetField(input.modifier, convert(input.getType(), generics.javaGenerics), input.getName());
    }

    private final Map<String, FunNGenerator.GenericParameters> usedFunN = new HashMap<>();
    private final Set<Integer> usedFunNSuperTypes = new HashSet<>();

    public Map<String, byte[]> auxiliaries = new HashMap<>();

    protected TargetType convert(RefTypeOrTPHOrWildcardOrGeneric input) {
        return convert(input, generics.javaGenerics);
    }

    private static void collectArguments(TargetSpecializedType tspec, List<TargetType> newParams) {
        for (var i = 0; i < tspec.params().size(); i++) {
            var param = tspec.params().get(i);
            if (param instanceof TargetSpecializedType fn) {
                collectArguments(tspec, newParams);
            } else {
                newParams.add(param);
            }
        }
    }

    static TargetType flattenFunNType(List<TargetType> params, FunNGenerator.GenericParameters gep) {
        var newParams = new ArrayList<TargetType>();
        for (TargetType param : params) {
            if (param instanceof TargetSpecializedType fn) {
                collectArguments(fn, newParams);
            } else {
                newParams.add(param);
            }
        }
        var filteredParams = new ArrayList<TargetType>();
        for (var i = 0; i < newParams.size(); i++) {
            if (gep.parameters.get(i) != null)
                filteredParams.add(newParams.get(i));
        }
        return TargetFunNType.fromParams(params, filteredParams);
    }

    protected TargetType convert(RefTypeOrTPHOrWildcardOrGeneric input, GenerateGenerics generics) {
        return input.acceptTV(new TypeVisitor<>() {
            @Override
            public TargetType visit(RefType refType) {
                var name = refType.getName().toString();
                if (name.equals("void"))
                    return null;
                if (refType.isPrimitive()) {
                    return TargetType.toPrimitive(refType);
                }

                var params = refType.getParaList().stream().map(type -> {
                    var res = convert(type, generics);
                    if (res == null) res = new TargetRefType("java.lang.Void");
                    return res;
                }).toList();

                if (name.matches("Fun\\d+\\$\\$")) { // TODO This seems like a bad idea
                    var className = FunNGenerator.getSpecializedClassName(FunNGenerator.getArguments(params), FunNGenerator.getReturnType(params));
                    if (!usedFunNSuperTypes.contains(params.size())) {
                        usedFunNSuperTypes.add(params.size());
                        var code = FunNGenerator.generateSuperBytecode(params.size() - 1);
                        var superClassName = FunNGenerator.getSuperClassName(params.size() - 1);
                        try {
                            classLoader.findClass(superClassName);
                        } catch (ClassNotFoundException e) {
                            try {
                                classLoader.loadClass(code);
                            } catch (LinkageError ignored) {}
                        }
                        auxiliaries.put(superClassName, code);
                    }
                    FunNGenerator.GenericParameters gep = null;
                    if (!usedFunN.containsKey(className)) {
                        gep = new FunNGenerator.GenericParameters();
                        var code = FunNGenerator.generateSpecializedBytecode(FunNGenerator.getArguments(params), FunNGenerator.getReturnType(params), gep);
                        try {
                            classLoader.findClass(className);
                        } catch (ClassNotFoundException e) {
                            try {
                                classLoader.loadClass(code);
                            } catch (LinkageError ignored) {}
                        }
                        usedFunN.put(className, gep);
                        auxiliaries.put(className, code);
                    } else {
                        gep = usedFunN.get(className);
                    }
                    return flattenFunNType(params, gep);
                }
                return new TargetRefType(name, params);
            }

            @Override
            public TargetType visit(SuperWildcardType superWildcardType) {
                return new TargetSuperWildcard(convert(superWildcardType.getInnerType(), generics));
            }

            @Override
            public TargetType visit(TypePlaceholder typePlaceholder) {
                return generics.getTargetType(typePlaceholder);
            }

            @Override
            public TargetType visit(ExtendsWildcardType extendsWildcardType) {
                return new TargetExtendsWildcard(convert(extendsWildcardType.getInnerType(), generics));
            }

            @Override
            public TargetType visit(GenericRefType genericRefType) {
                return new TargetGenericType(genericRefType.getParsedName());
            }
        });
    }
}