package org.apache.lucene.search.join;

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF 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.
 */

import org.apache.lucene.index.AtomicReaderContext;
import org.apache.lucene.search.DocIdSet;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.FieldComparator;
import org.apache.lucene.search.Filter;
import org.apache.lucene.util.FixedBitSet;

import java.io.IOException;

/**
 * A field comparator that allows parent documents to be sorted by fields
 * from the nested / child documents.
 *
 * @lucene.experimental
 */
public abstract class ToParentBlockJoinFieldComparator extends FieldComparator<Object> {

  private final Filter parentFilter;
  private final Filter childFilter;
  final int spareSlot;

  FieldComparator<Object> wrappedComparator;
  FixedBitSet parentDocuments;
  FixedBitSet childDocuments;

  ToParentBlockJoinFieldComparator(FieldComparator<Object> wrappedComparator, Filter parentFilter, Filter childFilter, int spareSlot) {
    this.wrappedComparator = wrappedComparator;
    this.parentFilter = parentFilter;
    this.childFilter = childFilter;
    this.spareSlot = spareSlot;
  }

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

  @Override
  public void setBottom(int slot) {
    wrappedComparator.setBottom(slot);
  }

  @Override
  public void setTopValue(Object value) {
    wrappedComparator.setTopValue(value);
  }

  @Override
  public FieldComparator<Object> setNextReader(AtomicReaderContext context) throws IOException {
    DocIdSet innerDocuments = childFilter.getDocIdSet(context, null);
    if (isEmpty(innerDocuments)) {
      this.childDocuments = null;
    } else if (innerDocuments instanceof FixedBitSet) {
      this.childDocuments = (FixedBitSet) innerDocuments;
    } else {
      DocIdSetIterator iterator = innerDocuments.iterator();
      if (iterator != null) {
        this.childDocuments = toFixedBitSet(iterator, context.reader().maxDoc());
      } else {
        childDocuments = null;
      }
    }
    DocIdSet rootDocuments = parentFilter.getDocIdSet(context, null);
    if (isEmpty(rootDocuments)) {
      this.parentDocuments = null;
    } else if (rootDocuments instanceof FixedBitSet) {
      this.parentDocuments = (FixedBitSet) rootDocuments;
    } else {
      DocIdSetIterator iterator = rootDocuments.iterator();
      if (iterator != null) {
        this.parentDocuments = toFixedBitSet(iterator, context.reader().maxDoc());
      } else {
        this.parentDocuments = null;
      }
    }

    wrappedComparator = wrappedComparator.setNextReader(context);
    return this;
  }

  private static boolean isEmpty(DocIdSet set) {
    return set == null;
  }

  private static FixedBitSet toFixedBitSet(DocIdSetIterator iterator, int numBits) throws IOException {
    FixedBitSet set = new FixedBitSet(numBits);
    int doc;
    while ((doc = iterator.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {
      set.set(doc);
    }
    return set;
  }

  @Override
  public Object value(int slot) {
    return wrappedComparator.value(slot);
  }

  /**
   * Concrete implementation of {@link ToParentBlockJoinSortField} to sorts the parent docs with the lowest values
   * in the child / nested docs first.
   */
  public static final class Lowest extends ToParentBlockJoinFieldComparator {

    /**
     * Create ToParentBlockJoinFieldComparator.Lowest
     *
     * @param wrappedComparator The {@link FieldComparator} on the child / nested level.
     * @param parentFilter Filter (must produce FixedBitSet per-segment) that identifies the parent documents.
     * @param childFilter Filter that defines which child / nested documents participates in sorting.
     * @param spareSlot The extra slot inside the wrapped comparator that is used to compare which nested document
     *                  inside the parent document scope is most competitive.
     */
    public Lowest(FieldComparator<Object> wrappedComparator, Filter parentFilter, Filter childFilter, int spareSlot) {
      super(wrappedComparator, parentFilter, childFilter, spareSlot);
    }

    @Override
    public int compareBottom(int parentDoc) throws IOException {
      if (parentDoc == 0 || parentDocuments == null || childDocuments == null) {
        return 0;
      }

      // We need to copy the lowest value from all child docs into slot.
      int prevParentDoc = parentDocuments.prevSetBit(parentDoc - 1);
      int childDoc = childDocuments.nextSetBit(prevParentDoc + 1);
      if (childDoc >= parentDoc || childDoc == -1) {
        return 0;
      }

      // We only need to emit a single cmp value for any matching child doc
      int cmp = wrappedComparator.compareBottom(childDoc);
      if (cmp > 0) {
        return cmp;
      }

      while (true) {
        childDoc = childDocuments.nextSetBit(childDoc + 1);
        if (childDoc >= parentDoc || childDoc == -1) {
          return cmp;
        }
        int cmp1 = wrappedComparator.compareBottom(childDoc);
        if (cmp1 > 0) {
          return cmp1;
        } else {
          if (cmp1 == 0) {
            cmp = 0;
          }
        }
      }
    }

    @Override
    public void copy(int slot, int parentDoc) throws IOException {
      if (parentDoc == 0 || parentDocuments == null || childDocuments == null) {
        return;
      }

      // We need to copy the lowest value from all child docs into slot.
      int prevParentDoc = parentDocuments.prevSetBit(parentDoc - 1);
      int childDoc = childDocuments.nextSetBit(prevParentDoc + 1);
      if (childDoc >= parentDoc || childDoc == -1) {
        return;
      }
      wrappedComparator.copy(spareSlot, childDoc);
      wrappedComparator.copy(slot, childDoc);

      while (true) {
        childDoc = childDocuments.nextSetBit(childDoc + 1);
        if (childDoc >= parentDoc || childDoc == -1) {
          return;
        }
        wrappedComparator.copy(spareSlot, childDoc);
        if (wrappedComparator.compare(spareSlot, slot) < 0) {
          wrappedComparator.copy(slot, childDoc);
        }
      }
    }

    @Override
    @SuppressWarnings("unchecked")
    public int compareTop(int parentDoc) throws IOException {
      if (parentDoc == 0 || parentDocuments == null || childDocuments == null) {
        return 0;
      }

      // We need to copy the lowest value from all nested docs into slot.
      int prevParentDoc = parentDocuments.prevSetBit(parentDoc - 1);
      int childDoc = childDocuments.nextSetBit(prevParentDoc + 1);
      if (childDoc >= parentDoc || childDoc == -1) {
        return 0;
      }

      // We only need to emit a single cmp value for any matching child doc
      int cmp = wrappedComparator.compareBottom(childDoc);
      if (cmp > 0) {
        return cmp;
      }

      while (true) {
        childDoc = childDocuments.nextSetBit(childDoc + 1);
        if (childDoc >= parentDoc || childDoc == -1) {
          return cmp;
        }
        int cmp1 = wrappedComparator.compareTop(childDoc);
        if (cmp1 > 0) {
          return cmp1;
        } else {
          if (cmp1 == 0) {
            cmp = 0;
          }
        }
      }
    }

  }

  /**
   * Concrete implementation of {@link ToParentBlockJoinSortField} to sorts the parent docs with the highest values
   * in the child / nested docs first.
   */
  public static final class Highest extends ToParentBlockJoinFieldComparator {

    /**
     * Create ToParentBlockJoinFieldComparator.Highest
     *
     * @param wrappedComparator The {@link FieldComparator} on the child / nested level.
     * @param parentFilter Filter (must produce FixedBitSet per-segment) that identifies the parent documents.
     * @param childFilter Filter that defines which child / nested documents participates in sorting.
     * @param spareSlot The extra slot inside the wrapped comparator that is used to compare which nested document
     *                  inside the parent document scope is most competitive.
     */
    public Highest(FieldComparator<Object> wrappedComparator, Filter parentFilter, Filter childFilter, int spareSlot) {
      super(wrappedComparator, parentFilter, childFilter, spareSlot);
    }

    @Override
    public int compareBottom(int parentDoc) throws IOException {
      if (parentDoc == 0 || parentDocuments == null || childDocuments == null) {
        return 0;
      }

      int prevParentDoc = parentDocuments.prevSetBit(parentDoc - 1);
      int childDoc = childDocuments.nextSetBit(prevParentDoc + 1);
      if (childDoc >= parentDoc || childDoc == -1) {
        return 0;
      }

      int cmp = wrappedComparator.compareBottom(childDoc);
      if (cmp < 0) {
        return cmp;
      }

      while (true) {
        childDoc = childDocuments.nextSetBit(childDoc + 1);
        if (childDoc >= parentDoc || childDoc == -1) {
          return cmp;
        }
        int cmp1 = wrappedComparator.compareBottom(childDoc);
        if (cmp1 < 0) {
          return cmp1;
        } else {
          if (cmp1 == 0) {
            cmp = 0;
          }
        }
      }
    }

    @Override
    public void copy(int slot, int parentDoc) throws IOException {
      if (parentDoc == 0 || parentDocuments == null || childDocuments == null) {
        return;
      }

      int prevParentDoc = parentDocuments.prevSetBit(parentDoc - 1);
      int childDoc = childDocuments.nextSetBit(prevParentDoc + 1);
      if (childDoc >= parentDoc || childDoc == -1) {
        return;
      }
      wrappedComparator.copy(spareSlot, childDoc);
      wrappedComparator.copy(slot, childDoc);

      while (true) {
        childDoc = childDocuments.nextSetBit(childDoc + 1);
        if (childDoc >= parentDoc || childDoc == -1) {
          return;
        }
        wrappedComparator.copy(spareSlot, childDoc);
        if (wrappedComparator.compare(spareSlot, slot) > 0) {
          wrappedComparator.copy(slot, childDoc);
        }
      }
    }

    @Override
    @SuppressWarnings("unchecked")
    public int compareTop(int parentDoc) throws IOException {
      if (parentDoc == 0 || parentDocuments == null || childDocuments == null) {
        return 0;
      }

      int prevParentDoc = parentDocuments.prevSetBit(parentDoc - 1);
      int childDoc = childDocuments.nextSetBit(prevParentDoc + 1);
      if (childDoc >= parentDoc || childDoc == -1) {
        return 0;
      }

      int cmp = wrappedComparator.compareBottom(childDoc);
      if (cmp < 0) {
        return cmp;
      }

      while (true) {
        childDoc = childDocuments.nextSetBit(childDoc + 1);
        if (childDoc >= parentDoc || childDoc == -1) {
          return cmp;
        }
        int cmp1 = wrappedComparator.compareTop(childDoc);
        if (cmp1 < 0) {
          return cmp1;
        } else {
          if (cmp1 == 0) {
            cmp = 0;
          }
        }
      }
    }

  }

}
