/*
 * 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.beam.sdk.extensions.kryo;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.KryoException;
import com.esotericsoftware.kryo.io.InputChunked;
import com.esotericsoftware.kryo.io.OutputChunked;
import java.io.EOFException;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import org.apache.beam.sdk.coders.CoderException;
import org.apache.beam.sdk.coders.CustomCoder;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.checkerframework.checker.nullness.qual.Nullable;

/**
 * Coder using Kryo as (de)serialization mechanism. See {@link KryoCoderProvider} to get more
 * details about usage.
 *
 * @param <T> type of element coder can handle
 */
public class KryoCoder<T> extends CustomCoder<T> {

  /**
   * Create a new {@link KryoCoder} with default {@link KryoOptions}.
   *
   * @param <T> type of element this class should decode/encode {@link Kryo} instance used by
   *     returned {@link KryoCoder}
   * @return Newly created a {@link KryoCoder}
   */
  public static <T> KryoCoder<T> of() {
    return of(PipelineOptionsFactory.create(), Collections.emptyList());
  }

  /**
   * Create a new {@link KryoCoder} with default {@link KryoOptions}.
   *
   * @param registrars {@link KryoRegistrar}s which are used to register classes with underlying
   *     kryo instance
   * @param <T> type of element this class should decode/encode {@link Kryo} instance used by
   *     returned {@link KryoCoder}
   * @return Newly created a {@link KryoCoder}
   */
  public static <T> KryoCoder<T> of(KryoRegistrar... registrars) {
    return of(PipelineOptionsFactory.create(), registrars);
  }

  /**
   * Create a new {@link KryoCoder} with default {@link KryoOptions}.
   *
   * @param registrars {@link KryoRegistrar}s which are used to register classes with underlying
   *     kryo instance
   * @param <T> type of element this class should decode/encode {@link Kryo} instance used by
   *     returned {@link KryoCoder}
   * @return Newly created a {@link KryoCoder}
   */
  public static <T> KryoCoder<T> of(List<KryoRegistrar> registrars) {
    return of(PipelineOptionsFactory.create(), registrars);
  }

  /**
   * Create a new {@link KryoCoder}.
   *
   * @param pipelineOptions Options used for coder setup. See {@link KryoOptions} for more details.
   * @param <T> type of element this class should decode/encode {@link Kryo} instance used by
   *     returned {@link KryoCoder}
   * @return Newly created a {@link KryoCoder}
   */
  public static <T> KryoCoder<T> of(PipelineOptions pipelineOptions) {
    return of(pipelineOptions, Collections.emptyList());
  }

  /**
   * Create a new {@link KryoCoder}.
   *
   * @param pipelineOptions Options used for coder setup. See {@link KryoOptions} for more details.
   * @param registrars {@link KryoRegistrar}s which are used to register classes with underlying
   *     kryo instance
   * @param <T> type of element this class should decode/encode {@link Kryo} instance used by
   *     returned {@link KryoCoder}
   * @return Newly created a {@link KryoCoder}
   */
  public static <T> KryoCoder<T> of(PipelineOptions pipelineOptions, KryoRegistrar... registrars) {
    return of(pipelineOptions, Arrays.asList(registrars));
  }

  /**
   * Create a new {@link KryoCoder}.
   *
   * @param pipelineOptions Options used for coder setup. See {@link KryoOptions} for more details.
   * @param registrars {@link KryoRegistrar}s which are used to register classes with underlying
   *     kryo instance
   * @param <T> type of element this class should decode/encode {@link Kryo} instance used by
   *     returned {@link KryoCoder}
   * @return Newly created a {@link KryoCoder}
   */
  public static <T> KryoCoder<T> of(
      PipelineOptions pipelineOptions, List<KryoRegistrar> registrars) {
    final KryoOptions kryoOptions = pipelineOptions.as(KryoOptions.class);
    return new KryoCoder<>(
        new SerializableOptions(
            kryoOptions.getKryoBufferSize(),
            kryoOptions.getKryoReferences(),
            kryoOptions.getKryoRegistrationRequired()),
        registrars);
  }

  /** Serializable wrapper for {@link KryoOptions}. */
  static class SerializableOptions implements Serializable {

    /** Size of input and output buffer. */
    private final int bufferSize;

    /** Enables kryo reference tracking. */
    private final boolean references;

    /** Enables kryo required registration. */
    private final boolean registrationRequired;

    private SerializableOptions(int bufferSize, boolean references, boolean registrationRequired) {
      this.bufferSize = bufferSize;
      this.references = references;
      this.registrationRequired = registrationRequired;
    }

    /**
     * {@link SerializableOptions#bufferSize}.
     *
     * @return buffer size
     */
    int getBufferSize() {
      return bufferSize;
    }

    /**
     * {@link SerializableOptions#references}.
     *
     * @return boolean flag
     */
    boolean getReferences() {
      return references;
    }

    /**
     * {@link SerializableOptions#registrationRequired}.
     *
     * @return boolean flag
     */
    boolean getRegistrationRequired() {
      return registrationRequired;
    }
  }

  /** Unique id of the {@link KryoCoder} instance. */
  private final String instanceId = UUID.randomUUID().toString();

  /** Options for underlying kryo instance. */
  private final SerializableOptions options;

  /** Client-defined class registrations to {@link Kryo}. */
  private final List<KryoRegistrar> registrars;

  private KryoCoder(SerializableOptions options, List<KryoRegistrar> registrars) {
    this.options = options;
    this.registrars = registrars;
  }

  @Override
  public void encode(T value, OutputStream outStream) throws IOException {
    final KryoState kryoState = KryoState.get(this);
    if (value == null) {
      throw new CoderException("Cannot encode a null value.");
    }
    final OutputChunked outputChunked = kryoState.getOutputChunked();
    outputChunked.setOutputStream(outStream);
    try {
      kryoState.getKryo().writeClassAndObject(outputChunked, value);
      outputChunked.endChunk();
      outputChunked.flush();
    } catch (KryoException e) {
      outputChunked.reset();
      if (e.getCause() instanceof EOFException) {
        throw (EOFException) e.getCause();
      }
      throw new CoderException("Cannot encode given object of type [" + value.getClass() + "].", e);
    } catch (IllegalArgumentException e) {
      String message = e.getMessage();
      if (message != null) {
        if (message.startsWith("Class is not registered")) {
          throw new CoderException(message);
        }
      }
      throw e;
    }
  }

  @Override
  public T decode(InputStream inStream) throws IOException {
    final KryoState kryoState = KryoState.get(this);
    final InputChunked inputChunked = kryoState.getInputChunked();
    inputChunked.setInputStream(inStream);
    try {
      @SuppressWarnings("unchecked")
      final T instance = (T) kryoState.getKryo().readClassAndObject(inputChunked);
      return instance;
    } catch (KryoException e) {
      throw new CoderException("Cannot decode object from input stream.", e);
    }
  }

  @Override
  public void verifyDeterministic() throws NonDeterministicException {
    // noop
  }

  /**
   * Create a new {@link KryoCoder} instance with the user provided registrar.
   *
   * @param registrar registrar to append to list of already registered registrars.
   * @return new kryo coder
   */
  public KryoCoder<T> withRegistrar(KryoRegistrar registrar) {
    final List<KryoRegistrar> newRegistrars = new ArrayList<>(registrars);
    registrars.add(registrar);
    return new KryoCoder<>(options, newRegistrars);
  }

  /**
   * {@link KryoCoder#instanceId}.
   *
   * @return instance id
   */
  String getInstanceId() {
    return instanceId;
  }

  /**
   * {@link KryoCoder#options}.
   *
   * @return options
   */
  SerializableOptions getOptions() {
    return options;
  }

  /**
   * {@link KryoCoder#registrars}.
   *
   * @return registrars
   */
  List<KryoRegistrar> getRegistrars() {
    return registrars;
  }

  @Override
  public int hashCode() {
    return instanceId.hashCode();
  }

  @Override
  public boolean equals(@Nullable Object other) {
    if (other != null && getClass().equals(other.getClass())) {
      return instanceId.equals(((KryoCoder<?>) other).instanceId);
    }
    return false;
  }
}
