D-Net

package eu.dnetlib.index.solr.client;

import java.io.IOException;

import java.util.ArrayList;

import java.util.HashMap;

import java.util.List;

import java.util.Map;

import com.google.common.collect.BiMap;

import com.google.common.collect.Maps;

import eu.dnetlib.clients.index.client.AbstractIndexClient;

import eu.dnetlib.clients.index.client.IndexClient;

import eu.dnetlib.clients.index.client.IndexClientException;

import eu.dnetlib.clients.index.client.response.BrowseEntry;

import eu.dnetlib.clients.index.client.response.BrowseValueEntry;

import eu.dnetlib.clients.index.client.response.LookupResponse;

import eu.dnetlib.clients.index.model.Any.ValueType;

import eu.dnetlib.clients.index.query.IndexQueryFactory;

import eu.dnetlib.clients.index.query.QueryLanguage;

import eu.dnetlib.clients.index.query.QueryResponseFactory;

import eu.dnetlib.clients.index.query.QueryResponseParser;

import eu.dnetlib.utils.MetadataReference;

import eu.dnetlib.index.cql.CqlValueTransformerMap;

import eu.dnetlib.index.query.SolrIndexQuery;

import eu.dnetlib.index.query.SolrIndexQueryFactory;

import eu.dnetlib.index.query.SolrIndexQueryResponse;

import eu.dnetlib.index.solr.cql.SolrTypeBasedCqlValueTransformerMapFactory;

import eu.dnetlib.rmi.provision.BrowsingRow;

import eu.dnetlib.rmi.provision.GroupResult;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import org.apache.solr.client.solrj.SolrServerException;

import org.apache.solr.client.solrj.impl.CloudSolrClient;

import org.apache.solr.client.solrj.request.LukeRequest;

import org.apache.solr.client.solrj.response.LukeResponse;

import org.apache.solr.client.solrj.response.LukeResponse.FieldInfo;

import org.apache.solr.client.solrj.response.LukeResponse.FieldTypeInfo;

import org.apache.solr.client.solrj.response.QueryResponse;

/**

 * The Class SolrIndexClient.

 */

public class SolrIndexClient extends AbstractIndexClient implements IndexClient {

	private static final Log log = LogFactory.getLog(SolrIndexClient.class);

	/**

	 * The server.

	 */

	private CloudSolrClient client;

	private SolrIndexQueryFactory solrIndexQueryFactory;

	/**

	 * The query response factory.

	 */

	private QueryResponseFactory<QueryResponse> queryResponseFactory;

	private SolrTypeBasedCqlValueTransformerMapFactory tMapFactory;

	/**

	 * The Constructor.

	 *

	 * @param format            the format

	 * @param layout            the layout

	 * @param interpretation    the interpretation

	 * @param serviceProperties the service properties

	 * @param tMapFactory

	 */

	public SolrIndexClient(final String format, final String layout, final String interpretation, final Map<String, String> serviceProperties,

			final SolrIndexQueryFactory indexQueryFactory, final QueryResponseFactory<QueryResponse> queryResponseFactory,

			final SolrTypeBasedCqlValueTransformerMapFactory tMapFactory) {

		super(format, layout, interpretation, serviceProperties);

		log.debug(String.format("Created a new instance of the index of type %s_%s_%s", format, layout, interpretation));

		this.solrIndexQueryFactory = indexQueryFactory;

		this.queryResponseFactory = queryResponseFactory;

		this.tMapFactory = tMapFactory;

	}

	public static Map<String, ValueType> getStringValueTypeMap(final Map<String, FieldInfo> fieldInfos, final Map<String, FieldTypeInfo> fieldTypeInfos) {

		final Map<String, ValueType> result = Maps.newHashMap();

		for (FieldInfo fieldInfo : fieldInfos.values()) {

			FieldTypeInfo fieldTypeInfo = fieldTypeInfos.get(fieldInfo.getType());

			final String fieldName = fieldTypeInfo.getName().toLowerCase();

			final ValueType fieldType = resolveSolrTypeClassName(fieldName);

			result.put(fieldInfo.getName(), fieldType);

		}

		return result;

	}

	private static ValueType resolveSolrTypeClassName(final String solrTypeName) {

		if (solrTypeName.contains("LongField")) {

			return ValueType.LONG;

		} else if (solrTypeName.contains("IntField")) {

			return ValueType.LONG;

		} else if (solrTypeName.contains("short")) {

			return ValueType.LONG;

		} else if (solrTypeName.contains("float")) {

			return ValueType.DOUBLE;

		} else if (solrTypeName.contains("double")) {

			return ValueType.DOUBLE;

		} else if (solrTypeName.contains("date")) {

			return ValueType.DATETIME;

		} else {

			return ValueType.STRING;

		}

	}

	/**

	 * Do delete.

	 *

	 * @param query the CQL query

	 * @return true, if do delete

	 * @throws IndexClientException the index service exception

	 */

	@Override

	public long delete(final String query) throws IndexClientException {

		try {

			log.debug("delete by query: " + query);

			MetadataReference mdRef = new MetadataReference(getFormat(), getLayout(), getInterpretation());

			SolrIndexQuery translatedQuery = (SolrIndexQuery) solrIndexQueryFactory.getIndexQuery(QueryLanguage.CQL, query, this, mdRef);

			String tquery = translatedQuery.getQuery();

			translatedQuery.setQueryLimit(0);

			SolrIndexQueryResponse rsp = new SolrIndexQueryResponse(client.query(translatedQuery));

			QueryResponseParser responseParser = queryResponseFactory.getQueryResponseParser(rsp, mdRef);

			long total = responseParser.getNumFound();

			client.deleteByQuery(tquery);

			client.commit();

			return total;

		} catch (Exception e) {

			throw new IndexClientException("unable to run delete by query: " + query, e);

		}

	}

	/**

	 * {@inheritDoc}

	 *

	 * @throws IndexClientException

	 */

	@Override

	public List<BrowseEntry> browse(final String query, final List<String> browseFields, final int max) throws IndexClientException {

		MetadataReference mdRef = new MetadataReference(getFormat(), getLayout(), getInterpretation());

		SolrIndexQuery translatedQuery = buildBrowseQuery(query, browseFields, max, mdRef);

		return executeBrowseQuery(query, translatedQuery, mdRef, browseFields);

	}

	/**

	 * {@inheritDoc}

	 *

	 * @throws IndexClientException

	 */

	@Override

	public List<BrowseEntry> browse(final String query, final List<String> browseFields, final int max, final List<String> filterQuery)

			throws IndexClientException {

		MetadataReference mdRef = new MetadataReference(getFormat(), getLayout(), getInterpretation());

		SolrIndexQuery translatedQuery = buildBrowseQuery(query, browseFields, max, mdRef);

		if (filterQuery != null) {

			log.debug("Filter Query:");

			for (String fq : filterQuery) {

				translatedQuery.addFilterQuery(fq);

				log.debug("- " + fq);

			}

		}

		return executeBrowseQuery(query, translatedQuery, mdRef, browseFields);

	}

	public SolrIndexQuery buildBrowseQuery(final String query, final List<String> browseFields, final int max, final MetadataReference mdRef)

			throws IndexClientException {

		log.debug("Browse request for the index collection for query:" + query);

		SolrIndexQuery translatedQuery = (SolrIndexQuery) solrIndexQueryFactory.getIndexQuery(QueryLanguage.CQL, query, this, mdRef);

		translatedQuery.setFacet(true);

		if (browseFields != null) {

			List<String> browsableFields = solrIndexQueryFactory.getBrowsableFields(browseFields, mdRef);

			log.debug("Browsing fields:");

			for (String field : browsableFields) {

				translatedQuery.addFacetField(field);

				log.debug("- " + field);

			}

			translatedQuery.setFacetLimit(max);

			log.debug("max number of browsing field :" + max);

		}

		return translatedQuery;

	}

	private List<BrowseEntry> executeBrowseQuery(final String originalQuery,

			final SolrIndexQuery query,

			final MetadataReference mdRef,

			final List<String> browseFields) throws IndexClientException {

		try {

			SolrIndexQueryResponse response = new SolrIndexQueryResponse(client.query(query));

			QueryResponseParser responseParser = queryResponseFactory.getQueryResponseParser(response, mdRef);

			List<BrowsingRow> results = responseParser.getBrowsingResults();

			List<BrowseEntry> out = convertBrowseEntry(browseFields, results, responseParser.getAliases());

			return out;

		} catch (Throwable e) {

			throw new IndexClientException("Error on executing a query " + originalQuery, e);

		}

	}

	/**

	 * Creates the connection.

	 *

	 * @return the string

	 * @throws IndexClientException the index client exception

	 */

	private String getUrl() throws IndexClientException {

		String address = serviceProperties.get(ZK_ADDRESS);

		if (address == null) { throw new IndexClientException("Unable to load a solr client, missing zk address"); }

		return address;

	}

	/**

	 * Gets the server.

	 *

	 * @return the server

	 * @throws IndexClientException the index client exception

	 */

	public CloudSolrClient getClient() throws IndexClientException {

		if (this.client == null) {

			String url = getUrl();

			log.debug("create new Client " + url);

			client = new CloudSolrClient(url);

			client.setDefaultCollection(String.format("%s_%s_%s", getFormat(), getLayout(), getInterpretation()));

			try {

				client.ping();

			} catch (Exception e) {

				throw new IndexClientException("oops something went wrong", e);

			}

		}

		return client;

	}

	/**

	 * Sets the server.

	 *

	 * @param client the Client

	 */

	public void setClient(final CloudSolrClient client) {

		this.client = client;

	}

	@Override

	public LookupResponse lookup(final String query, final List<String> filterQuery, final int from, final int to) throws IndexClientException {

		log.debug("lookup request for the index collection for query:" + query);

		MetadataReference mdRef = new MetadataReference(getFormat(), getLayout(), getInterpretation());

		SolrIndexQuery translatedQuery = (SolrIndexQuery) solrIndexQueryFactory.getIndexQuery(QueryLanguage.CQL, query, this, mdRef);

		translatedQuery.setQueryOffset(from);

		translatedQuery.setQueryLimit(to - from + 1);

		if (filterQuery != null) {

			for (String fq : filterQuery) {

				translatedQuery.addFilterQuery(fq);

			}

		}

		try {

			SolrIndexQueryResponse response = new SolrIndexQueryResponse(client.query(translatedQuery));

			QueryResponseParser responseParser = queryResponseFactory.getQueryResponseParser(response, mdRef);

			return new LookupResponse(responseParser);

		} catch (Throwable e) {

			throw new IndexClientException("Error on executing a query " + query, e);

		}

	}

	@Override

	public CqlValueTransformerMap getCqlValueTransformerMap(final MetadataReference mdRef) throws IndexClientException {

		try {

			return this.tMapFactory.getIt(readFieldNamesAndTypes(mdRef.toString()));

		} catch (Exception e) {

			throw new IndexClientException(e);

		}

	}

	@Override

	public IndexQueryFactory getIndexQueryFactory() {

		return solrIndexQueryFactory;

	}

	@Override

	public void stop() throws IndexClientException {

		try {

			log.debug("shutdown client: " + serviceProperties.get(ZK_ADDRESS));

			client.close();

		} catch (Throwable e) {

			throw new IndexClientException(e);

		}

	}

	private List<BrowseEntry> convertBrowseEntry(final List<String> browseFields, final List<BrowsingRow> results, final BiMap<String, String> aliases) {

		Map<String, BrowseEntry> mapResult = new HashMap<>();

		for (BrowsingRow row : results) {

			for (GroupResult groupResult : row.getGroupResult()) {

				String name = groupResult.getName();

				List<BrowseValueEntry> valuesEntry;

				BrowseEntry entry;

				if (mapResult.containsKey(name)) {

					entry = mapResult.get(name);

					valuesEntry = entry.getValues();

					if (valuesEntry == null) {

						valuesEntry = new ArrayList<>();

						entry.setValues(valuesEntry);

					}

				} else {

					entry = new BrowseEntry();

					entry.setField(name);

					entry.setLabel(name);

					valuesEntry = new ArrayList<>();

					entry.setValues(valuesEntry);

					mapResult.put(name, entry);

				}

				String value = groupResult.getValue();

				int count = groupResult.getCount();

				BrowseValueEntry entryValue = new BrowseValueEntry(value, count);

				valuesEntry.add(entryValue);

			}

		}

		List<BrowseEntry> out = new ArrayList<>();

		for (String b : browseFields) {

			String inverse = null;

			if (aliases != null) {

				inverse = aliases.get(b) != null ? aliases.get(b) : aliases.inverse().get(b);

			}

			if (mapResult.containsKey(b)) {

				out.add(mapResult.get(b));

			} else if (mapResult.containsKey(inverse) == true) {

				BrowseEntry data = mapResult.get(inverse);

				data.setField(b);

				out.add(data);

			}

		}

		return out;

	}

	private Map<String, ValueType> readFieldNamesAndTypes(final String coreName) throws SolrServerException, IOException, IndexClientException {

		// final SolrServer server = cloudServer.getSolrServer(coreName);

		final LukeRequest request = new LukeRequest();

		request.setShowSchema(true);

		// cloudServer.setDefaultCollection(coreName);

		request.setNumTerms(0);

		final LukeResponse response = request.process(getClient());

		final Map<String, FieldInfo> fieldInfos = response.getFieldInfo();

		final Map<String, LukeResponse.FieldTypeInfo> fieldTypeInfos = response.getFieldTypeInfo();

		return getStringValueTypeMap(fieldInfos, fieldTypeInfos);

	}

}