/*
    * 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.serializers;
  
  import java.io.IOException;
  import java.io.ObjectInputStream;
  import java.io.ObjectOutputStream;
  import java.lang.reflect.InvocationTargetException;
  import java.lang.reflect.Method;
  import java.util.ArrayList;
  
  import com.esotericsoftware.kryo.Kryo;
  import com.esotericsoftware.kryo.Serializer;
  import com.esotericsoftware.kryo.io.Input;
  import com.esotericsoftware.kryo.io.Output;
  import com.esotericsoftware.kryo.serializers.FieldSerializer;
  
  /**
   * Kryo Serializer for Fastutil collection class.
   * By default, because they extend boxed collections, are being serialized very
   * inefficiently through a lot of temporary object creation.
   *
   * We are relying that fastutil classes are written to be correctly serialized
   * with Java serialization, and have put transient on all array fields and are
   * doing custom efficient serialization in writeObject/readObject methods.
   * This Serializer then swaps ObjectOutputStream for all default fields with
   * FieldSerializer, and then calls appropriate writeObject/readObject methods.
   * We are also relying on defaultWriteObject/defaultReadObject being
   * effectively called first within those methods
   *
   * @param <T> Object type
   */
  public class FastUtilSerializer<T> extends Serializer<T> {
    /** List of all types generated by fastutil */
    private static final String[] PRIMITIVE_TYPES = new String[] {
      "Boolean", "Byte", "Short", "Int", "Long", "Float", "Double", "Char",
      "Object"};
    /** List of all types used as keys in fastutil */
    private static final String[] PRIMITIVE_KEY_TYPES = new String[] {
      "Byte", "Short", "Int", "Long", "Float", "Double", "Char", "Object"};
  
    /** Field serializer for this fastutil class */
    private final FieldSerializer<T> fieldSerializer;
  
    /** Handle to writeObject Method on this fastutil class*/
    private final Method writeMethod;
    /** Handle to readObject Method on this fastutil class*/
    private final Method readMethod;
    /** Reusable output stream wrapper */
    private final FastUtilSerializer.FastutilKryoObjectOutputStream outputWrapper;
    /** Reusable input stream wrapper */
    private final FastUtilSerializer.FastutilKryoObjectInputStream inputWrapper;
  
    /**
     * Creates and initializes new serializer for a given fastutil class.
     * @param kryo Kryo instance
     * @param type Fastutil class
     */
    public FastUtilSerializer(Kryo kryo, Class<T> type) {
      fieldSerializer = new FieldSerializer<>(kryo, type);
      fieldSerializer.setIgnoreSyntheticFields(false);
  
      try {
        writeMethod = type.getDeclaredMethod(
            "writeObject", ObjectOutputStream.class);
        writeMethod.setAccessible(true);
        readMethod = type.getDeclaredMethod(
            "readObject", ObjectInputStream.class);
        readMethod.setAccessible(true);
      } catch (NoSuchMethodException e) {
        throw new RuntimeException(
            "Fastutil class " + type +
            " doesn't have readObject/writeObject methods", e);
      }
  
      try {
        outputWrapper = new FastutilKryoObjectOutputStream();
        inputWrapper = new FastutilKryoObjectInputStream();
      } catch (IOException e) {
        throw new RuntimeException(e);
      }
    }
  
    /**
    * Register serializer for a given fastutil class.
    * @param kryo Kryo instance
    * @param fastutilClass Fastutil class
    */
   public static void register(Kryo kryo, Class<?> fastutilClass) {
     kryo.register(fastutilClass, new FastUtilSerializer<>(kryo, fastutilClass));
   }
 
   /**
    * Registers serializers for all possible fastutil classes.
    *
    * There are many fastutil classes, so it is recommended to call this
    * function at the end, so they fastutil classes don't use up small
    * registration numbers.
    *
    * @param kryo Kryo instance
    */
   public static void registerAll(Kryo kryo) {
     registerArrayLists(kryo);
     registerArrayBigList(kryo);
     registerOpenHashSets(kryo);
     registerArraySets(kryo);
     registerRBTreeSets(kryo);
     registerAVLTreeSets(kryo);
     registerOpenHashMaps(kryo);
     registerRBTreeMaps(kryo);
     registerAVLTreeMaps(kryo);
 
     // Note - HeapPriorityQueues don't extend boxed collection,
     // and so they work out of the box correctly
   }
 
   /**
    * Register all Fastutil ArrayLists.
    *
    * @param kryo Kryo instance
    */
   public static void registerArrayLists(Kryo kryo) {
     registerAll(kryo, singleTypes(
         "it.unimi.dsi.fastutil._t1_s._T1_ArrayList", PRIMITIVE_TYPES));
   }
 
   /**
    * Register all Fastutil ArrayBigLists.
    *
    * @param kryo Kryo instance
    */
   public static void registerArrayBigList(Kryo kryo) {
     registerAll(kryo, singleTypes(
         "it.unimi.dsi.fastutil._t1_s._T1_BigArrayBigList", PRIMITIVE_TYPES));
   }
 
   /**
    * Register all Fastutil OpenHashSets.
    *
    * @param kryo Kryo instance
    */
   public static void registerOpenHashSets(Kryo kryo) {
     registerAll(kryo, singleTypes(
         "it.unimi.dsi.fastutil._t1_s._T1_OpenHashSet", PRIMITIVE_TYPES));
   }
 
   /**
    * Register all Fastutil ArraySets.
    *
    * @param kryo Kryo instance
    */
   public static void registerArraySets(Kryo kryo) {
     registerAll(kryo, singleTypes(
         "it.unimi.dsi.fastutil._t1_s._T1_ArraySet", PRIMITIVE_TYPES));
   }
 
   /**
    * Register all Fastutil RBTreeSets.
    *
    * @param kryo Kryo instance
    */
   public static void registerRBTreeSets(Kryo kryo) {
     registerAll(kryo, singleTypes(
         "it.unimi.dsi.fastutil._t1_s._T1_RBTreeSet", PRIMITIVE_KEY_TYPES));
   }
 
   /**
    * Register all Fastutil AVLTreeSets.
    *
    * @param kryo Kryo instance
    */
   public static void registerAVLTreeSets(Kryo kryo) {
     registerAll(kryo, singleTypes(
         "it.unimi.dsi.fastutil._t1_s._T1_AVLTreeSet", PRIMITIVE_KEY_TYPES));
   }
 
   /**
    * Register all Fastutil OpenHashMaps.
    *
    * @param kryo Kryo instance
    */
   public static void registerOpenHashMaps(Kryo kryo) {
     registerAll(kryo, doubleTypes(
         "it.unimi.dsi.fastutil._t1_s._T1_2_T2_OpenHashMap",
         PRIMITIVE_KEY_TYPES, PRIMITIVE_TYPES));
   }
 
   /**
    * Register all Fastutil RBTreeMaps.
    *
    * @param kryo Kryo instance
    */
   public static void registerRBTreeMaps(Kryo kryo) {
     registerAll(kryo, doubleTypes(
         "it.unimi.dsi.fastutil._t1_s._T1_2_T2_RBTreeMap",
         PRIMITIVE_KEY_TYPES, PRIMITIVE_TYPES));
   }
 
   /**
    * Register all Fastutil AVLTreeMaps.
    *
    * @param kryo Kryo instance
    */
   public static void registerAVLTreeMaps(Kryo kryo) {
     registerAll(kryo, doubleTypes(
         "it.unimi.dsi.fastutil._t1_s._T1_2_T2_AVLTreeMap",
         PRIMITIVE_KEY_TYPES, PRIMITIVE_TYPES));
   }
 
   /**
    * Register all class from the list of classes.
    *
    * @param kryo Kryo instance
    * @param types List of classes
    */
   private static void registerAll(Kryo kryo, ArrayList<Class<?>> types) {
     for (Class<?> type : types) {
       register(kryo, type);
     }
   }
 
   /**
    * Returns list of all classes that are generated by using given
    * pattern, and replacing it with passed list of types.
    * Pattern contains _t1_ and _T1_, for lowercase and actual name.
    *
    * @param pattern Given pattern
    * @param types Given list of strings to replace into pattern
    * @return List of all classes
    */
   private static ArrayList<Class<?>> singleTypes(
       String pattern, String[] types) {
     ArrayList<Class<?>> result = new ArrayList<>();
 
     for (String type : types) {
       try {
         result.add(Class.forName(
             pattern.replaceAll("_T1_", type).replaceAll(
                 "_t1_", type.toLowerCase())));
       } catch (ClassNotFoundException e) {
         throw new RuntimeException(pattern + " " + type, e);
       }
     }
     return result;
   }
 
   /**
    * Returns list of all classes that are generated by using given
    * pattern, and replacing it with passed list of types.
    * Pattern contains two variable pairs: _t1_, _T1_ and _t2_, _T2_,
    * in each pair one for lowercase and one for actual name.
    *
    * @param pattern Given pattern
    * @param types1 Given list of strings to replace t1 into pattern
    * @param types2 Given list of strings to replace t2 into pattern
    * @return List of all classes
    */
   private static ArrayList<Class<?>> doubleTypes(
       String pattern, String[] types1, String[] types2) {
     ArrayList<Class<?>> result = new ArrayList<>();
 
     for (String type1 : types1) {
       for (String type2 : types2) {
         try {
           result.add(Class.forName(
               pattern.replaceAll("_T1_", type1).replaceAll(
                   "_t1_", type1.toLowerCase())
                 .replaceAll("_T2_", type2).replaceAll(
                     "_t2_", type2.toLowerCase())));
         } catch (ClassNotFoundException e) {
           throw new RuntimeException(pattern + " " + type1 + " " + type2, e);
         }
       }
     }
     return result;
   }
 
   @Override
   public void write(Kryo kryo, Output output, T object) {
     fieldSerializer.write(kryo, output, object);
 
     outputWrapper.set(output, kryo);
     try {
       writeMethod.invoke(object, outputWrapper);
     } catch (IllegalAccessException | InvocationTargetException e) {
       throw new RuntimeException("writeObject failed", e);
     }
   }
 
   @Override
   public T read(Kryo kryo, Input input, Class<T> type) {
     T result = fieldSerializer.read(kryo, input, type);
 
     if (result != null) {
       inputWrapper.set(input, kryo);
       try {
         readMethod.invoke(result, inputWrapper);
       } catch (IllegalAccessException | InvocationTargetException e) {
         throw new RuntimeException("readObject failed", e);
       }
     }
 
     return result;
   }
 
   /**
    * Wrapper around ObjectOutputStream that ignores defaultWriteObject (assumes
    * that needed logic was already executed before), and passes all other calls
    * to Output
    */
   private static class FastutilKryoObjectOutputStream
       extends ObjectOutputStream {
     /** Output */
     private Output output;
     /** Kryo */
     private Kryo kryo;
 
     /** Constructor */
     FastutilKryoObjectOutputStream() throws IOException {
       super();
     }
 
     /**
      * Setter
      *
      * @param output Output
      * @param kryo kryo
      */
     public void set(Output output, Kryo kryo) {
       this.output = output;
       this.kryo = kryo;
     }
 
     @Override
     public void defaultWriteObject() throws IOException {
     }
 
     @Override
     public void writeBoolean(boolean val) throws IOException {
       output.writeBoolean(val);
     }
 
     @Override
     public void writeByte(int val) throws IOException  {
       output.writeByte(val);
     }
 
     @Override
     public void writeShort(int val)  throws IOException {
       output.writeShort(val);
     }
 
     @Override
     public void writeChar(int val)  throws IOException {
       output.writeChar((char) val);
     }
 
     @Override
     public void writeInt(int val)  throws IOException {
       output.writeInt(val, false);
     }
 
     @Override
     public void writeLong(long val)  throws IOException {
       output.writeLong(val, false);
     }
 
     @Override
     public void writeFloat(float val) throws IOException {
       output.writeFloat(val);
     }
 
     @Override
     public void writeDouble(double val) throws IOException {
       output.writeDouble(val);
     }
 
     @Override
     protected void writeObjectOverride(Object obj) throws IOException {
       kryo.writeClassAndObject(output, obj);
     }
   }
 
   /**
    * Wrapper around ObjectOutputStream that ignores defaultReadObject
    * (assumes that needed logic was already executed before), and passes
    * all other calls to Output
    */
   private static class FastutilKryoObjectInputStream extends ObjectInputStream {
     /** Input */
     private Input input;
     /** Kryo */
     private Kryo kryo;
 
     /** Constructor */
     FastutilKryoObjectInputStream() throws IOException {
       super();
     }
 
     /**
      * Setter
      *
      * @param input Input
      * @param kryo Kryo
      */
     public void set(Input input, Kryo kryo) {
       this.input = input;
       this.kryo = kryo;
     }
 
     @Override
     public void defaultReadObject() throws IOException, ClassNotFoundException {
     }
 
     @Override
     public boolean readBoolean() throws IOException {
       return input.readBoolean();
     }
 
     @Override
     public byte readByte() throws IOException {
       return input.readByte();
     }
 
     @Override
     public char readChar() throws IOException {
       return input.readChar();
     }
 
     @Override
     public short readShort() throws IOException {
       return input.readShort();
     }
 
     @Override
     public int readInt() throws IOException {
       return input.readInt(false);
     }
 
     @Override
     public long readLong() throws IOException {
       return input.readLong(false);
     }
 
     @Override
     public float readFloat() throws IOException {
       return input.readFloat();
     }
 
     @Override
     public double readDouble() throws IOException {
       return input.readDouble();
     }
 
     @Override
     protected Object readObjectOverride()
         throws IOException, ClassNotFoundException {
       return kryo.readClassAndObject(input);
     }
   }
 }