/**
 * Copyright 2014 Thomas Feng
 *
 * 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 me.tfeng.play.mongodb;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import org.apache.avro.Schema;
import org.apache.avro.Schema.Field;
import org.apache.avro.generic.GenericData.Record;
import org.apache.avro.generic.GenericDatumReader;
import org.apache.avro.generic.IndexedRecord;
import org.apache.avro.io.Decoder;
import org.apache.avro.specific.SpecificDatumReader;
import org.bson.types.Binary;
import org.mortbay.util.ajax.JSON;

import com.mongodb.BasicDBObject;
import com.mongodb.DBObject;

/**
 * @author Thomas Feng (huining.feng@gmail.com)
 */
public class RecordConverter {

  public static final String MONGO_CLASS_PROPERTY = "mongo-class";

  public static final String MONGO_NAME_PROPERTY = "mongo-name";

  public static final String MONGO_TYPE_PROPERTY = "mongo-type";

  public static DBObject toDbObject(IndexedRecord record) {
    DBObject dbObject = new BasicDBObject();
    Schema schema = record.getSchema();
    for (Field field : schema.getFields()) {
      Object value = record.get(field.pos());
      if (value != null) {
        dbObject.put(getFieldName(field), getDbObject(field.schema(), value));
      }
    }
    return dbObject;
  }

  public static <T extends IndexedRecord> T toRecord(Class<T> recordClass, DBObject dbObject) {
    SpecificDatumReader<T> reader = new SpecificDatumReader<T>(recordClass);
    try {
      Decoder decoder = new DBObjectDecoder(recordClass, dbObject);
      return reader.read(null, decoder);
    } catch (IOException e) {
      throw new RuntimeException("Unable to convert MongoDB object " + dbObject
          + " into Avro record", e);
    }
  }

  public static Record toRecord(Schema schema, DBObject object, ClassLoader classLoader)
      throws IOException {
    GenericDatumReader<Record> reader = new GenericDatumReader<>(schema);
    return reader.read(null, new DBObjectDecoder(schema, object, classLoader));
  }

  protected static String getFieldName(Field field) {
    String mongoName = field.getProp(MONGO_NAME_PROPERTY);
    if (mongoName != null) {
      return mongoName;
    } else {
      return field.name();
    }
  }

  @SuppressWarnings("unchecked")
  private static Object getDbObject(Schema schema, Object object) {
    if (object == null) {
      return null;
    } else if (object instanceof IndexedRecord) {
      return toDbObject((IndexedRecord) object);
    } else if (object instanceof Collection) {
      return getDbObjects(schema, (Collection<Object>) object);
    } else if (object instanceof Map) {
      return getDbObjects(schema, (Map<String, Object>) object);
    } else if (object instanceof ByteBuffer) {
      return new Binary(((ByteBuffer) object).array());
    } else if (object.getClass().isEnum()) {
      return ((Enum<?>) object).name();
    } else {
      String mongoClassName = schema.getProp(MONGO_CLASS_PROPERTY);
      String mongoType = schema.getProp(MONGO_TYPE_PROPERTY);
      if (object instanceof CharSequence) {
        object = ((CharSequence) object).toString();
      }
      if (mongoClassName == null && mongoType == null) {
        return object;
      } else if (mongoClassName != null && mongoType != null) {
        throw new RuntimeException("mongo-class and mongo-type should not be both specified: "
            + schema);
      } else {
        try {
          Class<?> mongoClass;
          if (mongoClassName != null) {
            mongoClass = schema.getClass().getClassLoader().loadClass(mongoClassName);
          } else {
            mongoClass = MongoType.valueOf(mongoType).getMongoClass();
          }
          if (object instanceof String && mongoClass.isAssignableFrom(Object.class)) {
            return JSON.parse((String) object);
          } else {
            return MongoDbTypeConverter.convertToMongoDbType(mongoClass, object);
          }
        } catch (ClassNotFoundException e) {
          throw new RuntimeException("Unable to load mongo-class " + mongoClassName, e);
        }
      }
    }
  }

  private static List<Object> getDbObjects(Schema schema, Collection<Object> collection) {
    return collection.stream().map(object -> getDbObject(schema.getElementType(), object))
        .collect(Collectors.toList());
  }

  private static Map<String, Object> getDbObjects(Schema schema, Map<String, Object> map) {
    Map<String, Object> newMap = new HashMap<>(map.size());
    map.entrySet().forEach(
        entry -> newMap.put(entry.getKey(), getDbObject(schema.getValueType(), entry.getValue())));
    return newMap;
  }
}

@@ -1127,16 +1127,52 @@
 .Field;%0A
+import org.apache.avro.Schema.Type;%0A
 import o
@@ -3245,24 +3245,583 @@
 t) %7B%0A    if 
+(schema.getType() == Type.UNION) %7B%0A      List%3CSchema%3E types = schema.getTypes();%0A      if (types.size() != 2%0A          && types.get(0).getType() != Type.NULL%0A          && types.get(1).getType() != Type.NULL) %7B%0A        throw new RuntimeException(%0A            %22In a union type, only null unioned with exactly one other type is supported: %22%0A                + schema);%0A      %7D%0A      if (types.get(0).getType() == Type.NULL) %7B%0A        return getDbObject(types.get(1), object);%0A      %7D else %7B%0A        return getDbObject(types.get(0), object);%0A      %7D%0A    %7D else if 
 (object == n

Fix an issue with union types in mongodb record conversion.