import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.ClassNode;
import org.objectweb.asm.tree.InsnList;
import org.objectweb.asm.tree.MethodInsnNode;
import org.objectweb.asm.tree.MethodNode;
import org.objectweb.asm.tree.VarInsnNode;

import java.nio.file.Files;
import java.nio.file.Paths;

/**
 * Makes all of TDWCM's payload registrations optional.
 *
 * Upstream bug: net.dumonde.tdwcm.networking.NetworkRegistry.register() does
 *   registrar = event.registrar("1.0");
 *   MESSAGES.forEach((t, m) -> registrar.playBidirectional(t, ...));
 * — every channel (entity_rotation_sync + ~30 more) is REQUIRED, so clients
 * carrying the mod cannot join any server without it ("channel missing").
 *
 * Fix: insert PayloadRegistrar.optional() right after registrar creation:
 *   registrar = event.registrar("1.0").optional();
 * Optional payloads are skipped during login negotiation when absent on the
 * other side, and client→server messages are dropped server-side instead of
 * failing handshake. Exactly what a client-side-only mod wants.
 */
public class TdwcmPatcher {
    public static void main(String[] args) throws Exception {
        if (args.length != 2) {
            System.err.println("usage: TdwcmPatcher <NetworkRegistry.class-in> <class-out>");
            System.exit(2);
        }
        byte[] bytes = Files.readAllBytes(Paths.get(args[0]));
        ClassNode cn = new ClassNode();
        new ClassReader(bytes).accept(cn, 0);

        boolean patched = false;
        for (MethodNode mn : cn.methods) {
            if (!"register".equals(mn.name)) continue;
            for (AbstractInsnNode n : mn.instructions) {
                if (!(n instanceof MethodInsnNode mi)) continue;
                if (!"registrar".equals(mi.name)) continue;
                if (!(mi.getNext() instanceof VarInsnNode store) || store.getOpcode() != Opcodes.ASTORE) continue;

                InsnList ins = new InsnList();
                ins.add(new MethodInsnNode(
                        Opcodes.INVOKEVIRTUAL,
                        "net/neoforged/neoforge/network/registration/PayloadRegistrar",
                        "optional",
                        "()Lnet/neoforged/neoforge/network/registration/PayloadRegistrar;",
                        false));
                mn.instructions.insert(mi, ins);
                patched = true;
                break;
            }
        }
        if (!patched) throw new IllegalStateException("registrar(...) call site not found");

        // COMPUTE_FRAMES: offsets shifted by the insertion; classpath must contain
        // the original jar + neoforge universal jar for hierarchy resolution.
        ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
        cn.accept(cw);
        Files.write(Paths.get(args[1]), cw.toByteArray());
        System.out.println("patched " + cn.name + ".register: registrar(...).optional() inserted");
    }
}