//
// Licensed 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.apache.couchdb.nouveau.lucene;

import jakarta.ws.rs.WebApplicationException;
import jakarta.ws.rs.core.Response.Status;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.CharBuffer;
import java.nio.charset.CharacterCodingException;
import java.nio.charset.CharsetDecoder;
import java.nio.charset.StandardCharsets;
import java.nio.file.NoSuchFileException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import org.apache.couchdb.nouveau.api.DocumentDeleteRequest;
import org.apache.couchdb.nouveau.api.DocumentUpdateRequest;
import org.apache.couchdb.nouveau.api.DoubleField;
import org.apache.couchdb.nouveau.api.DoubleRange;
import org.apache.couchdb.nouveau.api.Field;
import org.apache.couchdb.nouveau.api.SearchHit;
import org.apache.couchdb.nouveau.api.SearchRequest;
import org.apache.couchdb.nouveau.api.SearchResults;
import org.apache.couchdb.nouveau.api.StoredField;
import org.apache.couchdb.nouveau.api.StringField;
import org.apache.couchdb.nouveau.api.TextField;
import org.apache.couchdb.nouveau.core.IOUtils;
import org.apache.couchdb.nouveau.core.Index;
import org.apache.couchdb.nouveau.core.ser.ByteArrayWrapper;
import org.apache.couchdb.nouveau.core.ser.DoubleWrapper;
import org.apache.couchdb.nouveau.core.ser.FloatWrapper;
import org.apache.couchdb.nouveau.core.ser.IntWrapper;
import org.apache.couchdb.nouveau.core.ser.LongWrapper;
import org.apache.couchdb.nouveau.core.ser.NullWrapper;
import org.apache.couchdb.nouveau.core.ser.PrimitiveWrapper;
import org.apache.couchdb.nouveau.core.ser.StringWrapper;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field.Store;
import org.apache.lucene.document.SortedDocValuesField;
import org.apache.lucene.document.StoredValue;
import org.apache.lucene.facet.FacetResult;
import org.apache.lucene.facet.Facets;
import org.apache.lucene.facet.FacetsCollector;
import org.apache.lucene.facet.FacetsCollectorManager;
import org.apache.lucene.facet.LabelAndValue;
import org.apache.lucene.facet.StringDocValuesReaderState;
import org.apache.lucene.facet.StringValueFacetCounts;
import org.apache.lucene.facet.range.DoubleRangeFacetCounts;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexableField;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.StoredFields;
import org.apache.lucene.index.Term;
import org.apache.lucene.queryparser.flexible.core.QueryNodeException;
import org.apache.lucene.search.BooleanClause.Occur;
import org.apache.lucene.search.BooleanQuery;
import org.apache.lucene.search.CollectorManager;
import org.apache.lucene.search.FieldComparator;
import org.apache.lucene.search.FieldComparatorSource;
import org.apache.lucene.search.FieldDoc;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.LeafFieldComparator;
import org.apache.lucene.search.MultiCollectorManager;
import org.apache.lucene.search.Pruning;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Scorable;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.SearcherManager;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.SortField;
import org.apache.lucene.search.SortedNumericSortField;
import org.apache.lucene.search.SortedSetSortField;
import org.apache.lucene.search.TermQuery;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.TopFieldCollectorManager;
import org.apache.lucene.store.Directory;
import org.apache.lucene.util.BytesRef;

public class LuceneIndex extends Index {

    private static final Sort DEFAULT_SORT =
            new Sort(SortField.FIELD_SCORE, new SortField("_id", SortField.Type.STRING));
    private static final Pattern SORT_FIELD_RE = Pattern.compile("^([-+])?([\\.\\w]+)$");

    private final Analyzer analyzer;
    private final IndexWriter writer;
    private final SearcherManager searcherManager;
    private final LuceneIndexSchema schema;

    public LuceneIndex(
            final Analyzer analyzer,
            final IndexWriter writer,
            final long updateSeq,
            final long purgeSeq,
            final SearcherManager searcherManager) {
        super(updateSeq, purgeSeq);
        this.analyzer = Objects.requireNonNull(analyzer);
        this.writer = Objects.requireNonNull(writer);
        this.searcherManager = Objects.requireNonNull(searcherManager);
        this.schema = initSchema(writer);
    }

    @Override
    public int doNumDocs() throws IOException {
        return writer.getDocStats().numDocs;
    }

    @Override
    public long doDiskSize() throws IOException {
        final Directory dir = writer.getDirectory();
        long result = 0;
        for (final String name : dir.listAll()) {
            try {
                result += dir.fileLength(name);
            } catch (final FileNotFoundException | NoSuchFileException e) {
                // deleted while we were looping.
            }
        }
        return result;
    }

    @Override
    public void doUpdate(final String docId, final DocumentUpdateRequest request) throws IOException {
        final Term docIdTerm = docIdTerm(docId);
        final Document doc = toDocument(docId, request);
        schema.update(request.fields());
        writer.updateDocument(docIdTerm, doc);
    }

    @Override
    public void doDelete(final String docId, final DocumentDeleteRequest request) throws IOException {
        final Query query = docIdQuery(docId);
        writer.deleteDocuments(query);
    }

    @Override
    public boolean doCommit(final long updateSeq, final long purgeSeq) throws IOException {
        if (!writer.hasUncommittedChanges()) {
            return false;
        }
        writer.setLiveCommitData(Map.of(
                        "update_seq",
                        Long.toString(updateSeq),
                        "purge_seq",
                        Long.toString(purgeSeq),
                        "_schema",
                        schema.toString())
                .entrySet());
        writer.commit();
        return true;
    }

    @Override
    public void doClose() throws IOException {
        IOUtils.runAll(
                () -> {
                    searcherManager.close();
                },
                () -> {
                    writer.rollback();
                },
                () -> {
                    if (isDeleteOnClose()) {
                        var dir = writer.getDirectory();
                        for (final String name : dir.listAll()) {
                            dir.deleteFile(name);
                        }
                    }
                });
    }

    @Override
    public boolean doIsOpen() {
        return writer.isOpen();
    }

    @Override
    public SearchResults doSearch(final SearchRequest request) throws IOException {
        final Query query = parse(request);

        // Construct CollectorManagers.
        final MultiCollectorManager cm;
        final CollectorManager<?, ? extends TopDocs> hits = hitCollector(request);
        if (request.hasCounts() || request.hasRanges()) {
            cm = new MultiCollectorManager(hits, new FacetsCollectorManager());
        } else {
            cm = new MultiCollectorManager(hits);
        }

        searcherManager.maybeRefreshBlocking();

        final IndexSearcher searcher = searcherManager.acquire();
        try {
            final Object[] reduces = searcher.search(query, cm);
            return toSearchResults(request, searcher, reduces);
        } catch (final IllegalStateException e) {
            throw new WebApplicationException(e.getMessage(), e, Status.BAD_REQUEST);
        } finally {
            searcherManager.release(searcher);
        }
    }

    private CollectorManager<?, ? extends TopDocs> hitCollector(final SearchRequest searchRequest) {
        final Sort sort = toSort(searchRequest);
        final FieldDoc fieldDoc = searchRequest
                .after()
                .map(after -> {
                    var result = toFieldDoc(after);
                    if (getLastSortField(sort).getReverse()) {
                        result.doc = 0;
                    } else {
                        result.doc = Integer.MAX_VALUE;
                    }
                    return result;
                })
                .orElse(null);

        return new TopFieldCollectorManager(sort, searchRequest.limitAsInt(), fieldDoc, 1000);
    }

    private SortField getLastSortField(final Sort sort) {
        final SortField[] sortFields = sort.getSort();
        return sortFields[sortFields.length - 1];
    }

    private SearchResults toSearchResults(
            final SearchRequest searchRequest, final IndexSearcher searcher, final Object[] reduces)
            throws IOException {
        final TopDocs topDocs = (TopDocs) reduces[0];
        final var hits = collectHits(searcher, topDocs);
        Map<String, Map<String, Number>> counts = Collections.emptyMap();
        Map<String, Map<String, Number>> ranges = Collections.emptyMap();

        if (reduces.length == 2) {
            counts = collectCounts(searchRequest, searcher, (FacetsCollector) reduces[1]);
            ranges = collectRanges(searchRequest, searcher, (FacetsCollector) reduces[1]);
        }
        return new SearchResults(topDocs.totalHits.value(), topDocs.totalHits.relation(), hits, counts, ranges);
    }

    private List<SearchHit> collectHits(final IndexSearcher searcher, final TopDocs topDocs) throws IOException {
        final List<SearchHit> result = new ArrayList<SearchHit>(topDocs.scoreDocs.length);
        final StoredFields storedFields = searcher.storedFields();

        for (final ScoreDoc scoreDoc : topDocs.scoreDocs) {
            final Document doc = storedFields.document(scoreDoc.doc);

            final List<StoredField> fields =
                    new ArrayList<StoredField>(doc.getFields().size());
            for (IndexableField field : doc.getFields()) {
                if (field.name().equals("_id")) {
                    continue;
                }
                final StoredValue storedValue = field.storedValue();
                switch (storedValue.getType()) {
                    case DOUBLE:
                        fields.add(new StoredField(field.name(), storedValue.getDoubleValue()));
                        break;
                    case STRING:
                        fields.add(new StoredField(field.name(), storedValue.getStringValue()));
                        break;
                    case BINARY:
                        fields.add(new StoredField(field.name(), storedValue.getBinaryValue()));
                        break;
                    default:
                        throw new IllegalArgumentException(storedValue.getType() + " not supported");
                }
            }

            final PrimitiveWrapper<?>[] after = toAfter(((FieldDoc) scoreDoc));
            result.add(new SearchHit(doc.get("_id"), after, fields));
        }
        return result;
    }

    private Map<String, Map<String, Number>> collectCounts(
            final SearchRequest searchRequest, final IndexSearcher searcher, final FacetsCollector fc)
            throws IOException {
        if (!searchRequest.hasCounts()) {
            return Collections.emptyMap();
        }
        var c = searchRequest.counts().get();
        final Map<String, Map<String, Number>> result = new HashMap<String, Map<String, Number>>(c.size());
        for (final String field : c) {
            final StringDocValuesReaderState state = new StringDocValuesReaderState(searcher.getIndexReader(), field);
            final StringValueFacetCounts counts = new StringValueFacetCounts(state, fc);
            result.put(field, collectFacets(counts, searchRequest.topNAsInt(), field));
        }
        return result;
    }

    private Map<String, Map<String, Number>> collectRanges(
            final SearchRequest searchRequest, final IndexSearcher searcher, final FacetsCollector fc)
            throws IOException {
        if (!searchRequest.hasRanges()) {
            return Collections.emptyMap();
        }
        var r = searchRequest.ranges().get();
        final Map<String, Map<String, Number>> result = new HashMap<String, Map<String, Number>>(r.size());
        for (final Entry<String, List<DoubleRange>> entry : r.entrySet()) {
            final DoubleRangeFacetCounts counts = toDoubleRangeFacetCounts(fc, entry.getKey(), entry.getValue());
            result.put(entry.getKey(), collectFacets(counts, searchRequest.topNAsInt(), entry.getKey()));
        }
        return result;
    }

    private DoubleRangeFacetCounts toDoubleRangeFacetCounts(
            final FacetsCollector fc, final String field, final List<DoubleRange> ranges) throws IOException {
        final org.apache.lucene.facet.range.DoubleRange[] luceneRanges =
                new org.apache.lucene.facet.range.DoubleRange[ranges.size()];
        for (int i = 0; i < luceneRanges.length; i++) {
            final DoubleRange range = ranges.get(i);
            luceneRanges[i] = new org.apache.lucene.facet.range.DoubleRange(
                    range.getLabel(),
                    range.getMin() != null ? range.getMin() : Double.NEGATIVE_INFINITY,
                    range.isMinInclusive(),
                    range.getMax() != null ? range.getMax() : Double.POSITIVE_INFINITY,
                    range.isMaxInclusive());
        }
        return new DoubleRangeFacetCounts(field, fc, luceneRanges);
    }

    private Map<String, Number> collectFacets(final Facets facets, final int topN, final String dim)
            throws IOException {
        final FacetResult topChildren = facets.getTopChildren(topN, dim);
        final Map<String, Number> result = new HashMap<String, Number>(topChildren.childCount);
        for (final LabelAndValue lv : topChildren.labelValues) {
            result.put(lv.label, lv.value);
        }
        return result;
    }

    // Ensure _id is final sort field so we can paginate.
    private Sort toSort(final SearchRequest searchRequest) {
        return searchRequest
                .sort()
                .map(s -> {
                    final List<String> sort = new ArrayList<String>(s);
                    final String last = sort.get(sort.size() - 1);
                    // Append _id field if not already present.
                    switch (last) {
                        case "-_id":
                        case "_id":
                            break;
                        default:
                            sort.add("_id");
                    }
                    return convertSort(sort);
                })
                .orElse(DEFAULT_SORT);
    }

    private Sort convertSort(final List<String> sort) {
        final SortField[] fields = new SortField[sort.size()];
        for (int i = 0; i < sort.size(); i++) {
            fields[i] = convertSortField(sort.get(i));
        }
        return new Sort(fields);
    }

    private SortField convertSortField(final String sortString) {
        if ("relevance".equals(sortString)) {
            return SortField.FIELD_SCORE;
        }
        final Matcher m = SORT_FIELD_RE.matcher(sortString);
        if (!m.matches()) {
            throw new WebApplicationException(sortString + " is not a valid sort parameter", Status.BAD_REQUEST);
        }
        final boolean reverse = "-".equals(m.group(1));

        var type = schema.getType(m.group(2));
        if (type == null) {
            return new UnknownSortField(m.group(2), reverse);
        }
        switch (type) {
            case STRING:
                return new SortedSetSortField(m.group(2), reverse);
            case DOUBLE:
                return new SortedNumericSortField(m.group(2), SortField.Type.DOUBLE, reverse);
            default:
                throw new WebApplicationException("can't sort on field " + m.group(2), Status.BAD_REQUEST);
        }
    }

    private static Document toDocument(final String docId, final DocumentUpdateRequest request) throws IOException {
        final Document result = new Document();

        // id
        result.add(new org.apache.lucene.document.StringField("_id", docId, Store.YES));
        result.add(new SortedDocValuesField("_id", new BytesRef(docId)));

        // partition (optional)
        request.partition()
                .ifPresent(partition ->
                        result.add(new org.apache.lucene.document.StringField("_partition", partition, Store.NO)));

        final CharsetDecoder utf8Decoder = StandardCharsets.UTF_8.newDecoder();

        for (Field field : request.fields()) {
            // Underscore-prefix is reserved.
            if (field.name().startsWith("_")) {
                continue;
            }
            if (field instanceof TextField) {
                var f = (TextField) field;
                result.add(new org.apache.lucene.document.TextField(
                        f.name(), f.value(), f.store() ? Store.YES : Store.NO));
            } else if (field instanceof StringField) {
                var f = (StringField) field;
                result.add(new org.apache.lucene.document.KeywordField(
                        f.name(), f.value(), f.store() ? Store.YES : Store.NO));
            } else if (field instanceof DoubleField) {
                var f = (DoubleField) field;
                result.add(new org.apache.lucene.document.DoubleField(
                        f.name(), f.value(), f.store() ? Store.YES : Store.NO));
            } else if (field instanceof StoredField) {
                var f = (StoredField) field;
                var val = f.value();
                if (val instanceof String) {
                    result.add(new org.apache.lucene.document.StoredField(f.name(), (String) val));
                } else if (val instanceof Number) {
                    result.add(new org.apache.lucene.document.StoredField(f.name(), ((Number) val).doubleValue()));
                } else if (val instanceof byte[]) {
                    try {
                        final CharBuffer buf = utf8Decoder.decode(ByteBuffer.wrap((byte[]) val));
                        result.add(new org.apache.lucene.document.StoredField(f.name(), buf.toString()));
                    } catch (final CharacterCodingException e) {
                        result.add(new org.apache.lucene.document.StoredField(f.name(), (byte[]) val));
                    }
                } else {
                    throw new WebApplicationException(field + " is not valid", Status.BAD_REQUEST);
                }
            } else {
                throw new WebApplicationException(field + " is not valid", Status.BAD_REQUEST);
            }
        }

        return result;
    }

    private FieldDoc toFieldDoc(final Object[] after) {
        final Object[] fields = new Object[after.length];
        for (int i = 0; i < after.length; i++) {
            if (after[i] instanceof PrimitiveWrapper<?>) {
                fields[i] = ((PrimitiveWrapper<?>) after[i]).getValue();
            }
            if (fields[i] instanceof byte[]) {
                fields[i] = new BytesRef((byte[]) fields[i]);
            }
            if (fields[i] instanceof String) {
                fields[i] = new BytesRef((String) fields[i]);
            }
        }
        return new FieldDoc(0, Float.NaN, fields);
    }

    private PrimitiveWrapper<?>[] toAfter(final FieldDoc fieldDoc) {
        final CharsetDecoder utf8Decoder = StandardCharsets.UTF_8.newDecoder();
        final PrimitiveWrapper<?>[] fields = new PrimitiveWrapper<?>[fieldDoc.fields.length];
        for (int i = 0; i < fields.length; i++) {
            if (fieldDoc.fields[i] instanceof String) {
                fields[i] = new StringWrapper((String) fieldDoc.fields[i]);
            } else if (fieldDoc.fields[i] instanceof BytesRef) {
                var bytes = toBytes((BytesRef) fieldDoc.fields[i]);
                try {
                    final CharBuffer buf = utf8Decoder.decode(ByteBuffer.wrap(bytes));
                    fields[i] = new StringWrapper(buf.toString());
                } catch (final CharacterCodingException e) {
                    fields[i] = new ByteArrayWrapper(bytes);
                }
            } else if (fieldDoc.fields[i] instanceof Double) {
                fields[i] = new DoubleWrapper((double) fieldDoc.fields[i]);
            } else if (fieldDoc.fields[i] instanceof Integer) {
                fields[i] = new IntWrapper((int) fieldDoc.fields[i]);
            } else if (fieldDoc.fields[i] instanceof Long) {
                fields[i] = new LongWrapper((long) fieldDoc.fields[i]);
            } else if (fieldDoc.fields[i] instanceof Float) {
                fields[i] = new FloatWrapper((float) fieldDoc.fields[i]);
            } else if (fieldDoc.fields[i] == null) {
                fields[i] = new NullWrapper();
            } else {
                throw new WebApplicationException(fieldDoc.fields[i].getClass() + " is not valid", Status.BAD_REQUEST);
            }
        }
        return fields;
    }

    private static byte[] toBytes(final BytesRef bytesRef) {
        return Arrays.copyOfRange(bytesRef.bytes, bytesRef.offset, bytesRef.offset + bytesRef.length);
    }

    private static Query docIdQuery(final String docId) {
        return new TermQuery(docIdTerm(docId));
    }

    private static Term docIdTerm(final String docId) {
        return new Term("_id", docId);
    }

    private Query parse(final SearchRequest request) {
        var pointsConfigMap = schema.toPointsConfigMap(request.getLocale());
        var queryParser = new NouveauQueryParser(analyzer, pointsConfigMap);

        try {
            var result = queryParser.parse(request.query(), "default");
            return request.partition()
                    .map(partition -> {
                        final BooleanQuery.Builder builder = new BooleanQuery.Builder();
                        builder.add(new TermQuery(new Term("_partition", partition)), Occur.MUST);
                        builder.add(result, Occur.MUST);
                        return (Query) builder.build();
                    })
                    .orElse(result);
        } catch (final QueryNodeException e) {
            throw new WebApplicationException(e.getMessage(), e, Status.BAD_REQUEST);
        }
    }

    private LuceneIndexSchema initSchema(IndexWriter writer) {
        var commitData = writer.getLiveCommitData();
        if (commitData == null) {
            return LuceneIndexSchema.emptySchema();
        }
        for (var entry : commitData) {
            if (entry.getKey().equals("_schema")) {
                return LuceneIndexSchema.fromString(entry.getValue());
            }
        }
        return LuceneIndexSchema.emptySchema();
    }

    @Override
    public String toString() {
        return "LuceneIndex [analyzer=" + analyzer + ", writer=" + writer + ", searcherManager=" + searcherManager
                + "]";
    }

    /**
     * This shard is unaware of the type of a sort field as no document within it currently has that field. This
     * custom sort field therefore assumes all documents have null for this field and therefore all hits are equal.
     */
    private static final class UnknownSortField extends SortField {

        private static final FieldComparatorSource COMPARATOR = new FieldComparatorSource() {

            @Override
            public FieldComparator<?> newComparator(String fieldname, int numHits, Pruning pruning, boolean reversed) {
                return new FieldComparator<Void>() {

                    @Override
                    public int compare(int slot1, int slot2) {
                        return 0;
                    }

                    @Override
                    public void setTopValue(Void value) {
                        // empty
                    }

                    @Override
                    public Void value(int slot) {
                        return null;
                    }

                    @Override
                    public LeafFieldComparator getLeafComparator(LeafReaderContext context) throws IOException {
                        return new LeafFieldComparator() {

                            @Override
                            public void setBottom(int slot) throws IOException {
                                // empty
                            }

                            @Override
                            public int compareBottom(int doc) throws IOException {
                                return 0;
                            }

                            @Override
                            public int compareTop(int doc) throws IOException {
                                return 0;
                            }

                            @Override
                            public void copy(int slot, int doc) throws IOException {
                                // empty
                            }

                            @Override
                            public void setScorer(Scorable scorer) throws IOException {
                                // empty
                            }
                        };
                    }
                };
            }
        };

        private UnknownSortField(final String fieldName, final boolean reverse) {
            super(fieldName, COMPARATOR, reverse);
        }
    }
}
