package org.apache.lucene.index;

import java.io.IOException;
import java.util.*;

import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.MockTokenizer;
import org.apache.lucene.analysis.TokenFilter;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.Tokenizer;
import org.apache.lucene.analysis.payloads.PayloadHelper;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.apache.lucene.analysis.tokenattributes.PayloadAttribute;
import org.apache.lucene.codecs.Codec;
import org.apache.lucene.codecs.lucene62.Lucene62Codec;
import org.apache.lucene.codecs.simpletext.SimpleTextCodec;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.index.IndexWriterConfig.OpenMode;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.store.Directory;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.LuceneTestCase;

public class TestPayloads extends LuceneTestCase {

    public void testPayloadsEncodingWithSimpleTextCodec() throws Exception {
        Directory dir = newDirectory();
        performTest(dir, true);
        dir.close();
    }

    public void testPayloadsEncodingWithLucene62Codec() throws Exception {
        Directory dir = newDirectory();
        performTest(dir, false);
        dir.close();
    }


    private void performTest(Directory dir, boolean isSimpleTextCodec) throws Exception {
        PayloadAnalyzer analyzer = new PayloadAnalyzer();

        final IndexWriterConfig config = new IndexWriterConfig(analyzer);
        Codec codec = new Lucene62Codec();
        if(isSimpleTextCodec)
            codec = new SimpleTextCodec();

        config.setCodec(codec);
        config.setOpenMode(OpenMode.CREATE);
        config.setMergePolicy(newLogMergePolicy());

        IndexWriter writer = new IndexWriter(dir, config);

        int numTerms = 1;
        final String fieldName = "f1";

        Document d1 = new Document();
        d1.add(newTextField(fieldName, "a", Field.Store.YES));
        byte[] scorePayload = PayloadHelper.encodeFloat(1f);
        analyzer.setPayloadData(fieldName, scorePayload, 0, 4);
        writer.addDocument(d1);


        Document d2 = new Document();
        d2.add(newTextField(fieldName, "a", Field.Store.YES));
        scorePayload = PayloadHelper.encodeFloat(2f);
        analyzer.setPayloadData(fieldName, scorePayload, 0, 4);
        writer.addDocument(d2);


        // make sure we create more than one segment to test merging
        writer.commit();

        writer.forceMerge(1);
        // flush
        writer.close();


        IndexReader reader = DirectoryReader.open(dir);


        Term[] terms = new Term[numTerms];
        terms[0] = new Term(fieldName, "a");

        float previousDecodedPayload = 0f;

        PostingsEnum[] tps = new PostingsEnum[numTerms];
        for (int termIndex = 0; termIndex < numTerms; termIndex++) {
            tps[termIndex] = MultiFields.getTermPositionsEnum(reader,
                    terms[termIndex].field(),
                    new BytesRef(terms[termIndex].text()));

            int docCounter = 0;
            while (tps[termIndex].nextDoc() != DocIdSetIterator.NO_MORE_DOCS) {
                for (int position = 0; position < tps[termIndex].freq(); position++) {
                    System.out.print("termIndex " + termIndex + " in doc: " + (docCounter) + " position: " + position + " term: `" + terms[termIndex].text() + "` freq: " + tps[termIndex].freq() + " payload: ");
                    tps[termIndex].nextPosition();
                    BytesRef payload = tps[termIndex].getPayload();
                    if (payload != null) {
                        float decodedPayload = PayloadHelper.decodeFloat(payload.bytes);
                        System.out.println("len:" + decodedPayload);
                        if (previousDecodedPayload == decodedPayload) {
                            fail("Payload in the same position of a different document should be different values: " + decodedPayload + ", " + previousDecodedPayload);
                        }
                        previousDecodedPayload = decodedPayload;
                    } else
                        System.out.println("null");
                }
                docCounter++;
            }
        }
        reader.close();

    }


    private static class PayloadData {
        byte[] data;
        int offset;
        int length;

        PayloadData(byte[] data, int offset, int length) {
            this.data = data;
            this.offset = offset;
            this.length = length;
        }
    }

    /**
     * This Analyzer uses an MockTokenizer and PayloadFilter.
     */
    private static class PayloadAnalyzer extends Analyzer {
        Map<String, PayloadData> fieldToData = new HashMap<>();

        public PayloadAnalyzer() {
            super(PER_FIELD_REUSE_STRATEGY);
        }

        public PayloadAnalyzer(String field, byte[] data, int offset, int length) {
            super(PER_FIELD_REUSE_STRATEGY);
            setPayloadData(field, data, offset, length);
        }

        void setPayloadData(String field, byte[] data, int offset, int length) {
            fieldToData.put(field, new PayloadData(data, offset, length));
        }

        @Override
        public TokenStreamComponents createComponents(String fieldName) {
            PayloadData payload = fieldToData.get(fieldName);
            Tokenizer ts = new MockTokenizer(MockTokenizer.WHITESPACE, false);
            TokenStream tokenStream = (payload != null) ?
                    new PayloadFilter(ts, fieldName, fieldToData) : ts;
            return new TokenStreamComponents(ts, tokenStream);
        }
    }


    /**
     * This Filter adds payloads to the tokens.
     */
    private static class PayloadFilter extends TokenFilter {
        PayloadAttribute payloadAtt;
        CharTermAttribute termAttribute;
        private Map<String, PayloadData> fieldToData;
        private String fieldName;
        private PayloadData payloadData;
        private int offset;

        public PayloadFilter(TokenStream in, String fieldName, Map<String, PayloadData> fieldToData) {
            super(in);
            this.fieldToData = fieldToData;
            this.fieldName = fieldName;
            payloadAtt = addAttribute(PayloadAttribute.class);
            termAttribute = addAttribute(CharTermAttribute.class);
        }

        @Override
        public boolean incrementToken() throws IOException {
            boolean hasNext = input.incrementToken();
            if (!hasNext) {
                return false;
            }

            // Some values of the same field are to have payloads and others not
            if (offset + payloadData.length <= payloadData.data.length && !termAttribute.toString().endsWith("NO PAYLOAD")) {
                BytesRef p = new BytesRef(payloadData.data, offset, payloadData.length);
                payloadAtt.setPayload(p);
                offset += payloadData.length;
            } else {
                payloadAtt.setPayload(null);
            }

            return true;
        }

        @Override
        public void reset() throws IOException {
            super.reset();
            this.payloadData = fieldToData.get(fieldName);
            this.offset = payloadData.offset;
        }
    }


}