package study.seata.loopback;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.MessageDigest;
import java.security.ProtectionDomain;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;

/**
 * 实验隔离：只修改 Seata 2.5.0 TC 的 RPC 绑定地址，不修改事务逻辑。
 * 与原版服务的差异必须在实验文档和证据中记录。
 */
public final class LoopbackAgent {
    private static final String TARGET = "org/apache/seata/core/rpc/netty/NettyServerBootstrap";
    private static final String SHA256 = "83fb0319743b67ad27b73c1589cbf22a2ad3a8c4005e4ffb6d8fc7a0237177ae";
    private static final String BOOTSTRAP = "io/netty/bootstrap/ServerBootstrap";
    private static final String OLD_DESCRIPTOR = "(I)Lio/netty/channel/ChannelFuture;";
    private static final String NEW_DESCRIPTOR = "(Ljava/lang/String;I)Lio/netty/channel/ChannelFuture;";

    private LoopbackAgent() {}

    public static void premain(String argument, Instrumentation instrumentation) {
        instrumentation.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String name, Class<?> type,
                                    ProtectionDomain domain, byte[] bytes) {
                if (!TARGET.equals(name)) return null;
                try {
                    byte[] changed = rewrite(bytes);
                    System.err.println("U38 loopback agent: Seata RPC bind 已限定为 127.0.0.1");
                    return changed;
                } catch (Throwable failure) {
                    System.err.println("U38 loopback agent 失败，禁止启动: " + failure);
                    Runtime.getRuntime().halt(78);
                    return null;
                }
            }
        });
    }

    private static byte[] rewrite(byte[] bytes) throws Exception {
        if (!SHA256.equals(hex(MessageDigest.getInstance("SHA-256").digest(bytes)))) {
            throw new IllegalArgumentException("Seata RPC 类 SHA-256 不匹配固定 2.5.0 发行包");
        }
        final int[] changed = {0};
        ClassReader reader = new ClassReader(bytes);
        ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_MAXS);
        reader.accept(new ClassVisitor(Opcodes.ASM6, writer) {
            @Override
            public MethodVisitor visitMethod(int access, String name, String descriptor,
                                             String signature, String[] exceptions) {
                MethodVisitor delegate = super.visitMethod(access, name, descriptor, signature, exceptions);
                if (!"start".equals(name) || !"()V".equals(descriptor)) return delegate;
                return new MethodVisitor(Opcodes.ASM6, delegate) {
                    @Override
                    public void visitMethodInsn(int opcode, String owner, String name,
                                                String descriptor, boolean isInterface) {
                        if (opcode == Opcodes.INVOKEVIRTUAL && BOOTSTRAP.equals(owner)
                                && "bind".equals(name) && OLD_DESCRIPTOR.equals(descriptor)) {
                            // 原栈 [bootstrap, port]；换为 [bootstrap, host, port]。
                            super.visitLdcInsn("127.0.0.1");
                            super.visitInsn(Opcodes.SWAP);
                            super.visitMethodInsn(opcode, owner, name, NEW_DESCRIPTOR, isInterface);
                            changed[0]++;
                        } else {
                            super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
                        }
                    }
                };
            }
        }, 0);
        if (changed[0] != 1) throw new IllegalStateException("预期恰好改写一个 bind(int)，实际=" + changed[0]);
        return writer.toByteArray();
    }

    private static String hex(byte[] bytes) {
        StringBuilder result = new StringBuilder(bytes.length * 2);
        for (byte b : bytes) result.append(String.format("%02x", b & 0xff));
        return result.toString();
    }

    /** 离线检查输入和替换结果；不启动网络服务。 */
    public static void main(String[] args) throws Exception {
        if (args.length != 1) throw new IllegalArgumentException("需要固定版本的目标 class 文件");
        byte[] rewritten = rewrite(Files.readAllBytes(Paths.get(args[0])));
        final int[] calls = {0};
        new ClassReader(rewritten).accept(new ClassVisitor(Opcodes.ASM6) {
            @Override
            public MethodVisitor visitMethod(int access, String name, String descriptor,
                                             String signature, String[] exceptions) {
                if (!"start".equals(name) || !"()V".equals(descriptor)) return null;
                return new MethodVisitor(Opcodes.ASM6) {
                    @Override
                    public void visitMethodInsn(int opcode, String owner, String name,
                                                String descriptor, boolean isInterface) {
                        if (opcode == Opcodes.INVOKEVIRTUAL && BOOTSTRAP.equals(owner) && "bind".equals(name)) {
                            if (!NEW_DESCRIPTOR.equals(descriptor)) throw new IllegalStateException("仍存在其他 bind 调用");
                            calls[0]++;
                        }
                    }
                };
            }
        }, 0);
        if (calls[0] != 1) throw new IllegalStateException("替换结果的 bind 调用数=" + calls[0]);
        System.out.println("U38 loopback agent 静态校验通过");
    }
}