/*
 * 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.
 */
package org.apache.lucene.search;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;

import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;

A Scorer for queries with a required part and an optional part.
Delays skipTo() on the optional part until a score() is needed.
/** A Scorer for queries with a required part and an optional part.
 * Delays skipTo() on the optional part until a score() is needed.
 */
class ReqOptSumScorer extends Scorer {
  private final Scorer reqScorer;
  private final Scorer optScorer;
  private final DocIdSetIterator reqApproximation;
  private final DocIdSetIterator optApproximation;
  private final TwoPhaseIterator optTwoPhase;
  private final DocIdSetIterator approximation;
  private final TwoPhaseIterator twoPhase;

  private final MaxScoreSumPropagator maxScorePropagator;
  private float minScore = 0;
  private float reqMaxScore;
  private boolean optIsRequired;

  
Construct a ReqOptScorer.
Params:
reqScorer – The required scorer. This must match.
optScorer – The optional scorer. This is used for scoring only.
scoreMode –  How the produced scorers will be consumed./**
   * Construct a <code>ReqOptScorer</code>.
   *
   * @param reqScorer The required scorer. This must match.
   * @param optScorer The optional scorer. This is used for scoring only.
   * @param scoreMode  How the produced scorers will be consumed.
   */
  public ReqOptSumScorer(Scorer reqScorer, Scorer optScorer, ScoreMode scoreMode) throws IOException {
    super(reqScorer.weight);
    assert reqScorer != null;
    assert optScorer != null;
    if (scoreMode == ScoreMode.TOP_SCORES) {
      this.maxScorePropagator = new MaxScoreSumPropagator(Arrays.asList(reqScorer, optScorer));
    } else {
      this.maxScorePropagator = null;
    }
    this.reqScorer = reqScorer;
    this.optScorer = optScorer;

    final TwoPhaseIterator reqTwoPhase = reqScorer.twoPhaseIterator();
    this.optTwoPhase = optScorer.twoPhaseIterator();
    if (reqTwoPhase == null) {
      reqApproximation = reqScorer.iterator();
    } else {
      reqApproximation = reqTwoPhase.approximation();
    }
    if (optTwoPhase == null) {
      optApproximation = optScorer.iterator();
    } else {
      optApproximation = optTwoPhase.approximation();
    }
    if (scoreMode != ScoreMode.TOP_SCORES) {
      approximation = reqApproximation;
      this.reqMaxScore = Float.POSITIVE_INFINITY;
    } else {
      reqScorer.advanceShallow(0);
      optScorer.advanceShallow(0);
      this.reqMaxScore = reqScorer.getMaxScore(NO_MORE_DOCS);
      this.approximation = new DocIdSetIterator() {
        int upTo = -1;
        float maxScore;

        private void moveToNextBlock(int target) throws IOException {
          upTo = advanceShallow(target);
          float reqMaxScoreBlock = reqScorer.getMaxScore(upTo);
          maxScore = getMaxScore(upTo);

          // Potentially move to a conjunction
          optIsRequired = reqMaxScoreBlock < minScore;
        }

        private int advanceImpacts(int target) throws IOException {
          if (target > upTo) {
            moveToNextBlock(target);
          }

          while (true) {
            if (maxScore >= minScore) {
              return target;
            }

            if (upTo == NO_MORE_DOCS) {
              return NO_MORE_DOCS;
            }

            target = upTo + 1;

            moveToNextBlock(target);
          }
        }

        @Override
        public int nextDoc() throws IOException {
          return advanceInternal(reqApproximation.docID()+1);
        }

        @Override
        public int advance(int target) throws IOException {
          return advanceInternal(target);
        }

        private int advanceInternal(int target) throws IOException {
          if (target == NO_MORE_DOCS) {
            reqApproximation.advance(target);
            return NO_MORE_DOCS;
          }
          int reqDoc = target;
          advanceHead: for (;;) {
            if (minScore != 0) {
              reqDoc = advanceImpacts(reqDoc);
            }
            if (reqApproximation.docID() < reqDoc) {
              reqDoc = reqApproximation.advance(reqDoc);
            }
            if (reqDoc == NO_MORE_DOCS || optIsRequired == false) {
              return reqDoc;
            }

            int upperBound = reqMaxScore < minScore ? NO_MORE_DOCS : upTo;
            if (reqDoc > upperBound) {
              continue;
            }

            // Find the next common doc within the current block
            for (;;) { // invariant: reqDoc >= optDoc
              int optDoc = optApproximation.docID();
              if (optDoc < reqDoc) {
                optDoc = optApproximation.advance(reqDoc);
              }
              if (optDoc > upperBound) {
                reqDoc = upperBound + 1;
                continue advanceHead;
              }

              if (optDoc != reqDoc) {
                reqDoc = reqApproximation.advance(optDoc);
                if (reqDoc > upperBound) {
                  continue advanceHead;
                }
              }

              if (reqDoc == NO_MORE_DOCS || optDoc == reqDoc) {
                return reqDoc;
              }
            }
          }
        }

        @Override
        public int docID() {
          return reqApproximation.docID();
        }

        @Override
        public long cost() {
          return reqApproximation.cost();
        }
      };
    }

    if (reqTwoPhase == null && optTwoPhase == null) {
      this.twoPhase = null;
    } else {
      this.twoPhase = new TwoPhaseIterator(approximation) {

        @Override
        public boolean matches() throws IOException {
          if (reqTwoPhase != null && reqTwoPhase.matches() == false) {
            return false;
          }
          if (optTwoPhase != null) {
            if (optIsRequired) {
              // The below condition is rare and can only happen if we transitioned to optIsRequired=true
              // after the opt approximation was advanced and before it was confirmed.
              if (reqScorer.docID() != optApproximation.docID()) {
                if (optApproximation.docID() < reqScorer.docID()) {
                  optApproximation.advance(reqScorer.docID());
                }
                if (reqScorer.docID() != optApproximation.docID()) {
                  return false;
                }
              }
              if (optTwoPhase.matches() == false) {
                // Advance the iterator to make it clear it doesn't match the current doc id
                optApproximation.nextDoc();
                return false;
              }
            } else if (optApproximation.docID() == reqScorer.docID() && optTwoPhase.matches() == false) {
              // Advance the iterator to make it clear it doesn't match the current doc id
              optApproximation.nextDoc();
            }
          }
          return true;
        }

        @Override
        public float matchCost() {
          float matchCost = 1;
          if (reqTwoPhase != null) {
            matchCost += reqTwoPhase.matchCost();
          }
          if (optTwoPhase != null) {
            matchCost += optTwoPhase.matchCost();
          }
          return matchCost;
        }
      };
    }
  }

  @Override
  public TwoPhaseIterator twoPhaseIterator() {
    return twoPhase;
  }

  @Override
  public DocIdSetIterator iterator() {
    if (twoPhase == null) {
      return approximation;
    } else {
      return TwoPhaseIterator.asDocIdSetIterator(twoPhase);
    }
  }

  @Override
  public int docID() {
    return reqScorer.docID();
  }

  @Override
  public float score() throws IOException {
    // TODO: sum into a double and cast to float if we ever send required clauses to BS1
    int curDoc = reqScorer.docID();
    float score = reqScorer.score();

    int optScorerDoc = optApproximation.docID();
    if (optScorerDoc < curDoc) {
      optScorerDoc = optApproximation.advance(curDoc);
      if (optTwoPhase != null && optScorerDoc == curDoc && optTwoPhase.matches() == false) {
        optScorerDoc = optApproximation.nextDoc();
      }
    }
    if (optScorerDoc == curDoc) {
      score += optScorer.score();
    }

    return score;
  }

  @Override
  public int advanceShallow(int target) throws IOException {
    int upTo = reqScorer.advanceShallow(target);
    if (optScorer.docID() <= target) {
      upTo = Math.min(upTo, optScorer.advanceShallow(target));
    } else if (optScorer.docID() != NO_MORE_DOCS) {
      upTo = Math.min(upTo, optScorer.docID() - 1);
    }
    return upTo;
  }

  @Override
  public float getMaxScore(int upTo) throws IOException {
    float maxScore = reqScorer.getMaxScore(upTo);
    if (optScorer.docID() <= upTo) {
      maxScore += optScorer.getMaxScore(upTo);
    }
    return maxScore;
  }

  @Override
  public void setMinCompetitiveScore(float minScore) throws IOException {
    this.minScore = minScore;
    // Potentially move to a conjunction
    if (reqMaxScore < minScore) {
      optIsRequired = true;
    }
    maxScorePropagator.setMinCompetitiveScore(minScore);
  }

  @Override
  public Collection<ChildScorable> getChildren() {
    ArrayList<ChildScorable> children = new ArrayList<>(2);
    children.add(new ChildScorable(reqScorer, "MUST"));
    children.add(new ChildScorable(optScorer, "SHOULD"));
    return children;
  }
}