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

package io.airbyte.scheduler.client;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.when;

import com.fasterxml.jackson.databind.JsonNode;
import com.google.common.collect.ImmutableMap;
import io.airbyte.commons.json.Jsons;
import io.airbyte.config.ConnectorJobOutput;
import io.airbyte.config.DestinationConnection;
import io.airbyte.config.JobCheckConnectionConfig;
import io.airbyte.config.JobConfig.ConfigType;
import io.airbyte.config.JobDiscoverCatalogConfig;
import io.airbyte.config.JobGetSpecConfig;
import io.airbyte.config.SourceConnection;
import io.airbyte.config.StandardCheckConnectionOutput;
import io.airbyte.protocol.models.AirbyteCatalog;
import io.airbyte.protocol.models.ConnectorSpecification;
import io.airbyte.scheduler.persistence.job_error_reporter.ConnectorJobReportingContext;
import io.airbyte.scheduler.persistence.job_error_reporter.JobErrorReporter;
import io.airbyte.scheduler.persistence.job_factory.OAuthConfigSupplier;
import io.airbyte.scheduler.persistence.job_tracker.JobTracker;
import io.airbyte.scheduler.persistence.job_tracker.JobTracker.JobState;
import io.airbyte.workers.temporal.JobMetadata;
import io.airbyte.workers.temporal.TemporalClient;
import io.airbyte.workers.temporal.TemporalResponse;
import java.io.IOException;
import java.nio.file.Path;
import java.util.UUID;
import java.util.function.Function;
import java.util.function.Supplier;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Nested;
import org.junit.jupiter.api.Test;

// the goal here is to test the "execute" part of this class and all of the various exceptional
// cases. then separately test submission of each job type without having to re-test all of the
// execution exception cases again.
class DefaultSynchronousSchedulerClientTest {

  private static final Path LOG_PATH = Path.of("/tmp");
  private static final String DOCKER_IMAGE = "foo/bar";
  private static final UUID WORKSPACE_ID = UUID.randomUUID();
  private static final UUID UUID1 = UUID.randomUUID();
  private static final UUID UUID2 = UUID.randomUUID();
  private static final String UNCHECKED = "unchecked";
  private static final JsonNode CONFIGURATION = Jsons.jsonNode(ImmutableMap.builder()
      .put("username", "airbyte")
      .put("password", "abc")
      .build());
  private static final SourceConnection SOURCE_CONNECTION = new SourceConnection()
      .withSourceId(UUID1)
      .withSourceDefinitionId(UUID2)
      .withConfiguration(CONFIGURATION);
  private static final DestinationConnection DESTINATION_CONNECTION = new DestinationConnection()
      .withDestinationId(UUID1)
      .withDestinationDefinitionId(UUID2)
      .withConfiguration(CONFIGURATION);
  private static final String SOURCE_DOCKER_IMAGE = "source-airbyte:1.2.3";

  private TemporalClient temporalClient;
  private JobTracker jobTracker;
  private JobErrorReporter jobErrorReporter;
  private OAuthConfigSupplier oAuthConfigSupplier;
  private DefaultSynchronousSchedulerClient schedulerClient;

  @BeforeEach
  void setup() throws IOException {
    temporalClient = mock(TemporalClient.class);
    jobTracker = mock(JobTracker.class);
    jobErrorReporter = mock(JobErrorReporter.class);
    oAuthConfigSupplier = mock(OAuthConfigSupplier.class);
    schedulerClient = new DefaultSynchronousSchedulerClient(temporalClient, jobTracker, jobErrorReporter, oAuthConfigSupplier);

    when(oAuthConfigSupplier.injectSourceOAuthParameters(any(), any(), eq(CONFIGURATION))).thenReturn(CONFIGURATION);
    when(oAuthConfigSupplier.injectDestinationOAuthParameters(any(), any(), eq(CONFIGURATION))).thenReturn(CONFIGURATION);
  }

  private static JobMetadata createMetadata(final boolean succeeded) {
    return new JobMetadata(
        succeeded,
        LOG_PATH);
  }

  @Nested
  @DisplayName("Test execute method.")
  class ExecuteSynchronousJob {

    @SuppressWarnings(UNCHECKED)
    @Test
    void testExecuteJobSuccess() {
      final UUID sourceDefinitionId = UUID.randomUUID();
      final Supplier<TemporalResponse<String>> function = mock(Supplier.class);
      final Function<String, String> mapperFunction = output -> output;
      when(function.get()).thenReturn(new TemporalResponse<>("hello", createMetadata(true)));

      final ConnectorJobReportingContext jobContext = new ConnectorJobReportingContext(UUID.randomUUID(), SOURCE_DOCKER_IMAGE);
      final SynchronousResponse<String> response = schedulerClient
          .execute(ConfigType.DISCOVER_SCHEMA, jobContext, sourceDefinitionId, function, mapperFunction, WORKSPACE_ID);

      assertNotNull(response);
      assertEquals("hello", response.getOutput());
      assertEquals(ConfigType.DISCOVER_SCHEMA, response.getMetadata().getConfigType());
      assertTrue(response.getMetadata().getConfigId().isPresent());
      assertEquals(sourceDefinitionId, response.getMetadata().getConfigId().get());
      assertTrue(response.getMetadata().isSucceeded());
      assertEquals(LOG_PATH, response.getMetadata().getLogPath());

      verify(jobTracker).trackDiscover(any(UUID.class), eq(sourceDefinitionId), eq(WORKSPACE_ID), eq(JobState.STARTED));
      verify(jobTracker).trackDiscover(any(UUID.class), eq(sourceDefinitionId), eq(WORKSPACE_ID), eq(JobState.SUCCEEDED));
      verifyNoInteractions(jobErrorReporter);
    }

    @SuppressWarnings(UNCHECKED)
    @Test
    void testExecuteMappedOutput() {
      final UUID sourceDefinitionId = UUID.randomUUID();
      final Supplier<TemporalResponse<Integer>> function = mock(Supplier.class);
      final Function<Integer, String> mapperFunction = Object::toString;
      when(function.get()).thenReturn(new TemporalResponse<>(42, createMetadata(true)));

      final ConnectorJobReportingContext jobContext = new ConnectorJobReportingContext(UUID.randomUUID(), SOURCE_DOCKER_IMAGE);
      final SynchronousResponse<String> response = schedulerClient
          .execute(ConfigType.DISCOVER_SCHEMA, jobContext, sourceDefinitionId, function, mapperFunction, WORKSPACE_ID);

      assertNotNull(response);
      assertEquals("42", response.getOutput());
      assertEquals(ConfigType.DISCOVER_SCHEMA, response.getMetadata().getConfigType());
      assertTrue(response.getMetadata().getConfigId().isPresent());
      assertEquals(sourceDefinitionId, response.getMetadata().getConfigId().get());
      assertTrue(response.getMetadata().isSucceeded());
      assertEquals(LOG_PATH, response.getMetadata().getLogPath());
    }

    @SuppressWarnings(UNCHECKED)
    @Test
    void testExecuteJobFailure() {
      final UUID sourceDefinitionId = UUID.randomUUID();
      final Supplier<TemporalResponse<String>> function = mock(Supplier.class);
      final Function<String, String> mapperFunction = output -> output;
      when(function.get()).thenReturn(new TemporalResponse<>(null, createMetadata(false)));

      final ConnectorJobReportingContext jobContext = new ConnectorJobReportingContext(UUID.randomUUID(), SOURCE_DOCKER_IMAGE);
      final SynchronousResponse<String> response = schedulerClient
          .execute(ConfigType.DISCOVER_SCHEMA, jobContext, sourceDefinitionId, function, mapperFunction, WORKSPACE_ID);

      assertNotNull(response);
      assertNull(response.getOutput());
      assertEquals(ConfigType.DISCOVER_SCHEMA, response.getMetadata().getConfigType());
      assertTrue(response.getMetadata().getConfigId().isPresent());
      assertEquals(sourceDefinitionId, response.getMetadata().getConfigId().get());
      assertFalse(response.getMetadata().isSucceeded());
      assertEquals(LOG_PATH, response.getMetadata().getLogPath());

      verify(jobTracker).trackDiscover(any(UUID.class), eq(sourceDefinitionId), eq(WORKSPACE_ID), eq(JobState.STARTED));
      verify(jobTracker).trackDiscover(any(UUID.class), eq(sourceDefinitionId), eq(WORKSPACE_ID), eq(JobState.FAILED));
      verifyNoInteractions(jobErrorReporter);
    }

    @SuppressWarnings(UNCHECKED)
    @Test
    void testExecuteRuntimeException() {
      final UUID sourceDefinitionId = UUID.randomUUID();
      final Supplier<TemporalResponse<String>> function = mock(Supplier.class);
      final Function<String, String> mapperFunction = output -> output;
      when(function.get()).thenThrow(new RuntimeException());

      final ConnectorJobReportingContext jobContext = new ConnectorJobReportingContext(UUID.randomUUID(), SOURCE_DOCKER_IMAGE);
      assertThrows(
          RuntimeException.class,
          () -> schedulerClient.execute(ConfigType.DISCOVER_SCHEMA, jobContext, sourceDefinitionId, function,
              mapperFunction, WORKSPACE_ID));

      verify(jobTracker).trackDiscover(any(UUID.class), eq(sourceDefinitionId), eq(WORKSPACE_ID), eq(JobState.STARTED));
      verify(jobTracker).trackDiscover(any(UUID.class), eq(sourceDefinitionId), eq(WORKSPACE_ID), eq(JobState.FAILED));
      verifyNoInteractions(jobErrorReporter);
    }

  }

  @Nested
  @DisplayName("Test job creation for each configuration type.")
  class TestJobCreation {

    @Test
    void testCreateSourceCheckConnectionJob() throws IOException {
      final JobCheckConnectionConfig jobCheckConnectionConfig = new JobCheckConnectionConfig()
          .withConnectionConfiguration(SOURCE_CONNECTION.getConfiguration())
          .withDockerImage(DOCKER_IMAGE);

      final StandardCheckConnectionOutput mockOutput = mock(StandardCheckConnectionOutput.class);
      final ConnectorJobOutput jobOutput = new ConnectorJobOutput().withCheckConnection(mockOutput);
      when(temporalClient.submitCheckConnection(any(UUID.class), eq(0), eq(jobCheckConnectionConfig)))
          .thenReturn(new TemporalResponse<>(jobOutput, createMetadata(true)));
      final SynchronousResponse<StandardCheckConnectionOutput> response =
          schedulerClient.createSourceCheckConnectionJob(SOURCE_CONNECTION, DOCKER_IMAGE);
      assertEquals(mockOutput, response.getOutput());
    }

    @Test
    void testCreateDestinationCheckConnectionJob() throws IOException {
      final JobCheckConnectionConfig jobCheckConnectionConfig = new JobCheckConnectionConfig()
          .withConnectionConfiguration(DESTINATION_CONNECTION.getConfiguration())
          .withDockerImage(DOCKER_IMAGE);

      final StandardCheckConnectionOutput mockOutput = mock(StandardCheckConnectionOutput.class);
      final ConnectorJobOutput jobOutput = new ConnectorJobOutput().withCheckConnection(mockOutput);
      when(temporalClient.submitCheckConnection(any(UUID.class), eq(0), eq(jobCheckConnectionConfig)))
          .thenReturn(new TemporalResponse<>(jobOutput, createMetadata(true)));
      final SynchronousResponse<StandardCheckConnectionOutput> response =
          schedulerClient.createDestinationCheckConnectionJob(DESTINATION_CONNECTION, DOCKER_IMAGE);
      assertEquals(mockOutput, response.getOutput());
    }

    @Test
    void testCreateDiscoverSchemaJob() throws IOException {
      final JobDiscoverCatalogConfig jobDiscoverCatalogConfig = new JobDiscoverCatalogConfig()
          .withConnectionConfiguration(SOURCE_CONNECTION.getConfiguration())
          .withDockerImage(DOCKER_IMAGE);

      final AirbyteCatalog mockOutput = mock(AirbyteCatalog.class);
      final ConnectorJobOutput jobOutput = new ConnectorJobOutput().withDiscoverCatalog(mockOutput);
      when(temporalClient.submitDiscoverSchema(any(UUID.class), eq(0), eq(jobDiscoverCatalogConfig)))
          .thenReturn(new TemporalResponse<>(jobOutput, createMetadata(true)));
      final SynchronousResponse<AirbyteCatalog> response = schedulerClient.createDiscoverSchemaJob(SOURCE_CONNECTION, DOCKER_IMAGE);
      assertEquals(mockOutput, response.getOutput());
    }

    @Test
    void testCreateGetSpecJob() throws IOException {
      final JobGetSpecConfig jobSpecConfig = new JobGetSpecConfig().withDockerImage(DOCKER_IMAGE);

      final ConnectorSpecification mockOutput = mock(ConnectorSpecification.class);
      final ConnectorJobOutput jobOutput = new ConnectorJobOutput().withSpec(mockOutput);
      when(temporalClient.submitGetSpec(any(UUID.class), eq(0), eq(jobSpecConfig)))
          .thenReturn(new TemporalResponse<>(jobOutput, createMetadata(true)));
      final SynchronousResponse<ConnectorSpecification> response = schedulerClient.createGetSpecJob(DOCKER_IMAGE);
      assertEquals(mockOutput, response.getOutput());
    }

  }

}
