package com.firstblick.elasticsearch.action.export;

import org.elasticsearch.common.xcontent.ToXContent;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.rest.action.support.RestXContentBuilder;
import org.elasticsearch.search.SearchHitField;
import org.elasticsearch.search.internal.InternalSearchHit;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class ExportFields implements ToXContent {

    private final List<String> fields;
    private InternalSearchHit hit;
    private final List<FieldExtractor> fieldExtractors;

    abstract class FieldExtractor implements ToXContent {
    }

    class HitFieldExtractor extends FieldExtractor {

        private final String fieldName;

        public HitFieldExtractor(String fieldName) {
            this.fieldName = fieldName;
        }

        @Override
        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
            SearchHitField field = hit.getFields().get(fieldName);
            if (field != null && !field.values().isEmpty()) {
                if (field.values().size() == 1) {
                    builder.field(field.name(), field.values().get(0));
                } else {
                    builder.field(field.name());
                    builder.startArray();
                    for (Object value : field.values()) {
                        builder.value(value);
                    }
                    builder.endArray();
                }
            }
            return builder;
        }
    }

    public void hit(InternalSearchHit hit) {
        this.hit = hit;
    }

    public ExportFields(List<String> fields) {
        this.fields = fields;
        this.fieldExtractors = getFieldExtractors();
    }

    private List<FieldExtractor> getFieldExtractors() {
        List<FieldExtractor> extractors = new ArrayList<FieldExtractor>(fields
                .size());
        for (String fn : fields) {
            FieldExtractor fc = null;
            if (fn.startsWith("_")) {
                if (fn.equals("_source")) {
                    fc = new FieldExtractor() {
                        @Override
                        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
                            RestXContentBuilder.restDocumentSource(hit.sourceRef(), builder, params);
                            return builder;
                        }
                    };
                } else if (fn.equals("_id")) {
                    fc = new FieldExtractor() {
                        @Override
                        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
                            return builder.field("_id", hit.getId());
                        }
                    };
                } else if (fn.equals("_type")) {
                    fc = new FieldExtractor() {
                        @Override
                        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
                            return builder.field("_type", hit.getType());
                        }
                    };
                }
            } else {
                fc = new HitFieldExtractor(fn);
            }
            extractors.add(fc);
        }
        return extractors;
    }

    @Override
    public XContentBuilder toXContent(XContentBuilder builder, Params
            params) throws IOException {
        builder.startObject();
        for (FieldExtractor fc : fieldExtractors) {
            fc.toXContent(builder, params);
        }
        builder.endObject();
        return builder;
    }
}


@@ -82,27 +82,28 @@
 mon.
-xcontent.ToXContent
+bytes.BytesReference
 ;%0Aim
@@ -145,89 +145,9 @@
 ent.
-XContentBuilder;%0Aimport org.elasticsearch.rest.action.support.RestXContentBuilder
+*
 ;%0Aim
@@ -515,16 +515,16 @@
 ctors;%0A%0A
-
     abst
@@ -576,24 +576,714 @@
 nt %7B%0A    %7D%0A%0A
+    class SourceFieldExtractor extends FieldExtractor %7B%0A%0A        @Override%0A        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException %7B%0A            BytesReference source = hit.sourceRef();%0A            XContentType contentType = XContentFactory.xContentType(source);%0A            XContentParser parser = XContentFactory.xContent(contentType)%0A                    .createParser(source);%0A            try %7B%0A                parser.nextToken();%0A                builder.field(%22_source%22);%0A                builder.copyCurrentStructure(parser);%0A            %7D finally %7B%0A                parser.close();%0A            %7D%0A            return builder;%0A        %7D%0A    %7D%0A%0A
     class Hi
@@ -2750,32 +2750,38 @@
        fc = new 
+Source
 FieldExtractor()
@@ -2784,357 +2784,8 @@
 or()
- %7B%0A                        @Override%0A                        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException %7B%0A                            RestXContentBuilder.restDocumentSource(hit.sourceRef(), builder, params);%0A                            return builder;%0A                        %7D%0A                    %7D
 ;%0A  

reformat json of source upon export