import java.io.*;
import java.util.*;
import java.util.jar.Attributes;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.jar.Manifest;

/**
 * @author yuwenbo1
 * @since 2021/3/24 15:23
 * 參考鏈接 {@see https://juejin.cn/post/6869589995066556430}
 */
public class CustomerLoadJarClassLoader extends ClassLoader implements Closeable {
    /**
     * META-INF.MF定義的鏈碼基礎掃描路徑
     */
    private static final String CHAINCODE_BASE_PACKAGE = "Chaincode-Base-Package";
    /**
     * 解壓的臨時包名
     */
    private static final String JAR_TMP_PACKAGE = "chaincode";

    private static final int BUFFER_SIZE = 1024;

    private Map<String, byte[]> map;
    /**
     * 如果用戶不配置,那么默認掃描com.xxx包下面的注解包
     */
    private List<String> packageScanList = Arrays.asList("com.jd.icity");

    private List<String> classNameList = new ArrayList<>(256);
    /**
     * 由頂級加載器和擴展加載器的列表信息
     */
    private List<String> bootAndExtLoaderList = Arrays.asList("sun", "java", "javax", "jdk", "javassist");

    public CustomerLoadJarClassLoader(String jarPath) throws FileNotFoundException {
        /**
         * 設置當前類加載器的父類為當前線程的類加載器,因為默認是使用的是appclassloader,
         * 但是springboot項目的classloader是springboot自定義的classloader(具體可看打包好后的jar文件的META-INF.MF文件,指定的啟動類為springboot的loader),
         * 這就會導致加載對應文件不直接屬于classloader,相關校驗就會失敗
         */
        super(Thread.currentThread().getContextClassLoader());
        if (!jarPath.endsWith(".jar")) {
            throw new IllegalArgumentException("jarFile is not a jar" + jarPath);
        }
        dealChaincodeScanList(jarPath);
        map = new HashMap<>(64);
        unzipJarAndRead(jarPath);
        dealLibJar(jarPath);
    }

    /**
     * 獲取jar名稱路徑,以及對應的需要掃描的類路徑信息
     *
     * @param jarPath 包含jar名稱的文件
     * @return
     * @throws FileNotFoundException
     */
    private void dealChaincodeScanList(String jarPath) throws FileNotFoundException {
        try {
            JarFile jarFile = new JarFile(jarPath);
            Manifest manifest = jarFile.getManifest();

            Attributes mainAttributes = manifest.getMainAttributes();
            //得到需要進行掃描的類信息
            String mainPackage = mainAttributes.getValue(CHAINCODE_BASE_PACKAGE);
            if (null != mainPackage) {
                this.packageScanList = Arrays.asList(mainPackage.split(","));
            }
        } catch (IOException e) {
            throw new FileNotFoundException("jar:" + jarPath + " not exist");
        }
    }

    /**
     * 解壓jar文件,并且進行讀取加載,
     * 如果讀取到類信息了,那么則進行讀取,
     * 如果讀取到的是jar內容,那么則進行解壓,然后進行文件處理,并讀取相應的類文件進行然后進行加載
     *
     * @param jarPath
     */
    private void unzipJarAndRead(String jarPath) {
        //直接解壓jar文件,然后遞歸查詢class文件,以及對應的.jar文件,并加載jar文件進行讀取
        try {
            /**
             * 直接在對應的jarPath路徑下進行解壓即可,因為本來路徑就是唯一的,而且清理的時候也比較方便
             */
            File tmpFolder = new File(jarPath);
            String folderParent = tmpFolder.getParent();
            File folder = new File(folderParent, JAR_TMP_PACKAGE);
            if (!folder.exists()) {
                folder.mkdir();
            }
            // 設置jvm關閉的時候自動刪除
            folder.deleteOnExit();

            JarFile jarFile = new JarFile(jarPath);
            Enumeration<JarEntry> entries = jarFile.entries();
            while (entries.hasMoreElements()) {
                JarEntry jarEntry = entries.nextElement();
                String name = jarEntry.getName();
                if (name.endsWith(".jar")) {
                    dealJarFile(name, folder, jarFile, jarEntry);
                } else if (name.endsWith(".class")) {
                    dealClass(name, jarFile, jarEntry);
                }
                //其他類型的暫時不關心,忽略即可
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 對文件是jar的類進行處理
     *
     * @param name
     * @param libJarPath 依賴的jar包解壓路徑
     */
    private static void dealJarFile(String name, File libJarPath, JarFile jarFile, JarEntry jarEntry) {
        System.out.println("unzip jar file:" + name);

        //對jar文件進行解壓
        String[] split = name.split("/");
        if (split.length > 1) {
            File file = new File(libJarPath, split[split.length - 1]);
            try {
                if (!file.exists()) {
                    boolean success = file.createNewFile();
                    if (!success) {
                        throw new RuntimeException("create file:" + file.getPath() + " exception");
                    }
                }
                unpack(jarFile, jarEntry, file);
                file.deleteOnExit();
            } catch (IOException e) {
                System.out.println("unzip jar exception:" + e.getMessage());
            }

        }
    }

    /**
     * 對jar包進行解壓
     *
     * @param jarFile
     * @param entry
     * @param file
     * @throws IOException
     */
    private static void unpack(JarFile jarFile, JarEntry entry, File file) throws IOException {
        try (InputStream inputStream = jarFile.getInputStream(entry)) {
            try (OutputStream outputStream = new FileOutputStream(file)) {
                byte[] buffer = new byte[BUFFER_SIZE];
                int bytesRead;
                while ((bytesRead = inputStream.read(buffer)) != -1) {
                    outputStream.write(buffer, 0, bytesRead);
                }
                outputStream.flush();
            }
        }
    }

    /**
     * 處理加載類信息的邏輯
     * 如果在map中已經存在了,那么則直接不進行處理,
     * 如果不存在,那么則進行讀取類文件,然后放入map中,并將掃描包列表信息進行判斷,是否是設置的開頭參數,如果有才放入類列表中
     *
     * @param name
     * @param jarFile
     * @param jarEntry
     */
    private void dealClass(String name, JarFile jarFile, JarEntry jarEntry) {
        String className = name.replace(".class", "").replaceAll("/", ".");
        if (!map.containsKey(className)) {
            byte[] b = getClassByte(jarFile, jarEntry);
            if (null != b) {
                map.put(className, b);
                //需要判斷是否是需要進行掃描的包路徑,然后進行添加
                packageScanList.stream().forEach(v -> {
                    if (className.startsWith(v)) {
                        classNameList.add(className);
                    }
                });
            }
        }
    }

    /**
     * 對有依賴的jar包進行處理,加載到內存中
     */
    private void dealLibJar(String jarPath) {
        File tmpFolder = new File(jarPath);
        String folderParent = tmpFolder.getParent();
        File folder = new File(folderParent, JAR_TMP_PACKAGE);
        if (!folder.exists()) {
            System.out.println("no lib jar exist,will be return");
            return;
        }

        File[] files = folder.listFiles();
        if (files.length == 0) {
//            System.out.println("jar list is empty,will be return");
            return;
        }
        for (File file : files) {
            JarFile jarFile = null;
            try {
                jarFile = new JarFile(file);
                Enumeration<JarEntry> entries = jarFile.entries();
                while (entries.hasMoreElements()) {
                    JarEntry jarEntry = entries.nextElement();
                    String name = jarEntry.getName();
                    /**
                     * 針對lib里面的jar包,只會解析.class文件,對其他不關心
                     */
                    if (name.endsWith(".class")) {
                        dealClass(name, jarFile, jarEntry);
                    }
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    public List<String> getClassNameList() {
        return classNameList;
    }

    @Override
    protected Class<?> findClass(String name) throws ClassNotFoundException {
        byte[] b;
        if (map.containsKey(name)) {
            b = map.get(name);
        } else {
            throw new ClassNotFoundException(name);
        }
        return defineClass(name, b, 0, b.length);
    }

    /**
     * 獲取某個jar里面的信息
     *
     * @param jarFile
     * @param jarEntry
     * @return
     */
    private byte[] getClassByte(JarFile jarFile, JarEntry jarEntry) {
        try (InputStream input = jarFile.getInputStream(jarEntry)) {
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            int bufferSize = 4096;
            byte[] buffer = new byte[bufferSize];
            int bytesNumRead = 0;
            while ((bytesNumRead = input.read(buffer)) != -1) {
                baos.write(buffer, 0, bytesNumRead);
            }
            return baos.toByteArray();
        } catch (FileNotFoundException e) {
        } catch (IOException e) {
        }
        return null;
    }

    @Override
    public void close() throws IOException {
        if (!map.isEmpty()) {
            map.clear();
        }
        if (!classNameList.isEmpty()) {
            classNameList.clear();
        }
    }
}