package top.r3944realms.lib39.core.lang;

import java.io.*;
import java.net.URL;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.HashMap;
import java.util.Map;

/**
 * The type Encrypted class loader.
 */
public class EncryptedClassLoader extends ClassLoader {
    static {
        System.loadLibrary("ClassEncrypt");
    }

    private native byte[] decryptClass(byte[] encryptedData, String key);

    private final String encryptedClassPath;
    private final String decryptionKey;

    // 缓存已加载的类字节码，避免重复加载
    private final Map<String, byte[]> classBytesCache = new HashMap<>();

    /**
     * Instantiates a new Encrypted class loader.
     *
     * @param encryptedClassPath the encrypted class path
     * @param key                the key
     */
    public EncryptedClassLoader(String encryptedClassPath, String key) {
        this.encryptedClassPath = encryptedClassPath;
        this.decryptionKey = key;
    }

    /**
     * Instantiates a new Encrypted class loader.
     *
     * @param encryptedClassPath the encrypted class path
     * @param key                the key
     * @param parent             the parent
     */
    public EncryptedClassLoader(String encryptedClassPath, String key, ClassLoader parent) {
        super(parent);
        this.encryptedClassPath = encryptedClassPath;
        this.decryptionKey = key;
    }

    @Override
    protected Class<?> findClass(String name) throws ClassNotFoundException {
        try {
            // 从缓存获取或加载类字节码
            byte[] classData;
            synchronized (classBytesCache) {
                classData = classBytesCache.get(name);
                if (classData == null) {
                    // 1. 读取加密的class文件
                    byte[] encryptedData = loadEncryptedClass(name);

                    // 2. 使用JNI解密
                    classData = decryptClass(encryptedData, decryptionKey);

                    // 3. 验证解密后的数据是否是有效的class文件
                    if (!isValidClass(classData)) {
                        throw new ClassNotFoundException("Invalid class data after decryption");
                    }

                    // 缓存类字节码
                    classBytesCache.put(name, classData);
                }
            }

            // 4. 定义类
            return defineClass(name, classData, 0, classData.length);

        } catch (Exception e) {
            throw new ClassNotFoundException("Failed to load encrypted class: " + name, e);
        }
    }

    private byte [] loadEncryptedClass(String className) throws IOException {
        String path = className.replace('.', File.separatorChar) + ".class";
        Path fullPath = Paths.get(encryptedClassPath, path);

        if (!Files.exists(fullPath)) {
            // 尝试寻找内部类
            int dollarIndex = className.lastIndexOf('$');
            if (dollarIndex != -1) {
                String outerClass = className.substring(0, dollarIndex);
                path = outerClass.replace('.', File.separatorChar) +
                        "$" + className.substring(dollarIndex + 1) + ".class";
                fullPath = Paths.get(encryptedClassPath, path);
            }
        }

        if (!Files.exists(fullPath)) {
            // 尝试其他可能的文件扩展名
            fullPath = Paths.get(encryptedClassPath, className.replace('.', File.separatorChar) + ".enc");
            if (!Files.exists(fullPath)) {
                throw new FileNotFoundException("Encrypted class not found: " + className);
            }
        }

        return Files.readAllBytes(fullPath);
    }

    private boolean isValidClass(byte [] data) {
        // Java class文件的魔数是0xCAFEBABE
        return data.length >= 4 &&
                data[0] == (byte)0xCA &&
                data[1] == (byte)0xFE &&
                data[2] == (byte)0xBA &&
                data[3] == (byte)0xBE;
    }

    @Override
    protected Class<?> loadClass(String name, boolean resolve)
            throws ClassNotFoundException {
        // 优先检查是否已加载
        Class<?> clazz = findLoadedClass(name);
        if (clazz != null) {
            return clazz;
        }

        // Java核心类库使用父加载器
        if (name.startsWith("java.") || name.startsWith("javax.") ||
                name.startsWith("sun.") || name.startsWith("jdk.")) {
            return super.loadClass(name, resolve);
        }

        try {
            // 尝试用自定义ClassLoader加载
            clazz = findClass(name);
        } catch (ClassNotFoundException e) {
            // 如果找不到，委托给父加载器
            clazz = super.loadClass(name, resolve);
        }

        if (resolve) {
            resolveClass(clazz);
        }

        return clazz;
    }

    /**
     * Get class bytes byte [ ].
     *
     * @param className the class name
     * @return the byte [ ]
     * @throws ClassNotFoundException the class not found exception
     */
// 添加获取类字节码的方法
    public byte[] getClassBytes(String className) throws ClassNotFoundException {
        synchronized (classBytesCache) {
            byte[] bytes = classBytesCache.get(className);
            if (bytes == null) {
                // 触发类加载以填充缓存
                loadClass(className);
                bytes = classBytesCache.get(className);
            }
            return bytes != null ? bytes.clone() : null; // 返回副本
        }
    }

    @Override
    public InputStream getResourceAsStream(String name) {
        try {
            // 处理.class资源请求
            if (name.endsWith(".class")) {
                String className = name.substring(0, name.length() - 6)
                        .replace('/', '.');
                byte[] classData = getClassBytes(className);
                if (classData != null) {
                    return new ByteArrayInputStream(classData);
                }
            }

            // 处理其他资源文件
            Path resourcePath = Paths.get(encryptedClassPath, name);
            if (Files.exists(resourcePath)) {
                return Files.newInputStream(resourcePath);
            }

        } catch (Exception e) {
            // 忽略异常，返回null让父加载器处理
        }

        // 委托给父加载器
        return super.getResourceAsStream(name);
    }

    @Override
    public URL getResource(String name) {
        try {
            Path resourcePath = Paths.get(encryptedClassPath, name);
            if (Files.exists(resourcePath)) {
                return resourcePath.toUri().toURL();
            }
        } catch (Exception e) {
            // 忽略异常
        }
        return super.getResource(name);
    }

    /**
     * Clear cache.
     */
    public void clearCache() {
        synchronized (classBytesCache) {
            classBytesCache.clear();
        }
    }

    /**
     * Clear cache.
     *
     * @param className the class name
     */
    public void clearCache(String className) {
        synchronized (classBytesCache) {
            classBytesCache.remove(className);
        }
    }
}