/**
 * Copyright 2019 Airbnb. Licensed under Apache-2.0. See LICENSE in the project root for license
 * information.
 */
package com.airbnb.dynein.scheduler.modules;

import com.airbnb.conveyor.async.AsyncSqsClient;
import com.airbnb.conveyor.async.AsyncSqsClientFactory;
import com.airbnb.conveyor.async.config.AsyncSqsClientConfiguration;
import com.airbnb.conveyor.async.config.OverrideConfiguration;
import com.airbnb.conveyor.async.config.RetryPolicyConfiguration;
import com.airbnb.conveyor.async.metrics.AsyncConveyorMetrics;
import com.airbnb.conveyor.async.metrics.NoOpConveyorMetrics;
import com.airbnb.dynein.common.job.JacksonJobSpecTransformer;
import com.airbnb.dynein.common.job.JobSpecTransformer;
import com.airbnb.dynein.common.token.JacksonTokenManager;
import com.airbnb.dynein.common.token.TokenManager;
import com.airbnb.dynein.scheduler.Constants;
import com.airbnb.dynein.scheduler.ManagedInboundJobQueueConsumer;
import com.airbnb.dynein.scheduler.ScheduleManager;
import com.airbnb.dynein.scheduler.ScheduleManagerFactory;
import com.airbnb.dynein.scheduler.config.DynamoDBConfiguration;
import com.airbnb.dynein.scheduler.config.SQSConfiguration;
import com.airbnb.dynein.scheduler.config.SchedulerConfiguration;
import com.airbnb.dynein.scheduler.config.WorkersConfiguration;
import com.airbnb.dynein.scheduler.dynamodb.DynamoDBScheduleManagerFactory;
import com.airbnb.dynein.scheduler.heartbeat.HeartbeatConfiguration;
import com.airbnb.dynein.scheduler.heartbeat.HeartbeatManager;
import com.airbnb.dynein.scheduler.metrics.Metrics;
import com.airbnb.dynein.scheduler.metrics.NoOpMetricsImpl;
import com.airbnb.dynein.scheduler.partition.K8SConsecutiveAllocationPolicy;
import com.airbnb.dynein.scheduler.partition.PartitionPolicy;
import com.airbnb.dynein.scheduler.partition.StaticAllocationPolicy;
import com.airbnb.dynein.scheduler.worker.PartitionWorker;
import com.airbnb.dynein.scheduler.worker.PartitionWorkerFactory;
import com.airbnb.dynein.scheduler.worker.WorkerManager;
import com.google.common.eventbus.EventBus;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.google.inject.AbstractModule;
import com.google.inject.Provides;
import com.google.inject.Singleton;
import com.google.inject.name.Names;
import java.time.Clock;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadPoolExecutor;
import javax.inject.Named;
import lombok.extern.slf4j.Slf4j;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;

@Slf4j
public class SchedulerModule extends AbstractModule {
  private final SchedulerConfiguration schedulerConfiguration;
  private final DynamoDBConfiguration dynamoDBConfiguration;
  private final SQSConfiguration sqsConfiguration;
  private final HeartbeatConfiguration heartbeatConfiguration;
  private final WorkersConfiguration workersConfiguration;

  public SchedulerModule(SchedulerConfiguration schedulerConfiguration) {
    this.schedulerConfiguration = schedulerConfiguration;
    this.dynamoDBConfiguration = schedulerConfiguration.getDynamoDb();
    this.sqsConfiguration = schedulerConfiguration.getSqs();
    this.heartbeatConfiguration = schedulerConfiguration.getHeartbeat();
    this.workersConfiguration = schedulerConfiguration.getWorkers();
  }

  @Override
  protected void configure() {
    bind(Clock.class).toInstance(Clock.systemUTC());
    bind(WorkersConfiguration.class).toInstance(workersConfiguration);
    bind(DynamoDBConfiguration.class).toInstance(dynamoDBConfiguration);
    bind(Integer.class)
        .annotatedWith(Names.named(Constants.MAX_PARTITIONS))
        .toInstance(schedulerConfiguration.getMaxPartitions());
    bind(String.class)
        .annotatedWith(Names.named(Constants.INBOUND_QUEUE_NAME))
        .toInstance(sqsConfiguration.getInboundQueueName());

    bind(DynamoDBScheduleManagerFactory.class).asEagerSingleton();
    bind(ScheduleManagerFactory.class).to(DynamoDBScheduleManagerFactory.class).asEagerSingleton();
    bind(ScheduleManager.class).toProvider(ScheduleManagerFactory.class).asEagerSingleton();
    bind(EventBus.class).toInstance(new EventBus());

    switch (workersConfiguration.getPartitionPolicy()) {
      case K8S:
        {
          bind(K8SConsecutiveAllocationPolicy.class).asEagerSingleton();
          bind(PartitionPolicy.class).to(K8SConsecutiveAllocationPolicy.class);
          break;
        }
      case STATIC:
        {
          bind(PartitionPolicy.class)
              .toInstance(
                  new StaticAllocationPolicy(workersConfiguration.getStaticPartitionList()));
          break;
        }
    }
    bind(Metrics.class).to(NoOpMetricsImpl.class).asEagerSingleton();
    bind(AsyncConveyorMetrics.class).to(NoOpConveyorMetrics.class).asEagerSingleton();
    bind(JobSpecTransformer.class).to(JacksonJobSpecTransformer.class).asEagerSingleton();
    bind(TokenManager.class).to(JacksonTokenManager.class).asEagerSingleton();
  }

  @Provides
  public DynamoDbAsyncClient providesDdbAsyncClient() {
    return DynamoDbAsyncClient.builder()
        .endpointOverride(dynamoDBConfiguration.getEndpoint())
        .build();
  }

  @Provides
  @Singleton
  public ManagedInboundJobQueueConsumer providesManagedInboundJobQueueConsumer(
      @Named(Constants.SQS_CONSUMER) AsyncSqsClient asyncClient, ScheduleManager scheduleManager) {
    return new ManagedInboundJobQueueConsumer(
        asyncClient,
        MoreExecutors.getExitingExecutorService(
            (ThreadPoolExecutor)
                Executors.newFixedThreadPool(
                    1, new ThreadFactoryBuilder().setNameFormat("inbound-consumer-%d").build())),
        scheduleManager,
        sqsConfiguration.getInboundQueueName());
  }

  private AsyncSqsClientConfiguration buildBaseSqsConfig() {
    AsyncSqsClientConfiguration asyncSqsConfig =
        new AsyncSqsClientConfiguration(AsyncSqsClientConfiguration.Type.PRODUCTION);

    asyncSqsConfig.setRegion(sqsConfiguration.getRegion());
    asyncSqsConfig.setEndpoint(sqsConfiguration.getEndpoint());
    return asyncSqsConfig;
  }

  @Provides
  @Named(Constants.SQS_CONSUMER)
  AsyncSqsClientConfiguration provideConsumerSqsConfig() {
    return buildBaseSqsConfig();
  }

  @Provides
  @Named(Constants.SQS_PRODUCER)
  AsyncSqsClientConfiguration provideProducerSqsConfig() {
    AsyncSqsClientConfiguration asyncSqsConfig = buildBaseSqsConfig();

    RetryPolicyConfiguration retryConfig = new RetryPolicyConfiguration();
    retryConfig.setPolicy(RetryPolicyConfiguration.Policy.CUSTOM);
    retryConfig.setCondition(RetryPolicyConfiguration.Condition.DEFAULT);
    retryConfig.setBackOff(RetryPolicyConfiguration.BackOff.EQUAL_JITTER);
    retryConfig.setMaximumBackoffTime(sqsConfiguration.getTimeouts().getMaxBackoffTime());
    retryConfig.setBaseDelay(sqsConfiguration.getTimeouts().getBaseDelay());

    OverrideConfiguration overrideConfig = new OverrideConfiguration();
    overrideConfig.setApiCallAttemptTimeout(
        sqsConfiguration.getTimeouts().getApiCallAttemptTimeout());
    overrideConfig.setApiCallTimeout(sqsConfiguration.getTimeouts().getApiCallTimeout());
    overrideConfig.setRetryPolicyConfiguration(retryConfig);
    asyncSqsConfig.setOverrideConfiguration(overrideConfig);
    return asyncSqsConfig;
  }

  @Provides
  @Singleton
  @Named(Constants.SQS_CONSUMER)
  public AsyncSqsClient provideAsyncConsumer(
      AsyncSqsClientFactory asyncSqsClientFactory,
      @Named(Constants.SQS_CONSUMER) AsyncSqsClientConfiguration asyncSqsClientConfiguration) {
    return asyncSqsClientFactory.create(asyncSqsClientConfiguration);
  }

  @Provides
  @Singleton
  @Named(Constants.SQS_PRODUCER)
  public AsyncSqsClient provideAsyncProducer(
      AsyncSqsClientFactory asyncSqsClientFactory,
      @Named(Constants.SQS_PRODUCER) AsyncSqsClientConfiguration asyncSqsClientConfiguration) {
    return asyncSqsClientFactory.create(asyncSqsClientConfiguration);
  }

  @Provides
  @Singleton
  public HeartbeatManager provideHeartBeat(EventBus eventBus, Clock clock) {
    HeartbeatManager heartbeatManager =
        new HeartbeatManager(
            new ConcurrentHashMap<>(),
            MoreExecutors.getExitingScheduledExecutorService(
                (ScheduledThreadPoolExecutor)
                    Executors.newScheduledThreadPool(
                        1,
                        new ThreadFactoryBuilder().setNameFormat("heartbeat-manager-%d").build())),
            eventBus,
            clock,
            heartbeatConfiguration);
    eventBus.register(heartbeatManager);
    return heartbeatManager;
  }

  @Provides
  @Singleton
  public WorkerManager provideWorkerManager(
      PartitionWorkerFactory factory, EventBus eventBus, Metrics metrics, PartitionPolicy policy) {

    List<PartitionWorker> partitionWorkers = new ArrayList<>();
    for (int i = 0; i < workersConfiguration.getNumberOfWorkers(); i++) {
      partitionWorkers.add(factory.get(i));
    }

    WorkerManager manager =
        new WorkerManager(
            partitionWorkers,
            factory,
            MoreExecutors.getExitingScheduledExecutorService(
                (ScheduledThreadPoolExecutor)
                    Executors.newScheduledThreadPool(
                        1, new ThreadFactoryBuilder().setNameFormat("worker-manager-%d").build())),
            new ConcurrentHashMap<>(),
            policy,
            metrics);

    eventBus.register(manager);
    return manager;
  }
}