/**
 * Copyright (c) 2025 SoftCorporation LLC. All rights reserved.
 *
 * The Software License, Version 1.0
 *
 * SoftCorporation LLC. grants you ("Licensee") a non-exclusive, royalty free,
 * license to use, modify and redistribute this software in source and binary
 * code form, provided that the following conditions are met:
 *
 * 1. The end-user documentation included with the redistribution,
 *    if any, must include the following acknowledgment:
 *       "This product includes software developed by the
 *        SoftCorporation LLC. (http://www.softcorporation.com)."
 *    Alternately, this acknowledgment may appear in the software itself,
 *    if and wherever such third-party acknowledgments normally appear.
 *
 * 2. This software and SoftCorporation name must not be used to
 *    promote products derived from the software without prior
 *    written permission. For written permission, please contact
 *    info@softcorporation.com.
 *
 * This software is provided "AS IS," without a warranty of any kind.
 * ALL EXPRESS OR IMPLIED CONDITIONS, REPRESENTATIONS AND WARRANTIES,
 * INCLUDING ANY IMPLIED WARRANTY OF MERCHANTABILITY, FITNESS FOR A
 * PARTICULAR PURPOSE OR NON-INFRINGEMENT, ARE HEREBY EXCLUDED.
 * IN NO EVENT SHALL THE SOFTCORPORATION BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
 * USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION).
 *
 */
package com.softcorporation.sandi.search.service.solr;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.apache.solr.client.solrj.SolrQuery;
import org.apache.solr.client.solrj.SolrQuery.ORDER;
import org.apache.solr.client.solrj.SolrRequest.METHOD;
import org.apache.solr.client.solrj.SolrServerException;
import org.apache.solr.client.solrj.response.Group;
import org.apache.solr.client.solrj.response.GroupCommand;
import org.apache.solr.client.solrj.response.GroupResponse;
import org.apache.solr.client.solrj.response.QueryResponse;
import org.apache.solr.common.SolrDocument;
import org.apache.solr.common.SolrDocumentList;
import org.json.JSONArray;
import org.json.JSONObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import com.softcorporation.sandi.base.model.Client;
import com.softcorporation.sandi.base.model.TextChunk;
import com.softcorporation.sandi.base.service.solr.SolrService;
import com.softcorporation.sandi.base.util.ClientException;
import com.softcorporation.sandi.base.util.Constants;
import com.softcorporation.sandi.base.util.Utils;
import com.softcorporation.sandi.search.model.SearchRequest;
import com.softcorporation.sandi.search.model.SearchResponse;
import com.softcorporation.sandi.search.model.SearchResult;
import com.softcorporation.sandi.search.service.SearchService;
import com.softcorporation.sandi.search.util.SearchException;
import com.softcorporation.util.StringUtil;

@Service
public class SolrSearchService extends SearchService {

    private static final Logger LOGGER = LoggerFactory.getLogger(SolrSearchService.class);
    
    @Autowired
    private SolrService solrService;

    public SearchResponse search(SearchRequest searchRequest) throws SearchException {
    	long timeStart = System.currentTimeMillis();
    	String requestId = searchRequest.getRequestId();
        String query = searchRequest.getQuery();
        String searchQuery = searchRequest.getSearchQuery();
        String filterQuery = searchRequest.getFilterQuery();

        init(searchRequest);
        
        try {
        	List<String> searchCollections = Utils.stringToList(searchRequest.getSearchCollections(), false);
            List<String> searchClientIds = Utils.stringToList(searchRequest.getSearchClients(), false);

        	String collections = getCollections(searchRequest, searchCollections, searchClientIds);
        	List<String> filterList = getFilters(searchRequest, filterQuery, searchClientIds);
        	
            LOGGER.debug("[{}] Searching for: '{}' with filter query: '{}' in collections: '{}' with filters: '{}'", 
            		requestId, searchQuery, filterQuery, collections, filterList);
        
            // default search query
            if (searchQuery == null) {
            	searchQuery = "*";
            } else {
                searchQuery = cleanSearchQuery(searchQuery);
                if (searchQuery.isEmpty()) {
                	searchQuery = "*";
                }
            }
            searchRequest.setSearchQuery(searchQuery);
            searchRequest.setCollections(collections);
            searchRequest.setFilterList(filterList);
            
            SearchResponse searchResponse;
            if (query == null) {
              	searchResponse = sandiSearch(searchRequest);
            } else {
            	searchResponse = customSearch(searchRequest);
            }

        	long timeFinish = System.currentTimeMillis();
            searchResponse.setTook((int) (timeFinish - timeStart));
            return searchResponse;
        } catch (Exception e) {
            LOGGER.error("[{}] Search failed for query '{}': {}", requestId, searchQuery, e.getMessage());
            throw new SearchException(e.getMessage());
        }
    }

    public SearchResponse sandiSearch(final SearchRequest searchRequest) 
    		throws SolrServerException, IOException, SearchException {
    	
    	String searchQuery = searchRequest.getSearchQuery();
    	String requestId = searchRequest.getRequestId();
        LOGGER.debug("[{}] Search query: {}", requestId, searchQuery);

        SolrQuery solrQuery = new SolrQuery();

        // fields
        boolean needChunks = setFields(requestId, searchRequest, solrQuery);
        
        // query
        if ("*".equals(searchQuery)) {
        	searchRequest.setLegacy(true);
        	searchRequest.setExpand(false);
        	searchRequest.setRerank(false);
        	searchRequest.setRag(false);
        }
        if (searchRequest.isLegacy()) {
            setLegacySearchQuery(searchRequest, solrQuery);
        } else {
            if (searchRequest.isExpand()) {
            	String expandedQuery = languageModelService.expandSearchQuery(searchRequest.getClient(), searchQuery);
            	searchRequest.setExpandedQuery(expandedQuery);
            }
        	setVectorSearchQuery(searchRequest, solrQuery);
        }
        solrQuery.set("distrib.singlePass", false);

        // pagination
        setPagination(requestId, searchRequest, solrQuery);

        // sorting
        setSorting(requestId, searchRequest, solrQuery);
        
        // grouping
        setGroups(requestId, searchRequest, solrQuery);

        // faceting
        setFacets(requestId, searchRequest, solrQuery);
        
        // highlighting
        setHighlight(requestId, searchRequest, solrQuery);

        if (searchRequest.isDebug()) {
        	solrQuery.setParam("debugQuery", Constants.TRUE);
//        	solrQuery.setParam("debug", Constants.TRUE);
        }
        LOGGER.debug("[{}] Final Solr query: {}", requestId, solrQuery);
        
        QueryResponse response = solrService.getSolrClient().query(searchRequest.getCollections(), solrQuery, METHOD.POST);
        LOGGER.debug("[{}] Searching collections: {} took {} ms", requestId, searchRequest.getCollections(), response.getQTime());

        Map<String, List<SolrDocument>> chunkMap = null;
        SearchResponse searchResponse;
        GroupResponse groupResponse = response.getGroupResponse();
        if (groupResponse == null) {
            SolrDocumentList solrDocuments = response.getResults();
    		if (solrDocuments.size() > 0 && needChunks) {
                chunkMap = searchChunks(searchRequest, solrDocuments);
            }
            searchResponse = processSearchResults(searchRequest, response, solrDocuments, chunkMap);
        } else {
    		if (groupResponse.getValues().size() > 0 && needChunks) {
                chunkMap = searchChunks(searchRequest, groupResponse);
            }
        	searchResponse = processGroupResults(searchRequest, response, groupResponse, chunkMap);
        }
        searchResponse.setStart((long) solrQuery.getStart());
        
        LOGGER.debug("[{}] Search returned {} document(s) out of {} found", requestId, 
        		searchResponse.getResults().size(), searchResponse.getFoundResults());

        return searchResponse;
    }

    /**
     * Custom search //TODO
     * 
     * @param searchRequest
     * @return
     * @throws SolrServerException
     * @throws IOException
     */
    public SearchResponse customSearch(final SearchRequest searchRequest) throws SolrServerException, IOException {
    	
    	String requestId = searchRequest.getRequestId();
        String query = searchRequest.getQuery();
        String collections = searchRequest.getCollections();
        List<String> filterList = searchRequest.getFilterList();
        LOGGER.debug("[{}] Custom search query: {}", requestId, query);

        SolrQuery solrQuery = new SolrQuery();
        solrQuery.setQuery(query);
        
        LOGGER.debug("[{}] Final solr query: {}", requestId, solrQuery);
        
        QueryResponse response = solrService.getSolrClient().query(collections, solrQuery);
        SearchResponse searchResponse = new SearchResponse(requestId, "SUCCESS", "Search completed. Solr query time: " + response.getQTime() + " ms");
        GroupResponse groupResponse = response.getGroupResponse();
        if (groupResponse != null) {
            searchResponse.setResults(groupResponse.getValues());
//            searchResponse.setTotalResults(groupResponse.getNumFound());
        } else {
            SolrDocumentList solrResults = response.getResults();
            searchResponse.setFoundResults(solrResults.getNumFound());
            searchResponse.setResults(solrResults);
        }
        searchResponse.setStart((long) solrQuery.getStart());
        
//        LOGGER.debug("[{}] Search found {} document(s) out of {} total", requestId, 
//        		solrResults.size(), solrResults.getNumFound());

        return searchResponse;
    }
    
    /**
     * Process search results
     * 
     * @param searchRequest
     * @param response
     * @param solrDocuments
     * @param chunkDocuments
     * @return
     */
    SearchResponse processSearchResults(final SearchRequest searchRequest, final QueryResponse response, 
    		final SolrDocumentList solrDocuments, final Map<String, List<SolrDocument>> chunkMap) {
        SearchResponse searchResponse = new SearchResponse(searchRequest.getRequestId(), "SUCCESS", "Solr took "
    		+ response.getQTime() +	" ms searching for query: " + searchRequest.getSearchQuery());
        
        searchResponse.setFoundResults(solrDocuments.getNumFound());
        LOGGER.debug("[{}] Processing flat results. Found docs: {}", searchRequest.getRequestId(), searchResponse.getFoundResults());

		List<SearchResult> searchResults = new ArrayList<>(solrDocuments.size());

        if (solrDocuments.size() > 0) {
            boolean noGreatMatches = false;
            // search result: _id, rank, document (solrDocument), result (final result), client
        	for (int i = 0; i < solrDocuments.size(); i++) {
        		SolrDocument solrDocument = solrDocuments.get(i);
        		String id = (String) solrDocument.getFieldValue(Constants._ID);
    			SearchResult searchResult = new SearchResult(id, i + 1, solrDocument);
    			searchResults.add(searchResult);
    			if (i == 0) {
    				Float score = (Float) solrDocument.getFieldValue(Constants.SCORE);
    			    if (score < searchScoreResultsBad) {
    				    noGreatMatches = true;
    	    	    }
        	    }
            }
            
            // TODO run in parallel to search
            if (searchRequest.isDym()) {
          	    String correctedQuery = spellCheck(searchRequest.getClient(), searchRequest.getSearchQuery());
          	    if (correctedQuery != null) {
    				searchResponse.setDymQuery(correctedQuery);
          	    }
            }
            
    		if (chunkMap != null) {
    			addTextToSearchResults(searchRequest, searchResults, chunkMap);
    		}
    		
    		if (searchRequest.isRerank()) {
    			String rerankQuery;
    			if (searchRequest.getExpandedQuery() == null) {
    				rerankQuery = searchRequest.getSearchQuery();
    			} else {
    				rerankQuery = searchRequest.getExpandedQuery();
    			}
    			rerankSearchResults(searchRequest.getClient(), rerankQuery, searchResults);
    			for (SearchResult searchResult : searchResults) {
    				SolrDocument solrDocument = searchResult.getSolrDocument();
    				solrDocument.put(Constants.RERANK_SCORE, searchResult.getScore());
    			}
    		}
    		
    		if (noGreatMatches) {
    		    searchResponse.setRagAnswer(Constants.NO_GREAT_MATCHES);
    		} else if (searchRequest.isRag() && searchRequest.getPageNumber() == 1) {
    		    String ragAnswer = getRagAnswer(searchRequest.getClient(), searchRequest.getSearchQuery(), searchResults);
    		    searchResponse.setRagAnswer(ragAnswer);
    		}
    		
            if (searchRequest.getHighlightFields() != null) {
                Map<String, Map<String, List<String>>> highlighting = response.getHighlighting();
            	for (SearchResult searchResult : searchResults) {
            		SolrDocument solrDocument = (SolrDocument) searchResult.getSolrDocument();
            		Object idObject = solrDocument.getFieldValue(Constants._ID);
            		if (idObject != null) {
            			String id = idObject.toString();
                        if (highlighting != null) {
                    		Map<String, List<String>> highlightDocument = highlighting.get(id);
                    		if (highlightDocument != null && highlightDocument.size() > 0) {
                        		solrDocument.put(Constants.HIGHLIGHT, highlightDocument);
                    		}
                            
//                    		for (String key : highlightDocument.keySet()) {
//                      		Object value = highlightDocument.get(key);
//                      		highlightDocument.put(key, value);
//                      	}
                    		
                  		}
                    }
                }
            }
            
            // cleanup
            List<String> fieldList = searchRequest.getFieldList();
    		if (fieldList != null && !fieldList.contains("*")) {
    			for (SolrDocument solrDocument : solrDocuments) {
    	    		for (Object field : solrDocument.getFieldNames().toArray()) {
    	                if (!fieldList.contains(field) && !Constants.HIGHLIGHT.equals(field)) {
    	                	solrDocument.remove(field);
    	                }
    	    		}
    			}
    		}
		}

		// select list of results
		List<SolrDocument> results = new ArrayList<>(searchResults.size());
		for (SearchResult searchResult : searchResults) {
			results.add(searchResult.getSolrDocument());
		}
        searchResponse.setResults(results);

        if (searchRequest.isDebug()) {
        	Map<String, Object> debugInfo = response.getDebugMap();
        	if (debugInfo != null) {
            	searchResponse.setDebug(debugInfo.toString());
        	}
        }

        return searchResponse;
    }

    /**
     * Process group results
     * 
     * @param requestId
     * @param searchRequest
     * @param response
     * @param solrDocuments
     * @return
     */
    SearchResponse processGroupResults(final SearchRequest searchRequest, final QueryResponse response,
    		final GroupResponse groupResponse, Map<String, List<SolrDocument>> chunkMap) {
        SearchResponse searchResponse = new SearchResponse(searchRequest.getRequestId(), "SUCCESS", "Search completed. (" + response.getQTime() + " ms)");
        
        List<GroupCommand> groupCommands = groupResponse.getValues();
        
        // search result: _id, rank, document (solrDocument), result (final result)
		ArrayList<JSONObject> searchResults = new ArrayList<>();
    	for (int i = 0; i < groupCommands.size(); i++) {
    		GroupCommand groupCommand = groupCommands.get(i);
    		List<Group> group = groupCommand.getValues();
    		
    		// TODO
    		
    		Object value = group.get(i);    		
    		SolrDocument solrDocument = (SolrDocument) value;
    		
    		Object idObject = solrDocument.getFieldValue(Constants._ID);
    		if (idObject != null) {
    			String id = idObject.toString();
    			JSONObject searchResult = new JSONObject();
    			searchResult.put(Constants._ID, id);
    			searchResult.put(Constants.RANK, i + 1);
    			searchResult.put(Constants.RESULT, solrDocument);
    			searchResults.add(searchResult);
            }
        }
        
        searchResponse.setResults(groupResponse.getValues());
        //searchResponse.setFoundResults(groupResponse.getNumFound());

        LOGGER.debug("[{}] Processing group results. Found docs: {}", searchRequest.getRequestId(), searchResponse.getFoundResults());

        return searchResponse;
    }
    
    /**
     * Create Legacy Search Query by precision
     * 
     * -   high: all words
     * - medium: main words + quoted
     * -    low: quoted
     * 
     * @param requestId
     * @param searchRequest
     * @param solrQuery
     * @throws IOException
     */
    void setLegacySearchQuery(final SearchRequest searchRequest, final SolrQuery solrQuery) throws IOException {
    	String searchQuery = searchRequest.getSearchQuery();
    	
        setFilters(searchRequest, solrQuery);

    	String query;
        if ("*".equals(searchQuery)) {
        	query = "*:*";
        } else {
        	String cleanSearchQuery = cleanSearchQuery(searchQuery);
        	
        	JSONArray queryKeywords = languageModelService.getQueryKeywords(searchRequest.getClient(), cleanSearchQuery);
            if (searchRequest.isSynonyms()) {
            	cleanSearchQuery = StringUtil.cleanText(cleanSearchQuery);
                processSynonyms(searchRequest, cleanSearchQuery, queryKeywords);
            }

        	float wPhrase = 0.01F * queryKeywords.length() * queryKeywords.length();
        	
        	String matchQuery = "\"ʡʡ " + Utils.removeQuotes(searchQuery) + " ʡʡ\"";
        	String phraseQuery = "\"" + Utils.removeQuotes(searchQuery) + "\"";
        	
    		query = "_text_high:(" + matchQuery + ")^" + weightLegacyHigh
                    + " _text_high:(" + phraseQuery + ")^" + weightLegacyHigh * wPhrase
                    + " _text_low:(" + phraseQuery + ")^" + weightLegacyLow * wPhrase
                    + " _text_content:(" + phraseQuery + ")^" + weightLegacyContent * wPhrase;
            
            // filter query
        	List<String> filters = searchRequest.getFilterList();
            String filterQuery = getFilterQuery(filters);
            solrQuery.setFilterQueries(filterQuery);

            String mainQuery = "";
            String quotedQuery = getSubQuery(queryKeywords, "_text_content", new String[] {Constants.QUOTED}, Constants.AND);
            if (quotedQuery.length() > 0) {
                if (mainQuery.length() > 0) {
                	mainQuery += " AND ";
                }
            	mainQuery += quotedQuery;
            }
            String includedQuery = getSubQuery(queryKeywords, "_text_content", new String[] {Constants.INCLUDED}, Constants.AND);
            if (includedQuery.length() > 0) {
                if (mainQuery.length() > 0) {
                	mainQuery += " AND ";
                }
                mainQuery += includedQuery;
            }
            String excludedQuery = getSubQuery(queryKeywords, "_text_content", new String[] {Constants.EXCLUDED}, Constants.OR);
            if (excludedQuery.length() > 0) {
                if (mainQuery.length() > 0) {
                	mainQuery += " AND ";
                }
            	mainQuery += " NOT (" + excludedQuery + ")";
            }
            if (mainQuery.length() > 0) {
            	filterQuery += " AND " + mainQuery;
            }

            String andQuery;
            String orQuery;
            if (Constants.PRECISION_HIGH.equals(searchRequest.getPrecision())) {
            	andQuery = getSubQuery(queryKeywords, "_text_content", new String[] {Constants.ENTITY, Constants.KEYWORD, Constants.NUMBER}, Constants.AND);
            	orQuery = "";
            } else if (Constants.PRECISION_LOW.equals(searchRequest.getPrecision())) {
            	andQuery = "";
                orQuery = getSubQuery(queryKeywords, "_text_content", new String[] {Constants.ENTITY, Constants.KEYWORD, Constants.NUMBER}, Constants.OR);
            } else {
            	andQuery = getSubQuery(queryKeywords, "_text_content", new String[] {Constants.ENTITY}, Constants.AND);
                orQuery = getSubQuery(queryKeywords, "_text_content", new String[] {Constants.ENTITY, Constants.KEYWORD, Constants.NUMBER}, Constants.OR);
            }

            if (andQuery.length() > 0 || orQuery.length() > 0) {
            	query += " (";
            	if (andQuery.length() > 0) {
            		query += " (" + andQuery + ")";
            	}
            	if (orQuery.length() > 0) {
            		query += " (" + orQuery + ")^0.5";
            	}
            	query += " )";
            }
            
        	String proxQuery = buildProximityQuery("_text_content", queryKeywords);
            if (proxQuery.length() > 0) {
                query += " (" + proxQuery + ")^0.5";
            }

            if (mainQuery.length() > 0) {
            	query = "(" + query + ") AND " +  mainQuery;
            }
        
        }
        solrQuery.setQuery(query);
        LOGGER.debug("[{}] Solr legacy query: {}", searchRequest.getRequestId(), query);
    }
    
    /**
     * Create Vector Search Query by precision
     * 
     * -   high: high weight + main words + quoted
     * - medium: medium weight + quoted
     * -    low: low weight + quoted
     * 
     * @param searchRequest
     * @param solrQuery
     * @throws IOException
     */
    protected void setVectorSearchQuery(final SearchRequest searchRequest, final SolrQuery solrQuery) throws IOException {
    	String searchQuery = searchRequest.getSearchQuery();
    	
    	String query = "";
        if ("*".equals(searchQuery)) {
        	query = "*:*";
        } else {
        	String cleanSearchQuery = cleanSearchQuery(searchQuery);

        	JSONArray queryKeywords = languageModelService.getQueryKeywords(searchRequest.getClient(), cleanSearchQuery);
            if (searchRequest.isSynonyms()) {
            	cleanSearchQuery = StringUtil.cleanText(cleanSearchQuery);
                processSynonyms(searchRequest, cleanSearchQuery, queryKeywords);
            }
            
        	String embedQuery;
			if (searchRequest.getExpandedQuery() == null) {
				embedQuery = getKeywordsQuery(queryKeywords);
			} else {
				embedQuery = searchRequest.getExpandedQuery();
			}
            String queryEmbedding = languageModelService.getQueryEmbeddingString(searchRequest.getClient(), embedQuery);
            searchRequest.setQueryEmbedding(queryEmbedding);
            solrQuery.set("vector", queryEmbedding);

            // filter query
        	List<String> filters = searchRequest.getFilterList();
            String filterQuery = getFilterQuery(filters);
            solrQuery.setFilterQueries(filterQuery);

            String mainQuery = "";
            String quotedQuery = getSubQuery(queryKeywords, "_text_content", new String[] {Constants.QUOTED}, Constants.AND);
            if (quotedQuery.length() > 0) {
                if (mainQuery.length() > 0) {
                	mainQuery += " AND ";
                }
            	mainQuery += quotedQuery;
            }
            String includedQuery = getSubQuery(queryKeywords, "_text_content", new String[] {Constants.INCLUDED}, Constants.AND);
            if (includedQuery.length() > 0) {
                if (mainQuery.length() > 0) {
                	mainQuery += " AND ";
                }
                mainQuery += includedQuery;
            }
            String excludedQuery = getSubQuery(queryKeywords, "_text_content", new String[] {Constants.EXCLUDED}, Constants.OR);
            if (excludedQuery.length() > 0) {
                if (mainQuery.length() > 0) {
                	mainQuery += " AND ";
                }
            	mainQuery += " NOT (" + excludedQuery + ")";
            }
            if (mainQuery.length() > 0) {
            	filterQuery += " AND " + mainQuery;
            }

            String andQuery, orQuery;
            if (Constants.PRECISION_HIGH.equals(searchRequest.getPrecision())) {
            	andQuery = getSubQuery(queryKeywords, "_text_content", new String[] {Constants.ENTITY, Constants.KEYWORD, Constants.NUMBER}, Constants.AND);
            	orQuery = getSubQuery(queryKeywords, "_text_content", new String[] {Constants.ENTITY}, Constants.AND);
            } else if (Constants.PRECISION_LOW.equals(searchRequest.getPrecision())) {
                andQuery = getSubQuery(queryKeywords, "_text_content", new String[] {Constants.ENTITY}, Constants.OR);
            	orQuery = "*:*";
            } else {
            	andQuery = getSubQuery(queryKeywords, "_text_content", new String[] {Constants.ENTITY}, Constants.AND);
                orQuery = getSubQuery(queryKeywords, "_text_content", new String[] {Constants.ENTITY, Constants.KEYWORD, Constants.NUMBER}, Constants.OR);
            }

            String childQuery = "";
            if (andQuery.length() > 0) {
                String firstQuery = filterQuery + " AND (" + andQuery + ")";
                String prefilter1 = "{!child of='_uuid:*'} " + firstQuery;
                solrQuery.set("prefilter1", prefilter1);
            	String vectorQuery1 = "{!knn f=_emb topK=" + searchResultsTopK + " preFilter=$prefilter1 v=$vector}";
                solrQuery.set("knn1", vectorQuery1);
                childQuery = " {!parent which=\"_uuid:*\" score=max v=$knn1}";
            }

            if (orQuery.length() > 0 && !andQuery.equals(orQuery)) {
            	String secondQuery = filterQuery + " AND (" + orQuery + ")";
                String prefilter2 = "{!child of='_uuid:*'} " + secondQuery;
                solrQuery.set("prefilter2", prefilter2);
            	String vectorQuery2 = "{!knn f=_emb topK=" + searchResultsTopK + " preFilter=$prefilter2 v=$vector}";
                solrQuery.set("knn2", vectorQuery2);
                childQuery += " {!parent which=\"_uuid:*\" score=max v=$knn2}";
        	}

            if (childQuery.length() == 0) {
                String prefilter = "{!child of='_uuid:*'} " + filterQuery;
                solrQuery.set("prefilter", prefilter);
            	String vectorQuery = "{!knn f=_emb topK=" + searchResultsTopK + " preFilter=$prefilter v=$vector}";
                solrQuery.set("knn", vectorQuery);
                childQuery = "{!parent which=\"_uuid:*\" score=max v=$knn}";
            }
            if (weightVectorQuery != 1.0) {
            	childQuery = "(" + childQuery + ")^" + weightVectorQuery;
            }

            // parent query
        	String parentQuery = getVectorParentQuery(mainQuery, queryKeywords);
        	query = parentQuery + " OR " + childQuery;
        }
        solrQuery.setQuery(query);
    }

    /**
     * Create Keywords Query
     * 
     * @param queryKeywords
     * @return keywords query
     */
    private String getKeywordsQuery(final JSONArray queryKeywords) {
    	StringBuilder sb = new StringBuilder();
        for (int i = 0; i < queryKeywords.length(); i++) {
        	JSONObject keyword = queryKeywords.getJSONObject(i);
        	String word = keyword.getString(Constants.TEXT);
        	String type = keyword.getString(Constants.TYPE);
        	if (!Constants.EXCLUDED.equals(type)) {
        		if (sb.length() > 0) {
        			sb.append(" ");
        		}
        		sb.append(word);
    		}
        }
        return sb.toString();
    }

    /**
     * Create Vector Parent Query
     * 
     * @param mainQuery
     * @param queryKeywords
     */
    private String getVectorParentQuery(final String mainQuery, final JSONArray queryKeywords) {
        String searchQuery = getKeywordsQuery(queryKeywords);

        // parent query
    	String matchQuery = "\"ʡʡ " + Utils.removeQuotes(searchQuery) + " ʡʡ\"";
    	String parentQuery = "_text_high:(" + matchQuery + ")^" + weightVectorHigh;

    	float wPhrase = 0.01F * queryKeywords.length() * queryKeywords.length();

    	String exactQuery = "\"" + Utils.removeQuotes(searchQuery) + "\"";
	    if (exactQuery.length() > 2) {
	    	parentQuery += " OR _text_content:(" + exactQuery + ")^" + weightVectorContent * wPhrase;
	    }

    	String proxQuery = buildProximityQuery("_text_content", queryKeywords);
        if (proxQuery.length() > 0) {
            parentQuery += " OR (" + proxQuery + ")^" + weightVectorContent * wPhrase;
        }
        
        if (mainQuery.length() > 0) {
        	parentQuery = "(" + mainQuery + ")^" + weightVectorContent + " AND (" + parentQuery + ")";
        }
	    parentQuery = "(" + parentQuery + ")";
        if (weightLegacyQuery != 1.0) {
        	parentQuery += "^" + weightLegacyQuery;
        }
        LOGGER.debug("Solr vector parent query: {}", parentQuery);
        return parentQuery;
    }

    /**
     * Get filter query
     * 
     * @param searchWords
     * @return sub query
     */
    private String getFilterQuery(final List<String> filters) {
        StringBuilder filterQuery = new StringBuilder();
    	for (String filter : filters) {
    		if (filter.length() > 0) {
    			if (filterQuery.length() > 0) {
    				filterQuery.append(" AND ");
    			}
    			filterQuery.append(filter);
    		}
    	}
        return filterQuery.toString();
    }

    /**
     * Get sub query of specified type
     * 
     * @param queryKeywords JSONArray
     * @param keywordType String
     * @return sub query
     */
    private String getSubQuery(final JSONArray queryKeywords, final String field, final String[] keywordTypes, final String operator) {
    	List<String> types = Arrays.asList(keywordTypes);
        StringBuilder subQuery = new StringBuilder();
        for (int i = 0; i < queryKeywords.length(); i++) {
        	JSONObject keyword = queryKeywords.getJSONObject(i);
        	String word = keyword.getString(Constants.TEXT);
        	JSONArray synonyms = keyword.optJSONArray(Constants.SYNONYMS, null);
        	String type = keyword.getString(Constants.TYPE);
        	if (types.contains(type)) {
    			if (subQuery.length() > 0) {
    				subQuery.append(" ");
    				subQuery.append(operator);
    				subQuery.append(" ");
    			}
    			subQuery.append(field);
    			subQuery.append(":(");
    			subQuery.append(getQueryWord(word));
    			if (synonyms != null) {
        			for (Object object : synonyms) {
        				String synonym = (String) object;
            			subQuery.append(" ");
            			subQuery.append(getQueryWord(synonym));
        			}
    			}
       			subQuery.append(")");
        	}
        }
        return subQuery.toString();
    }

    /**
     * Get span sub query
     * 
     * @param queryKeywords JSONArray
     * @param keywordType String
     * @return sub query
     */
    private String getSpanQuery(final JSONArray queryKeywords) {
        StringBuilder subQuery = new StringBuilder();
        for (int i = 0; i < queryKeywords.length(); i++) {
        	JSONObject keyword = queryKeywords.getJSONObject(i);
        	String word = keyword.getString(Constants.TEXT);
        	String type = keyword.getString(Constants.TYPE);
        	if (Constants.ENTITY.equals(type) || Constants.KEYWORD.equals(type) || Constants.NUMBER.equals(type)) {
				subQuery.append(" ");
    			subQuery.append(word);
    		}
        }
        if (subQuery.length() == 0) {
        	return "";
        }
		return "{!edismax qf=_text_content mm=60%}" + subQuery.toString();
    }

    /**
     * Get query word
     * 
     * @param text the word
     * @return query word
     */
    private String getQueryWord(final String text) {
    	String output = text.trim();
    	if ("AND".equals(output) || "OR".equals(output) || "NOT".equals(output)) {
    		output = output.toLowerCase();
    	} else if (output.contains(" ")) {
        	output = "\"" + Utils.removeQuotes(output) + "\"";
        }
        return output;
    }
    
    /**
     * Search chunks
     * 
     * @param searchRequest
     * @param solrDocuments
     * @return map of chunks by id
     * @throws SolrServerException
     * @throws IOException
     */
    public Map<String, List<SolrDocument>> searchChunks(final SearchRequest searchRequest, final SolrDocumentList solrDocuments) 
			throws SolrServerException, IOException {
    	StringBuilder sb = new StringBuilder();
    	for (SolrDocument solrDocument : solrDocuments) {
    		if (sb.length() > 0) {
    			sb.append(" ");
    		}
			sb.append("\"");
    		sb.append(solrDocument.getFieldValue(Constants._ID));
			sb.append("\"");
    	}
    	String ids = sb.toString();
        LOGGER.debug("[{}] Searching chunks for ids: ({})", searchRequest.getRequestId(), ids);

        SolrQuery solrQuery = new SolrQuery();
        String fields = Utils.listToString(searchRequest.getFieldList());
        fields += ",_id,_clientId,_field,_begin,_end,_type,_root_";
        solrQuery.set("fl", fields);

    	String vectorSimilarityQuery = "({!knn f=_emb topK=" + searchResultsTopK + " preFilter=$knnPreFilter v=$knnEmb})";
    	
        solrQuery.set("q", vectorSimilarityQuery);
        solrQuery.set("knnPreFilter", "_root_:(" + ids + ")");
        solrQuery.set("knnEmb", searchRequest.getQueryEmbedding());
    	solrQuery.set("rows", searchResultsTopK);
        LOGGER.debug("[{}] Final chunk query: {}", searchRequest.getRequestId(), solrQuery);
        
        QueryResponse response = solrService.getSolrClient().query(searchRequest.getCollections(), solrQuery, METHOD.POST);
        SolrDocumentList chunkDocuments = response.getResults();
        LOGGER.debug("[{}] Found {} chunks in collections: {} Search took {} ms", searchRequest.getRequestId(), 
        		chunkDocuments.getNumFound(), searchRequest.getCollections(), response.getQTime());

        Map<String, List<SolrDocument>> chunkMap = new HashMap<>();
        for (SolrDocument chunkDocument : chunkDocuments) {
        	String id = (String) chunkDocument.getFieldValue(Constants.ROOT);
        	List<SolrDocument> chunks = chunkMap.get(id);
        	if (chunks == null) {
        		chunks = new ArrayList<>();
        		chunks.add(chunkDocument);
        		chunkMap.put(id, chunks);
        	} else {
        		chunks.add(chunkDocument);
        	}
        }
        return chunkMap;
    }
    
    /**
     * Search chunks for groups
     * 
     * @param searchRequest
     * @param solrDocuments
     * @return map of chunks by id
     * @throws SolrServerException
     * @throws IOException
     */
    public Map<String, List<SolrDocument>> searchChunks(final SearchRequest searchRequest, final GroupResponse groupResponse) 
			throws SolrServerException, IOException {
    	StringBuilder sb = new StringBuilder();
    	
    	//TODO for groups
        
    	return null;	
    }
    
    /**
     * Set pagination
     * 
     * @param requestId
     * @param searchRequest
     * @param solrQuery
     */
    protected void setPagination(final String requestId, final SearchRequest searchRequest, final SolrQuery solrQuery) {
        int start = (searchRequest.getPageNumber() - 1) * searchRequest.getPageSize();
        solrQuery.setStart(start);
        solrQuery.setRows(searchRequest.getPageSize());
        LOGGER.debug("[{}] Search pagination: start={}, rows={}", requestId, start, searchRequest.getPageSize());
    }
    
    /**
     * Set Fields
     * 
     * @param requestId
     * @param searchRequest
     * @param solrQuery
     * @param needChunks
     */
    protected boolean setFields(final String requestId, final SearchRequest searchRequest, 
    		final SolrQuery solrQuery) {
    	List<String> queryFieldList = Utils.stringToList(searchRequest.getResultFields(), false);
    	searchRequest.setFieldList(queryFieldList);
    	
    	List<String> fieldList = new ArrayList<>();
    	if (queryFieldList != null) {
    	    fieldList.addAll(queryFieldList);
    	}
    	if (fieldList.size() == 0) {
    		fieldList = new ArrayList<>();
    		fieldList.add("score");
    		fieldList.add("_id");
    	} else if (fieldList.contains("*")) {
    		fieldList = new ArrayList<>();
    		fieldList.add("score");
    		fieldList.add("*");
    	}

    	boolean needChunks;
		if (searchRequest.isRerank() || fieldList.contains(Constants.CHUNKS) ||
				(searchRequest.isRag() && searchRequest.getPageNumber() == 1)) {
			needChunks = true;
		} else {
			needChunks = false;
		}
		if (needChunks) {
    		fieldList.add("score");
        	if (!fieldList.contains(Constants.ROOT)) {
    		    fieldList.add(Constants.ROOT);
        	}
        	if (!fieldList.contains(Constants._ID)) {
    		    fieldList.add(Constants._ID);
        	}
	    }
        solrQuery.set("fl", Utils.listToString(fieldList));
        LOGGER.debug("[{}] Search result fields: {}", requestId, fieldList);
        return needChunks;
    }

    /**
     * Set filters
     * 
     * @param searchRequest
     * @param solrQuery
     */
    protected void setFilters(final SearchRequest searchRequest, final SolrQuery solrQuery) {
    	List<String> filters = searchRequest.getFilterList();
        if (filters != null) {
            for (String filter : filters) {
            	solrQuery.addFilterQuery(filter);
                LOGGER.debug("[{}] Added filter: {}", searchRequest.getRequestId(), filter);
            }
        }
    }
    
    public String getCollections(SearchRequest searchRequest, List<String> collectionList, List<String> clientIds) throws ClientException {
    	if (collectionList == null) {
    		collectionList = new ArrayList<>();
        }
    	if (collectionList.size() == 0 && clientIds != null) {
    		for (String clientId : clientIds) {
    			Client client = clientService.getClient(clientId);
    			if (client == null) {
    	            throw new ClientException("Unknown client " + clientId);
    			} else {    				
        			String collection = client.getCollection();
        			if (!collectionList.contains(collection)) {
        				collectionList.add(collection);
        			}
    			}
    		}
    	}
    	if (collectionList.size() == 0) {
    		collectionList.add(searchRequest.getClient().getCollection());
    	}
    	
        StringBuilder sb = new StringBuilder();
        for (String collection : collectionList) {
        	if (sb.length() > 0) {
        		sb.append(",");
        	}
        	sb.append(collection);
        }
        String collections = sb.toString();
		return collections;
    }
    
    public List<String> getFilters(SearchRequest searchRequest, String filterQuery, List<String> clientIds) {
    	List<String> filters = new ArrayList<>();
    	if (filterQuery != null && filterQuery.trim().length() > 0) {
    		filters.add(filterQuery);
    	}
    	if (clientIds == null || clientIds.size() == 0) {
    		filters.add("_clientId:\"" + searchRequest.getClientId() + "\"");
    	} else {
            StringBuilder sb = new StringBuilder();
    		for (String clientId : clientIds) {
    			if (sb.length() > 0) {
    				sb.append(" ");
    			}
    	        sb.append("\"");
    	        sb.append(clientId);
    	        sb.append("\"");
    		}
    		if (sb.length() > 0) {
        		filters.add("_clientId:(" + sb.toString() + ")");
    		}    		
		}    		
		return filters;
    }

    /**
     * Set sorting
     * 
     * @param requestId
     * @param searchRequest
     * @param solrQuery
     * @throws SearchException 
     */
    protected void setSorting(String requestId, SearchRequest searchRequest, SolrQuery solrQuery) throws SearchException {
    	String fields = searchRequest.getSortFields();
    	if (fields != null) {
    		fields = StringUtil.cleanSpaces(fields);
    		List<String> fieldList = Utils.stringToList(fields, false);
            for (final String sortField : fieldList) {
        		String field;
        		String order;
            	String[] array = sortField.split(" ");
            	if (array.length == 1) {
            		field = array[0];
            		if ("score".equals(field)) {
            			order = "desc";
            		} else {
            			order = "asc";
            		}
            	} else if (array.length == 2) {
            		field = array[0];
            		order = array[1];
            	} else {
            		throw new SearchException("Invalid sorting parameter");
            	}
            	if ("asc".equals(order)) {
        			solrQuery.addSort(field, ORDER.asc);
        		} else if ("desc".equals(order)) {
        			solrQuery.addSort(field, ORDER.desc);
        		} else {
            		throw new SearchException("Invalid sort order parameter");
        		}
            }
    	}
    }
        
    /**
     * Set facets
     * 
     * @param requestId
     * @param searchRequest
     * @param solrQuery
     * @throws SearchException 
     */
    protected void setFacets(final String requestId, final SearchRequest searchRequest, final SolrQuery solrQuery) 
    		throws SearchException {
    	if (searchRequest.getFacetFields() != null) {
    		String fields = searchRequest.getFacetFields();
    		JSONObject jsonFacet;
        	if (fields.startsWith("{") && fields.endsWith("}")) {
        		jsonFacet = new JSONObject(fields);
    	    } else {
    	    	jsonFacet = new JSONObject();
        		List<String> fieldList = Utils.stringToList(fields, false);
                for (final String facetField : fieldList) {
            		String field;
            		int limit;
                	String[] array = facetField.split(" ");
                	if (array.length == 1) {
                		field = array[0];
                		limit = -1;
                	} else if (array.length == 2) {
                		field = array[0];
                		try {
                    		limit = Integer.parseInt(array[1]);
                        } catch (NumberFormatException e) {
                            throw new SearchException("Invalid facet limit parameter");
                        }
                	} else {
                	    throw new SearchException("Invalid facet parameter");
                	}
                	JSONObject facet = new JSONObject();
                	facet.put("type", "terms");
                	facet.put("field", field);
                	if (limit >= 0) {
                	    facet.put("limit", limit);
                    }
                	jsonFacet.put(field, facet);
                }
    	    }                
            solrQuery.set("json.facet", jsonFacet.toString());
        }
    }

    /**
     * Set groups
     * 
     * @param requestId
     * @param searchRequest
     * @param solrQuery
     */
	protected void setGroups(final String requestId, final SearchRequest searchRequest, final SolrQuery solrQuery) {
		if (searchRequest.getGroupFields() != null) {
			List<String> fieldList = Utils.stringToList(searchRequest.getGroupFields(), false);
	    	if (fieldList.size() > 0) {
	    		solrQuery.set("group", "true");
	    		for (final String field : fieldList) {
	    		    solrQuery.add("group.field", field);
	    		}
	    		if (searchRequest.isCollapse()) {
	    		    solrQuery.add("group.main", "true");
	    		    // TODO add facet count
	    		}
	    		solrQuery.add("group.ngroups", "true");
	            LOGGER.debug("[{}] Group results by: {}", requestId, searchRequest.getGroupFields());
	    	}
	    	
	    	// TODO
//			if (null != searchRequest.getGroupOffset()) {
//			    solrQuery.add("group.offset", searchRequest.getGroupOffset());
//			}
//			if (null != searchRequest.getGroupLimit()) {
//			    solrQuery.add("group.limit", searchRequest.getGroupLimit());
//			}

		}
	}

	/**
	 * Set highlight fields
	 * 
	 * @param requestId
	 * @param searchRequest
	 * @param solrQuery
	 */
	protected void setHighlight(final String requestId, final SearchRequest searchRequest, final SolrQuery solrQuery) {
    	if (searchRequest.getHighlightFields() != null) {
        	String highlightFields = searchRequest.getHighlightFields();
        	String highlightTags = searchRequest.getHighlightTags();
            LOGGER.debug("[{}] Highlighting fields: {}", requestId, highlightFields);
            
            solrQuery.set("hl", true);
            solrQuery.set("hl.fl", highlightFields);

            solrQuery.set("hl.method", "unified");
//            solrQuery.set("hl.q", cleanQuery);
//            solrQuery.set("hl.qparser", "lucene");
////          solrQuery.set("hl.mergeContiguous", true); // joins separate hl words
////          solrQuery.set("hl.usePhraseHighlighter", true);
//            solrQuery.set("hl.preserveMulti", true);
//            solrQuery.set("hl.fragsize", highlightTags);
            
            solrQuery.set("hl.formatter", "simple");
            if (highlightTags != null) {
                String[] tags = highlightTags.split(",");
                if (tags.length == 2) {
                	solrQuery.set("hl.tag.pre", tags[0]);
                	solrQuery.set("hl.tag.post", tags[1]);
                }
            }
            
        	if (!searchRequest.getFieldList().contains(Constants._ID)) {
        		String fields = solrQuery.get("fl");
        		solrQuery.set("fl", fields + "," + Constants._ID);
        	}
        }
    }

	/**
	 * Add text to search results
	 * 
	 * @param List of search results
	 */
	private void addTextToSearchResults(final SearchRequest searchRequest, final List<SearchResult> searchResults,
			final Map<String, List<SolrDocument>> chunkMap) {
        boolean addTextToChunk = searchRequest.getFieldList().contains(Constants.CHUNKS);
		for (SearchResult searchResult : searchResults) {
			SolrDocument solrDocument = searchResult.getSolrDocument();
			// TODO
    		String id = (String) solrDocument.getFieldValue(Constants._ID);
	        
			List<SolrDocument> chunkList = chunkMap.get(searchResult.getId());
			if (chunkList != null) {
				List<TextChunk> chunks = new ArrayList<>(chunkList.size());
				for (SolrDocument solrChunk : chunkList) {
					int begin = (int) solrChunk.get(Constants.CHUNK_BEGIN);
					int end = (int) solrChunk.get(Constants.CHUNK_END);
					String field = (String) solrChunk.get(Constants.CHUNK_FIELD);
					String type = (String) solrChunk.get(Constants.CHUNK_TYPE);
					TextChunk textChunk = new TextChunk(begin, end, field);
					textChunk.setType(type);

		    		// return chunks text
    				Object object = solrDocument.get(field);
					if (object != null) {
					    String text = Utils.getTextOfObject(object);
						if (text != null) {
							text = Utils.getSubText(text, begin - 200, end + 200); //TODO
						    textChunk.setText(text);
						    if (addTextToChunk) {
						    	solrChunk.put(Constants.CHUNK_TEXT, text);
						    }
						    // TODO need?
//						    if (solrChunk.containsKey(Constants.EMB_HIGH)) {
//						    	textChunk.setType(Constants.EMB_HIGH);
//						    } else if (solrChunk.containsKey(Constants.EMB_LOW)) {
//							    textChunk.setType(Constants.EMB_LOW);
//						    } else if (solrChunk.containsKey(Constants.EMB_CONTENT)) {
//							    textChunk.setType(Constants.EMB_CONTENT);
//						    }
						    chunks.add(textChunk);
						}
					}
				}
				searchResult.setChunks(chunks);
			}
		}
	}

    /**
     * Builds a proximity query with field specification.
     */
    public static String buildProximityQuery(String field, final JSONArray queryKeywords) {
    	StringBuilder sb = new StringBuilder();
    	List<String> queryWords = new ArrayList<>();
        for (int i = 0; i < queryKeywords.length(); i++) {
        	JSONObject keyword = queryKeywords.getJSONObject(i);
        	String word = keyword.getString(Constants.TEXT);
        	String type = keyword.getString(Constants.TYPE);
        	if (Constants.ENTITY.equals(type) || Constants.KEYWORD.equals(type) || Constants.NUMBER.equals(type)
        			|| Constants.INCLUDED.equals(type) || Constants.QUOTED.equals(type)) {
        		queryWords.add(word);
    		}
        }
    	
    	if (queryWords.size() < 2) {
        	return "";
    	}
        int proximity = (queryWords.size() - 1) * 4;
        String phrase = String.join(" ", queryWords);
    	return String.format("%s:\"%s\"~%d", field, phrase, proximity);
    }

}

