/*
 * Copyright (C) 2014-2016 LinkedIn Corp. All rights reserved.
 *
 * Licensed 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.
 */

package gobblin.metrics.reporter.util;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.util.Map;

import org.apache.avro.Schema;
import org.apache.commons.codec.DecoderException;
import org.apache.commons.codec.binary.Hex;

import com.google.common.base.Optional;
import com.google.common.base.Throwables;
import com.google.common.collect.Maps;
import com.typesafe.config.Config;

import gobblin.metrics.kafka.KafkaAvroSchemaRegistry;
import gobblin.metrics.kafka.SchemaRegistryException;
import gobblin.util.ConfigUtils;


/**
 * Implementation of {@link gobblin.metrics.reporter.util.SchemaVersionWriter} that uses a
 * {@link gobblin.metrics.kafka.KafkaAvroSchemaRegistry} to get Schema version identifier and write it to
 * {@link java.io.DataOutputStream}.
 */
public class SchemaRegistryVersionWriter implements SchemaVersionWriter<Schema> {

  private final KafkaAvroSchemaRegistry registry;
  private Map<Schema, String> registrySchemaIds;
  private final Optional<String> overrideName;
  private final Optional<Schema> schema;
  private final Optional<String> schemaId;
  private final int schemaIdLengthBytes;

  public SchemaRegistryVersionWriter(Config config) throws IOException {
    this(new KafkaAvroSchemaRegistry(ConfigUtils.configToProperties(config)), Optional.<String>absent(), Optional.<Schema>absent());
  }

  public SchemaRegistryVersionWriter(KafkaAvroSchemaRegistry registry, String overrideName, Optional<Schema> singleSchema)
      throws IOException {
    this(registry, Optional.of(overrideName), singleSchema);
  }

  public SchemaRegistryVersionWriter(KafkaAvroSchemaRegistry registry, Optional<String> overrideName, Optional<Schema> singleSchema)
      throws IOException {
    this.registry = registry;
    this.registrySchemaIds = Maps.newConcurrentMap();
    this.overrideName = overrideName;
    this.schema = singleSchema;
    this.schemaIdLengthBytes = registry.getSchemaIdLengthByte();
    if (this.schema.isPresent()) {
      try {
        this.schemaId = this.overrideName.isPresent()
            ? Optional.of(this.registry.register(this.schema.get(), this.overrideName.get()))
            : Optional.of(this.registry.register(this.schema.get()));
      } catch (SchemaRegistryException e) {
        throw Throwables.propagate(e);
      }
    } else {
      this.schemaId = Optional.absent();
    }
  }

  @Override
  public void writeSchemaVersioningInformation(Schema schema, DataOutputStream outputStream) throws IOException {

    String schemaId = this.schemaId.isPresent() ? this.schemaId.get() : this.getIdForSchema(schema);

    outputStream.writeByte(KafkaAvroSchemaRegistry.MAGIC_BYTE);
    try {
      outputStream.write(Hex.decodeHex(schemaId.toCharArray()));
    } catch (DecoderException exception) {
      throw new IOException(exception);
    }
  }

  private String getIdForSchema(Schema schema) {
    if (!this.registrySchemaIds.containsKey(schema)) {
      try {
        String schemaId = this.overrideName.isPresent()
            ? this.registry.register(schema, this.overrideName.get())
            : this.registry.register(schema);
        this.registrySchemaIds.put(schema, schemaId);
      } catch (SchemaRegistryException e) {
        throw Throwables.propagate(e);
      }
    }
    return this.registrySchemaIds.get(schema);
  }

  @Override
  public Schema readSchemaVersioningInformation(DataInputStream inputStream) throws IOException {
    if (inputStream.readByte() != KafkaAvroSchemaRegistry.MAGIC_BYTE) {
      throw new IOException("MAGIC_BYTE not found in Avro message.");
    }

    byte[] byteKey = new byte[schemaIdLengthBytes];
    int bytesRead = inputStream.read(byteKey, 0, schemaIdLengthBytes);
    if (bytesRead != schemaIdLengthBytes) {
      throw new IOException(String.format("Could not read enough bytes for schema id. Expected: %d, found: %d.",
          schemaIdLengthBytes, bytesRead));
    }
    String hexKey = Hex.encodeHexString(byteKey);

    try {
      return this.registry.getSchemaByKey(hexKey);
    } catch (SchemaRegistryException sre) {
      throw new IOException("Failed to retrieve schema for key " + hexKey, sre);
    }
  }
}

@@ -1738,16 +1738,22 @@
  config)
+%0A     
  throws 
@@ -1870,16 +1870,24 @@
 bsent(),
+%0A       
  Optiona
@@ -1994,33 +1994,224 @@
 ing overrideName
-,
+)%0A      throws IOException %7B%0A    this(registry, overrideName, Optional.%3CSchema%3Eabsent());%0A  %7D%0A%0A  public SchemaRegistryVersionWriter(KafkaAvroSchemaRegistry registry, String overrideName,%0A     
  Optional%3CSchema
@@ -2412,32 +2412,38 @@
 g%3E overrideName,
+%0A     
  Optional%3CSchema
@@ -2796,32 +2796,43 @@
 Name.isPresent()
+ ? Optional
 %0A            ? O
@@ -2828,26 +2828,16 @@
         
-? Optional
 .of(this
@@ -3222,32 +3222,38 @@
 am outputStream)
+%0A     
  throws IOExcept
@@ -3765,28 +3765,16 @@
 resent()
-%0A           
  ? this.
@@ -4161,16 +4161,22 @@
 tStream)
+%0A     
  throws 
@@ -4539,16 +4539,27 @@
 n(String
+%0A          
 .format(
@@ -4628,26 +4628,16 @@
 d: %25d.%22,
-%0A         
  schemaI
@@ -4649,16 +4649,30 @@
 thBytes,
+%0A             
  bytesRe

Add backward compatible constructor