package cn.reghao.tnb.sb.util;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.net.URL;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.jar.JarEntry;
import java.util.jar.JarInputStream;

/**
 * 扫描指定包路径下的所有 class 文件
 *
 * @author reghao
 * @date 2020-09-24 15:41:07
 */
public class PackageScanner {
    public List<Class<?>> doScan(Class<?> clazz, String basePackage) throws IOException {
        List<Class<?>> classList = new ArrayList<>();
        String pkgPath = basePackage.replace(".", "/");
        ClassLoader cl = Thread.currentThread().getContextClassLoader();

        URL url = cl.getResource(pkgPath);
        assert url != null;
        File file = new File(rootPath(url));
        String protocol = url.getProtocol();
        if ("jar".equals(protocol)) {
            readFromJarFile(file, pkgPath, classList);
        } else if ("file".equals(protocol)) {
            if (file.isDirectory()) {
                readFromDir(file, clazz, classList);
            } else {
                addClazz(file.getAbsolutePath(), clazz, classList);
            }
        }
        return classList;
    }

    private String rootPath(URL url) {
        String path = url.getPath();
        int pos = path.indexOf("!");
        if (pos == -1) {
            return path.replace("file:", "");
        } else {
            return path.substring(0, pos).replace("file:", "");
        }
    }

    private void readFromDir(File dir, Class<?> clazz, List<Class<?>> classList) {
        File[] files = dir.listFiles();
        assert files != null;
        Arrays.asList(files).forEach(file -> {
            if (file.isDirectory()) {
                readFromDir(file, clazz, classList);
            } else {
                String path = file.getPath();
                addClazz(path, clazz, classList);
            }
        });
    }

    private void addClazz(String filepath, Class<?> clazz, List<Class<?>> classList) {
        String tmp = filepath.replace(clazz.getResource("/").getPath(), "");
        String className = tmp.replace("/", ".")
            .replace(".class", "");
        try {
            classList.add(Class.forName(className));
        } catch (ClassNotFoundException e) {
            e.printStackTrace();
        }
    }

    private void readFromJarFile(File jarFile, String pkgPath, List<Class<?>> classList) throws IOException {
        if (!jarFile.exists()) {
            System.out.println("jar 文件 " + jarFile.getAbsolutePath() + " 不存在...");
        }

        JarInputStream in = new JarInputStream(new FileInputStream(jarFile));
        JarEntry entry = in.getNextJarEntry();
        while (entry != null) {
            String name = entry.getName();
            if (name.startsWith(pkgPath) && name.endsWith(".class")) {
                String pkg = name.replace("/", ".").replace(".class", "");
                try {
                    Class<?> clazz = Class.forName(pkg);
                    classList.add(clazz);
                } catch (ClassNotFoundException e) {
                    e.printStackTrace();
                }
            }
            entry = in.getNextJarEntry();
        }
    }
}
