// Dstl (c) Crown Copyright 2017
package uk.gov.dstl.baleen.resources.gazetteer;

import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import org.apache.uima.resource.Resource;
import org.bson.Document;
import org.bson.types.ObjectId;

import com.mongodb.client.MongoCollection;

import uk.gov.dstl.baleen.exceptions.BaleenException;
import uk.gov.dstl.baleen.exceptions.InvalidParameterException;
import uk.gov.dstl.baleen.resources.SharedMongoResource;

/** Connect to MongoDB and use as the backend for a Gazetteer */
public class MongoGazetteer extends AbstractMultiMapGazetteer<ObjectId> {
  public static final String CONFIG_COLLECTION = "collection";
  public static final String CONFIG_VALUE_FIELD = "valueField";

  public static final String DEFAULT_COLLECTION = "gazetteer";
  public static final String DEFAULT_VALUE_FIELD = "value";

  private MongoCollection<Document> coll;

  private String valueField;

  /**
   * Configure a new instance of MongoGazetteer. The following config parameters are
   * expected/allowed:
   *
   * <ul>
   *   <li><b>collection</b> - The collection containing the gazetteer; defaults to gazetteer
   *   <li><b>valueField</b> - The field containing the gazetteer values; defaults to value
   * </ul>
   *
   * @param connection A SharedMongoResource object to use to connect to Mongo
   * @param config A map of additional configuration options
   * @throws InvalidParameterException if the passed connection parameter is not a
   *     SharedMongoResource
   */
  @Override
  public void init(Resource connection, Map<String, Object> config) throws BaleenException {
    SharedMongoResource mongo;
    try {
      mongo = (SharedMongoResource) connection;
    } catch (ClassCastException cce) {
      throw new InvalidParameterException(
          "Unable to cast connection parameter to SharedMongoResource", cce);
    }

    valueField = DEFAULT_VALUE_FIELD;
    if (config.containsKey(CONFIG_VALUE_FIELD)) {
      valueField = config.get(CONFIG_VALUE_FIELD).toString();
    }

    String collection = DEFAULT_COLLECTION;
    if (config.containsKey(CONFIG_COLLECTION)) {
      collection = config.get(CONFIG_COLLECTION).toString();
    }

    coll = mongo.getDB().getCollection(collection);

    super.init(connection, config);
  }

  @Override
  public Map<String, Object> getAdditionalData(String key) {
    ObjectId id = getId(key);
    if (id == null) {
      return Collections.emptyMap();
    }

    Document doc =
        coll.find(new Document("_id", id))
            .projection(new Document(valueField, 0).append("_id", 0))
            .first();
    if (doc == null) {
      return Collections.emptyMap();
    }

    Map<String, Object> ret = new HashMap<>();

    for (Entry<String, Object> mongoEntry : doc.entrySet()) {
      ret.put(mongoEntry.getKey(), mongoEntry.getValue());
    }

    return ret;
  }

  @Override
  public void destroy() {
    coll = null;
    super.destroy();
  }

  @SuppressWarnings("unchecked")
  @Override
  public void reloadValues() {
    reset();

    for (Document doc : coll.find()) {
      ObjectId id = (ObjectId) doc.get("_id");

      Object val = doc.get(valueField);
      if (val instanceof String) {
        String s = (String) val;

        addTerm(id, s);
      } else if (val instanceof List) {
        for (String s : (List<String>) val) {
          addTerm(id, s);
        }
      }
    }
  }
}
