import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.ClassNode;
import org.objectweb.asm.tree.InsnList;
import org.objectweb.asm.tree.InsnNode;
import org.objectweb.asm.tree.LdcInsnNode;
import org.objectweb.asm.tree.MethodInsnNode;
import org.objectweb.asm.tree.MethodNode;
import org.objectweb.asm.tree.VarInsnNode;

/**
 * Patches zombie.vehicles.BaseVehicle for Landtrain chain support:
 * 1) remove forced breakConstraint() in addPointConstraint()
 * 2) route constraintChanged() driver lookups through helper that handles chain middle vehicles
 */
public final class BaseVehicleConstraintPatch {
    private static final String TARGET_NAME = "addPointConstraint";
    private static final String UPDATE_NAME = "update";
    private static final String CONSTRAINT_CHANGED_NAME = "constraintChanged";
    private static final String AUTHORIZATION_CHANGED_NAME = "authorizationChanged";
    private static final String AUTHORIZATION_CLIENT_COLLIDE_NAME = "authorizationClientCollide";
    private static final String AUTHORIZATION_SERVER_COLLIDE_NAME = "authorizationServerCollide";
    private static final String AUTHORIZATION_SERVER_ON_SEAT_NAME = "authorizationServerOnSeat";
    private static final String IS_ENTER_BLOCKED_NAME = "isEnterBlocked";
    private static final String IS_ENTER_BLOCKED2_NAME = "isEnterBlocked2";
    private static final String CLINIT_NAME = "<clinit>";
    private static final String VOID_NOARG_DESC = "()V";
    private static final String AUTHORIZATION_CHANGED_DESC =
            "(Lzombie/characters/IsoGameCharacter;)V";
    private static final String AUTHORIZATION_CLIENT_COLLIDE_DESC =
            "(Lzombie/characters/IsoPlayer;)V";
    private static final String AUTHORIZATION_SERVER_COLLIDE_DESC = "(SZ)V";
    private static final String AUTHORIZATION_SERVER_ON_SEAT_DESC =
            "(Lzombie/characters/IsoPlayer;Z)V";
    private static final String ENTER_BLOCKED_DESC = "(Lzombie/characters/IsoGameCharacter;I)Z";
    private static final String EXIT_BLOCKED_DESC = "(Lzombie/characters/IsoGameCharacter;I)Z";
    private static final String EXIT_BLOCKED2_DESC = "(I)Z";
    private static final String PATCH_LOG_LINE = "[Landtrain][BaseVehiclePatch] BaseVehicle override enabled";
    private static final String BREAK_DESC_OBJECT_BOOL = "(ZLjava/lang/Boolean;)V";
    private static final String BREAK_DESC_PRIMITIVE_BOOL = "(ZZ)V";
    private static final String BASE_VEHICLE_OWNER = "zombie/vehicles/BaseVehicle";
    private static final String BULLET_OWNER = "zombie/core/physics/Bullet";
    private static final String BULLET_ADD_ROPE = "addRopeConstraint";
    private static final String BULLET_ADD_POINT = "addPointConstraint";
    private static final String BULLET_ADD_ROPE_DESC = "(IIFFFFFFF)I";
    private static final String BULLET_ADD_POINT_DESC = "(IIFFFFFF)I";
    private static final String CAR_CONTROLLER_OWNER = "zombie/core/physics/CarController";
    private static final String CAR_CONTROLLER_PARK = "park";
    private static final String CAR_CONTROLLER_PARK_DESC = "()V";
    private static final String GET_DRIVER_DESC = "()Lzombie/characters/IsoGameCharacter;";
    private static final String HELPER_OWNER = "zombie/vehicles/LandtrainConstraintAuthHelper";
    private static final String HELPER_METHOD = "resolveConstraintDriver";
    private static final String HELPER_DESC =
            "(Lzombie/vehicles/BaseVehicle;)Lzombie/characters/IsoGameCharacter;";
    private static final String HELPER_ENTER_BLOCKED = "isEnterBlockedLandtrain";
    private static final String HELPER_ENTER_BLOCKED_DESC =
            "(Lzombie/vehicles/BaseVehicle;Lzombie/characters/IsoGameCharacter;I)Z";
    private static final String HELPER_ENTER_BLOCKED2 = "isEnterBlocked2Landtrain";
    private static final String HELPER_ENTER_BLOCKED2_DESC = "(Lzombie/vehicles/BaseVehicle;I)Z";
    private static final String HELPER_AUTHORIZATION_CHANGED =
            "authorizationChangedLandtrain";
    private static final String HELPER_AUTHORIZATION_CHANGED_DESC =
            "(Lzombie/vehicles/BaseVehicle;Lzombie/characters/IsoGameCharacter;)V";
    private static final String HELPER_AUTHORIZATION_CLIENT_COLLIDE =
            "authorizationClientCollideLandtrain";
    private static final String HELPER_AUTHORIZATION_CLIENT_COLLIDE_DESC =
            "(Lzombie/vehicles/BaseVehicle;Lzombie/characters/IsoPlayer;)V";
    private static final String HELPER_AUTHORIZATION_SERVER_COLLIDE =
            "authorizationServerCollideLandtrain";
    private static final String HELPER_AUTHORIZATION_SERVER_COLLIDE_DESC =
            "(Lzombie/vehicles/BaseVehicle;SZ)V";
    private static final String HELPER_AUTHORIZATION_SERVER_ON_SEAT =
            "authorizationServerOnSeatLandtrain";
    private static final String HELPER_AUTHORIZATION_SERVER_ON_SEAT_DESC =
            "(Lzombie/vehicles/BaseVehicle;Lzombie/characters/IsoPlayer;Z)V";
    private static final String HELPER_SAFE_PARK = "safeParkControllerLandtrain";
    private static final String HELPER_SAFE_PARK_DESC =
            "(Lzombie/core/physics/CarController;)V";

    private static final class AddPointPatchStats {
        int removedBreakCalls;
        int forcedRigidCalls;
    }

    private BaseVehicleConstraintPatch() {
    }

    public static void main(String[] args) throws Exception {
        if (args.length != 2) {
            System.err.println("Usage: BaseVehicleConstraintPatch <input BaseVehicle.class> <output BaseVehicle.class>");
            System.exit(2);
        }

        Path input = Paths.get(args[0]);
        Path output = Paths.get(args[1]);
        byte[] original = Files.readAllBytes(input);

        ClassNode classNode = new ClassNode();
        new ClassReader(original).accept(classNode, 0);

        int removedCalls = 0;
        int forcedRigidCalls = 0;
        int inspectedAddPointMethods = 0;
        int patchedConstraintDriverCalls = 0;
        int patchedEnterBlockedCalls = 0;
        int patchedEnterBlocked2Calls = 0;
        int patchedAuthorizationChangedMethods = 0;
        int patchedAuthorizationClientCollideMethods = 0;
        int patchedAuthorizationServerCollideMethods = 0;
        int patchedAuthorizationServerOnSeatMethods = 0;
        int patchedSafeParkCalls = 0;

        for (MethodNode method : classNode.methods) {
            if (TARGET_NAME.equals(method.name) && isTargetAddPointConstraint(method.desc)) {
                inspectedAddPointMethods++;
                AddPointPatchStats stats = patchAddPointConstraint(method);
                removedCalls += stats.removedBreakCalls;
                forcedRigidCalls += stats.forcedRigidCalls;
            } else if (UPDATE_NAME.equals(method.name) && VOID_NOARG_DESC.equals(method.desc)) {
                patchedSafeParkCalls += patchNullSafeParkCalls(method);
            } else if (CONSTRAINT_CHANGED_NAME.equals(method.name)
                    && VOID_NOARG_DESC.equals(method.desc)) {
                patchedConstraintDriverCalls += patchConstraintChangedDriverCalls(method);
            } else if (AUTHORIZATION_CHANGED_NAME.equals(method.name)
                    && AUTHORIZATION_CHANGED_DESC.equals(method.desc)) {
                patchedAuthorizationChangedMethods += patchMethodDelegateToHelper(
                        method,
                        HELPER_AUTHORIZATION_CHANGED,
                        HELPER_AUTHORIZATION_CHANGED_DESC);
            } else if (AUTHORIZATION_CLIENT_COLLIDE_NAME.equals(method.name)
                    && AUTHORIZATION_CLIENT_COLLIDE_DESC.equals(method.desc)) {
                patchedAuthorizationClientCollideMethods += patchMethodDelegateToHelper(
                        method,
                        HELPER_AUTHORIZATION_CLIENT_COLLIDE,
                        HELPER_AUTHORIZATION_CLIENT_COLLIDE_DESC);
            } else if (AUTHORIZATION_SERVER_COLLIDE_NAME.equals(method.name)
                    && AUTHORIZATION_SERVER_COLLIDE_DESC.equals(method.desc)) {
                patchedAuthorizationServerCollideMethods += patchMethodDelegateToHelper(
                        method,
                        HELPER_AUTHORIZATION_SERVER_COLLIDE,
                        HELPER_AUTHORIZATION_SERVER_COLLIDE_DESC);
            } else if (AUTHORIZATION_SERVER_ON_SEAT_NAME.equals(method.name)
                    && AUTHORIZATION_SERVER_ON_SEAT_DESC.equals(method.desc)) {
                patchedAuthorizationServerOnSeatMethods += patchMethodDelegateToHelper(
                        method,
                        HELPER_AUTHORIZATION_SERVER_ON_SEAT,
                        HELPER_AUTHORIZATION_SERVER_ON_SEAT_DESC);
            } else if (IS_ENTER_BLOCKED_NAME.equals(method.name)
                    && ENTER_BLOCKED_DESC.equals(method.desc)) {
                patchedEnterBlockedCalls += patchEnterBlockedCall(
                        method,
                        "isExitBlocked",
                        EXIT_BLOCKED_DESC,
                        HELPER_ENTER_BLOCKED,
                        HELPER_ENTER_BLOCKED_DESC);
            } else if (IS_ENTER_BLOCKED2_NAME.equals(method.name)
                    && ENTER_BLOCKED_DESC.equals(method.desc)) {
                patchedEnterBlocked2Calls += patchEnterBlockedCall(
                        method,
                        "isExitBlocked2",
                        EXIT_BLOCKED2_DESC,
                        HELPER_ENTER_BLOCKED2,
                        HELPER_ENTER_BLOCKED2_DESC);
            }
        }

        if (removedCalls < 2) {
            throw new IllegalStateException(
                    "Expected to remove 2 breakConstraint calls, removed "
                            + removedCalls
                            + " (inspected addPoint methods: "
                            + inspectedAddPointMethods
                            + ")");
        }
        if (forcedRigidCalls < 1) {
            throw new IllegalStateException(
                    "Expected to force at least 1 rope->point constraint call, patched "
                            + forcedRigidCalls);
        }
        if (patchedConstraintDriverCalls < 1) {
            throw new IllegalStateException(
                    "Expected to patch at least 1 constraintChanged getDriver call, patched "
                            + patchedConstraintDriverCalls);
        }
        if (patchedAuthorizationChangedMethods < 1) {
            throw new IllegalStateException(
                    "Expected to patch authorizationChanged, patched "
                            + patchedAuthorizationChangedMethods);
        }
        if (patchedAuthorizationClientCollideMethods < 1) {
            throw new IllegalStateException(
                    "Expected to patch authorizationClientCollide, patched "
                            + patchedAuthorizationClientCollideMethods);
        }
        if (patchedAuthorizationServerCollideMethods < 1) {
            throw new IllegalStateException(
                    "Expected to patch authorizationServerCollide, patched "
                            + patchedAuthorizationServerCollideMethods);
        }
        if (patchedAuthorizationServerOnSeatMethods < 1) {
            throw new IllegalStateException(
                    "Expected to patch authorizationServerOnSeat, patched "
                            + patchedAuthorizationServerOnSeatMethods);
        }
        if (patchedSafeParkCalls < 1) {
            throw new IllegalStateException(
                    "Expected to patch at least 1 CarController.park call, patched "
                            + patchedSafeParkCalls);
        }
        if (patchedEnterBlockedCalls < 1) {
            throw new IllegalStateException(
                    "Expected to patch isEnterBlocked call, patched " + patchedEnterBlockedCalls);
        }
        if (patchedEnterBlocked2Calls < 1) {
            throw new IllegalStateException(
                    "Expected to patch isEnterBlocked2 call, patched " + patchedEnterBlocked2Calls);
        }
        if (!ensureClassInitLog(classNode)) {
            throw new IllegalStateException("Failed to inject BaseVehicle class-init debug log");
        }

        ClassWriter writer = new ClassWriter(0);
        classNode.accept(writer);

        Files.createDirectories(output.getParent());
        Files.write(output, writer.toByteArray());
        System.out.println(
                "Patched BaseVehicle.class; removed breakConstraint calls: "
                        + removedCalls
                        + ", forced rigid constraints: "
                        + forcedRigidCalls
                        + ", constraint driver hooks: "
                        + patchedConstraintDriverCalls
                        + ", auth hooks (changed/client/server/seat): "
                        + patchedAuthorizationChangedMethods
                        + "/"
                        + patchedAuthorizationClientCollideMethods
                        + "/"
                        + patchedAuthorizationServerCollideMethods
                        + "/"
                        + patchedAuthorizationServerOnSeatMethods
                        + ", safe park hooks: "
                        + patchedSafeParkCalls
                        + ", enter-block hooks: "
                        + patchedEnterBlockedCalls
                        + "/"
                        + patchedEnterBlocked2Calls
                        + ", class-init debug log: enabled");
    }

    private static boolean isTargetAddPointConstraint(String methodDesc) {
        return "(Lzombie/characters/IsoPlayer;Lzombie/vehicles/BaseVehicle;Ljava/lang/String;Ljava/lang/String;Z)V"
                        .equals(methodDesc)
                || "(Lzombie/characters/IsoPlayer;Lzombie/vehicles/BaseVehicle;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Boolean;)V"
                        .equals(methodDesc);
    }

    private static AddPointPatchStats patchAddPointConstraint(MethodNode method) {
        AddPointPatchStats stats = new AddPointPatchStats();
        InsnList insns = method.instructions;

        for (AbstractInsnNode node = insns.getFirst(); node != null; ) {
            AbstractInsnNode next = node.getNext();
            if (node instanceof MethodInsnNode call
                    && BASE_VEHICLE_OWNER.equals(call.owner)
                    && "breakConstraint".equals(call.name)) {
                if (!(BREAK_DESC_OBJECT_BOOL.equals(call.desc)
                        || BREAK_DESC_PRIMITIVE_BOOL.equals(call.desc))) {
                    node = next;
                    continue;
                }

                // breakConstraint(...) consumes objectref + 2 args and returns void.
                InsnList replacement = new InsnList();
                replacement.add(new InsnNode(Opcodes.POP2));
                replacement.add(new InsnNode(Opcodes.POP));
                insns.insert(node, replacement);
                insns.remove(node);
                stats.removedBreakCalls++;
            } else if (node instanceof MethodInsnNode call
                    && BULLET_OWNER.equals(call.owner)
                    && BULLET_ADD_ROPE.equals(call.name)
                    && BULLET_ADD_ROPE_DESC.equals(call.desc)) {
                // Drop the rope-length float and call Bullet.addPointConstraint(...) instead.
                insns.insertBefore(call, new InsnNode(Opcodes.POP));
                MethodInsnNode replacement =
                        new MethodInsnNode(
                                Opcodes.INVOKESTATIC,
                                BULLET_OWNER,
                                BULLET_ADD_POINT,
                                BULLET_ADD_POINT_DESC,
                                false);
                insns.set(call, replacement);
                stats.forcedRigidCalls++;
            }
            node = next;
        }

        return stats;
    }

    private static int patchConstraintChangedDriverCalls(MethodNode method) {
        int patched = 0;
        InsnList insns = method.instructions;
        for (AbstractInsnNode node = insns.getFirst(); node != null; ) {
            AbstractInsnNode next = node.getNext();
            if (!(node instanceof MethodInsnNode call)) {
                node = next;
                continue;
            }
            if (!BASE_VEHICLE_OWNER.equals(call.owner)
                    || !"getDriver".equals(call.name)
                    || !GET_DRIVER_DESC.equals(call.desc)) {
                node = next;
                continue;
            }

            MethodInsnNode replacement =
                    new MethodInsnNode(
                            Opcodes.INVOKESTATIC,
                            HELPER_OWNER,
                            HELPER_METHOD,
                            HELPER_DESC,
                            false);
            insns.set(call, replacement);
            patched++;
            node = next;
        }
        return patched;
    }

    private static int patchNullSafeParkCalls(MethodNode method) {
        int patched = 0;
        InsnList insns = method.instructions;
        for (AbstractInsnNode node = insns.getFirst(); node != null; ) {
            AbstractInsnNode next = node.getNext();
            if (!(node instanceof MethodInsnNode call)) {
                node = next;
                continue;
            }
            if (!CAR_CONTROLLER_OWNER.equals(call.owner)
                    || !CAR_CONTROLLER_PARK.equals(call.name)
                    || !CAR_CONTROLLER_PARK_DESC.equals(call.desc)) {
                node = next;
                continue;
            }

            MethodInsnNode replacement =
                    new MethodInsnNode(
                            Opcodes.INVOKESTATIC,
                            HELPER_OWNER,
                            HELPER_SAFE_PARK,
                            HELPER_SAFE_PARK_DESC,
                            false);
            insns.set(call, replacement);
            patched++;
            node = next;
        }
        return patched;
    }

    private static int patchMethodDelegateToHelper(
            MethodNode method, String helperMethod, String helperDesc) {
        InsnList insns = new InsnList();

        int localIndex = 0;
        int maxStack = 0;
        if ((method.access & Opcodes.ACC_STATIC) == 0) {
            insns.add(new VarInsnNode(Opcodes.ALOAD, 0));
            localIndex = 1;
            maxStack = 1;
        }

        Type[] argumentTypes = Type.getArgumentTypes(method.desc);
        for (Type argumentType : argumentTypes) {
            insns.add(new VarInsnNode(argumentType.getOpcode(Opcodes.ILOAD), localIndex));
            localIndex += argumentType.getSize();
            maxStack += argumentType.getSize();
        }

        insns.add(new MethodInsnNode(Opcodes.INVOKESTATIC, HELPER_OWNER, helperMethod, helperDesc, false));

        Type returnType = Type.getReturnType(method.desc);
        if (returnType.getSort() == Type.VOID) {
            insns.add(new InsnNode(Opcodes.RETURN));
        } else {
            insns.add(new InsnNode(returnType.getOpcode(Opcodes.IRETURN)));
        }

        method.instructions.clear();
        method.instructions.add(insns);
        method.tryCatchBlocks.clear();
        if (method.localVariables != null) {
            method.localVariables.clear();
        }
        method.maxLocals = Math.max(method.maxLocals, localIndex);
        method.maxStack = Math.max(method.maxStack, maxStack);
        return 1;
    }

    private static int patchEnterBlockedCall(
            MethodNode method,
            String targetCallName,
            String targetCallDesc,
            String helperMethod,
            String helperDesc) {
        int patched = 0;
        InsnList insns = method.instructions;
        for (AbstractInsnNode node = insns.getFirst(); node != null; ) {
            AbstractInsnNode next = node.getNext();
            if (!(node instanceof MethodInsnNode call)) {
                node = next;
                continue;
            }
            if (!BASE_VEHICLE_OWNER.equals(call.owner)
                    || !targetCallName.equals(call.name)
                    || !targetCallDesc.equals(call.desc)) {
                node = next;
                continue;
            }

            MethodInsnNode replacement =
                    new MethodInsnNode(
                            Opcodes.INVOKESTATIC,
                            HELPER_OWNER,
                            helperMethod,
                            helperDesc,
                            false);
            insns.set(call, replacement);
            patched++;
            node = next;
        }
        return patched;
    }

    private static boolean ensureClassInitLog(ClassNode classNode) {
        MethodNode clinit = null;
        for (MethodNode method : classNode.methods) {
            if (CLINIT_NAME.equals(method.name) && VOID_NOARG_DESC.equals(method.desc)) {
                clinit = method;
                break;
            }
        }

        if (clinit == null) {
            clinit = new MethodNode(Opcodes.ACC_STATIC, CLINIT_NAME, VOID_NOARG_DESC, null, null);
            clinit.instructions = new InsnList();
            clinit.instructions.add(createPatchLogInstructions());
            clinit.instructions.add(new InsnNode(Opcodes.RETURN));
            clinit.maxStack = 2;
            clinit.maxLocals = 0;
            classNode.methods.add(clinit);
            return true;
        }

        if (hasPatchLog(clinit)) {
            return true;
        }

        boolean inserted = false;
        for (AbstractInsnNode node = clinit.instructions.getFirst(); node != null; ) {
            AbstractInsnNode next = node.getNext();
            if (node.getOpcode() == Opcodes.RETURN) {
                clinit.instructions.insertBefore(node, createPatchLogInstructions());
                inserted = true;
            }
            node = next;
        }

        if (inserted) {
            clinit.maxStack = Math.max(clinit.maxStack, 2);
        }
        return inserted;
    }

    private static boolean hasPatchLog(MethodNode method) {
        for (AbstractInsnNode node = method.instructions.getFirst(); node != null; node = node.getNext()) {
            if (node instanceof LdcInsnNode ldc
                    && ldc.cst instanceof String text
                    && PATCH_LOG_LINE.equals(text)) {
                return true;
            }
        }
        return false;
    }

    private static InsnList createPatchLogInstructions() {
        InsnList insns = new InsnList();
        insns.add(new org.objectweb.asm.tree.FieldInsnNode(
                Opcodes.GETSTATIC,
                "java/lang/System",
                "out",
                "Ljava/io/PrintStream;"));
        insns.add(new LdcInsnNode(PATCH_LOG_LINE));
        insns.add(new MethodInsnNode(
                Opcodes.INVOKEVIRTUAL,
                "java/io/PrintStream",
                "println",
                "(Ljava/lang/String;)V",
                false));
        return insns;
    }
}