package com.novitechie;

import com.janetfilter.core.commons.DebugInfo;
import com.janetfilter.core.plugin.MyTransformer;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.tree.*;

import java.util.ListIterator;

import static org.objectweb.asm.Opcodes.*;

public class LicensingFacadeTransformer implements MyTransformer {

    private static final int ASM_VERSION = ASM9;

    @Override
    public String getHookClassName() {
        return "com/intellij/ui/LicensingFacade";
    }

    @Override
    public byte[] transform(String className, byte[] classBytes, int order) throws Exception {
        if (classBytes == null || classBytes.length == 0) {
            return classBytes;
        }
        try {
            ClassReader reader = new ClassReader(classBytes);
            ClassNode node = new ClassNode(ASM_VERSION);
            reader.accept(node, 0);

            for (MethodNode m : node.methods) {
                if ("getLicenseExpirationDate".equals(m.name)) {
                    InsnList list = new InsnList();
                    LabelNode L0 = new LabelNode();
                    list.add(new MethodInsnNode(INVOKESTATIC, "com/novitechie/rules/StackTraceRule", "hook", "()Ljava/util/Date;", false));
                    list.add(new InsnNode(DUP));
                    list.add(new JumpInsnNode(IFNULL, L0));
                    list.add(new InsnNode(ARETURN));
                    list.add(L0);
                    list.add(new InsnNode(POP));
                    m.instructions.insert(list);
                } else if ("toJson".equals(m.name) && "()Ljava/lang/String;".equals(m.desc)) {
                    insertToJsonFilterSecure(m);
                } else if ("getInstance".equals(m.name) && "()Lcom/intellij/ui/LicensingFacade;".equals(m.desc)) {
                    insertGetInstanceGuard(m);
                }
            }

            ClassWriter writer = new SafeClassWriter(null, null, ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
            node.accept(writer);
            return writer.toByteArray();
        } catch (Exception e) {
            DebugInfo.warn("[PRIVACY] Failed to transform LicensingFacade, skipping", e);
            return classBytes;
        }
    }

    private void insertGetInstanceGuard(MethodNode m) {
        if (m.instructions == null || m.instructions.size() == 0) {
            return;
        }
        int storeIdx = m.maxLocals;
        m.maxLocals += 1;

        ListIterator<AbstractInsnNode> iterator = m.instructions.iterator();
        while (iterator.hasNext()) {
            AbstractInsnNode current = iterator.next();
            if (current.getOpcode() == ARETURN) {
                LabelNode returnOriginal = new LabelNode();
                InsnList patch = new InsnList();
                patch.add(new VarInsnNode(ASTORE, storeIdx));
                patch.add(new MethodInsnNode(INVOKESTATIC, "com/novitechie/rules/StackTraceRule", "check", "()Z", false));
                patch.add(new JumpInsnNode(IFEQ, returnOriginal));
                patch.add(new InsnNode(ACONST_NULL));
                patch.add(new InsnNode(ARETURN));
                patch.add(returnOriginal);
                patch.add(new VarInsnNode(ALOAD, storeIdx));
                m.instructions.insertBefore(current, patch);
            }
        }
    }

    private void insertToJsonFilterSecure(MethodNode m) {
        if (m.instructions == null || m.instructions.size() == 0) {
            return;
        }

        int originalJsonIdx = m.maxLocals;
        int filteredJsonIdx = originalJsonIdx + 1;
        m.maxLocals += 2;

        ListIterator<AbstractInsnNode> iterator = m.instructions.iterator();
        while (iterator.hasNext()) {
            AbstractInsnNode current = iterator.next();

            if (current.getOpcode() == ARETURN) {
                LabelNode useOriginalLabel = new LabelNode();
                LabelNode endLabel = new LabelNode();
                InsnList patch = new InsnList();
                patch.add(new VarInsnNode(ASTORE, originalJsonIdx));
                patch.add(new VarInsnNode(ALOAD, originalJsonIdx));
                patch.add(new MethodInsnNode(INVOKESTATIC, "com/novitechie/rules/LicenseRule", "filterJson", "(Ljava/lang/String;)Ljava/lang/String;", false));
                patch.add(new VarInsnNode(ASTORE, filteredJsonIdx));
                patch.add(new VarInsnNode(ALOAD, filteredJsonIdx));
                patch.add(new JumpInsnNode(IFNULL, useOriginalLabel));
                patch.add(new VarInsnNode(ALOAD, filteredJsonIdx));
                patch.add(new JumpInsnNode(GOTO, endLabel));
                patch.add(useOriginalLabel);
                patch.add(new VarInsnNode(ALOAD, originalJsonIdx));
                patch.add(endLabel);
                m.instructions.insertBefore(current, patch);
            }
        }
    }

}