/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.giraph.writable.kryo;

import java.io.DataInput;
import java.io.DataOutput;
import java.util.Arrays;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Random;

import com.esotericsoftware.kryo.util.DefaultClassResolver;
import de.javakaffee.kryoserializers.guava.ImmutableMapSerializer;
import org.apache.giraph.conf.GiraphConfigurationSettable;
import com.esotericsoftware.kryo.ClassResolver;
import com.esotericsoftware.kryo.ReferenceResolver;
import com.esotericsoftware.kryo.util.MapReferenceResolver;
import org.apache.giraph.types.ops.collections.Basic2ObjectMap;
import org.apache.giraph.types.ops.collections.BasicSet;
import org.apache.giraph.writable.kryo.markers.NonKryoWritable;
import org.apache.giraph.writable.kryo.markers.KryoIgnoreWritable;
import org.apache.giraph.writable.kryo.serializers.ArraysAsListSerializer;
import org.apache.giraph.writable.kryo.serializers.CollectionsNCopiesSerializer;
import org.apache.giraph.writable.kryo.serializers.DirectWritableSerializer;
import org.apache.giraph.writable.kryo.serializers.FastUtilSerializer;
import org.apache.giraph.writable.kryo.serializers.ImmutableBiMapSerializerUtils;
import org.apache.giraph.writable.kryo.serializers.ReusableFieldSerializer;
import org.apache.hadoop.conf.Configurable;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.Writable;
import org.apache.log4j.Logger;
import org.objenesis.strategy.StdInstantiatorStrategy;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.Serializer;
import com.esotericsoftware.kryo.factories.SerializerFactory;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.InputChunked;
import com.esotericsoftware.kryo.io.Output;
import com.esotericsoftware.kryo.io.OutputChunked;
import com.esotericsoftware.kryo.pool.KryoCallback;
import com.esotericsoftware.kryo.pool.KryoFactory;
import com.esotericsoftware.kryo.pool.KryoPool;
import com.esotericsoftware.kryo.serializers.ClosureSerializer;
import com.esotericsoftware.kryo.serializers.FieldSerializer;
import com.esotericsoftware.kryo.util.ObjectMap;
import com.google.common.base.Preconditions;

import de.javakaffee.kryoserializers.guava.ImmutableListSerializer;

/**
 * Kryo instance that provides serialization through DataInput/DataOutput
 * that org.apache.hadoop.io.Writable uses.
 *
 * All public APIs are static.
 *
 * It extends Kryo to reuse KryoPool functionality, but have additional needed
 * objects cached as well. If we move to ThreadLocal or other caching
 * technique, we can use composition, instead of inheritance here.
 *
 * TODO: Refactor this class into two separate classes depending on
 * whether the reference tracking is enabled or disabled.
 */
public class HadoopKryo extends Kryo {
  /** Pool of reusable Kryo objects, since they are expensive to create */
  private static final KryoPool KRYO_POOL = new KryoPool.Builder(
      new KryoFactory() {
        @Override
        public Kryo create() {
          return createKryo(true, true);
        }
      }).build();
  /** Thread local HadoopKryo object */
  private static final ThreadLocal<HadoopKryo> KRYO =
    new ThreadLocal<HadoopKryo>() {
      @Override protected HadoopKryo initialValue() {
        return  createKryo(false, false);
      }
    };

  /**
   * List of interfaces/parent classes that will not be allowed to be
   * serialized, together with explanation of why, that will be shown
   * when throwing such exception
   */
  private static final Map<Class<?>, String> NON_SERIALIZABLE;

  static {
    NON_SERIALIZABLE = new LinkedHashMap<>();
    NON_SERIALIZABLE.put(
        NonKryoWritable.class,
        "it is marked to not allow serialization, " +
        "look at the class for more details");
    NON_SERIALIZABLE.put(
        KryoWritableWrapper.class, "recursion is disallowed");
    NON_SERIALIZABLE.put(
        Configuration.class,
        "it cannot be supported since it contains ClassLoader");
    NON_SERIALIZABLE.put(
        GiraphConfigurationSettable.class, "configuration cannot be set");
    NON_SERIALIZABLE.put(
        Configurable.class, "configuration cannot be set");
    NON_SERIALIZABLE.put(
        Random.class,
        "it should be rarely serialized, since it would create same stream " +
        "of numbers everywhere, use TransientRandom instead");
    NON_SERIALIZABLE.put(
        Logger.class,
        "Logger must be a static field");
  }

  /** Reusable Input object */
  private InputChunked input;
  /** Reusable Output object */
  private OutputChunked output;

  /** Reusable DataInput wrapper stream */
  private DataInputWrapperStream dataInputWrapperStream;
  /** Reusable DataOutput wrapper stream */
  private DataOutputWrapperStream dataOutputWrapperStream;

  /**
   * Map of already initialized serializers used
   * for readIntoObject/writeOutOfObject pair of methods
   */
  private final ObjectMap<Class<?>, ReusableFieldSerializer<Object>>
  classToIntoSerializer = new ObjectMap<>();

  /** Hide constructor, so all access go through pool of cached objects */
  private HadoopKryo() {
  }

  /**
   * Constructor that takes custom class resolver and reference resolver.
   * @param classResolver Class resolver
   * @param referenceResolver Reference resolver
   */
  private HadoopKryo(ClassResolver classResolver,
                     ReferenceResolver referenceResolver) {
    super(classResolver, referenceResolver);
  }

  // Public API:

  /**
   * Write type of given object and the object itself to the output stream.
   * Inverse of readClassAndObj.
   *
   * @param out Output stream
   * @param object Object to write
   */
  public static void writeClassAndObj(
          final DataOutput out, final Object object) {
    writeInternal(out, object, false);
  }

  /**
   * Read object from the input stream, by reading first type of the object,
   * and then all of its fields.
   * Inverse of writeClassAndObject.
   *
   * @param in Input stream
   * @return Deserialized object
   * @param <T> Type of the object being read
   */
  public static <T> T readClassAndObj(DataInput in) {
    return readInternal(in, null, false);
  }

  /**
   * Write an object to output, in a way that can be read by readIntoObject.
   *
   * @param out Output stream
   * @param object Object to be written
   */
  public static void writeOutOfObject(
      final DataOutput out, final Object object) {
    writeInternal(out, object, true);
  }

  /**
   * Reads an object, from input, into a given object,
   * allowing object reuse.
   * Inverse of writeOutOfObject.
   *
   * @param in Input stream
   * @param object Object to fill from input
   */
  public static void readIntoObject(DataInput in, Object object) {
    readInternal(in, object, true);
  }

  /**
   * Writes class and object to specified output stream with specified
   * Kryo object. It does not use interim buffers.
   * @param kryo Kryo object
   * @param out Output stream
   * @param object Object
   */
  public static void writeWithKryo(
          final HadoopKryo kryo, final Output out,
          final Object object) {
    kryo.writeClassAndObject(out, object);
    out.close();
  }

  /**
   * Write out of object with given kryo
   * @param kryo Kryo object
   * @param out Output
   * @param object Object to write
   */
  public static void writeWithKryoOutOfObject(
          final HadoopKryo kryo, final Output out,
          final Object object) {
    kryo.writeOutOfObject(out, object);
    out.close();
  }

  /**
   * Reads class and object from specified input stream with
   * specified kryo object.
   * it does not use interim buffers.
   * @param kryo Kryo object
   * @param in Input buffer
   * @param <T> Object type parameter
   * @return Object
   */
  public static <T> T readWithKryo(
          final HadoopKryo kryo, final Input in) {
    T object;
    object = (T) kryo.readClassAndObject(in);
    in.close();
    return object;
  }

  /**
   * Read into object with given kryo.
   * @param kryo Kryo object
   * @param in Input
   * @param object Object to read into
   */
  public static void readWithKryoIntoObject(
          final HadoopKryo kryo, final Input in, Object object) {
    kryo.readIntoObject(in, object);
    in.close();
  }

  /**
   * Create copy of the object, by magically recursively copying
   * all of its fields, keeping reference structures (like cycles)
   *
   * @param object Object to be copied
   * @return Copy of the object.
   * @param <T> Type of the object
   */
  public static <T> T createCopy(final T object) {
    return KRYO_POOL.run(new KryoCallback<T>() {
      @Override
      public T execute(Kryo kryo) {
        return kryo.copy(object);
      }
    });
  }

  /**
   * Returns a kryo which doesn't track objects, hence
   * serialization of recursive/nested objects is not
   * supported.
   *
   * Reference tracking significantly degrades the performance
   * since kryo has to store all serialized objects and search
   * the history to check if an object has been already serialized.
   *
   * @return Hadoop kryo which doesn't track objects.
   */
  public static HadoopKryo getNontrackingKryo() {
    return KRYO.get();
  }

  // Private implementation:

  /**
   * Create new instance of HadoopKryo, properly initialized.
   * @param trackReferences if true, object references are tracked.
   * @param hasBuffer if true, an interim buffer is used.
   * @return new HadoopKryo instance
   */
  private static HadoopKryo createKryo(boolean trackReferences,
                                       boolean hasBuffer) {
    HadoopKryo kryo;
    if (trackReferences) {
      kryo = new HadoopKryo();
    } else {
      // Only use GiraphClassResolver if it is properly initialized.
      // This is to enable test cases which use KryoSimpleWrapper
      // but don't start ZK.
      kryo = new HadoopKryo(
              GiraphClassResolver.isInitialized() ? new GiraphClassResolver() :
                                                    new DefaultClassResolver(),
              new MapReferenceResolver());
    }

    try {
      kryo.register(Class.forName("java.lang.invoke.SerializedLambda"));
      kryo.register(Class.forName(
        "com.esotericsoftware.kryo.serializers.ClosureSerializer$Closure"),
        new ClosureSerializer());
    } catch (ClassNotFoundException e) {
      throw new IllegalStateException(
        "Trying to use Kryo on Java version " +
          System.getProperty("java.version") +
          ", but unable to find needed classes", e);
    }

    kryo.register(Arrays.asList().getClass(), new ArraysAsListSerializer());
    kryo.register(Collections.nCopies(1, new Object()).getClass(),
        new CollectionsNCopiesSerializer());

    ImmutableListSerializer.registerSerializers(kryo);
    ImmutableMapSerializer.registerSerializers(kryo);
    ImmutableBiMapSerializerUtils.registerSerializers(kryo);

    // There are many fastutil classes, register them at the end,
    // so they don't use up small registration numbers
    FastUtilSerializer.registerAll(kryo);

    kryo.setInstantiatorStrategy(new DefaultInstantiatorStrategy(
        new StdInstantiatorStrategy()));

    SerializerFactory customSerializerFactory = new SerializerFactory() {
      @SuppressWarnings("rawtypes")
      @Override
      public Serializer makeSerializer(Kryo kryo, final Class<?> type) {
        for (final Entry<Class<?>, String> entry :
            NON_SERIALIZABLE.entrySet()) {
          if (entry.getKey().isAssignableFrom(type)) {
            // Allow Class object to be serialized, but not a live instance.
            return new Serializer() {
              @Override
              public Object read(Kryo kryo, Input input, Class type) {
                throw new RuntimeException("Cannot serialize " + type +
                    ". Objects being serialized cannot capture " +
                    entry.getKey() + " because " + entry.getValue() +
                    ". Either remove field in question" +
                    ", or make it transient (so that it isn't serialized)");
              }

              @Override
              public void write(Kryo kryo, Output output, Object object) {
                throw new RuntimeException("Cannot serialize " + type +
                    ". Objects being serialized cannot capture " +
                    entry.getKey() + " because " + entry.getValue() +
                    ". Either remove field in question" +
                    ", or make it transient (so that it isn't serialized)");
              }
            };
          }
        }

        if (Writable.class.isAssignableFrom(type) &&
            !KryoIgnoreWritable.class.isAssignableFrom(type) &&
            // remove BasicSet, BasicArrayList and Basic2ObjectMap temporarily,
            // for lack of constructors
            !BasicSet.class.isAssignableFrom(type) &&
            !Basic2ObjectMap.class.isAssignableFrom(type)) {
          // use the Writable method defined by the type
          DirectWritableSerializer serializer = new DirectWritableSerializer();
          return serializer;
        } else {
          FieldSerializer serializer = new FieldSerializer<>(kryo, type);
          serializer.setIgnoreSyntheticFields(false);
          return serializer;
        }
      }
    };

    kryo.addDefaultSerializer(Writable.class, customSerializerFactory);
    kryo.setDefaultSerializer(customSerializerFactory);

    if (hasBuffer) {
      kryo.input = new InputChunked(4096);
      kryo.output = new OutputChunked(4096);
      kryo.dataInputWrapperStream = new DataInputWrapperStream();
      kryo.dataOutputWrapperStream = new DataOutputWrapperStream();
    }

    if (!trackReferences) {
      kryo.setReferences(false);

      // Auto reset can only be disabled if the GiraphClassResolver is
      // properly initialized.
      if (GiraphClassResolver.isInitialized()) {
        kryo.setAutoReset(false);
      }
    }
    return kryo;
  }

  /**
   * Register serializer for class with class name
   *
   * @param kryo HadoopKryo
   * @param className Name of the class for which to register serializer
   * @param serializer Serializer to use
   */
  private static void registerSerializer(HadoopKryo kryo, String className,
      Serializer serializer) {
    try {
      kryo.register(Class.forName(className), serializer);
    } catch (ClassNotFoundException e) {
      throw new IllegalStateException("Class " + className + " is missing", e);
    }
  }

  /**
   * Initialize reusable objects for reading from given DataInput.
   *
   * @param in Input stream
   */
  private void setDataInput(DataInput in) {
    dataInputWrapperStream.setDataInput(in);
    input.setInputStream(dataInputWrapperStream);
  }

  /**
   * Initialize reusable objects for writing into given DataOutput.
   *
   *  @param out Output stream
   */
  private void setDataOutput(DataOutput out) {
    dataOutputWrapperStream.setDataOutput(out);
    output.setOutputStream(dataOutputWrapperStream);
  }

  /**
   * Get or create reusable serializer for given class.
   *
   * @param type Type of the object
   * @return Serializer
   */
  private ReusableFieldSerializer<Object> getOrCreateReusableSerializer(
      Class<?> type) {
    ReusableFieldSerializer<Object> serializer =
        classToIntoSerializer.get(type);
    if (serializer == null) {
      serializer = new ReusableFieldSerializer<>(this, type);
      classToIntoSerializer.put(type, serializer);
    }
    return serializer;
  }

  /**
   * Internal write implementation, that reuses HadoopKryo objects
   * from the pool.
   *
   * @param out Output stream
   * @param object Object to be written
   * @param outOf whether we are writing reusable objects,
   *              or full objects with class name
   */
  private static void writeInternal(
      final DataOutput out, final Object object, final boolean outOf) {
    KRYO_POOL.run(new KryoCallback<Void>() {
      @Override
      public Void execute(Kryo kryo) {
        HadoopKryo hkryo = (HadoopKryo) kryo;
        hkryo.setDataOutput(out);

        if (outOf) {
          hkryo.writeOutOfObject(hkryo.output, object);
        } else {
          hkryo.writeClassAndObject(hkryo.output, object);
        }

        hkryo.output.endChunks();
        hkryo.output.close();

        return null;
      }
    });
  }

  /**
   * Internal read implementation, that reuses HadoopKryo objects
   * from the pool.
   *
   * @param in Input stream
   * @param outObject Object to fill from input (if not null)
   * @param into whether we are reading reusable objects,
   *             or full objects with class name
   * @return Read object (new one, or same passed in if we use reusable)
   * @param <T> Type of the object to read
   */
  @SuppressWarnings("unchecked")
  private static <T> T readInternal(
      final DataInput in, final T outObject, final boolean into) {
    return KRYO_POOL.run(new KryoCallback<T>() {
      @Override
      public T execute(Kryo kryo) {
        HadoopKryo hkryo = (HadoopKryo) kryo;
        hkryo.setDataInput(in);

        T object;
        if (into) {
          hkryo.readIntoObject(hkryo.input, outObject);
          object = outObject;
        } else {
          object = (T) hkryo.readClassAndObject(hkryo.input);
        }
        hkryo.input.nextChunks();

        hkryo.input.close();
        return object;
      }
    });
  }

  /**
   * Reads an object, from input, into a given object,
   * allowing object reuse.
   *
   * @param input Input stream
   * @param object Object to fill from input
   */
  private void readIntoObject(Input input, Object object) {
    Preconditions.checkNotNull(object);

    Class<?> type = object.getClass();
    ReusableFieldSerializer<Object> serializer =
        getOrCreateReusableSerializer(type);

    serializer.setReadIntoObject(object);
    Object result = readObject(input, type, serializer);

    Preconditions.checkState(result == object);
  }

  /**
   * Write an object to output, in a way that can be read
   * using readIntoObject.
   * @param output Output stream
   * @param object Object to be written
   */
  private void writeOutOfObject(Output output, Object object) {
    ReusableFieldSerializer<Object> serializer =
        getOrCreateReusableSerializer(object.getClass());
    writeObject(output, object, serializer);
  }

}
