package org.openjdk.gcbench.roots;

import org.objectweb.asm.*;

import java.net.URL;
import java.net.URLClassLoader;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

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

public class ClassGenerator {

    private static final int CLASSFILE_VERSION = 50;

    private static final AtomicInteger idx = new AtomicInteger();


    public static Class<?> generate(ByteClassLoader cl) throws Exception {
        ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_MAXS);

        String name = "Class" + idx.incrementAndGet();

        cw.visit(CLASSFILE_VERSION,
                ACC_PUBLIC + ACC_SUPER,
                name,
                null,
                Type.getInternalName(Object.class),
                new String[0]);

        cw.visitSource(name + ".java", null);

        {
            MethodVisitor mv = cw.visitMethod(ACC_PUBLIC, "<init>", "()V", null, null);
            mv.visitCode();
            mv.visitVarInsn(ALOAD, 0);
            mv.visitMethodInsn(INVOKESPECIAL, "java/lang/Object", "<init>", "()V", false);
            mv.visitInsn(RETURN);
            mv.visitMaxs(1, 1);
            mv.visitEnd();
        }

        cw.visitField(ACC_STATIC, "ref", "Ljava/lang/Object;", "Ljava/lang/Object;", null);

        cw.visitEnd();

        cl.put(name, cw.toByteArray());
        return cl.findClass(name);
    }

    public static class ByteClassLoader extends URLClassLoader {
        private final Map<String, byte[]> customClasses;

        public ByteClassLoader() {
            super(new URL[] {});
            customClasses = new ConcurrentHashMap<>();
        }

        @Override
        protected Class<?> findClass(final String name) throws ClassNotFoundException {
            if (customClasses.containsKey(name)) {
                byte[] bs = customClasses.get(name);
                return defineClass(name, bs, 0, bs.length);
            }
            return super.findClass(name);
        }

        public void put(String name, byte[] bytes) {
            customClasses.put(name, bytes);
        }
    }

}