/*
 * 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.hudi.keygen;

import org.apache.hudi.common.table.HoodieTableConfig;
import org.apache.hudi.common.util.Option;
import org.apache.hudi.exception.HoodieKeyException;
import org.apache.hudi.keygen.constant.KeyGeneratorType;

import org.apache.avro.Schema;
import org.apache.avro.generic.GenericData;
import org.apache.avro.generic.GenericRecord;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import static org.apache.hudi.common.table.HoodieTableConfig.KEY_GENERATOR_CLASS_NAME;
import static org.apache.hudi.common.table.HoodieTableConfig.RECORDKEY_FIELDS;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

public class TestKeyGenUtils {

  @Test
  public void testInferKeyGeneratorType() {
    assertEquals(
        KeyGeneratorType.SIMPLE,
        KeyGenUtils.inferKeyGeneratorType(Option.of("col1"), "partition1"));
    assertEquals(
        KeyGeneratorType.COMPLEX,
        KeyGenUtils.inferKeyGeneratorType(Option.of("col1"), "partition1,partition2"));
    assertEquals(
        KeyGeneratorType.COMPLEX,
        KeyGenUtils.inferKeyGeneratorType(Option.of("col1,col2"), "partition1"));
    assertEquals(
        KeyGeneratorType.COMPLEX,
        KeyGenUtils.inferKeyGeneratorType(Option.of("col1,col2"), "partition1,partition2"));
    assertEquals(
        KeyGeneratorType.NON_PARTITION,
        KeyGenUtils.inferKeyGeneratorType(Option.of("col1,col2"), ""));
    assertEquals(
        KeyGeneratorType.NON_PARTITION,
        KeyGenUtils.inferKeyGeneratorType(Option.of("col1,col2"), null));
  }

  @Test
  public void testExtractRecordKeys() {
    // test complex key form: field1:val1,field2:val2,...
    String[] s1 = KeyGenUtils.extractRecordKeys("id:1");
    Assertions.assertArrayEquals(new String[] {"1"}, s1);

    String[] s2 = KeyGenUtils.extractRecordKeys("id:1,id:2");
    Assertions.assertArrayEquals(new String[] {"1", "2"}, s2);

    String[] s3 = KeyGenUtils.extractRecordKeys("id:1,id2:__null__,id3:__empty__");
    Assertions.assertArrayEquals(new String[] {"1", null, ""}, s3);

    String[] s4 = KeyGenUtils.extractRecordKeys("id:ab:cd,id2:ef");
    Assertions.assertArrayEquals(new String[] {"ab:cd", "ef"}, s4);

    // test simple key form: val1
    String[] s5 = KeyGenUtils.extractRecordKeys("1");
    Assertions.assertArrayEquals(new String[] {"1"}, s5);

    String[] s6 = KeyGenUtils.extractRecordKeys("id:1,id2:2,2");
    Assertions.assertArrayEquals(new String[]{"1", "2", "2"}, s6);
  }

  @Test
  public void testExtractRecordKeysWithFields() {
    List<String> fields = new ArrayList<>(1);
    fields.add("id2");

    String[] s1 = KeyGenUtils.extractRecordKeysByFields("id1:1,id2:2,id3:3", fields);
    Assertions.assertArrayEquals(new String[] {"2"}, s1);

    String[] s2 = KeyGenUtils.extractRecordKeysByFields("id1:1,id2:2,2,id3:3", fields);
    Assertions.assertArrayEquals(new String[] {"2", "2"}, s2);
  }

  @Test
  void testGetRecordKey() {
    Schema nullableStringSchema = Schema.createUnion(Schema.create(Schema.Type.NULL), Schema.create(Schema.Type.STRING));
    Schema schema = Schema.createRecord("TestRecord", "doc", "test", false,
        Arrays.asList(
            new Schema.Field("key1", nullableStringSchema, "", null),
            new Schema.Field("key2", nullableStringSchema, "", null),
            new Schema.Field("key3", nullableStringSchema, "", null),
            new Schema.Field("key4", nullableStringSchema, "", null)
        ));
    GenericRecord avroRecord = new GenericData.Record(schema);
    avroRecord.put("key1", "value1");
    avroRecord.put("key2", "value2");
    avroRecord.put("key3", null);
    avroRecord.put("key4", "");

    assertEquals("key1:value1",
        KeyGenUtils.getRecordKey(avroRecord, Arrays.asList("key1"), true));
    assertThrows(HoodieKeyException.class,
        () -> KeyGenUtils.getRecordKey(avroRecord, Arrays.asList("key3"), true),
        "recordKey values: \"key3:__null__\" for fields: [key3] cannot be entirely null or empty.");
    assertThrows(HoodieKeyException.class,
        () -> KeyGenUtils.getRecordKey(avroRecord, Arrays.asList("key4"), true),
        "recordKey values: \"key4:__empty__\" for fields: [key4] cannot be entirely null or empty.");
    assertEquals("key1:value1,key2:value2",
        KeyGenUtils.getRecordKey(avroRecord, Arrays.asList("key1", "key2"), true));
    assertEquals("key1:value1,key3:__null__",
        KeyGenUtils.getRecordKey(avroRecord, Arrays.asList("key1", "key3"), true));
    assertEquals("key1:value1,key4:__empty__",
        KeyGenUtils.getRecordKey(avroRecord, Arrays.asList("key1", "key4"), true));

    assertEquals("value1",
        KeyGenUtils.getRecordKey(avroRecord, "key1", true));
    assertThrows(HoodieKeyException.class,
        () -> KeyGenUtils.getRecordKey(avroRecord, "key3", true),
        "recordKey value: \"null\" for field: \"key3\" cannot be null or empty.");
    assertThrows(HoodieKeyException.class,
        () -> KeyGenUtils.getRecordKey(avroRecord, "key4", true),
        "recordKey value: \"\" for field: \"key4\" cannot be null or empty.");
  }

  @Test
  void testIsComplexKeyGeneratorWithSingleRecordKeyField() {
    HoodieTableConfig tableConfig = new HoodieTableConfig();
    tableConfig.setValue(KEY_GENERATOR_CLASS_NAME, "org.apache.hudi.keygen.ComplexKeyGenerator");
    tableConfig.setValue(RECORDKEY_FIELDS, "id");
    assertTrue(KeyGenUtils.isComplexKeyGeneratorWithSingleRecordKeyField(tableConfig));

    tableConfig = new HoodieTableConfig();
    tableConfig.setValue(RECORDKEY_FIELDS, "userId");
    tableConfig.setValue(KEY_GENERATOR_CLASS_NAME, "org.apache.hudi.keygen.ComplexAvroKeyGenerator");
    assertTrue(KeyGenUtils.isComplexKeyGeneratorWithSingleRecordKeyField(tableConfig));
  }

  @Test
  void testIsComplexKeyGeneratorWithSingleRecordKeyFieldOnMultipleFields() {
    HoodieTableConfig tableConfig = new HoodieTableConfig();
    tableConfig.setValue(KEY_GENERATOR_CLASS_NAME, "org.apache.hudi.keygen.ComplexKeyGenerator");
    tableConfig.setValue(RECORDKEY_FIELDS, "id,userId");
    assertFalse(KeyGenUtils.isComplexKeyGeneratorWithSingleRecordKeyField(tableConfig));

    tableConfig = new HoodieTableConfig();
    tableConfig.setValue(KEY_GENERATOR_CLASS_NAME, "org.apache.hudi.keygen.ComplexAvroKeyGenerator");
    tableConfig.setValue(RECORDKEY_FIELDS, "id,userId,name");
    assertFalse(KeyGenUtils.isComplexKeyGeneratorWithSingleRecordKeyField(tableConfig));
  }

  @Test
  void testIsComplexKeyGeneratorWithSingleRecordKeyFieldOnNonComplexGenerator() {
    HoodieTableConfig tableConfig = new HoodieTableConfig();
    tableConfig.setValue(KEY_GENERATOR_CLASS_NAME, "org.apache.hudi.keygen.SimpleKeyGenerator");
    tableConfig.setValue(RECORDKEY_FIELDS, "id");
    assertFalse(KeyGenUtils.isComplexKeyGeneratorWithSingleRecordKeyField(tableConfig));

    tableConfig = new HoodieTableConfig();
    tableConfig.setValue(KEY_GENERATOR_CLASS_NAME, "org.apache.hudi.keygen.SimpleAvroKeyGenerator");
    tableConfig.setValue(RECORDKEY_FIELDS, "userId");
    assertFalse(KeyGenUtils.isComplexKeyGeneratorWithSingleRecordKeyField(tableConfig));

    tableConfig = new HoodieTableConfig();
    tableConfig.setValue(KEY_GENERATOR_CLASS_NAME, "org.apache.hudi.keygen.TimestampBasedKeyGenerator");
    tableConfig.setValue(RECORDKEY_FIELDS, "id");
    assertFalse(KeyGenUtils.isComplexKeyGeneratorWithSingleRecordKeyField(tableConfig));

    tableConfig = new HoodieTableConfig();
    tableConfig.setValue(KEY_GENERATOR_CLASS_NAME, "org.apache.hudi.keygen.CustomKeyGenerator");
    tableConfig.setValue(RECORDKEY_FIELDS, "id");
    assertFalse(KeyGenUtils.isComplexKeyGeneratorWithSingleRecordKeyField(tableConfig));
  }

  @Test
  void testIsComplexKeyGeneratorWithSingleRecordKeyFieldOnNoRecordKeyFields() {
    HoodieTableConfig tableConfig = new HoodieTableConfig();
    tableConfig.setValue(KEY_GENERATOR_CLASS_NAME, "org.apache.hudi.keygen.ComplexKeyGenerator");
    assertFalse(KeyGenUtils.isComplexKeyGeneratorWithSingleRecordKeyField(tableConfig));
  }

  @Test
  void testIsComplexKeyGeneratorWithSingleRecordKeyFieldEmptyRecordKeyFields() {
    HoodieTableConfig tableConfig = new HoodieTableConfig();
    tableConfig.setValue(KEY_GENERATOR_CLASS_NAME, "org.apache.hudi.keygen.ComplexKeyGenerator");
    tableConfig.setValue(RECORDKEY_FIELDS, "");
    assertFalse(KeyGenUtils.isComplexKeyGeneratorWithSingleRecordKeyField(tableConfig));
  }
}
