package liquibase.snapshot;

import liquibase.logging.LogFactory;
import liquibase.servicelocator.ServiceLocator;
import liquibase.structure.DatabaseObject;
import liquibase.structure.core.Catalog;
import liquibase.structure.core.Schema;
import liquibase.util.StringUtils;

import java.util.Arrays;
import java.util.HashSet;
import java.util.Set;

public class SnapshotControl {

    private Set<Class<? extends DatabaseObject>> types;
    private static Set<Class<? extends DatabaseObject>> defaultTypes;

    public SnapshotControl() {
        this.types = getDefaultTypes();
    }

    public SnapshotControl(Class<? extends DatabaseObject>... types) {
        if (types == null || types.length == 0) {
            this.types = getDefaultTypes();
        } else {
            this.types = new HashSet<Class<? extends DatabaseObject>>(Arrays.asList(types));
        }
    }

    public SnapshotControl(String types) {
        this.types = readTypesString(types);
    }

    private Set<Class<? extends DatabaseObject>> readTypesString(String typesString) {
        if (StringUtils.trimToNull(typesString) == null) {
            return getDefaultTypes();
        } else {
            Set<Class<? extends DatabaseObject>> returnSet = new HashSet<Class<? extends DatabaseObject>>();

            Set<String> typesToInclude = new HashSet<String>(Arrays.asList(typesString.toLowerCase().split("\\s*,\\s*")));

            Class<? extends DatabaseObject>[] classes = ServiceLocator.getInstance().findClasses(DatabaseObject.class);
            for (Class<? extends DatabaseObject> clazz : classes) {
                if (typesToInclude.contains(clazz.getSimpleName().toLowerCase()) || typesToInclude.contains(clazz.getSimpleName().toLowerCase()+"s")) {
                    returnSet.add(clazz);
                }
            }
            return returnSet;
        }
    }

    private Set<Class<? extends DatabaseObject>> getDefaultTypes() {
        if (defaultTypes == null) {
            Set<Class<? extends DatabaseObject>> set = new HashSet<Class<? extends DatabaseObject>>();

            Class<? extends DatabaseObject>[] classes = ServiceLocator.getInstance().findClasses(DatabaseObject.class);
            for (Class<? extends DatabaseObject> clazz : classes) {
                try {
                    if (clazz.newInstance().snapshotByDefault()) {
                        set.add(clazz);
                    }
                } catch (Exception e) {
                    LogFactory.getLogger().info("Cannot construct "+clazz.getName()+" to determine if it should be included in the snapshot by default");
                }
            }

            defaultTypes = set;
        }
        return defaultTypes;
    }

    public Set<Class<? extends DatabaseObject>> getTypesToInclude() {
        return types;
    }

    public boolean shouldInclude(Class<? extends DatabaseObject> type) {
        return type.equals(Catalog.class) || types.contains(type);
    }
}

@@ -22,16 +22,73 @@
 pshot;%0A%0A
+import liquibase.exception.UnexpectedLiquibaseException;%0A
 import l
@@ -1455,16 +1455,93 @@
 %5Cs*%22)));
+%0A            Set%3CString%3E typesNotFound = new HashSet%3CString%3E(typesToInclude);
 %0A%0A      
@@ -1924,33 +1924,376 @@
                 
-%7D
+    typesNotFound.remove(clazz.getSimpleName().toLowerCase());%0A                    typesNotFound.remove(clazz.getSimpleName().toLowerCase()+%22s%22);%0A                %7D%0A            %7D%0A            if (typesNotFound.size() %3E 0) %7B%0A                throw new UnexpectedLiquibaseException(%22Unknown snapshot type(s) %22+StringUtils.join(typesNotFound, %22, %22));
 %0A            %7D%0A 

Throw error if unknown snapshot type