/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch 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.suggest.completion2x;

import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.apache.lucene.analysis.tokenattributes.PayloadAttribute;
import org.apache.lucene.analysis.tokenattributes.PositionIncrementAttribute;
import org.apache.lucene.analysis.tokenattributes.TermToBytesRefAttribute;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.BytesRefBuilder;
import org.apache.lucene.util.IntsRef;
import org.apache.lucene.util.fst.Util;

import java.io.IOException;
import java.util.Iterator;
import java.util.Set;

/**
 *
 */
public final class CompletionTokenStream extends TokenStream {

    private final PayloadAttribute payloadAttr = addAttribute(PayloadAttribute.class);
    private final PositionIncrementAttribute posAttr = addAttribute(PositionIncrementAttribute.class);
    private final ByteTermAttribute bytesAtt = addAttribute(ByteTermAttribute.class);;


    private final TokenStream input;
    private BytesRef payload;
    private Iterator<IntsRef> finiteStrings;
    private ToFiniteStrings toFiniteStrings;
    private int posInc = -1;
    private static final int MAX_PATHS = 256;
    private CharTermAttribute charTermAttribute;

    public CompletionTokenStream(TokenStream input, BytesRef payload, ToFiniteStrings toFiniteStrings) {
        // Don't call the super(input) ctor - this is a true delegate and has a new attribute source since we consume
        // the input stream entirely in toFiniteStrings(input)
        this.input = input;
        this.payload = payload;
        this.toFiniteStrings = toFiniteStrings;
    }

    @Override
    public boolean incrementToken() throws IOException {
        clearAttributes();
        if (finiteStrings == null) {
            Set<IntsRef> strings = toFiniteStrings.toFiniteStrings(input);

            if (strings.size() > MAX_PATHS) {
                throw new IllegalArgumentException("TokenStream expanded to " + strings.size() + " finite strings. Only <= " + MAX_PATHS
                        + " finite strings are supported");
            }
            posInc = strings.size();
            finiteStrings = strings.iterator();
        }
        if (finiteStrings.hasNext()) {
            posAttr.setPositionIncrement(posInc);
            /*
             * this posInc encodes the number of paths that this surface form
             * produced. Multi Fields have the same surface form and therefore sum up
             */
            posInc = 0;
            Util.toBytesRef(finiteStrings.next(), bytesAtt.builder()); // now we have UTF-8
            if (charTermAttribute != null) {
                charTermAttribute.setLength(0);
                charTermAttribute.append(bytesAtt.toUTF16());
            }
            if (payload != null) {
                payloadAttr.setPayload(this.payload);
            }
            return true;
        }

        return false;
    }

    @Override
    public void end() throws IOException {
        super.end();
        if (posInc == -1) {
            input.end();
        }
    }

    @Override
    public void close() throws IOException {
        input.close();
    }

    public interface ToFiniteStrings {
        Set<IntsRef> toFiniteStrings(TokenStream stream) throws IOException;
    }

    @Override
    public void reset() throws IOException {
        super.reset();
        if (hasAttribute(CharTermAttribute.class)) {
            // we only create this if we really need it to safe the UTF-8 to UTF-16 conversion
            charTermAttribute = getAttribute(CharTermAttribute.class);
        }
        finiteStrings = null;
        posInc = -1;
    }

    public interface ByteTermAttribute extends TermToBytesRefAttribute {
        // marker interface

        /**
         * Return the builder from which the term is derived.
         */
        BytesRefBuilder builder();

        CharSequence toUTF16();
    }
}

@@ -1209,52 +1209,194 @@
 til.
-BytesRef;%0Aimport org.apache.lucene.util.Byte
+AttributeImpl;%0Aimport org.apache.lucene.util.AttributeReflector;%0Aimport org.apache.lucene.util.BytesRef;%0Aimport org.apache.lucene.util.BytesRefBuilder;%0Aimport org.apache.lucene.util.Char
 sRef
@@ -4139,32 +4139,39 @@
   %7D%0A%0A    public 
+static 
 interface ToFini
@@ -4182,32 +4182,39 @@
 rings %7B%0A        
+public 
 Set%3CIntsRef%3E toF
@@ -4862,24 +4862,24 @@
 builder();%0A%0A
-
         Char
@@ -4903,11 +4903,1089 @@
 );%0A    %7D
+%0A%0A    public static final class ByteTermAttributeImpl extends AttributeImpl implements ByteTermAttribute, TermToBytesRefAttribute %7B%0A        private final BytesRefBuilder bytes = new BytesRefBuilder();%0A        private CharsRefBuilder charsRef;%0A%0A        @Override%0A        public BytesRefBuilder builder() %7B%0A            return bytes;%0A        %7D%0A%0A        @Override%0A        public BytesRef getBytesRef() %7B%0A            return bytes.get();%0A        %7D%0A%0A        @Override%0A        public void clear() %7B%0A            bytes.clear();%0A        %7D%0A%0A        @Override%0A        public void reflectWith(AttributeReflector reflector) %7B%0A%0A        %7D%0A%0A        @Override%0A        public void copyTo(AttributeImpl target) %7B%0A            ByteTermAttributeImpl other = (ByteTermAttributeImpl) target;%0A            other.bytes.copyBytes(bytes);%0A        %7D%0A%0A        @Override%0A        public CharSequence toUTF16() %7B%0A            if (charsRef == null) %7B%0A                charsRef = new CharsRefBuilder();%0A            %7D%0A            charsRef.copyUTF8Bytes(getBytesRef());%0A            return charsRef.get();%0A        %7D%0A    %7D
 %0A%7D%0A

Fix CompletionSuggestSearchIT and CompletionSuggestSearch2xIT