/*
 * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
 */

package io.airbyte.integrations.destination.aws_datalake;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import software.amazon.awssdk.auth.credentials.AwsCredentials;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.athena.AthenaClient;
import software.amazon.awssdk.services.athena.model.AthenaException;
import software.amazon.awssdk.services.athena.model.GetQueryExecutionRequest;
import software.amazon.awssdk.services.athena.model.GetQueryExecutionResponse;
import software.amazon.awssdk.services.athena.model.GetQueryResultsRequest;
import software.amazon.awssdk.services.athena.model.QueryExecutionContext;
import software.amazon.awssdk.services.athena.model.QueryExecutionState;
import software.amazon.awssdk.services.athena.model.ResultConfiguration;
import software.amazon.awssdk.services.athena.model.StartQueryExecutionRequest;
import software.amazon.awssdk.services.athena.model.StartQueryExecutionResponse;
import software.amazon.awssdk.services.athena.paginators.GetQueryResultsIterable;

public class AthenaHelper {

  private AthenaClient athenaClient;
  private String outputBucket;
  private String workGroup;
  private static final Logger LOGGER = LoggerFactory.getLogger(AthenaHelper.class);

  public AthenaHelper(AwsCredentials credentials, Region region, String outputBucket, String workGroup) {
    LOGGER.debug(String.format("region = %s, outputBucket = %s, workGroup = %s", region, outputBucket, workGroup));
    var credProvider = StaticCredentialsProvider.create(credentials);
    this.athenaClient = AthenaClient.builder().region(region).credentialsProvider(credProvider).build();
    this.outputBucket = outputBucket;
    this.workGroup = workGroup;
  }

  public String submitAthenaQuery(String database, String query) {
    try {

      // The QueryExecutionContext allows us to set the database
      QueryExecutionContext queryExecutionContext = QueryExecutionContext.builder()
          .database(database).build();

      // The result configuration specifies where the results of the query should go
      ResultConfiguration resultConfiguration = ResultConfiguration.builder()
          .outputLocation(outputBucket)
          .build();

      StartQueryExecutionRequest startQueryExecutionRequest = StartQueryExecutionRequest.builder()
          .queryString(query)
          .queryExecutionContext(queryExecutionContext)
          .resultConfiguration(resultConfiguration)
          .workGroup(workGroup)
          .build();

      StartQueryExecutionResponse startQueryExecutionResponse = athenaClient.startQueryExecution(startQueryExecutionRequest);
      return startQueryExecutionResponse.queryExecutionId();
    } catch (AthenaException e) {
      e.printStackTrace();
      System.exit(1);
    }
    return "";
  }

  public void waitForQueryToComplete(String queryExecutionId) throws InterruptedException {
    GetQueryExecutionRequest getQueryExecutionRequest = GetQueryExecutionRequest.builder()
        .queryExecutionId(queryExecutionId).build();

    GetQueryExecutionResponse getQueryExecutionResponse;
    boolean isQueryStillRunning = true;
    while (isQueryStillRunning) {
      getQueryExecutionResponse = athenaClient.getQueryExecution(getQueryExecutionRequest);
      String queryState = getQueryExecutionResponse.queryExecution().status().state().toString();
      if (queryState.equals(QueryExecutionState.FAILED.toString())) {
        throw new RuntimeException("The Amazon Athena query failed to run with error message: " + getQueryExecutionResponse
            .queryExecution().status().stateChangeReason());
      } else if (queryState.equals(QueryExecutionState.CANCELLED.toString())) {
        throw new RuntimeException("The Amazon Athena query was cancelled.");
      } else if (queryState.equals(QueryExecutionState.SUCCEEDED.toString())) {
        isQueryStillRunning = false;
      } else {
        // Sleep an amount of time before retrying again
        Thread.sleep(1000);
      }
    }
  }

  public GetQueryResultsIterable getResults(String queryExecutionId) {

    try {

      // Max Results can be set but if its not set,
      // it will choose the maximum page size
      GetQueryResultsRequest getQueryResultsRequest = GetQueryResultsRequest.builder()
          .queryExecutionId(queryExecutionId)
          .build();

      GetQueryResultsIterable getQueryResultsResults = athenaClient.getQueryResultsPaginator(getQueryResultsRequest);
      return getQueryResultsResults;

    } catch (AthenaException e) {
      e.printStackTrace();
      System.exit(1);
    }
    return null;
  }

  public GetQueryResultsIterable runQuery(String database, String query) throws InterruptedException {
    int retryCount = 0;

    while (retryCount < 10) {
      var execId = submitAthenaQuery(database, query);
      try {
        waitForQueryToComplete(execId);
      } catch (RuntimeException e) {
        e.printStackTrace();
        LOGGER.info("Athena query failed once. Retrying.");
        retryCount++;
        continue;
      }
      return getResults(execId);
    }
    LOGGER.info("Athena query failed and we are out of retries.");
    return null;
  }

}
