/*
 * Licensed to De Bitmanager under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. De Bitmanager 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 nl.bitmanager.elasticsearch.search;

import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Map;

import org.apache.lucene.index.LeafReaderContext;
import org.elasticsearch.common.document.DocumentField;
import org.elasticsearch.index.fielddata.AtomicFieldData;
import org.elasticsearch.index.fielddata.IndexFieldData;
import org.elasticsearch.index.fielddata.IndexFieldDataService;
import org.elasticsearch.index.mapper.DocumentFieldMappers;
import org.elasticsearch.index.mapper.DocumentMapper;
import org.elasticsearch.index.mapper.FieldMapper;
import org.elasticsearch.index.mapper.MappedFieldType;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.fetch.FetchSubPhase;
import org.elasticsearch.search.internal.SearchContext;

import nl.bitmanager.elasticsearch.typehandlers.TypeHandler;

/**
 * This class optional fetches extra diagnostic information like segment-level info and docvalues.
 * It is controlled via the SearchParms (supplied via ext: {})
 */
public class FetchDiagnostics implements FetchSubPhase {
    private final static boolean DEBUG=false;

    @Override
    public void hitExecute(SearchContext context, HitContext hitContext) {
        try {
            if (DEBUG) System.out.println("HIT execute1");
            
            SearchParms x = (SearchParms)context.getSearchExt("_bm");
            if (x==null || !x.diagnostics) return;
            if (context.storedFieldsContext() != null && context.storedFieldsContext().fetchFields() == false) {
                return ;
            }
            SearchHit hit = hitContext.hit();
            Map<String, DocumentField> fields = hitContext.hit().fieldsOrNull();
            if (fields == null) {
                fields = new HashMap<>();
                hit.fields(fields);
            }
            LinkedHashMap<String, Object> map = new LinkedHashMap<String, Object>();
            LeafReaderContext ctx = hitContext.readerContext();
            int docid = hit.docId();
            int rel_docid = docid - ctx.docBase;
            map.put("shard",  context.getQueryShardContext().getShardId());
            map.put("segment",  ctx.ord);
            map.put("docid",  docid);
            map.put("docid_rel",  rel_docid);
            map.put("index_uuid",  context.getQueryShardContext().index().getUUID());
            fields.put("_bm", new DocumentField("_bm", Collections.singletonList(map)));
            
            DocumentMapper docMapper = context.mapperService().documentMapper(hitContext.hit().getType());
            DocumentFieldMappers mappers = docMapper.mappers();
            
            //FieldInfo[] segmentFields = DocInverter.getFields(ctx.reader());
//PW            
//            IndexFieldDataService fds = context.fieldData();
            LinkedHashMap<String, Object> dst = new LinkedHashMap<String, Object>();
            for (FieldMapper f: mappers) {
                if (DEBUG) System.out.println("HIT execute2 F=" + f.name() + ", dv=" + f.fieldType().hasDocValues());
                if (!f.fieldType().hasDocValues()) continue;
//PW                extractDocValues  (dst, f, fds, hitContext);
            }
            if (dst.size()>0) map.put("docvalues", dst);
            //.parentFieldMapper();
        } catch (Exception e) {
            throw new RuntimeException (e.getMessage(), e);
        }
        
        
    }
    
    private void extractDocValues(Map<String, Object> dst, FieldMapper field, IndexFieldDataService fieldDataService, HitContext hitContext) throws IOException {
        IndexFieldData<?> fd;
        MappedFieldType fieldType = field.fieldType();
        try {
            fd = fieldDataService.getForField(fieldType);
        } catch(Throwable th) {
            String x = th.toString();
            if (th instanceof IllegalArgumentException) {
                if (x.indexOf("not supported") > 0) return;
            }
            dst.put(field.name(), "error: " + x);
            return;
        }

        AtomicFieldData dv = fd.load(hitContext.readerContext());
        TypeHandler typeHandler = TypeHandler.create(fieldType);
        if (dv != null) {
            int reldoc = hitContext.docId();
            if (reldoc < 0) {
                dst.put(field.name(), String.format("negative reldoc=%d docid=%d, base=%d", reldoc, hitContext.docId() , hitContext.readerContext().docBase));
                return;
            }
            if (DEBUG) System.out.println("reldoc=" + reldoc);
            Object[] values = typeHandler.docValuesToObjects (dv, reldoc);
            if (values != null && values.length > 0)
                dst.put(field.name(), values);
        }
    }

    
}

@@ -1192,72 +1192,8 @@
 ta;%0A
-import org.elasticsearch.index.fielddata.IndexFieldDataService;%0A
 impo
@@ -1408,16 +1408,72 @@
 ldType;%0A
+import org.elasticsearch.index.query.QueryShardContext;%0A
 import o
@@ -3479,151 +3479,65 @@
     
-//FieldInfo%5B%5D segmentFields = DocInverter.getFields(ctx.reader());%0A//PW            %0A//            IndexFieldDataService fds = context.fieldData
+QueryShardContext shardCtx = context.getQueryShardContext
 ();%0A
@@ -3802,17 +3802,16 @@
     if (
-!
 f.fieldT
@@ -3836,22 +3836,13 @@
 ()) 
-continue;%0A//PW
+%0A    
     
@@ -3884,11 +3884,16 @@
  f, 
-fds
+shardCtx
 , hi
@@ -3978,44 +3978,8 @@
 t);%0A
-            //.parentFieldMapper();%0A
     
@@ -4187,46 +4187,34 @@
 ld, 
-IndexFieldDataService fieldDataService
+QueryShardContext shardCtx
 , Hi
@@ -4375,24 +4375,16 @@
 d = 
-fieldDataService
+shardCtx
 .get

Make docvalues appear again on the diagnostic fields.