package io.ray.runtime.serializer;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import java.io.IOException;
import java.lang.reflect.Array;
import java.math.BigInteger;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.apache.commons.lang3.mutable.MutableBoolean;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.apache.commons.lang3.tuple.Pair;
import org.msgpack.core.MessageBufferPacker;
import org.msgpack.core.MessagePack;
import org.msgpack.core.MessagePacker;
import org.msgpack.core.MessageUnpacker;
import org.msgpack.value.ArrayValue;
import org.msgpack.value.ExtensionValue;
import org.msgpack.value.IntegerValue;
import org.msgpack.value.Value;
import org.msgpack.value.ValueType;

// We can't pack List / Map by MessagePack, because we don't know the type class when unpacking.
public class MessagePackSerializer {

  private static final byte LANGUAGE_SPECIFIC_TYPE_EXTENSION_ID = 101;
  // MessagePack length is an int takes up to 9 bytes.
  // https://github.com/msgpack/msgpack/blob/master/spec.md#int-format-family
  private static final int MESSAGE_PACK_OFFSET = 9;

  // Pakcers indexed by its corresponding Java class object.
  private static Map<Class<?>, TypePacker> packers = new HashMap<>();
  // Unpackers indexed by its corresponding MessagePack ValueType.
  private static Map<ValueType, TypeUnpacker> unpackers = new HashMap<>();
  // Null and array don't have a corresponding class, so define them separately.
  private static final TypePacker NULL_PACKER;
  private static final TypePacker ARRAY_PACKER;
  private static final TypePacker EXTENSION_PACKER;

  static {
    // ===== Initialize packers =====
    // Null packer.
    NULL_PACKER = (object, packer, javaSerializer) -> packer.packNil();

    // Array packer.
    ARRAY_PACKER =
        ((object, packer, javaSerializer) -> {
          int length = Array.getLength(object);
          packer.packArrayHeader(length);
          for (int i = 0; i < length; ++i) {
            pack(Array.get(object, i), packer, javaSerializer);
          }
        });

    // Extension packer.
    EXTENSION_PACKER =
        ((object, packer, javaSerializer) -> {
          javaSerializer.serialize(object, packer);
        });

    packers.put(
        Boolean.class, ((object, packer, javaSerializer) -> packer.packBoolean((Boolean) object)));
    packers.put(Byte.class, ((object, packer, javaSerializer) -> packer.packByte((Byte) object)));
    packers.put(
        Short.class, ((object, packer, javaSerializer) -> packer.packShort((Short) object)));
    packers.put(
        Integer.class, ((object, packer, javaSerializer) -> packer.packInt((Integer) object)));
    packers.put(Long.class, ((object, packer, javaSerializer) -> packer.packLong((Long) object)));
    packers.put(
        BigInteger.class,
        ((object, packer, javaSerializer) -> packer.packBigInteger((BigInteger) object)));
    packers.put(
        Float.class, ((object, packer, javaSerializer) -> packer.packFloat((Float) object)));
    packers.put(
        Double.class, ((object, packer, javaSerializer) -> packer.packDouble((Double) object)));
    packers.put(
        String.class, ((object, packer, javaSerializer) -> packer.packString((String) object)));
    packers.put(
        byte[].class,
        ((object, packer, javaSerializer) -> {
          byte[] bytes = (byte[]) object;
          packer.packBinaryHeader(bytes.length);
          packer.writePayload(bytes);
        }));

    // ===== Initialize unpackers =====
    List<Class<?>> booleanClasses = ImmutableList.of(Boolean.class, boolean.class);
    List<Class<?>> byteClasses = ImmutableList.of(Byte.class, byte.class);
    List<Class<?>> shortClasses = ImmutableList.of(Short.class, short.class);
    List<Class<?>> intClasses = ImmutableList.of(Integer.class, int.class);
    List<Class<?>> longClasses = ImmutableList.of(Long.class, long.class);
    List<Class<?>> bigIntClasses = ImmutableList.of(BigInteger.class);
    List<Class<?>> floatClasses = ImmutableList.of(Float.class, float.class);
    List<Class<?>> doubleClasses = ImmutableList.of(Double.class, double.class);
    List<Class<?>> stringClasses = ImmutableList.of(String.class);
    List<Class<?>> binaryClasses = ImmutableList.of(byte[].class);

    // Null unpacker.
    unpackers.put(ValueType.NIL, (value, targetClass, javaDeserializer) -> null);
    // Boolean unpacker.
    unpackers.put(
        ValueType.BOOLEAN,
        (value, targetClass, javaDeserializer) -> {
          Preconditions.checkArgument(
              checkTypeCompatible(booleanClasses, targetClass),
              "Boolean can't be deserialized as {}.",
              targetClass);
          return value.asBooleanValue().getBoolean();
        });
    // Integer unpacker.
    unpackers.put(
        ValueType.INTEGER,
        ((value, targetClass, javaDeserializer) -> {
          IntegerValue iv = value.asIntegerValue();
          if (iv.isInByteRange() && checkTypeCompatible(byteClasses, targetClass)) {
            return iv.asByte();
          } else if (iv.isInShortRange() && checkTypeCompatible(shortClasses, targetClass)) {
            return iv.asShort();
          } else if (iv.isInIntRange() && checkTypeCompatible(intClasses, targetClass)) {
            return iv.asInt();
          } else if (iv.isInLongRange() && checkTypeCompatible(longClasses, targetClass)) {
            return iv.asLong();
          } else if (checkTypeCompatible(bigIntClasses, targetClass)) {
            return iv.asBigInteger();
          }
          throw new IllegalArgumentException(
              "Integer can't be deserialized as " + targetClass + ".");
        }));
    // Float unpacker.
    unpackers.put(
        ValueType.FLOAT,
        ((value, targetClass, javaDeserializer) -> {
          if (checkTypeCompatible(doubleClasses, targetClass)) {
            return value.asFloatValue().toDouble();
          } else if (checkTypeCompatible(floatClasses, targetClass)) {
            return value.asFloatValue().toFloat();
          }
          throw new IllegalArgumentException("Float can't be deserialized as " + targetClass + ".");
        }));
    // String unpacker.
    unpackers.put(
        ValueType.STRING,
        ((value, targetClass, javaDeserializer) -> {
          Preconditions.checkArgument(
              checkTypeCompatible(stringClasses, targetClass),
              "String can't be deserialized as {}.",
              targetClass);
          return value.asStringValue().asString();
        }));
    // Binary unpacker.
    unpackers.put(
        ValueType.BINARY,
        ((value, targetClass, javaDeserializer) -> {
          Preconditions.checkArgument(
              checkTypeCompatible(binaryClasses, targetClass),
              "Binary can't be deserialized as {}.",
              targetClass);
          return value.asBinaryValue().asByteArray();
        }));
    // Array unpacker.
    unpackers.put(
        ValueType.ARRAY,
        ((value, targetClass, javaDeserializer) -> {
          ArrayValue av = value.asArrayValue();
          Class<?> componentType =
              targetClass.isArray() ? targetClass.getComponentType() : Object.class;
          Object array = Array.newInstance(componentType, av.size());
          for (int i = 0; i < av.size(); ++i) {
            Array.set(array, i, unpack(av.get(i), componentType, javaDeserializer));
          }
          return array;
        }));
    // Extension unpacker.
    unpackers.put(
        ValueType.EXTENSION,
        ((value, targetClass, javaDeserializer) -> {
          ExtensionValue ev = value.asExtensionValue();
          byte extType = ev.getType();
          if (extType == LANGUAGE_SPECIFIC_TYPE_EXTENSION_ID) {
            return javaDeserializer.deserialize(ev);
          }
          throw new IllegalArgumentException("Unknown extension type id " + ev.getType() + ".");
        }));
  }

  interface JavaSerializer {

    void serialize(Object object, MessagePacker packer) throws IOException;
  }

  interface JavaDeserializer {

    Object deserialize(ExtensionValue v);
  }

  interface TypePacker {

    void pack(Object object, MessagePacker packer, JavaSerializer javaSerializer)
        throws IOException;
  }

  interface TypeUnpacker {

    Object unpack(Value value, Class<?> targetClass, JavaDeserializer javaDeserializer);
  }

  private static boolean checkTypeCompatible(List<Class<?>> expected, Class<?> actual) {
    for (Class<?> expectedClass : expected) {
      if (actual.isAssignableFrom(expectedClass)) {
        return true;
      }
    }
    return false;
  }

  private static void pack(Object object, MessagePacker packer, JavaSerializer javaSerializer)
      throws IOException {
    TypePacker typePacker;
    if (object == null) {
      typePacker = NULL_PACKER;
    } else {
      Class<?> type = object.getClass();
      typePacker = packers.get(type);
      if (typePacker == null) {
        if (type.isArray()) {
          typePacker = ARRAY_PACKER;
        } else {
          typePacker = EXTENSION_PACKER;
        }
      }
    }
    try {
      typePacker.pack(object, packer, javaSerializer);
    } catch (Exception e) {
      if (typePacker != EXTENSION_PACKER) {
        EXTENSION_PACKER.pack(object, packer, javaSerializer);
      } else {
        throw e;
      }
    }
  }

  private static Object unpack(Value v, Class<?> type, JavaDeserializer javaDeserializer) {
    return unpackers.get(v.getValueType()).unpack(v, type, javaDeserializer);
  }

  public static Pair<byte[], Boolean> encode(Object obj) {
    MessageBufferPacker packer = MessagePack.newDefaultBufferPacker();
    try {
      // Reserve MESSAGE_PACK_OFFSET bytes for MessagePack bytes length.
      packer.writePayload(new byte[MESSAGE_PACK_OFFSET]);
      // Serialize input object by MessagePack.
      MutableBoolean isCrossLanguage = new MutableBoolean(true);
      pack(
          obj,
          packer,
          ((object, packer1) -> {
            byte[] payload = FstSerializer.encode(object);
            packer1.packExtensionTypeHeader(LANGUAGE_SPECIFIC_TYPE_EXTENSION_ID, payload.length);
            packer1.addPayload(payload);
            isCrossLanguage.setFalse();
          }));
      byte[] msgpackBytes = packer.toByteArray();
      // Serialize MessagePack bytes length.
      MessageBufferPacker headerPacker = MessagePack.newDefaultBufferPacker();
      Preconditions.checkState(msgpackBytes.length >= MESSAGE_PACK_OFFSET);
      headerPacker.packLong(msgpackBytes.length - MESSAGE_PACK_OFFSET);
      byte[] msgpackBytesLength = headerPacker.toByteArray();
      // Check serialized MessagePack bytes length is valid.
      Preconditions.checkState(msgpackBytesLength.length <= MESSAGE_PACK_OFFSET);
      // Write MessagePack bytes length to reserved buffer.
      System.arraycopy(msgpackBytesLength, 0, msgpackBytes, 0, msgpackBytesLength.length);
      return ImmutablePair.of(msgpackBytes, isCrossLanguage.getValue());
    } catch (Exception e) {
      throw new RuntimeException(e);
    }
  }

  @SuppressWarnings("unchecked")
  public static <T> T decode(byte[] bs, Class<?> type) {
    try {
      // Read MessagePack bytes length.
      MessageUnpacker headerUnpacker = MessagePack.newDefaultUnpacker(bs, 0, MESSAGE_PACK_OFFSET);
      long msgpackBytesLength = headerUnpacker.unpackLong();
      headerUnpacker.close();
      // Check MessagePack bytes length is valid.
      Preconditions.checkState(MESSAGE_PACK_OFFSET + msgpackBytesLength <= bs.length);
      // Deserialize MessagePack bytes from MESSAGE_PACK_OFFSET.
      MessageUnpacker unpacker =
          MessagePack.newDefaultUnpacker(bs, MESSAGE_PACK_OFFSET, (int) msgpackBytesLength);
      Value v = unpacker.unpackValue();
      if (type == null) {
        type = Object.class;
      }
      return (T) unpack(v, type, ((ExtensionValue ev) -> FstSerializer.decode(ev.getData())));
    } catch (Exception e) {
      throw new RuntimeException(e);
    }
  }
}
