/*
 * @author Ranides Atterwim {@literal <ranides@gmail.com>}
 * @copyright Ranides Atterwim
 * @license WTFPL
 * @url http://ranides.net/projects/assira
 */
package net.ranides.assira.reflection.util;

import net.ranides.assira.collection.iterators.EnumerationUtils;
import net.ranides.assira.collection.iterators.ForwardSpliterator;
import net.ranides.assira.collection.iterators.IteratorUtils;
import net.ranides.assira.collection.query.CQuery;
import net.ranides.assira.generic.ValueUtils;
import net.ranides.assira.io.ZipReader;
import net.ranides.assira.io.uri.URIUtils;
import net.ranides.assira.text.StringUtils;
import net.ranides.assira.trace.ExceptionUtils;
import org.slf4j.Logger;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.net.URL;
import java.net.URLClassLoader;
import java.nio.charset.Charset;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.jar.JarEntry;
import java.util.jar.JarInputStream;

/**
 * This class can be used for discovery of all classes in classpath.
 *
 * Because of technical restrictions, by default you can't scan all packages, you have to
 * define root package. Scanner will look into it and all other subpackages.
 *
 * It is possible to scan all classpath, if application is loaded from single JAR.
 * Internally, scanner will deduce JAR filename and read it's content.
 *
 * Package scanner tries to load class names lazily. It does not store results in memory.
 * The same applies to classes: it will convert detected names to Class on demand.
 *
 * If discovered classname can't be converted to Class, configured fallback "resolver" will be called.
 */
public final class PackageScanner {

    private BiFunction<String, Throwable, Class<?>> handler;

    private String root;

    private ClassLoader loader;

    private PackageScanner() {
        // do nothing
    }

    /**
     * Creates new instance of PackageScanner.
     * PackageScanner objects behave in a way similar to streams or CQueries.
     *
     * PackageScanner must be configured before use.
     * If you don't define root package and application is not loaded from single JAR,
     * result will be empty.
     *
     * @return PackageScanner
     */
    public static PackageScanner scanner() {
        return new PackageScanner();
    }

    /**
     * Configures root package used as entry point for class discovery.
     *
     * @param root root
     * @return this
     */
    public PackageScanner scope(String root) {
        this.root = StringUtils.replace(root, '.', '/');
        return this;
    }

    /**
     * Configures root directory used as entry point for class discovery.
     * Root directory is essentially the same as "root package", but it uses "/" separators instead of "."
     *
     * @param root root
     * @return this
     */
    public PackageScanner directory(String root) {
        this.root = root;
        return this;
    }

    /**
     * Configures classloader used to convert class names to "Class" objects
     *
     * @param loader loader
     * @return this
     */
    public PackageScanner loader(ClassLoader loader) {
        this.loader = loader;
        return this;
    }

    /**
     * Configures fall-back handler, invoked if classloader can't convert class name.
     *
     * Handler can use some custom strategy (for example bytecode generation) to create requested class.
     * If handler returns "null", then scanner assumes it was impossible to load class anyway.
     *
     * @param handler handler
     * @return this
     */
    public PackageScanner resolver(BiFunction<String, Throwable, Class<?>> handler) {
        this.handler = handler;
        return this;
    }

    /**
     * Configures fall-back handler, invoked if classloader can't convert class name.
     *
     * This handler can't be used to create class. It receives exception from classloader, but
     * doesn't return any value. After exectution, scanner assumes it was impossible to load class anyway.
     *
     * @param handler handler
     * @return this
     */
    public PackageScanner logger(BiConsumer<String, Throwable> handler) {
        this.handler = (name,error) -> { handler.accept(name, error); return null; };
        return this;
    }

    /**
     * Configures fall-back handler, invoked if classloader can't convert class name.
     *
     * It logs information about class name, and passes control to the scanner which assumes
     * it was impossible to load class anyway.
     *
     * @param logger logger
     * @return this
     */
    public PackageScanner logger(Logger logger) {
        this.handler = (name,error) -> { logger.error("Can't load class: " + name, error); return null; };
        return this;
    }

    /**
     * Returns query with detected classes.
     * It uses configured classloader extensively, to resolve class names to "Class" objects.
     *
     * @return CQuery
     * @throws IOException IOException
     */
    public CQuery<Class<?>> types() throws IOException {
        return names().map(toClass(getLoader())).filter(Objects::nonNull);
    }

    /**
     * Returns query with detected class names.
     *
     * @return CQuery
     * @throws IOException IOException
     */
    public CQuery<String> names() throws IOException {
        return files().filter(this::isClass).map(PackageScanner::toClassName);
    }

    /**
     * Returns query with detected class files (".class").
     *
     * @return CQuery
     * @throws IOException IOException
     */
    public CQuery<String> files() throws IOException {
        if(root == null && loader instanceof URLClassLoader) {
            return CQuery.from().spliterator(() -> new PackageSpliterator(root, ((URLClassLoader) this.loader).getURLs()));
        }
        if(root == null && loader == null) {
            return CQuery.from().spliterator(() -> new PackageSpliterator());
        }
        return CQuery.from().spliterator(() -> new PackageSpliterator(root, getLoader()));
    }

    private ClassLoader getLoader() {
        return ValueUtils.or(loader, ClassLoaderUtils.getDefault());
    }

    private static final class PackageSpliterator extends ForwardSpliterator<String> {

        private final String nroot;
        private final String croot;
        private final Iterator<String> resources;

        private Spliterator<String> current;

        public PackageSpliterator() {
            this.nroot = "";
            this.croot = "";
            this.resources = Collections.singleton(PackageScanner.class.getProtectionDomain().getCodeSource().getLocation().getPath()+"!").iterator();
            this.current = Spliterators.emptySpliterator();
        }

        public PackageSpliterator(String root, ClassLoader loader) throws IOException {
            this.nroot = (root==null)? ":" : (root + "/");
            this.croot = root;
            this.resources = EnumerationUtils.map(loader.getResources(croot), URL::getPath);
            this.current = Spliterators.emptySpliterator();
        }

        public PackageSpliterator(String root, URL[] urls) {
            this.nroot = (root==null)? ":" : (root + "/");
            this.croot = root;
            this.resources =  IteratorUtils.map(Arrays.asList(urls).iterator(), URL::getPath);
            this.current = Spliterators.emptySpliterator();
        }

        @Override
        public boolean tryAdvance(Consumer<? super String> action) {
            return current.tryAdvance(action) || tryAdvanceResource(action);
        }

        private boolean tryAdvanceResource(Consumer<? super String> action) {
            if(!resources.hasNext()) {
                return false;
            }
            String uri = URIUtils.decodeURL(resources.next(), Charset.defaultCharset());
            File file = new File(uri);
            current = file.isDirectory() ? forDirectory(file.toPath()) : forJAR(uri);
            return current.tryAdvance(action);
        }

        private Spliterator<String> forJAR(String uri) {
            try {
                String jar = StringUtils.between(uri, ":", "!");
                if(jar.isEmpty()) {
                    jar = uri;
                }
                return new JARSpliterator(ZipReader.openJAR(new FileInputStream(jar)));
            } catch (IOException cause) {
                throw ExceptionUtils.rethrow(cause);
            }
        }

        private Spliterator<String> forDirectory(Path file) {
            try {
                return Files
                        .walk(file)
                        .filter(p -> !Files.isDirectory(p))
                        .map(s ->  StringUtils.replace(nroot + file.relativize(s), '\\','/'))
                        .spliterator();
            } catch(IOException cause) {
                throw ExceptionUtils.rethrow(cause);
            }
        }

        private final class JARSpliterator extends ForwardSpliterator<String> {

            private final JarInputStream istream;

            public JARSpliterator(JarInputStream istream) throws IOException {
                this.istream = istream;
            }

            @Override
            public boolean tryAdvance(Consumer<? super String> action) {
                try {
                    JarEntry entry;
                    while(null != (entry = istream.getNextJarEntry())) {
                        if((null == croot || entry.getName().startsWith(croot)) && !entry.isDirectory()) {
                            action.accept(entry.getName());
                            return true;
                        }
                    }
                    istream.close();
                    return false;
                } catch(IOException cause) {
                    throw ExceptionUtils.rethrow(cause);
                }
            }

        }

    }

    private boolean isClass(String path) {
        return path.endsWith(".class");
    }
    
    private Function<String, Class<?>> toClass(ClassLoader loader) {
        return name -> {
            try {
                return Class.forName(name, false, loader);
            } catch(ClassNotFoundException | NoClassDefFoundError cause) {
                if(handler != null) {
                    return handler.apply(name, cause);
                }
                return null;
            }
        };
    }

    private static String toClassName(String path) {
        String fname = StringUtils.removeSuffix(path, ".class");
        return StringUtils.replace(StringUtils.replace(fname, '/', '.'), '\\', '.');
    }

}
