/*
 * Copyright 2017 StreamSets Inc.
 *
 * 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.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.streamsets.pipeline.solr.impl;

import com.streamsets.pipeline.api.StageException;
import com.streamsets.pipeline.solr.api.Errors;
import com.streamsets.pipeline.solr.api.SdcSolrTarget;
import com.streamsets.pipeline.solr.api.SolrInstanceAPIType;
import org.apache.solr.client.solrj.SolrClient;
import org.apache.solr.client.solrj.SolrServerException;
import org.apache.solr.client.solrj.impl.CloudSolrClient;
import org.apache.solr.client.solrj.impl.HttpClientUtil;
import org.apache.solr.client.solrj.impl.HttpSolrClient;
import org.apache.solr.client.solrj.request.schema.SchemaRequest;
import org.apache.solr.client.solrj.response.schema.SchemaRepresentation;
import org.apache.solr.client.solrj.response.schema.SchemaResponse;
import org.apache.solr.common.SolrInputDocument;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.MalformedURLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

public class SolrTarget06 implements SdcSolrTarget {
  private final static Logger LOG = LoggerFactory.getLogger(SolrTarget06.class);
  private final static String VERSION = "6.1.0";
  private SolrClient solrClient;

  private final String solrURI;
  private final String zookeeperConnect;
  private final String defaultCollection;
  private final String instanceType;
  private final boolean kerberosAuth;
  private final boolean skipValidation;
  private final boolean ignoreOptionalFields;
  private final boolean waitFlush;
  private final boolean waitSearcher;
  private final boolean softCommit;
  private final int connectionTimeout;
  private final int socketTimeout;
  private List<String> requiredFieldNamesMap;
  private List<String> optionalFieldNamesMap;

  public SolrTarget06(
      String instanceType,
      String solrURI,
      String zookeeperConnect,
      String defaultCollection,
      boolean kerberosAuth,
      boolean skipValidation,
      boolean waitFlush,
      boolean waitSearcher,
      boolean softCommit,
      boolean ignoreOptionalFields,
      int connectionTimeout,
      int socketTimeout
  ) {
    this.instanceType = instanceType;
    this.solrURI = solrURI;
    this.zookeeperConnect = zookeeperConnect;
    this.defaultCollection = defaultCollection;
    this.kerberosAuth = kerberosAuth;
    this.skipValidation = skipValidation;
    this.waitFlush = waitFlush;
    this.waitSearcher = waitSearcher;
    this.softCommit = softCommit;
    this.ignoreOptionalFields = ignoreOptionalFields;
    this.requiredFieldNamesMap = new ArrayList<>();
    this.optionalFieldNamesMap = new ArrayList<>();
    this.connectionTimeout = connectionTimeout;
    this.socketTimeout = socketTimeout;
  }

  public void init() throws Exception {
    solrClient = getSolrClient();
    if (!skipValidation) {
      solrClient.ping();
    }
    getRequiredFieldNames();
    if (!ignoreOptionalFields) {
      getOptionalFieldNames();
    }
  }

  public List<String> getRequiredFieldNamesMap() {
    return requiredFieldNamesMap;
  }

  private void getRequiredFieldNames() throws SolrServerException, IOException {
    SchemaRequest schemaRequest = new SchemaRequest();
    SchemaResponse schemaResponse = schemaRequest.process(solrClient);
    SchemaRepresentation schemaRepresentation = schemaResponse.getSchemaRepresentation();
    List<Map<String, Object>> fields = schemaRepresentation.getFields();
    for (Map<String, Object> field : fields) {
      if (field.containsKey(REQUIRED) && field.get(REQUIRED).equals(true)) {
        requiredFieldNamesMap.add(field.get(NAME).toString());
      }
    }
  }

  public List<String> getOptionalFieldNamesMap() {
    return optionalFieldNamesMap;
  }

  private void getOptionalFieldNames() throws SolrServerException, IOException {
    SchemaRequest schemaRequest = new SchemaRequest();
    SchemaResponse schemaResponse = schemaRequest.process(solrClient);
    SchemaRepresentation schemaRepresentation = schemaResponse.getSchemaRepresentation();
    List<Map<String, Object>> fields = schemaRepresentation.getFields();
    for (Map<String, Object> field : fields) {
      if (!field.containsKey(REQUIRED) || field.get(REQUIRED).equals(false)) {
        optionalFieldNamesMap.add(field.get(NAME).toString());
      }
    }
  }

  private SolrClient getSolrClient() throws MalformedURLException {
    if(kerberosAuth) {
      // set kerberos before create SolrClient
      addSecurityProperties();
    }

    if (SolrInstanceAPIType.SINGLE_NODE.toString().equals(this.instanceType)) {
      HttpSolrClient solrClient = new HttpSolrClient.Builder(this.solrURI).build();
      solrClient.setConnectionTimeout(this.connectionTimeout);
      solrClient.setSoTimeout(this.socketTimeout);
      return solrClient;
    } else {
      CloudSolrClient cloudSolrClient = new CloudSolrClient.Builder().withZkHost(this.zookeeperConnect).build();
      cloudSolrClient.setDefaultCollection(this.defaultCollection);
      return cloudSolrClient;
    }
  }

  private void addSecurityProperties() {
    HttpClientUtil.setConfigurer(new SdcKrb5HttpClientConfigurer());
  }

  public void destroy() throws IOException{
    if(this.solrClient != null) {
      this.solrClient.close();
    }
  }

  @SuppressWarnings("unchecked")
  public void add(Map<String, Object> fieldMap) throws StageException {
    SolrInputDocument document = createDocument(fieldMap);
    try {
      this.solrClient.add(document);
    } catch (SolrServerException | IOException ex) {
      throw new StageException(Errors.SOLR_04, ex.toString(), ex);
    }
  }

  @SuppressWarnings("unchecked")
  public void add(List<Map<String, Object>> fieldMaps) throws StageException {
    List<SolrInputDocument> documents = new ArrayList();
    for(Map<String, Object> fieldMap : fieldMaps) {
      SolrInputDocument document = createDocument(fieldMap);
      documents.add(document);
    }

    try {
      this.solrClient.add(documents);
    } catch (SolrServerException | IOException | HttpSolrClient.RemoteSolrException ex) {
      throw new StageException(Errors.SOLR_04, ex.toString(), ex);
    }
  }

  private SolrInputDocument createDocument(Map<String, Object> fieldMap) {
    SolrInputDocument document = new SolrInputDocument();
    for (Map.Entry<String, Object> entry : fieldMap.entrySet()) {
      document.addField(entry.getKey(), entry.getValue());
    }
    return document;
  }


  public void commit() throws StageException {
    try {
      this.solrClient.commit(waitFlush, waitSearcher, softCommit);
    } catch (SolrServerException | IOException ex) {
      throw new StageException(Errors.SOLR_05, ex.toString(), ex);
    }
  }

  public String getVersion() {
    return this.VERSION;
  }
}