/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.beam.it.gcp.bigquery;

import static com.google.common.truth.Truth.assertThat;
import static org.junit.Assert.assertThrows;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import com.google.cloud.bigquery.BigQuery;
import com.google.cloud.bigquery.BigQueryException;
import com.google.cloud.bigquery.Dataset;
import com.google.cloud.bigquery.DatasetId;
import com.google.cloud.bigquery.DatasetInfo;
import com.google.cloud.bigquery.Field;
import com.google.cloud.bigquery.InsertAllRequest.RowToInsert;
import com.google.cloud.bigquery.QueryJobConfiguration;
import com.google.cloud.bigquery.Schema;
import com.google.cloud.bigquery.StandardSQLTypeName;
import com.google.cloud.bigquery.StandardTableDefinition;
import com.google.cloud.bigquery.Table;
import com.google.cloud.bigquery.TableId;
import com.google.cloud.bigquery.TableInfo;
import com.google.cloud.bigquery.TimePartitioning;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.Answers;
import org.mockito.ArgumentCaptor;
import org.mockito.Captor;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnit;
import org.mockito.junit.MockitoRule;

/** Unit tests for {@link BigQueryResourceManager}. */
@RunWith(JUnit4.class)
public class BigQueryResourceManagerTest {

  @Rule public final MockitoRule mockito = MockitoJUnit.rule();

  @Mock(answer = Answers.RETURNS_DEEP_STUBS)
  private BigQuery bigQuery;

  private Schema schema;
  private RowToInsert rowToInsert;
  private TimePartitioning timePartition;
  private static final String PARTITION_FIELD = "time"; //  name of column to use for partitioning
  private static final String TABLE_NAME = "table-name";
  private static final String DATASET_ID = "dataset-id";
  private static final String TEST_ID = "test-id";
  private static final String PROJECT_ID = "test-project";
  private static final String REGION = "us-central1";

  private BigQueryResourceManager testManager;

  @Captor private ArgumentCaptor<TableInfo> tableCaptor;

  @Before
  public void setUp() {
    schema = Schema.of(Field.of("name", StandardSQLTypeName.STRING));
    rowToInsert = RowToInsert.of("1", ImmutableMap.of("name", "Jake"));
    timePartition =
        TimePartitioning.newBuilder(TimePartitioning.Type.HOUR).setField(PARTITION_FIELD).build();
    testManager = new BigQueryResourceManager(TEST_ID, PROJECT_ID, bigQuery);
  }

  @Test
  public void testGetProjectIdReturnsCorrectValue() {
    assertThat(new BigQueryResourceManager(TEST_ID, PROJECT_ID, bigQuery).getProjectId())
        .isEqualTo(PROJECT_ID);
  }

  @Test
  public void testGetDatasetIdReturnsCorrectValue() {
    BigQueryResourceManager tm = BigQueryResourceManager.builder(TEST_ID, PROJECT_ID, null).build();

    assertThat(tm.getDatasetId()).matches(TEST_ID.replace('-', '_') + "_\\d{8}_\\d{6}_\\d{6}");
  }

  @Test
  public void testCreateDatasetShouldThrowErrorWhenDatasetCreateFails() {
    when(bigQuery.create(any(DatasetInfo.class))).thenThrow(RuntimeException.class);

    assertThrows(
        BigQueryResourceManagerException.class, () -> testManager.createDataset(DATASET_ID));
  }

  @Test
  public void testCreateDatasetShouldNotCreateDatasetWhenDatasetAlreadyExists() {
    testManager.createDataset(DATASET_ID);

    assertThrows(IllegalStateException.class, () -> testManager.createDataset(DATASET_ID));
  }

  @Test
  public void testCreateDatasetShouldWorkWhenBigQueryDoesNotThrowAnyError() {
    testManager.createDataset(DATASET_ID);
    verify(bigQuery).create(any(DatasetInfo.class));
  }

  @Test
  public void testCreateTableShouldThrowErrorWhenTableNameIsNotValid() {
    assertThrows(IllegalArgumentException.class, () -> testManager.createTable("", schema));
  }

  @Test
  public void testCreateTableShouldThrowErrorWhenSchemaIsNull() {
    assertThrows(IllegalArgumentException.class, () -> testManager.createTable(TABLE_NAME, null));
  }

  @Test
  public void testCreateTableShouldCreateDatasetWhenDatasetDoesNotExist() {
    when(bigQuery.create(any(DatasetInfo.class)).getDatasetId().getDataset())
        .thenReturn(DATASET_ID);
    when(bigQuery.getTable(any())).thenReturn(null);

    testManager.createTable(TABLE_NAME, schema);

    verify(bigQuery).create(any(DatasetInfo.class));
    verify(bigQuery).create(any(TableInfo.class));
  }

  @Test
  public void testCreateTableShouldThrowErrorWhenCreateFails() {
    testManager.createDataset(DATASET_ID);
    when(bigQuery.create(any(TableInfo.class))).thenThrow(BigQueryException.class);

    assertThrows(
        BigQueryResourceManagerException.class, () -> testManager.createTable(TABLE_NAME, schema));
  }

  @Test
  public void testCreateTableShouldThrowErrorWhenTableExists() {
    testManager.createDataset(DATASET_ID);

    when(bigQuery.getTable(any())).thenReturn(any());

    assertThrows(
        BigQueryResourceManagerException.class, () -> testManager.createTable(TABLE_NAME, schema));
  }

  @Test
  public void testCreateTableShouldWorkWhenBigQueryDoesNotThrowAnyError() {
    when(bigQuery.create(any(DatasetInfo.class)).getDatasetId().getDataset())
        .thenReturn(DATASET_ID);
    when(bigQuery.getTable(any())).thenReturn(null);

    testManager.createTable(TABLE_NAME, schema);

    verify(bigQuery).create(any(DatasetInfo.class));
  }

  @Test
  public void testCreateTimePartitionedTableShouldThrowErrorWhenTableNameIsNotValid() {
    assertThrows(
        IllegalArgumentException.class,
        () -> testManager.createTimePartitionedTable("", schema, timePartition));
  }

  @Test
  public void testCreateTimePartitionedTableShouldThrowErrorWhenSchemaIsNull() {
    assertThrows(
        IllegalArgumentException.class,
        () -> testManager.createTimePartitionedTable(TABLE_NAME, null, timePartition));
  }

  @Test
  public void testCreateTimePartitionedTableShouldThrowErrorWhenPartitionInfoIsNull() {
    assertThrows(
        IllegalArgumentException.class,
        () -> testManager.createTimePartitionedTable(TABLE_NAME, schema, null));
  }

  @Test
  public void testCreateTimePartitionedTableShouldCreateDatasetWhenDatasetDoesNotExist() {
    when(bigQuery.create(any(DatasetInfo.class)).getDatasetId().getDataset())
        .thenReturn(DATASET_ID);
    when(bigQuery.getTable(any())).thenReturn(null);

    testManager.createTimePartitionedTable(TABLE_NAME, schema, timePartition);

    verify(bigQuery).create(any(DatasetInfo.class));
  }

  @Test
  public void testCreateTimePartitionedTableShouldThrowErrorWhenCreateFails() {
    testManager.createDataset(DATASET_ID);
    when(bigQuery.create(any(TableInfo.class))).thenThrow(BigQueryException.class);

    assertThrows(
        BigQueryResourceManagerException.class,
        () -> testManager.createTimePartitionedTable(TABLE_NAME, schema, timePartition));
  }

  @Test
  public void testCreateTimePartitionedTableShouldThrowErrorWhenTableExists() {
    testManager.createDataset(DATASET_ID);

    when(bigQuery.getTable(any())).thenReturn(any());

    assertThrows(
        BigQueryResourceManagerException.class,
        () -> testManager.createTimePartitionedTable(TABLE_NAME, schema, timePartition));
  }

  @Test
  public void testCreateTimePartitionedTableShouldWorkWhenBigQueryDoesNotThrowAnyError() {
    when(bigQuery.create(any(DatasetInfo.class)).getDatasetId().getDataset())
        .thenReturn(DATASET_ID);
    when(bigQuery.getTable(any())).thenReturn(null);

    testManager.createTimePartitionedTable(TABLE_NAME, schema, timePartition);

    verify(bigQuery).create(any(TableInfo.class));
    verify(bigQuery).create(tableCaptor.capture());
    TableInfo capturedTableInfo = tableCaptor.getValue();
    StandardTableDefinition capturedTableDefinition = capturedTableInfo.getDefinition();
    TimePartitioning capturedTimePartitioning = capturedTableDefinition.getTimePartitioning();
    assertThat(capturedTimePartitioning).isEqualTo(timePartition);
  }

  @Test
  public void testWriteShouldThrowErrorWhenDatasetDoesNotExist() {
    assertThrows(IllegalStateException.class, () -> testManager.write(TABLE_NAME, rowToInsert));
    assertThrows(
        IllegalStateException.class,
        () -> testManager.write(TABLE_NAME, ImmutableList.of(rowToInsert)));
  }

  @Test
  public void testWriteShouldThrowErrorWhenTableDoesNotExist() {
    testManager.createDataset(DATASET_ID);

    when(bigQuery.create(any(DatasetInfo.class)).get(anyString())).thenReturn(null);

    assertThrows(IllegalStateException.class, () -> testManager.write(TABLE_NAME, rowToInsert));
    assertThrows(
        IllegalStateException.class,
        () -> testManager.write(TABLE_NAME, ImmutableList.of(rowToInsert)));
  }

  @Test
  public void testWriteShouldThrowErrorWhenInsertFails() {
    Dataset mockDataset = bigQuery.create(any(DatasetInfo.class));
    when(mockDataset.getDatasetId().getDataset()).thenReturn(DATASET_ID);

    testManager.createDataset(DATASET_ID);

    when(mockDataset.get(anyString()).insert(any())).thenThrow(BigQueryException.class);

    assertThrows(
        BigQueryResourceManagerException.class, () -> testManager.write(TABLE_NAME, rowToInsert));
    assertThrows(
        BigQueryResourceManagerException.class,
        () -> testManager.write(TABLE_NAME, ImmutableList.of(rowToInsert)));
  }

  @Test
  public void testWriteShouldWorkWhenBigQueryDoesNotThrowAnyError() {
    when(bigQuery.create(any(DatasetInfo.class)).getDatasetId().getDataset())
        .thenReturn(DATASET_ID);
    when(bigQuery.getTable(any())).thenReturn(null);

    testManager.createTable(TABLE_NAME, schema);

    testManager.write(TABLE_NAME, rowToInsert);
    testManager.write(TABLE_NAME, ImmutableList.of(rowToInsert));
  }

  @Test
  public void testReadTableShouldThrowErrorWhenDatasetDoesNotExist() {
    assertThrows(IllegalStateException.class, () -> testManager.readTable(TABLE_NAME));
  }

  @Test
  public void testReadTableShouldThrowErrorWhenTableDoesNotExist() {
    testManager.createDataset(REGION);
    when(bigQuery.create(any(DatasetInfo.class)).get(anyString())).thenReturn(null);
    assertThrows(IllegalStateException.class, () -> testManager.readTable(TABLE_NAME));
  }

  @Test
  public void testReadTableShouldThrowErrorWhenBigQueryFailsToReadRows()
      throws InterruptedException {
    when(bigQuery.create(any(DatasetInfo.class)).getDatasetId().getDataset())
        .thenReturn(DATASET_ID);
    when(bigQuery.getTable(any())).thenReturn(null);

    testManager.createTable(TABLE_NAME, schema);

    doThrow(BigQueryException.class).when(bigQuery).query(any(QueryJobConfiguration.class));

    assertThrows(BigQueryResourceManagerException.class, () -> testManager.readTable(TABLE_NAME));
  }

  @Test
  public void testReadTableShouldWorkWhenBigQueryDoesNotThrowAnyError()
      throws InterruptedException {
    when(bigQuery.create(any(DatasetInfo.class)).getDatasetId().getDataset())
        .thenReturn(DATASET_ID);
    when(bigQuery.getTable(any())).thenReturn(null);

    testManager.createTable(TABLE_NAME, schema);

    testManager.readTable(TABLE_NAME);

    verify(bigQuery).query(any(QueryJobConfiguration.class));
  }

  @Test
  public void testCleanupShouldThrowErrorWhenTableDeleteFails() {
    testManager.createDataset(DATASET_ID);

    Table mockTable = mock(Table.class);
    when(bigQuery.listTables(any(DatasetId.class)).iterateAll())
        .thenReturn(ImmutableList.of(mockTable));
    when(bigQuery.delete(any(TableId.class))).thenThrow(BigQueryException.class);

    assertThrows(BigQueryResourceManagerException.class, () -> testManager.cleanupAll());
  }

  @Test
  public void testCleanupShouldThrowErrorWhenDatasetDeleteFails() {
    testManager.createDataset(DATASET_ID);
    when(bigQuery.delete(any(DatasetId.class))).thenThrow(BigQueryException.class);

    assertThrows(BigQueryResourceManagerException.class, () -> testManager.cleanupAll());
  }

  @Test
  public void testCleanupShouldWorkWhenBigQueryDoesNotThrowAnyError() {
    when(bigQuery.create(any(DatasetInfo.class)).getDatasetId().getDataset())
        .thenReturn(DATASET_ID);

    testManager.createDataset(DATASET_ID);
    Table mockTable = mock(Table.class, Answers.RETURNS_DEEP_STUBS);
    when(bigQuery.listTables(any(DatasetId.class)).iterateAll())
        .thenReturn(ImmutableList.of(mockTable));
    when(mockTable.getTableId().getTable()).thenReturn(TABLE_NAME);

    testManager.cleanupAll();

    verify(bigQuery).delete(any(TableId.class));
    verify(bigQuery).delete(any(DatasetId.class));
  }
}
