package com.atlassian.vcache.marshallers;
   
   import java.io.ByteArrayInputStream;
   import java.io.ByteArrayOutputStream;
   import java.io.IOException;
   import java.io.InputStream;
   import java.io.ObjectInputStream;
   import java.io.ObjectOutputStream;
   import java.io.ObjectStreamClass;
  import java.io.Serializable;
  import java.util.Optional;
  import javax.annotation.Nonnull;
  
  import com.atlassian.annotations.Internal;
  import com.atlassian.vcache.Marshaller;
  import com.atlassian.vcache.MarshallerException;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * Implementation for {@link Serializable} objects that uses the standard Java Serialization mechanism. Any
   * IOExceptions are wrapped as {@link MarshallerException}s.
   *
   * @param <T> the type being marshalled
   *
   * @since 1.0
   */
  @Internal
  class JavaSerializationMarshaller<T extends Serializable> implements Marshaller<T>
  {
      private final Class<T> clazz;
      private final Optional<ClassLoader> loader;
  
      /**
       * Creates an instance that uses the default implementation of
       * {@link java.io.ObjectInputStream#resolveClass(ObjectStreamClass)} to resolve {@link Class} objects.
       */
      JavaSerializationMarshaller(Class<T> clazz)
      {
          this.clazz = requireNonNull(clazz);
          this.loader = Optional.empty();
      }
  
      /**
       * Creates an instance that uses the supplied {@link ClassLoader} to resolve
       * {@link Class} objects.
       */
      JavaSerializationMarshaller(Class<T> clazz, ClassLoader loader)
      {
          this.clazz = requireNonNull(clazz);
          this.loader = Optional.of(loader);
      }
  
      @Nonnull
      @Override
      public byte[] marshall(T obj) throws MarshallerException
      {
          final ByteArrayOutputStream baos = new ByteArrayOutputStream();
          try (ObjectOutputStream oos = new ObjectOutputStream(baos))
          {
              oos.writeObject(obj);
          }
          catch (IOException ioe)
          {
              throw new MarshallerException("Unable to marshall", ioe);
          }
  
          return baos.toByteArray();
      }
  
      @Nonnull
      @Override
      public T unmarshall(byte[] raw) throws MarshallerException
      {
          final ByteArrayInputStream bais = new ByteArrayInputStream(raw);
          try (ObjectInputStream ois = createObjectInputStream(bais))
          {
              return clazz.cast(ois.readObject());
          }
          catch (ClassCastException | ClassNotFoundException | IOException ex)
          {
              throw new MarshallerException("Unable to unmarshall", ex);
          }
      }
  
      private ObjectInputStream createObjectInputStream(InputStream istr) throws IOException
      {
          return loader.isPresent()
                  ? new ObjectInputStreamWithLoader(istr, loader.get())
                  : new ObjectInputStream(istr);
      }
  
      private static class ObjectInputStreamWithLoader extends ObjectInputStream
      {
          private final ClassLoader loader;
  
          public ObjectInputStreamWithLoader(InputStream in, ClassLoader loader) throws IOException
          {
              super(in);
             this.loader = requireNonNull(loader);
         }
 
         @Override
         protected Class<?> resolveClass(ObjectStreamClass desc) throws IOException, ClassNotFoundException
         {
             return Class.forName(desc.getName(), false, loader);
         }
     }
 }