package org.securegraph.elasticsearch;

import org.apache.lucene.index.*;
import org.apache.lucene.search.BitsFilteredDocIdSet;
import org.apache.lucene.search.DocIdSet;
import org.apache.lucene.search.Filter;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.OpenBitSet;
import org.securegraph.inmemory.security.Authorizations;
import org.securegraph.inmemory.security.ColumnVisibility;
import org.securegraph.inmemory.security.VisibilityEvaluator;
import org.securegraph.inmemory.security.VisibilityParseException;

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

public class AuthorizationsFilter extends Filter {
    public static String VISIBILITY_FIELD_NAME = "__visibility";
    private final VisibilityEvaluator visibilityEvaluator;
    private static final Map<BytesRef, ColumnVisibility> columnVisibilityCache = new HashMap<BytesRef, ColumnVisibility>();

    public AuthorizationsFilter(Authorizations authorizations) {
        this.visibilityEvaluator = new VisibilityEvaluator(authorizations);
    }

    @Override
    public DocIdSet getDocIdSet(AtomicReaderContext context, Bits acceptDocs) throws IOException {
        AtomicReader reader = context.reader();
        Fields fields = reader.fields();
        Terms terms = fields.terms(VISIBILITY_FIELD_NAME);
        if (terms == null) {
            return null;
        } else {
            OpenBitSet bitSet = new OpenBitSet(reader.maxDoc());
            TermsEnum iterator = terms.iterator(null);
            BytesRef bytesRef;
            while ((bytesRef = iterator.next()) != null) {
                makeVisible(iterator, bitSet, acceptDocs, isVisible(visibilityEvaluator, bytesRef));
            }
            return BitsFilteredDocIdSet.wrap(bitSet, acceptDocs);
        }
    }

    private void makeVisible(TermsEnum iterator, OpenBitSet bitSet, Bits liveDocs, boolean visible) throws IOException {
        DocsEnum docsEnum = iterator.docs(liveDocs, null);
        int doc;
        while ((doc = docsEnum.nextDoc()) != DocsEnum.NO_MORE_DOCS) {
            if (visible) {
                bitSet.set(doc);
            } else {
                bitSet.clear(doc);
            }
        }
    }

    private static boolean isVisible(VisibilityEvaluator visibilityEvaluator, BytesRef bytesRef) throws IOException {
        ColumnVisibility visibility = lookupColumnVisibility(bytesRef);
        if (visibility == null) {
            return true;
        }
        try {
            return visibilityEvaluator.evaluate(visibility);
        } catch (VisibilityParseException e) {
            throw new IOException(e);
        }
    }

    private static ColumnVisibility lookupColumnVisibility(BytesRef bytesRef) {
        ColumnVisibility visibility = columnVisibilityCache.get(bytesRef);
        if (visibility != null) {
            return visibility;
        }

        byte[] expression = trim(bytesRef);
        if (expression.length == 0) {
            return null;
        }
        visibility = new ColumnVisibility(expression);
        columnVisibilityCache.put(bytesRef, visibility);
        return visibility;
    }

    private static byte[] trim(BytesRef bytesRef) {
        byte[] buf = new byte[bytesRef.length];
        System.arraycopy(bytesRef.bytes, bytesRef.offset, buf, 0, bytesRef.length);
        return buf;
    }
}

@@ -642,16 +642,63 @@
 til.Map;
+%0Aimport java.util.concurrent.ConcurrentHashMap;
 %0A%0Apublic
@@ -955,42 +955,26 @@
 new 
-HashMap%3CBytesRef, ColumnVisibility
+ConcurrentHashMap%3C
 %3E();

change columnVisibilityCache in AuthorizationsFilter to ConcurrentHashMap