// Copyright 2022 The Bazel Authors. All rights reserved.
//
// 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 com.google.devtools.build.skyframe.state;

import static com.google.common.base.Preconditions.checkNotNull;

import com.google.common.collect.Lists;
import com.google.devtools.build.lib.events.ExtendedEventHandler;
import com.google.devtools.build.skyframe.SkyFunction;
import com.google.devtools.build.skyframe.SkyFunction.Environment;
import com.google.devtools.build.skyframe.SkyKey;
import com.google.devtools.build.skyframe.SkyValue;
import com.google.devtools.build.skyframe.SkyframeLookupResult;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.function.Consumer;

/**
 * This class drives a {@link StateMachine} instance.
 *
 * <p>One recommended usage pattern for this class is to embed an instance within a top level {@link
 * StateMachine} implementation and from there, re-export the {@link #drive} method. Then the
 * results from the {@link StateMachine} will be readily retrievable from the {@link SkyFunction}
 * state.
 */
// TODO(shahan); this is incompatible with partial re-evaluation, which causes the assumption that
// an unavailable previously requested dependency implies an error to no longer be true. This can be
// fixed by integrating with the partial re-evaluation mailbox.
public final class Driver {
  // TODO(shahan): tune this parameter.
  private static final int DEFAULT_LIST_CAPACITY = 128;

  private final ArrayDeque<TaskTreeNode> ready = new ArrayDeque<>();

  /** A Skyframe lookup has not yet been made for the key. */
  private final ArrayList<Lookup> newlyAdded = new ArrayList<>(DEFAULT_LIST_CAPACITY);

  /** A Skyframe lookup has already been made for the key, but it was not available. */
  private final ArrayList<Lookup> pending = new ArrayList<>(DEFAULT_LIST_CAPACITY);

  public Driver(StateMachine root) {
    ready.addLast(new TaskTreeNode(/* parent= */ null, root));
  }

  /**
   * Drives the machine as far as it can go without a Skyframe restart.
   *
   * @return true if execution is complete, false if a restart is needed.
   */
  public boolean drive(Environment env, ExtendedEventHandler listener) throws InterruptedException {
    if (!pending.isEmpty()) {
      // If pending is non-empty, it means there was a Skyframe restart. Either everything that was
      // pending is available now or we are in error bubbling. In the latter case, an error will be
      // present and this code will fail fast.
      //
      // NB: this assumption does not hold under partial re-evaluation and likewise the inference
      // below about unavailable values being errors.
      SkyframeLookupResult result = env.getLookupHandleForPreviouslyRequestedDeps();
      boolean hasException = false;
      for (var lookup : pending) {
        if (!result.queryDep(lookup.key(), lookup)) {
          // Since the key was previously requested, unavailability here implies an unhandled or
          // unavailable exception. Requests the error to ensure that this environment instance
          // knows the failure is due to a child error.
          var unusedNull = env.getValue(lookup.key());
          // Failing fast here would make behavior dependent on element ordering, so instead, flags
          // the exception and fails after all elements have been processed.
          hasException = true;
        }
      }
      if (hasException) {
        return false;
      }
      pending.clear();
    }

    while (true) {
      // Executes as many ready tasks as possible.
      while (!ready.isEmpty()) {
        TaskTreeNode next = ready.removeFirst();
        if (next.run(new TasksImpl(next), listener)) {
          TaskTreeNode parent = next.parent();
          if (parent == null) {
            return true; // Root is complete.
          }
          signalParentAndEnqueueIfReady(parent);
        }
      }

      if (newlyAdded.isEmpty()) {
        return false; // No tasks are ready and all remaining lookups are pending.
      }

      // Performs lookups for any newly added keys.
      SkyframeLookupResult result =
          env.getValuesAndExceptions(Lists.transform(newlyAdded, Lookup::key));
      for (var lookup : newlyAdded) {
        if (!result.queryDep(lookup.key(), lookup)) {
          pending.add(lookup); // Unhandled exceptions also end up here.
        }
      }
      newlyAdded.clear(); // Every entry is either done or has moved to pending.
    }
  }

  private void signalParentAndEnqueueIfReady(TaskTreeNode parent) {
    if (parent.decrementChildCount()) {
      ready.addLast(parent);
    }
  }

  private final class TasksImpl implements StateMachine.Tasks {
    private final TaskTreeNode parent;

    private TasksImpl(TaskTreeNode parent) {
      this.parent = parent;
    }

    @Override
    public void enqueue(StateMachine subtask) {
      parent.incrementChildCount();
      ready.addLast(new TaskTreeNode(parent, subtask));
    }

    @Override
    public void lookUp(SkyKey key, Consumer<SkyValue> sink) {
      lookUp(key, sink, exception -> false);
    }

    @Override
    public <E extends Exception> void lookUp(
        SkyKey key, Class<E> exceptionClass, StateMachine.ValueOrExceptionSink<E> sink) {
      lookUp(
          key,
          value -> sink.accept(value, /* exception= */ null),
          e -> {
            if (exceptionClass.isInstance(e)) {
              sink.accept(/* value= */ null, exceptionClass.cast(e));
              return true;
            }
            return false;
          });
    }

    @Override
    public <E1 extends Exception, E2 extends Exception> void lookUp(
        SkyKey key,
        Class<E1> exceptionClass1,
        Class<E2> exceptionClass2,
        StateMachine.ValueOrException2Sink<E1, E2> sink) {
      lookUp(
          key,
          value -> sink.accept(value, /* e1= */ null, /* e2= */ null),
          e -> {
            if (exceptionClass1.isInstance(e)) {
              sink.accept(/* value= */ null, exceptionClass1.cast(e), /* e2= */ null);
              return true;
            }
            if (exceptionClass2.isInstance(e)) {
              sink.accept(/* value= */ null, /* e1= */ null, exceptionClass2.cast(e));
              return true;
            }
            return false;
          });
    }

    @Override
    public <E1 extends Exception, E2 extends Exception, E3 extends Exception> void lookUp(
        SkyKey key,
        Class<E1> exceptionClass1,
        Class<E2> exceptionClass2,
        Class<E3> exceptionClass3,
        StateMachine.ValueOrException3Sink<E1, E2, E3> sink) {
      lookUp(
          key,
          value -> sink.accept(value, /* e1= */ null, /* e2= */ null, /* e3= */ null),
          e -> {
            if (exceptionClass1.isInstance(e)) {
              sink.accept(
                  /* value= */ null, exceptionClass1.cast(e), /* e2= */ null, /* e3= */ null);
              return true;
            }
            if (exceptionClass2.isInstance(e)) {
              sink.accept(
                  /* value= */ null, /* e1= */ null, exceptionClass2.cast(e), /* e3= */ null);
              return true;
            }
            if (exceptionClass3.isInstance(e)) {
              sink.accept(
                  /* value= */ null, /* e1= */ null, /* e2= */ null, exceptionClass3.cast(e));
              return true;
            }
            return false;
          });
    }

    /**
     * Adds a dependency to look up.
     *
     * <p>Exactly one of {@code sink} or {@code exceptionHandler} will be called. The callback could
     * be deferred until the next Skyframe restart if the queried key is not immediately available.
     */
    private void lookUp(SkyKey key, Consumer<SkyValue> sink, ExceptionHandler exceptionHandler) {
      parent.incrementChildCount();
      newlyAdded.add(new Lookup(parent, checkNotNull(key), sink, exceptionHandler));
    }
  }

  private final class Lookup implements SkyframeLookupResult.QueryDepCallback {
    private final TaskTreeNode parent;
    private final SkyKey key;
    private final Consumer<SkyValue> sink;
    private final ExceptionHandler exceptionHandler;

    private Lookup(
        TaskTreeNode parent,
        SkyKey key,
        Consumer<SkyValue> sink,
        ExceptionHandler exceptionHandler) {
      this.parent = parent;
      this.key = key;
      this.sink = sink;
      this.exceptionHandler = exceptionHandler;
    }

    private SkyKey key() {
      return key;
    }

    @Override
    public void acceptValue(SkyKey unusedKey, SkyValue value) {
      sink.accept(value);
      signalParentAndEnqueueIfReady(parent);
    }

    @Override
    public boolean tryHandleException(SkyKey unusedKey, Exception exception) {
      if (exceptionHandler.tryHandleException(exception)) {
        signalParentAndEnqueueIfReady(parent);
        return true;
      }
      return false;
    }
  }

  /**
   * Handles exceptions.
   *
   * <p>Receiving an exception and not handling it results in it bubbling up at the Skyframe level.
   *
   * <p>An exception may be offered to this method multiple times. Changing the return value results
   * in undefined behavior.
   */
  @FunctionalInterface
  interface ExceptionHandler {
    /** Returns true if an exception was handled and false otherwise. */
    boolean tryHandleException(Exception e);
  }
}