package com.speedment.common.codegen.internal.util;

import com.speedment.common.codegen.model.File;
import com.speedment.common.codegen.model.Import;
import com.speedment.common.codegen.model.trait.HasImports;

import java.util.ArrayList;
import java.util.List;
import java.util.function.Function;

/**
 * Utility class for operating on a tree of code models.
 *
 * @author Emil Forslund
 * @since  2.5
 */
public final class ModelTreeUtil {

    /**
     * Returns a list of all the imports of the specified model by traversing
     * the entire tree.
     *
     * @param root  the root to start at (usually the {@link File})
     * @return      the list of imports
     */
    public static List<Import> importsOf(HasImports<?> root) {
        return traverse(HasImports.class, root, HasImports::getImports);
    }

    /**
     * Traverses the tree of models searching for children of a particular type
     * by applying the getter method on each found element. Even if the root
     * itself is of the searched type, it will not be added.
     *
     * @param parentType  the parent type
     * @param root        the root node in the tree
     * @param getter      getter method used to obtain a list of children
     * @param <T>         the type of the model searched for
     * @param <P>         the type of the parent
     * @return            list of hits
     */
    public static <T, P> List<T>
    traverse(Class<P> parentType, P root, Function<? super P, List<? extends T>> getter) {
        final List<T> found = new ArrayList<>(getter.apply(root));

        for (final T element : found) {
            if (parentType.isInstance(element)) {
                @SuppressWarnings("unchecked")
                final P asParent = (P) element;
                found.addAll(getter.apply(asParent));
            }
        }
        return found;
    }

    /**
     * Utility classes should not be instantiated.
     */
    private ModelTreeUtil() {}
}

@@ -195,22 +195,391 @@
 rait.Has
-Import
+Classes;%0Aimport com.speedment.common.codegen.model.trait.HasConstructors;%0Aimport com.speedment.common.codegen.model.trait.HasFields;%0Aimport com.speedment.common.codegen.model.trait.HasImports;%0Aimport com.speedment.common.codegen.model.trait.HasInitializers;%0Aimport com.speedment.common.codegen.model.trait.HasJavadoc;%0Aimport com.speedment.common.codegen.model.trait.HasMethod
 s;%0A%0Aimpo
@@ -628,44 +628,8 @@
 ist;
-%0Aimport java.util.function.Function;
 %0A%0A/*
@@ -1078,381 +1078,476 @@
     
-return traverse(HasImports.class, root, HasImports::getImports);%0A    %7D%0A%0A    /**%0A     * Traverses the tree of models searching for children of a particular type%0A     * by applying the getter method on each found element. Even if the root%0A     * itself is of the searched type, it will not be added.
+final List%3CImport%3E found = new ArrayList%3C%3E(root.getImports());%0A%0A        if (root instanceof HasClasses) %7B%0A            final HasClasses%3C?%3E classes = (HasClasses%3C?%3E) root;%0A            classes.getClasses().forEach(clazz -%3E found.addAll(importsOf(clazz)));%0A        %7D%0A%0A        if (root instanceof HasFields) %7B%0A            final HasFields%3C?%3E classes = (HasFields%3C?%3E) root;
 %0A     
-*%0A
      
-* @param parentType  the parent type%0A     * @param root
+  classes.getFields().forEach(field -%3E found.addAll(importsOf(field)));%0A        %7D%0A%0A
         
 the 
@@ -1546,465 +1546,480 @@
     
-the 
+if (
 root 
-node in the tree%0A     * @param getter      getter method used to obtain a list of children%0A     * @param %3CT%3E         the type of the model searched for%0A     * @param %3CP%3E         the type of the parent%0A     * @return            list of hits%0A     */%0A    public static %3CT, P%3E List%3CT%3E%0A    traverse(Class%3CP%3E parentType, P root, Function%3C? super P, List%3C? extends T%3E%3E getter) %7B%0A        final List%3CT%3E found = new ArrayList%3C%3E(getter.apply(root));%0A%0A
+instanceof HasMethods) %7B%0A            final HasMethods%3C?%3E classes = (HasMethods%3C?%3E) root;%0A            classes.getMethods().forEach(field -%3E found.addAll(importsOf(field)));%0A        %7D%0A%0A        if (root instanceof HasConstructors) %7B%0A            final HasConstructors%3C?%3E classes = (HasConstructors%3C?%3E) root;%0A            classes.getConstructors().forEach(constr -%3E found.addAll(importsOf(constr)));%0A        %7D%0A%0A        if (root instanceof HasInitializers) %7B%0A   
         
 for 
@@ -2018,89 +2018,156 @@
     
-for (
+ 
 final 
-T element : found) %7B%0A            if (parentType.isInstance(element)) %7B
+HasInitializers%3C?%3E classes = (HasInitializers%3C?%3E) root;%0A            classes.getInitializers().forEach(init -%3E found.addAll(importsOf(init)));
 %0A   
@@ -2171,24 +2171,27 @@
         
+%7D%0A%0A
         
 @Suppres
@@ -2186,42 +2186,42 @@
     
-@SuppressWarnings(%22unchecked%22)%0A   
+if (root instanceof HasJavadoc) %7B%0A
     
@@ -2232,91 +2232,137 @@
     
- 
 final 
-P asParent = (P) element;%0A                found.addAll(getter.apply(asParent
+HasJavadoc%3C?%3E classes = (HasJavadoc%3C?%3E) root;%0A            classes.getJavadoc().ifPresent(doc -%3E found.addAll(importsOf(doc)
 ));%0A
@@ -2373,23 +2373,10 @@
     
-    
 %7D%0A
-        %7D
 %0A   

Fix bug with inherited import statements