package org.optaplanner.optapy;

import java.io.PrintStream;
import java.lang.annotation.RetentionPolicy;
import java.lang.reflect.Array;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiFunction;
import java.util.function.Function;
import java.util.function.Supplier;

import org.objectweb.asm.Type;
import org.optaplanner.core.api.domain.entity.PinningFilter;
import org.optaplanner.core.api.domain.entity.PlanningEntity;
import org.optaplanner.core.api.domain.lookup.PlanningId;
import org.optaplanner.core.api.domain.solution.PlanningEntityCollectionProperty;
import org.optaplanner.core.api.domain.solution.PlanningSolution;
import org.optaplanner.core.api.domain.solution.ProblemFactCollectionProperty;
import org.optaplanner.core.api.domain.valuerange.ValueRangeProvider;
import org.optaplanner.core.api.domain.variable.VariableListener;
import org.optaplanner.core.api.function.TriFunction;
import org.optaplanner.core.api.score.calculator.ConstraintMatchAwareIncrementalScoreCalculator;
import org.optaplanner.core.api.score.calculator.EasyScoreCalculator;
import org.optaplanner.core.api.score.calculator.IncrementalScoreCalculator;
import org.optaplanner.core.api.score.stream.ConstraintProvider;

import io.quarkus.gizmo.AnnotationCreator;
import io.quarkus.gizmo.BranchResult;
import io.quarkus.gizmo.BytecodeCreator;
import io.quarkus.gizmo.ClassCreator;
import io.quarkus.gizmo.ClassOutput;
import io.quarkus.gizmo.FieldDescriptor;
import io.quarkus.gizmo.MethodCreator;
import io.quarkus.gizmo.MethodDescriptor;
import io.quarkus.gizmo.ResultHandle;

public class PythonWrapperGenerator {
    /**
     * The Gizmo generated bytecode. Used by
     * gizmoClassLoader when not run in Quarkus
     * in order to create an instance of the Member
     * Accessor
     */
    private static final Map<String, byte[]> classNameToBytecode = new HashMap<>();

    /**
     * A custom classloader that looks for the class in
     * classNameToBytecode
     */
    static ClassLoader gizmoClassLoader = new ClassLoader() {
        // getName() is an abstract method in Java 11 but not in Java 8
        public String getName() {
            return "OptaPy Generated Classes ClassLoader";
        }

        @Override
        public Class<?> findClass(String name) throws ClassNotFoundException {
            if (classNameToBytecode.containsKey(name)) {
                // Gizmo generated class
                byte[] byteCode = classNameToBytecode.get(name);
                return defineClass(name, byteCode, 0, byteCode.length);
            } else {
                // Not a Gizmo generated class; load from parent class loader
                return PythonWrapperGenerator.class.getClassLoader().loadClass(name);
            }
        }
    };

    // These functions are set in Python code
    // Maps a OpaquePythonReference to a unique, numerical id
    static Function<OpaquePythonReference, Number> pythonObjectToId;

    private static Function<OpaquePythonReference, String> pythonObjectToString;

    private static Function<OpaquePythonReference, Class<?>> pythonGetJavaClass;

    // Maps a OpaquePythonReference that represents an array of objects to a list of its values
    private static Function<OpaquePythonReference, List<OpaquePythonReference>> pythonArrayIdToIdArray;

    // Maps a OptaquePythonReference that represents an array of primitive types to their values
    private static Function<OpaquePythonReference, List<Object>> pythonArrayToJavaList;

    // Reads an attribute on a OpaquePythonReference
    private static BiFunction<OpaquePythonReference, String, Object> pythonObjectIdAndAttributeNameToValue;

    // Sets an attribute on a OpaquePythonReference
    private static TriFunction<OpaquePythonReference, String, Object, Object> pythonObjectIdAndAttributeSetter;

    // These functions are used in Python to set fields to the corresponding Python function
    @SuppressWarnings("unused")
    public static void setPythonObjectToString(Function<OpaquePythonReference, String> pythonObjectToString) {
        PythonWrapperGenerator.pythonObjectToString = pythonObjectToString;
    }

    @SuppressWarnings("unused")
    public static void setPythonGetJavaClass(Function<OpaquePythonReference, Class<?>> pythonGetJavaClass) {
        PythonWrapperGenerator.pythonGetJavaClass = pythonGetJavaClass;
    }

    @SuppressWarnings("unused")
    public static void setPythonObjectToId(Function<OpaquePythonReference, Number> pythonObjectToId) {
        PythonWrapperGenerator.pythonObjectToId = pythonObjectToId;
    }

    @SuppressWarnings("unused")
    public static Object getValueFromPythonObject(OpaquePythonReference objectId, String attributeName) {
        return pythonObjectIdAndAttributeNameToValue.apply(objectId, attributeName);
    }

    @SuppressWarnings("unused")
    public static void setValueOnPythonObject(OpaquePythonReference objectId, String attributeName, Object value) {
        pythonObjectIdAndAttributeSetter.apply(objectId, attributeName, value);
    }

    @SuppressWarnings("unused")
    public static void setPythonArrayIdToIdArray(Function<OpaquePythonReference, List<OpaquePythonReference>> function) {
        pythonArrayIdToIdArray = function;
    }

    @SuppressWarnings("unused")
    public static void setPythonArrayToJavaList(Function<OpaquePythonReference, List<Object>> function) {
        pythonArrayToJavaList = function;
    }

    @SuppressWarnings("unused")
    public static void setPythonObjectIdAndAttributeNameToValue(BiFunction<OpaquePythonReference, String, Object> function) {
        pythonObjectIdAndAttributeNameToValue = function;
    }

    @SuppressWarnings("unused")
    public static void setPythonObjectIdAndAttributeSetter(
            TriFunction<OpaquePythonReference, String, Object, Object> setter) {
        pythonObjectIdAndAttributeSetter = setter;
    }

    @SuppressWarnings("unused")
    public static String getPythonObjectString(OpaquePythonReference pythonObject) {
        return pythonObjectToString.apply(pythonObject);
    }

    @SuppressWarnings("unused")
    public static OpaquePythonReference getPythonObject(PythonObject pythonObject) {
        return pythonObject.get__optapy_Id();
    }

    public static OpaquePythonReference getPythonObject(PythonComparable pythonObject) {
        return pythonObject.reference;
    }

    @SuppressWarnings("unused")
    public static ClassLoader getClassLoaderForAliasMap(Map<String, Class<?>> aliasMap) {
        return new ClassLoader() {
            // getName() is an abstract method in Java 11 but not in Java 8
            public String getName() {
                return "OptaPy Alias Map ClassLoader";
            }

            @Override
            public Class<?> findClass(String name) throws ClassNotFoundException {
                if (aliasMap.containsKey(name)) {
                    // Gizmo generated class
                    return aliasMap.get(name);
                } else {
                    // Not a Gizmo generated class; load from parent class loader
                    return gizmoClassLoader.loadClass(name);
                }
            }
        };
    }

    @SuppressWarnings("unused")
    public static Class<?> getJavaClass(OpaquePythonReference object) {
        return pythonGetJavaClass.apply(object);
    }

    @SuppressWarnings("unused")
    public static Boolean wrapBoolean(boolean value) {
        return value;
    }

    @SuppressWarnings("unused")
    public static Byte wrapByte(byte value) {
        return value;
    }

    @SuppressWarnings("unused")
    public static Short wrapShort(short value) {
        return value;
    }

    @SuppressWarnings("unused")
    public static Integer wrapInt(int value) {
        return value;
    }

    @SuppressWarnings("unused")
    public static Long wrapLong(long value) {
        return value;
    }

    // Used in Python to get the array type of a class; used in
    // determining what class a @ProblemFactCollection / @PlanningEntityCollection
    // should be
    @SuppressWarnings("unused")
    public static Class<?> getArrayClass(Class<?> elementClass) {
        return Array.newInstance(elementClass, 0).getClass();
    }

    @SuppressWarnings("unused")
    public static <T> String getCollectionSignature(Class<?> collectionClass, Class<T> elementClass) {
        StringBuilder out = new StringBuilder();
        out.append('L').append(Type.getInternalName(collectionClass)); // Return is of class Collection
        out.append("<"); // Collection is of generic type...
        out.append('L').append(Type.getInternalName(elementClass)); // The collection type
        out.append(";>;"); // end of signature
        String result = out.toString();
        return result;
    }

    // Holds the OpaquePythonReference
    static final String PYTHON_BINDING_FIELD_NAME = "__optaplannerPythonValue";
    static final String REFERENCE_MAP_FIELD_NAME = "__optaplannerReferenceMap";

    private static <T> T wrapArray(Class<T> javaClass, OpaquePythonReference object, Number id, Map<Number, Object> map) {
        // If the class is an array, we need to extract
        // its elements from the OpaquePythonReference
        if (Comparable.class.isAssignableFrom(javaClass.getComponentType()) ||
                Number.class.isAssignableFrom(javaClass.getComponentType())) {
            List<Object> items = pythonArrayToJavaList.apply(object);
            int length = items.size();
            Object out = Array.newInstance(javaClass.getComponentType(), length);

            // Put the array into the python id to java instance map
            map.put(id, out);

            // Set the elements of the array to the wrapped python items
            for (int i = 0; i < length; i++) {
                Object item = items.get(i);
                if (javaClass.getComponentType().equals(Integer.class) && item instanceof Long) {
                    item = ((Long) item).intValue();
                }
                Array.set(out, i, item);
            }
            return (T) out;
        }
        List<OpaquePythonReference> itemIds = pythonArrayIdToIdArray.apply(object);
        int length = itemIds.size();
        Object out = Array.newInstance(javaClass.getComponentType(), length);

        // Put the array into the python id to java instance map
        map.put(id, out);

        // Set the elements of the array to the wrapped python items
        for (int i = 0; i < length; i++) {
            Array.set(out, i, wrap(javaClass.getComponentType(), itemIds.get(i), map));
        }
        return (T) out;
    }

    public static <T> T wrapCollection(OpaquePythonReference object, Number id, Map<Number, Object> map) {
        PythonList out = new PythonList(object, id, map);
        map.put(id, out);
        return (T) out;
    }

    @SuppressWarnings("unchecked")
    public static <T> T wrap(Class<T> javaClass, OpaquePythonReference object, Map<Number, Object> map) {
        if (object == null) {
            return null;
        }

        // Check to see if we already created the object
        Number id = pythonObjectToId.apply(object);
        if (map.containsKey(id)) {
            return (T) map.get(id);
        }

        try {
            if (javaClass.isArray()) {
                return wrapArray(javaClass, object, id, map);
            } else if (javaClass.isAssignableFrom(PythonList.class)) {
                return wrapCollection(object, id, map);
            } else if (javaClass.isAssignableFrom(OpaquePythonReference.class)) {
                // Don't wrap OpaquePythonReference if it is a pointer to an OpaquePythonReference
                return (T) object;
            } else {
                // Create a new instance of the Java Class. Its constructor will put the instance into the map
                return javaClass.getConstructor(OpaquePythonReference.class, Number.class, Map.class).newInstance(object,
                        id, map);
            }
        } catch (IllegalAccessException | NoSuchMethodException | InstantiationException | InvocationTargetException e) {
            throw new IllegalStateException("Error occurred when wrapping object (" + getPythonObjectString(object) + ")", e);
        }
    }

    private static ClassOutput getClassOutput(AtomicReference<byte[]> bytesReference) {
        return (path, byteCode) -> {
            bytesReference.set(byteCode);
        };
    }

    /**
     * Creates a class that looks like this:
     *
     * class JavaWrapper implements NaryFunction<A0,A1,A2,...,AN> {
     * public static NaryFunction<A0,A1,A2,...,AN> delegate;
     *
     * #64;Override
     * public AN apply(A0 arg0, A1 arg1, ..., A(N-1) finalArg) {
     * return delegate.apply(arg0,arg1,...,finalArg);
     * }
     * }
     *
     * @param className The simple name of the generated class
     * @param baseInterface the base interface
     * @param delegate The Python function to delegate to
     * @return never null
     */
    @SuppressWarnings({ "unused", "unchecked" })
    public static <A> Class<? extends A> defineWrapperFunction(String className, Class<A> baseInterface,
            Object delegate) {
        Method[] interfaceMethods = baseInterface.getMethods();
        if (interfaceMethods.length != 1) {
            throw new IllegalArgumentException("Can only call this function for functional interfaces (only 1 method)");
        }
        className = "org.optaplanner.optapy.generated." + className + ".GeneratedClass";
        if (classNameToBytecode.containsKey(className)) {
            try {
                return (Class<? extends A>) gizmoClassLoader.loadClass(className);
            } catch (ClassNotFoundException e) {
                throw new IllegalStateException(
                        "Impossible State: the class (" + className + ") should exists since it was created");
            }
        }
        AtomicReference<byte[]> classBytecodeHolder = new AtomicReference<>();
        ClassOutput classOutput = getClassOutput(classBytecodeHolder);

        // holds the delegate (static; same one is reused; should be stateless)
        FieldDescriptor delegateField;
        try (ClassCreator classCreator = ClassCreator.builder()
                .className(className)
                .interfaces(baseInterface)
                .classOutput(classOutput)
                .build()) {
            delegateField = classCreator.getFieldCreator("delegate", baseInterface)
                    .setModifiers(Modifier.STATIC | Modifier.PUBLIC)
                    .getFieldDescriptor();
            MethodCreator methodCreator = classCreator.getMethodCreator(MethodDescriptor.ofMethod(interfaceMethods[0]));

            ResultHandle pythonProxy = methodCreator.readStaticField(delegateField);
            ResultHandle[] args = new ResultHandle[interfaceMethods[0].getParameterCount()];
            for (int i = 0; i < args.length; i++) {
                args[i] = methodCreator.getMethodParam(i);
            }
            ResultHandle constraints = methodCreator.invokeInterfaceMethod(
                    MethodDescriptor.ofMethod(interfaceMethods[0]),
                    pythonProxy, args);
            methodCreator.returnValue(constraints);
        } catch (Exception e) {
            throw new IllegalStateException(e);
        }
        classNameToBytecode.put(className, classBytecodeHolder.get());
        try {
            // Now that the class created, we need to set it static field to the delegate function
            Class<? extends A> out = (Class<? extends A>) gizmoClassLoader.loadClass(className);
            out.getField(delegateField.getName()).set(null, delegate);
            return out;
        } catch (Exception e) {
            throw new IllegalStateException(
                    "Impossible State: the class (" + className + ") should exists since it was just created");
        }
    }

    /**
     * Creates a class that looks like this:
     *
     * class JavaWrapper implements SomeInterface {
     * public static Supplier&lt;SomeInterface&gt; supplier;
     *
     * private SomeInterface delegate;
     *
     * public JavaWrapper() {
     * delegate = supplier.get();
     * }
     *
     * #64;Override
     * public Result interfaceMethod1(A0 arg0, A1 arg1, ..., A(N-1) finalArg) {
     * return delegate.interfaceMethod1(arg0,arg1,...,finalArg);
     * }
     *
     * #64;Override
     * public Result interfaceMethod2(A0 arg0, A1 arg1, ..., A(N-1) finalArg) {
     * return delegate.interfaceMethod2(arg0,arg1,...,finalArg);
     * }
     * }
     *
     * @param className The simple name of the generated class
     * @param baseInterface the base interface
     * @param delegateSupplier The Python class to delegate to
     * @return never null
     */
    @SuppressWarnings({ "unused", "unchecked" })
    public static <A> Class<? extends A> defineWrapperClass(String className, Class<? extends A> baseInterface,
            Supplier<? extends A> delegateSupplier) {
        Method[] interfaceMethods = baseInterface.getMethods();
        className = "org.optaplanner.optapy.generated." + className + ".GeneratedClass";
        if (classNameToBytecode.containsKey(className)) {
            try {
                return (Class<? extends A>) gizmoClassLoader.loadClass(className);
            } catch (ClassNotFoundException e) {
                throw new IllegalStateException(
                        "Impossible State: the class (" + className + ") should exists since it was created");
            }
        }
        AtomicReference<byte[]> classBytecodeHolder = new AtomicReference<>();
        ClassOutput classOutput = getClassOutput(classBytecodeHolder);

        // holds the supplier of the delegate (static)
        FieldDescriptor supplierField;

        // holds the delegate (instance; new one created for each instance)
        FieldDescriptor delegateField;
        try (ClassCreator classCreator = ClassCreator.builder()
                .className(className)
                .interfaces(baseInterface)
                .classOutput(classOutput)
                .build()) {
            supplierField = classCreator.getFieldCreator("delegateSupplier", Supplier.class)
                    .setModifiers(Modifier.STATIC | Modifier.PUBLIC)
                    .getFieldDescriptor();
            delegateField = classCreator.getFieldCreator("delegate", baseInterface)
                    .setModifiers(Modifier.PUBLIC | Modifier.FINAL)
                    .getFieldDescriptor();

            MethodCreator constructorCreator =
                    classCreator.getMethodCreator(MethodDescriptor.ofConstructor(classCreator.getClassName()));
            constructorCreator.invokeSpecialMethod(MethodDescriptor.ofConstructor(Object.class), constructorCreator.getThis());
            constructorCreator.writeInstanceField(delegateField, constructorCreator.getThis(),
                    constructorCreator.invokeInterfaceMethod(MethodDescriptor.ofMethod(Supplier.class, "get", Object.class),
                            constructorCreator.readStaticField(supplierField)));
            constructorCreator.returnValue(constructorCreator.getThis());

            for (Method method : interfaceMethods) {
                MethodCreator methodCreator = classCreator.getMethodCreator(MethodDescriptor.ofMethod(method));
                ResultHandle pythonProxy = methodCreator.readInstanceField(delegateField, methodCreator.getThis());
                ResultHandle[] args = new ResultHandle[method.getParameterCount()];
                for (int i = 0; i < args.length; i++) {
                    args[i] = methodCreator.getMethodParam(i);
                }
                ResultHandle result = methodCreator.invokeInterfaceMethod(
                        MethodDescriptor.ofMethod(method),
                        pythonProxy, args);
                methodCreator.returnValue(result);
            }
        } catch (Exception e) {
            throw new IllegalStateException(e);
        }
        classNameToBytecode.put(className, classBytecodeHolder.get());
        try {
            // Now that the class created, we need to set it static field to the supplier of the delegate
            Class<? extends A> out = (Class<? extends A>) gizmoClassLoader.loadClass(className);
            out.getField(supplierField.getName()).set(null, delegateSupplier);
            return out;
        } catch (Exception e) {
            throw new IllegalStateException(
                    "Impossible State: the class (" + className + ") should exists since it was just created");
        }
    }

    /**
     * Creates a class that looks like this:
     *
     * class PythonConstraintProvider implements ConstraintProvider {
     * public static Function<ConstraintFactory, Constraint[]> defineConstraintsImpl;
     *
     * &#64;Override
     * public Constraint[] defineConstraints(ConstraintFactory constraintFactory) {
     * return defineConstraintsImpl.apply(constraintFactory);
     * }
     * }
     *
     * @param className The simple name of the generated class
     * @param defineConstraintsImpl The Python function that return the list of constraints
     * @return never null
     */
    @SuppressWarnings("unused")
    public static Class<?> defineConstraintProviderClass(String className,
            ConstraintProvider defineConstraintsImpl) {
        return defineWrapperFunction(className, ConstraintProvider.class, defineConstraintsImpl);
    }

    /**
     * Creates a class that looks like this:
     *
     * class PythonEasyScoreCalculator implements EasyScoreCalculator {
     * public static EasyScoreCalculator easyScoreCalculatorImpl;
     *
     * &#64;Override
     * public Score calculateScore(Solution solution) {
     * return easyScoreCalculatorImpl.calculateScore(solution);
     * }
     * }
     *
     * @param className The simple name of the generated class
     * @param easyScoreCalculatorImpl The Python function that return the score for the solution
     * @return never null
     */
    @SuppressWarnings("unused")
    public static Class<?> defineEasyScoreCalculatorClass(String className,
            EasyScoreCalculator easyScoreCalculatorImpl) {
        return defineWrapperFunction(className, EasyScoreCalculator.class, easyScoreCalculatorImpl);
    }

    /**
     * Creates a class that looks like this:
     *
     * class PythonIncrementalScoreCalculator implements IncrementalScoreCalculator {
     * public static Supplier&lt;IncrementalScoreCalculator&gt; supplier;
     * public final IncrementalScoreCalculator delegate;
     *
     * public PythonIncrementalScoreCalculator() {
     * delegate = supplier.get();
     * }
     *
     * &#64;Override
     * public Score calculateScore(Solution solution) {
     * return delegate.calculateScore(solution);
     * }
     *
     * ...
     * }
     *
     * @param className The simple name of the generated class
     * @param incrementalScoreCalculatorSupplier A supplier that returns a new instance of the incremental score calculator on
     *        each call
     * @return never null
     */
    @SuppressWarnings("unused")
    public static Class<?> defineIncrementalScoreCalculatorClass(String className,
            Supplier<? extends IncrementalScoreCalculator> incrementalScoreCalculatorSupplier,
            boolean constraintMatchAware) {
        if (constraintMatchAware) {
            return defineWrapperClass(className, ConstraintMatchAwareIncrementalScoreCalculator.class,
                    (Supplier<ConstraintMatchAwareIncrementalScoreCalculator>) incrementalScoreCalculatorSupplier);
        }
        return defineWrapperClass(className, IncrementalScoreCalculator.class, incrementalScoreCalculatorSupplier);
    }

    /**
     * Creates a class that looks like this:
     *
     * class PythonVariableListener implements VariableListener {
     * public static Supplier&lt;VariableListener&gt; supplier;
     * public final VariableListener delegate;
     *
     * public PythonVariableListener() {
     * delegate = supplier.get();
     * }
     *
     * public void afterVariableChange(scoreDirector, entity) {
     * delegate.afterVariableChange(scoreDirector, entity);
     * }
     * ...
     * }
     *
     * @param className The simple name of the generated class
     * @param variableListenerSupplier A supplier that returns a new instance of the variable listener on
     *        each call
     * @return never null
     */
    @SuppressWarnings("unused")
    public static Class<?> defineVariableListenerClass(String className,
            Supplier<? extends VariableListener> variableListenerSupplier) {
        return defineWrapperClass(className, VariableListener.class, variableListenerSupplier);
    }

    /*
     * The Planning Entity, Problem Fact, and Planning Solution classes look similar, with the only
     * difference being their top-level annotation. They all look like this:
     *
     * (none or @PlanningEntity or @PlanningSolution)
     * public class PojoForPythonObject implements PythonObject {
     * OpaquePythonReference __optaplannerPythonValue;
     * String string$field;
     * AnotherPojoForPythonObject otherObject$field;
     *
     * public PojoForPythonObject(OpaquePythonReference reference, Number id, Map<Number, PythonObject>
     * pythonIdToPythonObjectMap) {
     * this.__optaplannerPythonValue = reference;
     * pythonIdToPythonObjectMap.put(id, this);
     * string$field = PythonWrapperGenerator.getValueFromPythonObject(reference, "string");
     * OpaquePythonReference otherObjectReference = PythonWrapperGenerator.getValueFromPythonObject(reference, "otherObject");
     * otherObject$field = PythonWrapperGenerator.wrap(otherObjectReference, id, pythonIdToPythonObjectMap);
     * }
     *
     * public OpaquePythonReference get__optapy_Id() {
     * return __optaplannerPythonValue;
     * }
     *
     * public String getStringField() {
     * return string$field;
     * }
     *
     * public void setStringField(String val) {
     * PythonWrapperGenerator.setValueOnPythonObject(__optaplannerPythonValue, "string", val);
     * this.string$field = val;
     * }
     * // Repeat for otherObject
     * }
     */
    @SuppressWarnings("unused")
    public static Class<?> definePlanningEntityClass(String className, Class<?> parentClass,
            boolean defineEqualsAndHashcode,
            List<List<Object>> optaplannerMethodAnnotations,
            Map<String, Object> planningEntityAnnotations) {
        className = "org.optaplanner.optapy.generated." + className + ".GeneratedClass";
        if (classNameToBytecode.containsKey(className)) {
            try {
                return gizmoClassLoader.loadClass(className);
            } catch (ClassNotFoundException e) {
                throw new IllegalStateException(
                        "Impossible State: the class (" + className + ") should exists since it was created");
            }
        }
        AtomicReference<byte[]> classBytecodeHolder = new AtomicReference<>();
        ClassOutput classOutput = getClassOutput(classBytecodeHolder);
        try (ClassCreator classCreator = ClassCreator.builder()
                .className(className)
                .superClass(parentClass != null ? parentClass : Object.class)
                .interfaces(PythonObject.class)
                .classOutput(classOutput)
                .build()) {
            AnnotationCreator annotationCreator = classCreator.addAnnotation(PlanningEntity.class);
            Object pinningFilter = planningEntityAnnotations.get("pinningFilter");
            if (pinningFilter != null) {
                Class<? extends PinningFilter> pinningFilterClass = defineWrapperFunction(className + "PinningFilter",
                        PinningFilter.class, pinningFilter);
                annotationCreator.addValue("pinningFilter", pinningFilterClass);
            }
            FieldDescriptor valueField;

            if (parentClass == null) {
                valueField = classCreator.getFieldCreator(PYTHON_BINDING_FIELD_NAME, OpaquePythonReference.class)
                        .setModifiers(Modifier.PUBLIC).getFieldDescriptor();
            } else {
                valueField = FieldDescriptor.of(parentClass, PYTHON_BINDING_FIELD_NAME, OpaquePythonReference.class);
            }
            FieldDescriptor referenceMapField = classCreator.getFieldCreator(REFERENCE_MAP_FIELD_NAME, Map.class)
                    .setModifiers(Modifier.PUBLIC).getFieldDescriptor();
            generateWrapperMethods(classCreator, parentClass, defineEqualsAndHashcode, valueField, referenceMapField,
                    optaplannerMethodAnnotations);
        }
        classNameToBytecode.put(className, classBytecodeHolder.get());
        try {
            return gizmoClassLoader.loadClass(className);
        } catch (ClassNotFoundException e) {
            throw new IllegalStateException(
                    "Impossible State: the class (" + className + ") should exists since it was just created");
        }
    }

    @SuppressWarnings("unused")
    public static Class<?> defineProblemFactClass(String className, Class<?> parentClass,
            boolean defineEqualsAndHashcode,
            List<List<Object>> optaplannerMethodAnnotations) {
        className = "org.optaplanner.optapy.generated." + className + ".GeneratedClass";
        if (classNameToBytecode.containsKey(className)) {
            try {
                return gizmoClassLoader.loadClass(className);
            } catch (ClassNotFoundException e) {
                throw new IllegalStateException(
                        "Impossible State: the class (" + className + ") should exists since it was created");
            }
        }
        AtomicReference<byte[]> classBytecodeHolder = new AtomicReference<>();
        ClassOutput classOutput = getClassOutput(classBytecodeHolder);
        try (ClassCreator classCreator = ClassCreator.builder()
                .className(className)
                .superClass(parentClass != null ? parentClass : Object.class)
                .interfaces(PythonObject.class)
                .classOutput(classOutput)
                .build()) {
            FieldDescriptor valueField = classCreator.getFieldCreator(PYTHON_BINDING_FIELD_NAME, OpaquePythonReference.class)
                    .setModifiers(Modifier.PUBLIC).getFieldDescriptor();
            FieldDescriptor referenceMapField = classCreator.getFieldCreator(REFERENCE_MAP_FIELD_NAME, Map.class)
                    .setModifiers(Modifier.PUBLIC).getFieldDescriptor();
            generateWrapperMethods(classCreator, parentClass, defineEqualsAndHashcode, valueField, referenceMapField,
                    optaplannerMethodAnnotations);
        }
        classNameToBytecode.put(className, classBytecodeHolder.get());
        try {
            return gizmoClassLoader.loadClass(className);
        } catch (ClassNotFoundException e) {
            throw new IllegalStateException(
                    "Impossible State: the class (" + className + ") should exists since it was just created");
        }
    }

    @SuppressWarnings("unused")
    public static Class<?> definePlanningSolutionClass(String className,
            boolean defineEqualsAndHashcode,
            List<List<Object>> optaplannerMethodAnnotations) {
        className = "org.optaplanner.optapy.generated." + className + ".GeneratedClass";
        if (classNameToBytecode.containsKey(className)) {
            try {
                return gizmoClassLoader.loadClass(className);
            } catch (ClassNotFoundException e) {
                throw new IllegalStateException(
                        "Impossible State: the class (" + className + ") should exists since it was created");
            }
        }
        AtomicReference<byte[]> classBytecodeHolder = new AtomicReference<>();
        ClassOutput classOutput = getClassOutput(classBytecodeHolder);
        try (ClassCreator classCreator = ClassCreator.builder()
                .className(className)
                .interfaces(PythonObject.class)
                .classOutput(classOutput)
                .build()) {
            classCreator.addAnnotation(PlanningSolution.class)
                    .addValue("solutionCloner", Type.getType(PythonPlanningSolutionCloner.class));
            FieldDescriptor valueField = classCreator.getFieldCreator(PYTHON_BINDING_FIELD_NAME, OpaquePythonReference.class)
                    .setModifiers(Modifier.PUBLIC).getFieldDescriptor();
            FieldDescriptor referenceMapField = classCreator.getFieldCreator(REFERENCE_MAP_FIELD_NAME, Map.class)
                    .setModifiers(Modifier.PUBLIC).getFieldDescriptor();
            generateWrapperMethods(classCreator, null, defineEqualsAndHashcode, valueField, referenceMapField,
                    optaplannerMethodAnnotations);
        }
        classNameToBytecode.put(className, classBytecodeHolder.get());
        try {
            return gizmoClassLoader.loadClass(className);
        } catch (ClassNotFoundException e) {
            throw new IllegalStateException(
                    "Impossible State: the class (" + className + ") should exists since it was just created");
        }
    }

    // Used for debugging; prints a result handle
    private static void print(MethodCreator methodCreator, ResultHandle toPrint) {
        ResultHandle out = methodCreator.readStaticField(FieldDescriptor.of(System.class, "out", PrintStream.class));
        methodCreator.invokeVirtualMethod(MethodDescriptor.ofMethod(PrintStream.class, "println", void.class, Object.class),
                out, toPrint);
    }

    // Generate PythonObject interface methods
    private static void generateAsPointer(ClassCreator classCreator, FieldDescriptor valueField,
            FieldDescriptor referenceMapField) {
        MethodCreator methodCreator = classCreator.getMethodCreator("get__optapy_Id", OpaquePythonReference.class);
        ResultHandle valueResultHandle = methodCreator.readInstanceField(valueField, methodCreator.getThis());
        methodCreator.returnValue(valueResultHandle);

        methodCreator = classCreator.getMethodCreator("get__optapy_reference_map", Map.class);
        ResultHandle referenceMapResultHandle = methodCreator.readInstanceField(referenceMapField, methodCreator.getThis());
        methodCreator.returnValue(referenceMapResultHandle);
    }

    // Create all methods in the class
    @SuppressWarnings("unchecked")
    private static void generateWrapperMethods(ClassCreator classCreator, Class<?> parentClass,
            boolean defineEqualsAndHashcode, FieldDescriptor valueField,
            FieldDescriptor referenceMapField,
            List<List<Object>> optaplannerMethodAnnotations) {
        if (parentClass == null) {
            generateAsPointer(classCreator, valueField, referenceMapField);
        }

        // We only need to create methods/fields for methods with OptaPlanner annotations
        // optaplannerMethodAnnotations: list of tuples (methodName, returnType, annotationList)
        // (Each annotation is represented by a Map)
        List<FieldDescriptor> fieldDescriptorList = new ArrayList<>(optaplannerMethodAnnotations.size());
        List<Object> returnTypeList = new ArrayList<>(optaplannerMethodAnnotations.size());
        for (List<Object> optaplannerMethodAnnotation : optaplannerMethodAnnotations) {
            String methodName = (String) (optaplannerMethodAnnotation.get(0));
            Class<?> returnType = (Class<?>) (optaplannerMethodAnnotation.get(1));
            String signature = (String) (optaplannerMethodAnnotation.get(2));
            if (returnType == null) {
                returnType = Object.class;
            }
            List<Map<String, Object>> annotations = (List<Map<String, Object>>) optaplannerMethodAnnotation.get(3);
            fieldDescriptorList
                    .add(generateWrapperMethod(classCreator, valueField, methodName, returnType, signature, annotations,
                            returnTypeList));
        }
        createConstructor(classCreator, valueField, referenceMapField, parentClass, fieldDescriptorList, returnTypeList);

        if (parentClass == null) {
            createToString(classCreator, valueField);
        }

        if (defineEqualsAndHashcode) {
            createEqualsAndHashcode(classCreator, valueField);
        }
    }

    private static void createToString(ClassCreator classCreator, FieldDescriptor valueField) {
        MethodCreator methodCreator =
                classCreator.getMethodCreator(MethodDescriptor.ofMethod(classCreator.getClassName(), "toString", String.class));
        methodCreator.returnValue(methodCreator.invokeStaticMethod(
                MethodDescriptor.ofMethod(PythonWrapperGenerator.class, "getPythonObjectString", String.class,
                        OpaquePythonReference.class),
                methodCreator.readInstanceField(valueField, methodCreator.getThis())));
    }

    private static void createEqualsAndHashcode(ClassCreator classCreator, FieldDescriptor valueField) {
        // equals
        MethodCreator methodCreator =
                classCreator.getMethodCreator(
                        MethodDescriptor.ofMethod(classCreator.getClassName(), "equals", boolean.class, Object.class));
        ResultHandle parameter = methodCreator.getMethodParam(0);
        ResultHandle isInstance = methodCreator.instanceOf(parameter, classCreator.getClassName());
        BranchResult branchResult = methodCreator.ifTrue(isInstance);
        BytecodeCreator bytecodeCreator = branchResult.trueBranch();
        bytecodeCreator.returnValue(bytecodeCreator.invokeStaticMethod(
                MethodDescriptor.ofMethod(PythonComparable.class, "isPythonObjectEqualToOther", boolean.class,
                        OpaquePythonReference.class, OpaquePythonReference.class),
                bytecodeCreator.readInstanceField(valueField, methodCreator.getThis()),
                bytecodeCreator.readInstanceField(valueField, parameter)));
        bytecodeCreator = branchResult.falseBranch();
        bytecodeCreator.returnValue(bytecodeCreator.load(false));

        // hashCode
        methodCreator =
                classCreator.getMethodCreator(MethodDescriptor.ofMethod(classCreator.getClassName(), "hashCode", int.class));
        methodCreator.returnValue(methodCreator.invokeStaticMethod(
                MethodDescriptor.ofMethod(PythonComparable.class, "getPythonObjectHash", int.class,
                        OpaquePythonReference.class),
                methodCreator.readInstanceField(valueField, methodCreator.getThis())));
    }

    private static void createConstructor(ClassCreator classCreator, FieldDescriptor valueField,
            FieldDescriptor referenceMapField, Class<?> parentClass,
            List<FieldDescriptor> fieldDescriptorList, List<Object> returnTypeList) {
        MethodCreator methodCreator = classCreator.getMethodCreator(MethodDescriptor.ofConstructor(classCreator.getClassName(),
                OpaquePythonReference.class, Number.class, Map.class));
        methodCreator.setModifiers(Modifier.PUBLIC);

        ResultHandle value = methodCreator.getMethodParam(0);
        if (parentClass != null) {
            methodCreator.invokeSpecialMethod(
                    MethodDescriptor.ofConstructor(parentClass, OpaquePythonReference.class, Number.class, Map.class),
                    methodCreator.getThis(), value, methodCreator.getMethodParam(1),
                    methodCreator.getMethodParam(2));
        } else {
            methodCreator.invokeSpecialMethod(MethodDescriptor.ofConstructor(Object.class), methodCreator.getThis());
            methodCreator.invokeInterfaceMethod(
                    MethodDescriptor.ofMethod(Map.class, "put", Object.class, Object.class, Object.class),
                    methodCreator.getMethodParam(2), methodCreator.getMethodParam(1), methodCreator.getThis());
            methodCreator.writeInstanceField(valueField, methodCreator.getThis(), value);
            methodCreator.writeInstanceField(referenceMapField, methodCreator.getThis(), methodCreator.getMethodParam(2));
        }

        for (int i = 0; i < fieldDescriptorList.size(); i++) {
            FieldDescriptor fieldDescriptor = fieldDescriptorList.get(i);
            Object returnType = returnTypeList.get(i);
            String methodName = fieldDescriptor.getName().substring(0, fieldDescriptor.getName().length() - 6);

            ResultHandle outResultHandle = methodCreator.invokeStaticMethod(
                    MethodDescriptor.ofMethod(PythonWrapperGenerator.class, "getValueFromPythonObject", Object.class,
                            OpaquePythonReference.class, String.class),
                    value, methodCreator.load(methodName));

            if (returnType instanceof Class) {
                Class<?> returnTypeClass = (Class<?>) returnType;
                if (Comparable.class.isAssignableFrom(returnTypeClass) || Number.class.isAssignableFrom(returnTypeClass)
                        || OpaquePythonReference.class.isAssignableFrom(returnTypeClass)) {
                    // It is a number/String, so it already translated to the corresponding Java type
                    if (Integer.class.equals(returnTypeClass)) {
                        ResultHandle isLong = methodCreator.instanceOf(outResultHandle, Long.class);
                        BranchResult ifLongBranchResult = methodCreator.ifTrue(isLong);
                        BytecodeCreator bytecodeCreator = ifLongBranchResult.trueBranch();
                        bytecodeCreator.writeInstanceField(fieldDescriptor, bytecodeCreator.getThis(),
                                bytecodeCreator.invokeVirtualMethod(MethodDescriptor.ofMethod(Long.class, "intValue",
                                        int.class), outResultHandle));
                        bytecodeCreator = ifLongBranchResult.falseBranch();
                        bytecodeCreator.writeInstanceField(fieldDescriptor, methodCreator.getThis(), outResultHandle);
                    } else {
                        methodCreator.writeInstanceField(fieldDescriptor, methodCreator.getThis(), outResultHandle);
                    }
                } else {
                    // We need to wrap it
                    ResultHandle actualClass;

                    if (returnTypeClass.isArray()) {
                        actualClass = methodCreator.loadClass(returnTypeClass);
                    } else if (Collection.class.isAssignableFrom(returnTypeClass)) {
                        actualClass = methodCreator.loadClass(PythonList.class);
                    } else {
                        actualClass = methodCreator.invokeStaticMethod(
                                MethodDescriptor.ofMethod(PythonWrapperGenerator.class, "getJavaClass", Class.class,
                                        OpaquePythonReference.class),
                                outResultHandle);
                    }
                    methodCreator.writeInstanceField(fieldDescriptor, methodCreator.getThis(),
                            methodCreator.invokeStaticMethod(MethodDescriptor.ofMethod(PythonWrapperGenerator.class,
                                    "wrap", Object.class, Class.class, OpaquePythonReference.class, Map.class),
                                    actualClass, outResultHandle, methodCreator.getMethodParam(2)));
                }
            } else {
                // It a reference to the current class; need to be wrapped
                ResultHandle actualClass = methodCreator.loadClass(classCreator.getClassName());
                methodCreator.writeInstanceField(fieldDescriptor, methodCreator.getThis(),
                        methodCreator.invokeStaticMethod(MethodDescriptor.ofMethod(PythonWrapperGenerator.class,
                                "wrap", Object.class, Class.class, OpaquePythonReference.class, Map.class),
                                actualClass, outResultHandle, methodCreator.getMethodParam(2)));
            }
        }
        methodCreator.returnValue(methodCreator.getThis());
    }

    private static FieldDescriptor generateWrapperMethod(ClassCreator classCreator, FieldDescriptor valueField,
            String methodName, Class<?> returnType, String signature, List<Map<String, Object>> annotations,
            List<Object> returnTypeList) {
        // Python types are not required, so we need to discover them. If the type is unknown, we default to Object,
        // but some annotations need something more specific than Object
        Object actualReturnType = returnType;
        for (Map<String, Object> annotation : annotations) {
            if (PlanningId.class.isAssignableFrom((Class<?>) annotation.get("annotationType"))
                    && !Comparable.class.isAssignableFrom(returnType)) {
                // A PlanningId MUST be comparable
                actualReturnType = Comparable.class;
            } else if ((ProblemFactCollectionProperty.class.isAssignableFrom((Class<?>) annotation.get("annotationType"))
                    || PlanningEntityCollectionProperty.class.isAssignableFrom((Class<?>) annotation.get("annotationType"))
                    || ValueRangeProvider.class.isAssignableFrom((Class<?>) annotation.get("annotationType")))
                    && !(Collection.class.isAssignableFrom(returnType) || returnType.isArray())) {
                // A ProblemFactCollection/PlanningEntityCollection MUST be a collection or array
                actualReturnType = List.class;
            } else if (SelfType.class.equals(returnType)) {
                actualReturnType = classCreator.getClassName();
            }
        }
        returnTypeList.add(actualReturnType);
        FieldDescriptor fieldDescriptor =
                classCreator.getFieldCreator(methodName + "$field", actualReturnType).getFieldDescriptor();

        String javaMethodName = methodName;
        if (javaMethodName.startsWith("get_") && javaMethodName.length() >= 5) {
            javaMethodName = "get" + Character.toUpperCase(javaMethodName.charAt(4)) + javaMethodName.substring(5);
        }

        MethodCreator methodCreator = classCreator.getMethodCreator(javaMethodName, actualReturnType);
        if (signature != null) {
            methodCreator.setSignature("()" + signature);
        }

        // Create method annotations for each annotation in the list
        for (Map<String, Object> annotation : annotations) {
            // The class representing the annotation is in the annotationType parameter.
            AnnotationCreator annotationCreator = methodCreator.addAnnotation((Class<?>) annotation.get("annotationType"));
            createAnnotation(annotationCreator, annotation);
        }

        // Getter is simply: return this.field
        methodCreator.returnValue(methodCreator.readInstanceField(fieldDescriptor, methodCreator.getThis()));

        // Assumption: all getters have a setter
        if (methodName.startsWith("get")) {
            String setterMethodName = "set" + methodName.substring(3);
            String javaSetterMethodName = "set" + javaMethodName.substring(3);
            MethodCreator setterMethodCreator =
                    classCreator.getMethodCreator(javaSetterMethodName, void.class, actualReturnType);
            if (signature != null) {
                setterMethodCreator.setSignature("(" + signature + ")V;");
            }

            // Use PythonWrapperGenerator.setValueOnPythonObject(obj, attribute, value)
            // to set the value on the Python Object
            setterMethodCreator.invokeStaticMethod(
                    MethodDescriptor.ofMethod(PythonWrapperGenerator.class, "setValueOnPythonObject", void.class,
                            OpaquePythonReference.class, String.class, Object.class),
                    setterMethodCreator.readInstanceField(valueField, setterMethodCreator.getThis()),
                    setterMethodCreator.load(setterMethodName),
                    setterMethodCreator.getMethodParam(0));
            // Update the field on the Pojo
            setterMethodCreator.writeInstanceField(fieldDescriptor, setterMethodCreator.getThis(),
                    setterMethodCreator.getMethodParam(0));
            setterMethodCreator.returnValue(null);
        }
        return fieldDescriptor;
    }

    private static void createAnnotation(AnnotationCreator annotationCreator, Map<String, Object> annotation) {
        Class<?> annotationType = (Class<?>) annotation.get("annotationType");
        for (Method method : annotationType.getMethods()) {
            if (method.getParameterCount() != 0
                    || !method.getDeclaringClass().equals(annotation.get("annotationType"))) {
                // skip if the parameter is not from the actual annotation (toString, hashCode, etc.)
                continue;
            }
            Object annotationValue = convertAnnotationValue(annotation.get(method.getName()));

            if (annotationValue != null) {
                annotationCreator.addValue(method.getName(), annotationValue);
            }
        }
    }

    private static Object convertAnnotationValue(Object annotationValue) {
        if (annotationValue == null) {
            return null;
        }
        if (annotationValue.getClass().isArray()) {
            int arrayLength = Array.getLength(annotationValue);
            Object[] out = new Object[arrayLength];
            for (int i = 0; i < out.length; i++) {
                out[i] = convertAnnotationValue(Array.get(annotationValue, i));
            }
            return out;
        }
        if (annotationValue instanceof List) {
            List<?> annotationValueList = (List<?>) annotationValue;
            Object[] out = new Object[annotationValueList.size()];
            for (int i = 0; i < out.length; i++) {
                out[i] = convertAnnotationValue(annotationValueList.get(i));
            }
            return out;
        } else if (annotationValue instanceof Map) {
            Map<String, Object> nestedAnnotation = (Map<String, Object>) annotationValue;
            Class<?> nestedAnnotationClass = (Class<?>) nestedAnnotation.get("annotationType");
            AnnotationCreator nestedAnnotationValue =
                    AnnotationCreator.of(nestedAnnotationClass.getName(), RetentionPolicy.RUNTIME);
            createAnnotation(nestedAnnotationValue, nestedAnnotation);
            return nestedAnnotationValue;
        } else {
            return annotationValue;
        }
    }
}
