import com.google.common.util.concurrent.MoreExecutors;
import com.lucidworks.apollo.component.ExecutorComponent;
import com.lucidworks.apollo.pipeline.*;
import com.lucidworks.apollo.pipeline.async.AbstractAsyncStageConfig;
import com.lucidworks.apollo.pipeline.async.AsyncStage;
import com.lucidworks.apollo.pipeline.async.AsyncStageConfig;

import java.util.concurrent.ExecutorService;

/**
 * Convenience base class for stages which need asynchronous processing.
 *
 * @param <C> The config class, subclass of {@link AbstractAsyncStageConfig}
 * @param <R> The intermediate result type, passed between {@link #doProcessing} and {@link #mergeResults}.
 * @see AsyncStage
 */
public abstract class AbstractAsyncQueryStage<C extends AbstractAsyncStageConfig, R> extends QueryStage<C>
    implements AsyncStage<R, QueryRequestAndResponse> {

  private final ExecutorService executorService;

  protected AbstractAsyncQueryStage(StageAssistFactoryParams params, ExecutorComponent executorComponent) {
    super(params);
    if(getConfiguration().isAsyncEnabled()) {
      int threadPoolSize = Math.max(0, params.getConfigurationComponent().getInt(
          "com.lucidworks.apollo.pipeline.query.async.threadPoolSize", 1000));
      if(threadPoolSize == 0) {
        executorService = executorComponent.getOrCreateCachedThreadPool("system-async-query-pipelines");
      } else {
        executorService = executorComponent.getOrCreateFixedThreadPool(threadPoolSize, "system-async-query-pipelines");
      }
    } else {
      executorService = MoreExecutors.newDirectExecutorService();
    }
  }

  @Override
  public final QueryRequestAndResponse process(QueryRequestAndResponse message, Context context) throws Exception {

    if(getConfiguration().isAsyncEnabled()) {
      AsyncStageConfig asyncConfig = getConfiguration().getAsyncConfig();
      AsyncStage.runInExecutor(this, message, context, asyncConfig.getAsyncId(), executorService);
    } else {
      AsyncStage.runInForeground(this, message, context);
    }
    return message;
  }
}