import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.Serializer;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import com.google.protobuf.Message;
import java.lang.reflect.Method;
import java.util.HashMap;

/**
 * Kryo serializer for Protobuf instances. Not thread-safe (neither is Kryo, so that is fine).
 *
 * <p>Registered as the default Kryo serializer for {@link Message} via {@code serialization-config}.
 */
public class ProtobufSerializer extends Serializer<Message> {

    // Forever cache, based on the number of classes seen which _should be_ finite
    // Cache for the `parseFrom(byte[] bytes)` method
    protected final HashMap<Class<?>, Method> methodCache = new HashMap<>();
    // Cache for the `getDefaultInstance()` method
    protected final HashMap<Class<?>, Method> defaultInstanceMethodCache = new HashMap<>();

    /**
     * This is slow, so we should cache to avoid killing perf.
     */
    private Method getMethodFromCache(
            Class<?> clazz,
            HashMap<Class<?>, Method> cache,
            String methodName,
            Class<?>... parameterTypes) throws Exception {

        Method method = cache.get(clazz);
        if (null == method) {
            method = clazz.getMethod(methodName, parameterTypes);
            cache.put(clazz, method);
        }
        return method;
    }

    protected Method getParse(Class<?> clazz) throws Exception {
        return getMethodFromCache(clazz, methodCache, "parseFrom", byte[].class);
    }

    protected Method getDefaultInstance(Class<?> clazz) throws Exception {
        return getMethodFromCache(clazz, defaultInstanceMethodCache, "getDefaultInstance");
    }

    @Override
    public void write(Kryo kryo, Output output, Message message) {
        output.writeBoolean(message != null);
        if (message != null) {
            byte[] byteArray = message.toByteArray();
            output.writeInt(byteArray.length, true);
            output.writeBytes(byteArray);
        }
    }

    @Override
    public Message read(Kryo kryo, Input input, Class<? extends Message> clazz) {
        try {
            if (!input.readBoolean()) {
                return null;
            }
            int size = input.readInt(true);
            if (size < 0) {
                throw new RuntimeException(String.format("Corrupted protobuf data: negative size %s for class: %s",
                        size, clazz.getName()));
            }
            if (size == 0) {
                return (Message) getDefaultInstance(clazz).invoke(null);
            }
            byte[] byteArray = new byte[size];
            input.readBytes(byteArray);

            //noinspection PrimitiveArrayArgumentToVarargsMethod
            return (Message) getParse(clazz).invoke(null, byteArray);
        } catch (Exception e) {
            throw new RuntimeException("Could not deserialize protobuf for class: " + clazz.getName(), e);
        }
    }

    @Override
    public Message copy(Kryo kryo, Message original) {
        try {
            byte[] bytes = original.toByteArray();

            //noinspection PrimitiveArrayArgumentToVarargsMethod
            return (Message) getParse(original.getClass()).invoke(null, bytes);
        } catch (Exception e) {
            throw new RuntimeException("Failed to copy message class: " + original.getClass().getName(), e);
        }
    }
}