D-Net

package eu.dnetlib.functionality.index.query;

import java.io.ByteArrayInputStream;

import java.io.IOException;

import java.util.Collection;

import java.util.List;

import java.util.Map;

import java.util.zip.GZIPInputStream;

import com.google.common.base.Predicate;

import com.google.common.collect.BiMap;

import com.google.common.collect.Iterables;

import com.google.common.collect.Lists;

import eu.dnetlib.data.provision.index.rmi.BrowsingRow;

import eu.dnetlib.data.provision.index.rmi.GroupResult;

import eu.dnetlib.functionality.index.utils.IndexFieldUtility;

import eu.dnetlib.miscutils.functional.UnaryFunction;

import org.apache.commons.io.IOUtils;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

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

import org.apache.solr.client.solrj.response.FacetField.Count;

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

import org.apache.solr.common.SolrDocument;

import org.apache.solr.common.SolrDocumentList;

import static eu.dnetlib.miscutils.collections.MappedCollection.listMap;

/**

 * The Class SolrResponseParser.

 */

public class SolrResponseParser extends QueryResponseParser {

	/**

	 * logger.

	 */

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

	/**

	 * Lower level response.

	 */

	private QueryResponse queryRsp = null;

	/** The wrapper rank. */

	protected final UnaryFunction<String, SolrDocument> wrapperRank =

			doc -> addRanking(

					getSingleField(doc, IndexFieldUtility.RESULT),

					getSingleField(doc, IndexFieldUtility.SCORE_FIELD));

	/** The wrapper no rank. */

	protected final UnaryFunction<String, SolrDocument> wrapperNoRank =

			doc -> wrap(getSingleField(doc, IndexFieldUtility.RESULT));

	/**

	 * The Constructor.

	 *

	 * @param highlightUtils

	 *            the highlight utils

	 * @param aliases

	 *            the aliases

	 * @param returnEmptyFields

	 *            the return empty fields

	 * @param includeRanking

	 *            the include ranking

	 * @param response

	 *            the response

	 */

	public SolrResponseParser(final UnaryFunction<String, String> highlightUtils, final BiMap<String, String> aliases, final boolean returnEmptyFields,

			final boolean includeRanking, final QueryResponse response) {

		super(highlightUtils, aliases, returnEmptyFields, includeRanking);

		this.queryRsp = response;

	}

	/**

	 * {@inheritDoc}

	 *

	 * @see eu.dnetlib.functionality.index.query.QueryResponseParser#getNumFound()

	 */

	@Override

	public long getNumFound() {

		return this.queryRsp.getResults().getNumFound();

	}

	/**

	 * {@inheritDoc}

	 *

	 * @see eu.dnetlib.functionality.index.query.QueryResponseParser#getQueryTime()

	 */

	@Override

	public int getQueryTime() {

		return queryRsp.getQTime();

	}

	/**

	 * {@inheritDoc}

	 *

	 * @see eu.dnetlib.functionality.index.query.QueryResponseParser#getElapsedTime()

	 */

	@Override

	public long getElapsedTime() {

		return queryRsp.getElapsedTime();

	}

	/**

	 * {@inheritDoc}

	 *

	 * @see eu.dnetlib.functionality.index.query.QueryResponseParser#getStatus()

	 */

	@Override

	public String getStatus() {

		return String.valueOf(queryRsp.getStatus());

	}

	/**

	 * {@inheritDoc}

	 *

	 * @see eu.dnetlib.functionality.index.query.QueryResponseParser#getCurrentSize()

	 */

	@Override

	public int getCurrentSize() {

		return queryRsp.getResults().size();

	}

	/**

	 * Gets the query response.

	 *

	 * @return the query response

	 */

	public QueryResponse getQueryResponse() {

		return queryRsp;

	}

	/**

	 * {@inheritDoc}

	 *

	 * @see eu.dnetlib.functionality.index.query.QueryResponseParser#getResults()

	 */

	@Override

	public List<String> getResults() {

		return asRankedList(queryRsp.getResults());

	}

	/**

	 * {@inheritDoc}

	 *

	 * @see eu.dnetlib.functionality.index.query.QueryResponseParser#getNumberOfBrowsingResults()

	 */

	@Override

	public Long getNumberOfBrowsingResults() {

		List<FacetField> ffList = queryRsp.getFacetFields();

		Long maxCount = 0L;

		if (ffList != null) {

			for (FacetField ff : ffList) {

				if (ff != null) {

					Long countFacets = countFacets(ff.getValues());

					if (countFacets > maxCount) {

						maxCount = countFacets;

					}

				}

			}

		}

		return maxCount;

	}

	/**

	 * {@inheritDoc}

	 *

	 * @see eu.dnetlib.functionality.index.query.QueryResponseParser#getBrowsingResults()

	 */

	@Override

	public List<BrowsingRow> getBrowsingResults() {

		List<BrowsingRow> bresList = Lists.newArrayList();

		List<GroupResult> facets = Lists.newArrayList();

		final List<FacetField> ffList = queryRsp.getFacetFields();

		Long numberOfBrowsingResults = getNumberOfBrowsingResults();

		for (int i = 0; (ffList != null) && (i < numberOfBrowsingResults); i++) {

			for (FacetField ff : ffList) {

				String name = null;

				if (aliases != null) {

					name = aliases.inverse().get(ff.getName());

				}

				// fix #1456

				if (name == null) {

					name = ff.getName();

				}

				final Count facet = getFacet(ff, i);

				if ((facet != null) && (facet.getCount() > 0)) {

					final String value = facet.getName();

					final int count = (int) facet.getCount();

					if (returnEmptyFields || !value.isEmpty()) {

						facets.add(new GroupResult(name, value, count));

					}

				}

			}

			if (facets.size() > 0) {

				bresList.add(new BrowsingRow(Lists.newArrayList(facets)));

				facets.clear();

			}

		}

		if (log.isDebugEnabled()) {

			log.debug("BrowsingResult size: " + bresList.size());

		}

		return bresList;

	}

	// /////////////// helpers

	/**

	 * Gets the facet.

	 *

	 * @param ff

	 *            the ff

	 * @param pos

	 *            the pos

	 * @return the facet

	 */

	private Count getFacet(final FacetField ff, final int pos) {

		if ((ff.getValues() == null) || (pos >= ff.getValues().size())) return null;

		return ff.getValues().get(pos);

	}

	/**

	 * Given SolrDocumentList, return a List of Strings, representing it.

	 *

	 * @param documentList

	 *            the document list

	 * @return the list< string>

	 */

	private List<String> asRankedList(final SolrDocumentList documentList) {

		UnaryFunction<String, SolrDocument> wrapper = includeRanking ? wrapperRank : wrapperNoRank;

		if (queryRsp.getHighlighting() != null) return listMap(listMap(documentList, doc -> {

			String score = getSingleField(doc, IndexFieldUtility.SCORE_FIELD);

			String hl = getHighlighting(getSingleField(doc, IndexFieldUtility.INDEX_RECORD_ID));

			String res = hl != null ? hl : getSingleField(doc, IndexFieldUtility.RESULT);

			return includeRanking ? addRanking(res, score) : wrap(res);

		}), highlightUtils);

		return listMap(documentList, wrapper);

	}

	/**

	 * Converts a String document to

	 *

	 * <record rank="score"> [document] </record>.

	 *

	 * @param doc

	 *            the doc

	 * @param score

	 *            the score

	 * @return the string

	 */

	private String addRanking(final String doc, final String score) {

		return new String("<record rank=\"" + score + "\">" + doc + "</record>");

	}

	/**

	 * Wraps the given document as <record> [document] </record>.

	 *

	 * @param doc

	 *            the doc

	 * @return the string

	 */

	private String wrap(final String doc) {

		return new String("<record>" + doc + "</record>");

	}

	/**

	 * Gets the single field.

	 *

	 * @param doc

	 *            the doc

	 * @param fieldName

	 *            the field name

	 * @return the single field

	 */

	@SuppressWarnings("unchecked")

	protected String getSingleField(final SolrDocument doc, final String fieldName) {

		Object value = doc.getFieldValue(fieldName);

		if (value instanceof Collection) return Iterables.getOnlyElement((Iterable<String>) value);

		return String.valueOf(value);

	}

	private byte[] base64Decode(final String s) {

		return org.apache.solr.common.util.Base64.base64ToByteArray(s);

	}

	private String unzip(final byte[] b) {

		try (ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(b)) {

			try (GZIPInputStream gzipInputStream = new GZIPInputStream(byteArrayInputStream)) {

				return new String(IOUtils.toByteArray(gzipInputStream));

			}

		} catch(IOException e) {

			throw new RuntimeException("Failed to unzip content", e);

		}

	}

	/**

	 * Gets the highlighting.

	 *

	 * @param docId

	 *            the doc id

	 * @return the highlighting

	 */

	private String getHighlighting(final String docId) {

		final Map<String, List<String>> highlight = queryRsp.getHighlighting().get(docId);

		String result = new String();

		if ((highlight != null) && (highlight.get(IndexFieldUtility.RESULT) != null)) {

			for (String s : highlight.get(IndexFieldUtility.RESULT)) {

				result = result.concat(s);

			}

			return result;

		}

		return null;

	}

	/**

	 * helper method.

	 *

	 * @param facets

	 *            the list of facets to analyze

	 * @return the number of non-empty facets in the list whose count is greater than zero

	 */

	private Long countFacets(final List<Count> facets) {

		if (facets == null) return 0L;

		return (long) Iterables.size(Iterables.filter(facets, new Predicate<Count>() {

			@Override

			public boolean apply(final Count c) {

				return (c != null) && (c.getName() != null) && !c.getName().isEmpty() && (c.getCount() > 0);

			}

		}));

	}

	@Override

	public long getStart() {

		// TODO Auto-generated method stub

		return queryRsp.getResults().getStart();

	}

}