/*
 * Licensed to Elastic Search and Shay Banon under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. Elastic Search licenses this
 * file to you 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 org.elasticsearch.search.highlight;

import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.vectorhighlight.*;
import org.elasticsearch.ElasticSearchException;
import org.elasticsearch.common.collect.ImmutableMap;
import org.elasticsearch.index.mapper.DocumentMapper;
import org.elasticsearch.index.mapper.FieldMapper;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchParseElement;
import org.elasticsearch.search.SearchPhase;
import org.elasticsearch.search.fetch.FetchPhaseExecutionException;
import org.elasticsearch.search.internal.InternalSearchHit;
import org.elasticsearch.search.internal.SearchContext;

import java.io.IOException;
import java.util.Map;

import static org.elasticsearch.common.collect.Maps.*;

/**
 * @author kimchy (shay.banon)
 */
public class HighlightPhase implements SearchPhase {

    @Override public Map<String, ? extends SearchParseElement> parseElements() {
        return ImmutableMap.of("highlight", new HighlighterParseElement());
    }

    @Override public void preProcess(SearchContext context) {
    }

    @Override public void execute(SearchContext context) throws ElasticSearchException {
        if (context.highlight() == null) {
            return;
        }

        for (SearchHit hit : context.fetchResult().hits().hits()) {
            InternalSearchHit internalHit = (InternalSearchHit) hit;

            DocumentMapper documentMapper = context.mapperService().type(internalHit.type());
            int docId = internalHit.docId();

            Map<String, HighlightField> highlightFields = newHashMap();
            for (SearchContextHighlight.Field field : context.highlight().fields()) {
                String fieldName = field.field();
                FieldMapper mapper = documentMapper.mappers().smartNameFieldMapper(field.field());
                if (mapper != null) {
                    fieldName = mapper.names().indexName();
                }

                FastVectorHighlighter highlighter = buildHighlighter(field);
                FieldQuery fieldQuery = buildFieldQuery(highlighter, context.query(), context.searcher().getIndexReader(), field);

                String[] fragments;
                try {
                    fragments = highlighter.getBestFragments(fieldQuery, context.searcher().getIndexReader(), docId, fieldName, field.fragmentCharSize(), field.numberOfFragments());
                } catch (IOException e) {
                    throw new FetchPhaseExecutionException(context, "Failed to highlight field [" + field.field() + "]", e);
                }
                HighlightField highlightField = new HighlightField(field.field(), fragments);
                highlightFields.put(highlightField.name(), highlightField);
            }

            internalHit.highlightFields(highlightFields);
        }
    }

    private FieldQuery buildFieldQuery(FastVectorHighlighter highlighter, Query query, IndexReader indexReader, SearchContextHighlight.Field field) {
        CustomFieldQuery.reader.set(indexReader);
        CustomFieldQuery.highlightFilters.set(field.highlightFilter());
        return new CustomFieldQuery(query, highlighter);
    }

    private FastVectorHighlighter buildHighlighter(SearchContextHighlight.Field field) {
        FragListBuilder fragListBuilder;
        FragmentsBuilder fragmentsBuilder;
        if (field.numberOfFragments() == 0) {
            fragListBuilder = new SingleFragListBuilder();
            fragmentsBuilder = new SimpleFragmentsBuilder(field.preTags(), field.postTags());
            // a HACK to make highlighter do highlighting, even though its using the single frag list builder
            field.numberOfFragments(1);
        } else {
            fragListBuilder = new SimpleFragListBuilder();
            if (field.scoreOrdered()) {
                fragmentsBuilder = new ScoreOrderFragmentsBuilder(field.preTags(), field.postTags());
            } else {
                fragmentsBuilder = new SimpleFragmentsBuilder(field.preTags(), field.postTags());
            }
        }

        return new FastVectorHighlighter(true, false, fragListBuilder, fragmentsBuilder);
    }
}

@@ -2102,32 +2102,50 @@
 urn;%0A        %7D%0A%0A
+        try %7B%0A    
         for (Sea
@@ -2208,16 +2208,20 @@
         
+    
 Internal
@@ -2282,16 +2282,20 @@
         
+  
+  
 Document
@@ -2380,16 +2380,20 @@
         
+    
 int docI
@@ -2430,16 +2430,20 @@
         
+ 
+   
 Map%3CStri
@@ -2490,24 +2490,28 @@
 wHashMap();%0A
+    
             
@@ -2596,24 +2596,28 @@
             
+ 
+   
 String field
@@ -2634,24 +2634,28 @@
 ld.field();%0A
+    
             
@@ -2753,24 +2753,28 @@
             
+    
 if (mapper !
@@ -2775,32 +2775,36 @@
 pper != null) %7B%0A
+    
                 
@@ -2855,35 +2855,42 @@
                 
-%7D%0A%0A
+    %7D%0A%0A   
                 
@@ -2889,16 +2889,17 @@
         
+ 
 FastVect
@@ -2947,24 +2947,28 @@
 ter(field);%0A
+    
             
@@ -3099,24 +3099,28 @@
             
+    
 String%5B%5D fra
@@ -3135,32 +3135,36 @@
                 
+    
 try %7B%0A          
@@ -3165,32 +3165,36 @@
                 
+  
+  
 fragments = high
@@ -3355,16 +3355,20 @@
         
+    
 %7D catch 
@@ -3405,16 +3405,20 @@
         
+  
+  
 throw ne
@@ -3522,32 +3522,36 @@
                 
+    
 %7D%0A              
@@ -3552,16 +3552,20 @@
         
+ 
+   
 Highligh
@@ -3650,16 +3650,20 @@
         
+    
 highligh
@@ -3718,35 +3718,43 @@
 d);%0A            
-%7D%0A%0A
+    %7D%0A%0A    
             inte
@@ -3791,24 +3791,160 @@
 ghtFields);%0A
+            %7D%0A        %7D finally %7B%0A            CustomFieldQuery.reader.remove();%0A            CustomFieldQuery.highlightFilters.remove();%0A
         %7D%0A  

clean thread local vars eagerly