package com.axis.innovators.box.decompilation.util;

import org.objectweb.asm.*;
import org.objectweb.asm.commons.*;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.jar.*;

/**
 * 混淆工具（未完工）
 * @author tzdwindows 7
 */
public class JarObfuscator {
    private final File jarFile;
    private final Map<String, byte[]> jarEntries = new LinkedHashMap<>();
    private final List<String> whitelistList = new ArrayList<>();
    private final Map<String, String> classMapping = new HashMap<>();
    private final Map<String, String> classNameMapping = new HashMap<>();
    private final Map<String, String> fieldMapping = new HashMap<>();
    private final Map<String, String> methodMapping = new HashMap<>();
    private final ClassNameGenerator classNameGenerator = new ClassNameGenerator();
    private final MemberNameGenerator fieldNameGenerator = new MemberNameGenerator();
    private final MemberNameGenerator methodNameGenerator = new MemberNameGenerator();


    public JarObfuscator(File file) {
        this.jarFile = file;
    }

    public void obfuscate() throws IOException {
        collectObfuscationNames();
        applyObfuscation();
    }

    public void addWhitelist(String classPath) {
        whitelistList.add(classPath.replace(".", "/"));
    }

    private void collectObfuscationNames() throws IOException {
        try (JarFile jar = new JarFile(jarFile)) {
            Enumeration<JarEntry> entries = jar.entries();
            while (entries.hasMoreElements()) {
                JarEntry entry = entries.nextElement();
                String entryName = entry.getName();

                if (entry.isDirectory() || !entryName.endsWith(".class")) {
                    continue;
                }

                try (InputStream is = jar.getInputStream(entry)) {
                    byte[] bytes = readAllBytes(is);
                    ClassReader cr = new ClassReader(bytes);
                    cr.accept(new MappingCollector(), ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);
                }
            }
        }
    }

    private class MappingCollector extends ClassVisitor {
        private String originalClassName;

        public MappingCollector() {
            super(Opcodes.ASM9);
        }

        @Override
        public void visit(int version, int access, String name, String signature,
                          String superName, String[] interfaces) {
            this.originalClassName = name;
            if (whitelistList.contains(originalClassName)){
                classMapping.put(originalClassName, originalClassName);
                return;
            }

            String originalPackage = name.contains("/") ?
                    name.substring(0, name.lastIndexOf('/')) : "";
            String className = name.contains("/") ?
                    name.substring(name.lastIndexOf('/') + 1) : name;

            String obfuscatedPackage = classNameGenerator.generatePackageName(originalPackage);
            String obfuscatedClassName = classNameGenerator.generateClassName(className);
            String newInternalName = obfuscatedPackage.isEmpty() ?
                    obfuscatedClassName : obfuscatedPackage + "/" + obfuscatedClassName;

            classMapping.put(originalClassName, newInternalName);
            classNameMapping.put(originalClassName, obfuscatedClassName);
            super.visit(version, access, name, signature, superName, interfaces);
        }

        @Override
        public FieldVisitor visitField(int access, String name, String descriptor,
                                       String signature, Object value) {
            if (whitelistList.contains(originalClassName)) {
                fieldMapping.put(originalClassName + "#" + name + "#" + descriptor, name);
                return super.visitField(access, name, descriptor, signature, value);
            }
            String newName = fieldNameGenerator.generate(name);
            fieldMapping.put(originalClassName + "#" + name + "#" + descriptor, newName);
            return super.visitField(access, name, descriptor, signature, value);
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor,
                                         String signature, String[] exceptions) {
            if (!"<init>".equals(name) && !"<clinit>".equals(name)
                    || whitelistList.contains(originalClassName)) {
                String newName = methodNameGenerator.generate(name);
                methodMapping.put(originalClassName + "#" + name + "#" + descriptor, newName);
            } else {
                methodMapping.put(originalClassName + "#" + name + "#" + descriptor, name);
            }
            return super.visitMethod(access, name, descriptor, signature, exceptions);
        }
    }

    private void applyObfuscation() throws IOException {
        try (JarFile jar = new JarFile(jarFile)) {
            Enumeration<JarEntry> entries = jar.entries();
            while (entries.hasMoreElements()) {
                JarEntry entry = entries.nextElement();
                String entryName = entry.getName();

                if (entry.isDirectory()) {
                    continue;
                }

                try (InputStream is = jar.getInputStream(entry)) {
                    byte[] bytes = readAllBytes(is);
                    if (entryName.endsWith(".class")) {
                        bytes = processClass(bytes);
                        String newName = classMapping.get(entryName.replace(".class", "")) + ".class";
                        jarEntries.put(newName, bytes);
                    } else {
                        jarEntries.put(entryName, bytes);
                    }
                }
            }
        }
    }

    private static class ClassNameGenerator {
        private final Map<String, String> packageMappings = new HashMap<>();
        private final Map<String, String> classMappings = new HashMap<>();
        private final Map<String, String> pathMappings = new HashMap<>();
        private int packageCounter = 0;
        private int classCounter = 0;

        public String generatePackageName(String originalPath) {
            return pathMappings.computeIfAbsent(originalPath, k -> {
                String[] parts = originalPath.split("/");
                StringBuilder newPath = new StringBuilder();
                for (String part : parts) {
                    String pkg = packageMappings.computeIfAbsent(part, p -> generateShortName(packageCounter++, 'a'));
                    newPath.append(pkg).append('/');
                }
                return newPath.substring(0, newPath.length() - 1);
            });
        }

        public String generateClassName(String originalName) {
            return classMappings.computeIfAbsent(originalName, k -> generateShortName(classCounter++, 'A'));
        }

        private String generateShortName(int index, char base) {
            StringBuilder sb = new StringBuilder();
            do {
                sb.append((char) (base + index % 26));
                index = index / 26 - 1;
            } while (index >= 0);
            return sb.reverse().toString();
        }
    }

    // 专用成员名生成器（处理字段和方法名）
    private static class MemberNameGenerator {
        private final Map<String, String> cache = new HashMap<>();
        private int counter = 0;

        public String generate(String original) {
            return cache.computeIfAbsent(original, k -> generateShortName(counter++));
        }

        private String generateShortName(int index) {
            StringBuilder sb = new StringBuilder();
            do {
                sb.append((char) ('a' + index % 26));
                index = index / 26 - 1;
            } while (index >= 0);
            return sb.reverse().toString();
        }
    }

    private byte[] processClass(byte[] original) {
        ClassReader cr = new ClassReader(original);
        ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_MAXS);

        Remapper remapper = new CustomRemapper();
        ClassVisitor cv = new ClassRemapper(cw, remapper);

        ClassVisitor chain = new ClassVisitor(Opcodes.ASM9, cv) {
            private String originalClassName;

            @Override
            public void visit(int version, int access, String name, String signature,
                              String superName, String[] interfaces) {
                this.originalClassName = name;
                String newInternalName = classMapping.get(name);

                super.visit(version, access, newInternalName,
                        remapper.mapSignature(signature, false),
                        remapper.mapType(superName),
                        interfaces != null ? Arrays.stream(interfaces)
                                .map(remapper::mapType)
                                .toArray(String[]::new) : null);
            }

            @Override
            public FieldVisitor visitField(int access, String name, String descriptor,
                                           String signature, Object value) {
                String newName = fieldMapping.get(originalClassName + "#" + name + "#" + descriptor);
                return super.visitField(access, newName,
                        remapper.mapDesc(descriptor),
                        remapper.mapSignature(signature, true),
                        value);
            }

            @Override
            public MethodVisitor visitMethod(int access, String name, String descriptor,
                                             String signature, String[] exceptions) {

                if ("<init>".equals(name) || "<clinit>".equals(name)) {
                    return super.visitMethod(access, name, descriptor, signature, exceptions);
                }

                String newName = methodMapping.get(originalClassName + "#" + name + "#" + descriptor);

                MethodVisitor mv = super.visitMethod(access, newName,
                        remapper.mapMethodDesc(descriptor),
                        remapper.mapSignature(signature, false),
                        exceptions != null ? Arrays.stream(exceptions)
                                .map(remapper::mapType)
                                .toArray(String[]::new) : null);

                return new MethodVisitor(Opcodes.ASM9, mv) {
                    @Override
                    public void visitMethodInsn(int opcode, String owner, String name,
                                                String descriptor, boolean isInterface) {
                        super.visitMethodInsn(opcode,
                                remapper.mapType(owner),
                                remapper.mapMethodName(owner, name, descriptor),
                                remapper.mapMethodDesc(descriptor),
                                isInterface);
                    }


                    @Override
                    public void visitFieldInsn(int opcode, String owner, String name,
                                               String descriptor) {
                        super.visitFieldInsn(opcode,
                                remapper.mapType(owner),
                                remapper.mapFieldName(owner, name, descriptor),
                                remapper.mapDesc(descriptor));
                    }

                    @Override
                    public void visitTypeInsn(int opcode, String type) {
                        String mappedType = remapper.mapType(type);
                        super.visitTypeInsn(opcode, mappedType);
                    }

                    @Override
                    public void visitLocalVariable(String name, String descriptor,
                                                   String signature, Label start, Label end, int index) {
                        super.visitLocalVariable(name,
                                remapper.mapDesc(descriptor),
                                remapper.mapSignature(signature, true),
                                start, end, index);
                    }

                    @Override
                    public void visitLdcInsn(Object value) {
                        if (value instanceof String) {
                            String encrypted = xorEncrypt((String) value);
                            super.visitLdcInsn(encrypted);
                            super.visitMethodInsn(
                                    Opcodes.INVOKESTATIC,
                                    "StringDecoder",
                                    "decode",
                                    "(Ljava/lang/String;)Ljava/lang/String;",
                                    false
                            );
                        } else {
                            super.visitLdcInsn(value);
                        }
                    }
                };
            }
        };

        cr.accept(chain, ClassReader.EXPAND_FRAMES);
        return cw.toByteArray();
    }

    private class CustomRemapper extends Remapper {
        @Override
        public String mapType(String internalName) {
            return classMapping.getOrDefault(internalName, internalName);
        }

        @Override
        public String mapMethodName(String owner, String name, String descriptor) {
            return methodMapping.getOrDefault(owner + "#" + name + "#" + descriptor, name);
        }

        @Override
        public String mapFieldName(String owner, String name, String descriptor) {
            return fieldMapping.getOrDefault(owner + "#" + name + "#" + descriptor, name);
        }

        @Override
        public String mapSignature(String signature, boolean typeSignature) {
            return super.mapSignature(signature, typeSignature);
        }

        @Override
        public String map(String internalName) {
            return classMapping.getOrDefault(internalName, internalName);
        }
    }

    public void save(String path) throws IOException {
        injectStringDecoder();

        try (JarOutputStream jos = new JarOutputStream(new FileOutputStream(path))) {
            for (Map.Entry<String, byte[]> entry : jarEntries.entrySet()) {
                String name = entry.getKey();
                jos.putNextEntry(new JarEntry(name));
                jos.write(entry.getValue());
                jos.closeEntry();
            }
        }
    }

    private void injectStringDecoder() {
        ClassWriter cw = new ClassWriter(0);
        cw.visit(Opcodes.V1_8, Opcodes.ACC_PUBLIC | Opcodes.ACC_SUPER,
                "StringDecoder", null, "java/lang/Object", null);

        MethodVisitor mv = cw.visitMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC,
                "decode", "(Ljava/lang/String;)Ljava/lang/String;", null, null);
        mv.visitCode();
        mv.visitVarInsn(Opcodes.ALOAD, 0);
        mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/String", "toCharArray", "()[C", false);
        mv.visitVarInsn(Opcodes.ASTORE, 1);
        mv.visitVarInsn(Opcodes.ALOAD, 1);
        mv.visitInsn(Opcodes.ARRAYLENGTH);
        mv.visitVarInsn(Opcodes.ISTORE, 2);
        mv.visitInsn(Opcodes.ICONST_0);
        mv.visitVarInsn(Opcodes.ISTORE, 3);

        Label loop = new Label();
        mv.visitLabel(loop);
        mv.visitVarInsn(Opcodes.ILOAD, 3);
        mv.visitVarInsn(Opcodes.ILOAD, 2);
        Label end = new Label();
        mv.visitJumpInsn(Opcodes.IF_ICMPGE, end);
        mv.visitVarInsn(Opcodes.ALOAD, 1);
        mv.visitVarInsn(Opcodes.ILOAD, 3);
        mv.visitInsn(Opcodes.DUP2);
        mv.visitInsn(Opcodes.CALOAD);
        mv.visitIntInsn(Opcodes.BIPUSH, 0x55);
        mv.visitInsn(Opcodes.IXOR);
        mv.visitInsn(Opcodes.I2C);
        mv.visitInsn(Opcodes.CASTORE);
        mv.visitIincInsn(3, 1);
        mv.visitJumpInsn(Opcodes.GOTO, loop);
        mv.visitLabel(end);
        mv.visitVarInsn(Opcodes.ALOAD, 1);
        mv.visitMethodInsn(Opcodes.INVOKESTATIC, "java/lang/String", "valueOf", "([C)Ljava/lang/String;", false);
        mv.visitInsn(Opcodes.ARETURN);
        mv.visitMaxs(4, 4);
        mv.visitEnd();

        cw.visitEnd();
        jarEntries.put("StringDecoder.class", cw.toByteArray());
    }

    public Map<String, String> getComparisonTable() {
        return Collections.unmodifiableMap(classMapping);
    }

    public void saveMappingTable(String path) throws IOException {
        Properties prop = new Properties();
        classMapping.forEach((k, v) -> {
            //System.out.println("[Class Mapping] " + k + " → " + v);
            prop.setProperty("class:" + k, v);
        });
        fieldMapping.forEach((k, v) -> {
            //System.out.println("[Field Mapping] " + k + " → " + v);
            prop.setProperty("field:" + k, v);
        });
        methodMapping.forEach((k, v) -> {
            //System.out.println("[Method Mapping] " + k + " → " + v);
            prop.setProperty("method:" + k, v);
        });

        try (Writer writer = new OutputStreamWriter(new FileOutputStream(path), StandardCharsets.UTF_8)) {
            prop.store(writer, "Obfuscation Mapping Table");
        }
    }

    private static String xorEncrypt(String input) {
        char[] chars = input.toCharArray();
        for (int i = 0; i < chars.length; i++) {
            chars[i] ^= 0x55;
        }
        return new String(chars);
    }

    private static byte[] readAllBytes(InputStream is) throws IOException {
        ByteArrayOutputStream buffer = new ByteArrayOutputStream();
        byte[] data = new byte[4096];
        int bytesRead;
        while ((bytesRead = is.read(data, 0, data.length)) != -1) {
            buffer.write(data, 0, bytesRead);
        }
        return buffer.toByteArray();
    }

    public static void main(String[] args) throws IOException {
        File inputJar = new File("C:\\Users\\Administrator\\MCreatorWorkspaces\\AxisInnovatorsBox\\build\\libs\\AxisInnovatorsBox-0.0.1.jar");
        JarObfuscator obfuscator = new JarObfuscator(inputJar);
        //obfuscator.addWhitelist("com/axis/innovators/box/Main");
        //obfuscator.addWhitelist("com/axis/innovators/box/plugins/BoxClassLoader");
        //obfuscator.addWhitelist("com/axis/innovators/box/plugins/IClassTransformer");
        //obfuscator.addWhitelist("com/axis/innovators/box/plugins/LoadingCorePlugin");
        //obfuscator.addWhitelist("com/axis/innovators/box/plugins/PluginDescriptor");
        //obfuscator.addWhitelist("com/axis/innovators/box/plugins/PluginDescriptor");
        obfuscator.obfuscate();
        obfuscator.saveMappingTable("C:\\Users\\Administrator\\Desktop\\mapping.properties");
        obfuscator.save("C:\\Users\\Administrator\\Desktop\\obfuscated.jar");
    }
}