/*
 * 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.celeborn.server.common.service.config;

import java.io.IOException;
import java.sql.SQLException;
import java.sql.Statement;
import java.text.DateFormat;
import java.text.ParseException;
import java.text.SimpleDateFormat;
import java.time.Instant;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.junit.After;
import org.junit.Assert;
import org.junit.Test;

import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.server.common.service.config.DynamicConfig.ConfigType;
import org.apache.celeborn.server.common.service.model.ClusterInfo;
import org.apache.celeborn.server.common.service.store.IServiceManager;
import org.apache.celeborn.server.common.service.store.db.DBSessionFactory;

public class ConfigServiceSuiteJ {
  private ConfigService configService;
  private static final ThreadLocal<DateFormat> DATE_FORMAT =
      ThreadLocal.withInitial(() -> new SimpleDateFormat("yyyy-MM-dd HH:mm:ss"));

  @Test
  public void testDbConfig() throws IOException, ParseException {
    CelebornConf celebornConf = new CelebornConf();
    celebornConf.set(
        CelebornConf.DYNAMIC_CONFIG_STORE_DB_HIKARI_JDBC_URL(),
        "jdbc:h2:mem:test;MODE=MYSQL;INIT=RUNSCRIPT FROM 'classpath:celeborn-0.6.0-h2.sql'\\;"
            + "RUNSCRIPT FROM 'classpath:celeborn-0.6.0-h2-ut-data.sql';DB_CLOSE_DELAY=-1;");
    celebornConf.set(
        CelebornConf.DYNAMIC_CONFIG_STORE_DB_HIKARI_DRIVER_CLASS_NAME(), "org.h2.Driver");
    celebornConf.set(CelebornConf.DYNAMIC_CONFIG_STORE_DB_HIKARI_MAXIMUM_POOL_SIZE(), "1");
    configService = new DbConfigServiceImpl(celebornConf);
    verifySystemConfig(configService);
    verifyTenantConfig(configService);
    verifyTenantUserConfig(configService);
    verifyTags(configService);

    configService.upsertSystemConfig(Collections.singletonMap("celeborn.test.int.only", "100"));
    SqlSessionFactory sqlSessionFactory = DBSessionFactory.get(celebornConf);
    try (SqlSession sqlSession = sqlSessionFactory.openSession(true)) {
      Statement statement = sqlSession.getConnection().createStatement();

      statement.execute("UPDATE celeborn_cluster_tags SET tag = 'tag3' WHERE tag='tag2'");
    } catch (SQLException e) {
      throw new RuntimeException(e);
    }

    configService.refreshCache();
    verifyConfigChanged(configService);
    verifyTagsChanged(configService);
    verifyServiceManager(
        ((DbConfigServiceImpl) configService).getServiceManager(),
        celebornConf,
        DATE_FORMAT.get().parse("2023-08-26 22:08:30").toInstant());

    configService.upsertSystemConfig(
        Collections.singletonMap("celeborn.test.system.upsert", "insert"));
    configService.refreshCache();
    SystemConfig systemConfig = configService.getSystemConfigFromCache();
    Assert.assertEquals(
        "insert",
        systemConfig.getValue(
            "celeborn.test.system.upsert", null, String.class, ConfigType.STRING));
    configService.deleteSystemConfigByKeys(Collections.singletonList("celeborn.test.int.only"));
    configService.refreshCache();
    systemConfig = configService.getSystemConfigFromCache();
    Assert.assertFalse(systemConfig.configs.containsKey("celeborn.test.int.only"));

    Map<String, String> tenantConfigs = new HashMap<>();
    tenantConfigs.put(CelebornConf.CLIENT_PUSH_BUFFER_INITIAL_SIZE().key(), "20480");
    tenantConfigs.put("celeborn.test.tenant.upsert", "insert");
    configService.upsertTenantConfig(ConfigLevel.TENANT, "tenant_id", null, tenantConfigs);
    configService.refreshCache();
    TenantConfig tenantConfig = configService.getRawTenantConfigFromCache("tenant_id");
    Assert.assertEquals(
        20480,
        tenantConfig
            .getValue(
                CelebornConf.CLIENT_PUSH_BUFFER_INITIAL_SIZE().key(),
                CelebornConf.CLIENT_PUSH_BUFFER_INITIAL_SIZE(),
                Long.TYPE,
                ConfigType.BYTES)
            .longValue());
    Assert.assertEquals(
        "insert",
        tenantConfig.getValue(
            "celeborn.test.tenant.upsert", null, String.class, ConfigType.STRING));
    configService.deleteTenantConfigByKeys(
        ConfigLevel.TENANT,
        "tenant_id",
        null,
        Arrays.asList(
            CelebornConf.CLIENT_PUSH_BUFFER_INITIAL_SIZE().key(), "celeborn.test.tenant.upsert"));
    configService.refreshCache();
    tenantConfig = configService.getRawTenantConfigFromCache("tenant_id");
    Assert.assertFalse(
        tenantConfig.configs.containsKey(CelebornConf.CLIENT_PUSH_BUFFER_INITIAL_SIZE().key()));
    Assert.assertFalse(tenantConfig.configs.containsKey("celeborn.test.tenant.upsert"));

    tenantConfigs = new HashMap<>();
    tenantConfigs.put(CelebornConf.CLIENT_PUSH_BUFFER_INITIAL_SIZE().key(), "2k");
    tenantConfigs.put("celeborn.test.tenant.user.upsert", "insert");
    configService.upsertTenantConfig(ConfigLevel.TENANT_USER, "tenant_id1", "Jerry", tenantConfigs);
    configService.refreshCache();
    tenantConfig = configService.getRawTenantUserConfigFromCache("tenant_id1", "Jerry");
    Assert.assertEquals(
        2048,
        tenantConfig
            .getValue(
                CelebornConf.CLIENT_PUSH_BUFFER_INITIAL_SIZE().key(),
                CelebornConf.CLIENT_PUSH_BUFFER_INITIAL_SIZE(),
                Long.TYPE,
                ConfigType.BYTES)
            .longValue());
    Assert.assertEquals(
        "insert",
        tenantConfig.getValue(
            "celeborn.test.tenant.user.upsert", null, String.class, ConfigType.STRING));
    configService.deleteTenantConfigByKeys(
        ConfigLevel.TENANT_USER,
        "tenant_id1",
        "Jerry",
        Arrays.asList(
            CelebornConf.CLIENT_PUSH_BUFFER_INITIAL_SIZE().key(),
            "celeborn.test.tenant.user.upsert"));
    configService.refreshCache();
    tenantConfig = configService.getRawTenantUserConfigFromCache("tenant_id1", "Jerry");
    Assert.assertFalse(
        tenantConfig.configs.containsKey(CelebornConf.CLIENT_PUSH_BUFFER_INITIAL_SIZE().key()));
    Assert.assertFalse(tenantConfig.configs.containsKey("celeborn.test.tenant.user.upsert"));
  }

  @Test
  public void testFsConfig() throws IOException {
    CelebornConf celebornConf = new CelebornConf();
    String file = getClass().getResource("/dynamicConfig.yaml").getFile();
    celebornConf.set(CelebornConf.DYNAMIC_CONFIG_STORE_FS_PATH(), file);
    celebornConf.set(CelebornConf.DYNAMIC_CONFIG_REFRESH_INTERVAL(), 5L);
    configService = new FsConfigServiceImpl(celebornConf);
    verifySystemConfig(configService);
    verifyTenantConfig(configService);
    verifyTenantUserConfig(configService);
    verifyTags(configService);

    // change -> refresh config
    file = getClass().getResource("/dynamicConfig_2.yaml").getFile();
    celebornConf.set(CelebornConf.DYNAMIC_CONFIG_STORE_FS_PATH(), file);
    configService.refreshCache();

    verifyConfigChanged(configService);
    verifyTagsChanged(configService);
  }

  @After
  public void teardown() {
    if (configService != null) {
      configService.shutdown();
    }
  }

  public void verifySystemConfig(ConfigService configService) {
    // ------------- Verify SystemConfig ----------------- //
    SystemConfig systemConfig = configService.getSystemConfigFromCache();
    // verify systemConfig's bytesConf -- use systemConfig
    Long value =
        systemConfig.getValue(
            CelebornConf.CLIENT_PUSH_BUFFER_INITIAL_SIZE().key(),
            CelebornConf.CLIENT_PUSH_BUFFER_INITIAL_SIZE(),
            Long.TYPE,
            ConfigType.BYTES);
    Assert.assertEquals(value.longValue(), 102400);

    // verify systemConfig's bytesConf -- defer to celebornConf
    value =
        systemConfig.getValue(
            CelebornConf.SHUFFLE_PARTITION_SPLIT_THRESHOLD().key(),
            CelebornConf.SHUFFLE_PARTITION_SPLIT_THRESHOLD(),
            Long.TYPE,
            ConfigType.BYTES);
    Assert.assertEquals(value.longValue(), 1073741824);

    // verify systemConfig's bytesConf only -- use systemConfig
    value =
        systemConfig.getValue(
            "celeborn.client.push.buffer.initial.size.only", null, Long.TYPE, ConfigType.BYTES);
    Assert.assertEquals(value.longValue(), 10240);

    // verify systemConfig's bytesConf with none
    value =
        systemConfig.getValue(
            "celeborn.client.push.buffer.initial.size.only.none",
            null,
            Long.TYPE,
            ConfigType.BYTES);
    Assert.assertNull(value);

    // verify systemConfig's timesConf
    value =
        systemConfig.getValue("celeborn.test.timeoutMs.only", null, Long.TYPE, ConfigType.TIME_MS);
    Assert.assertEquals(value.longValue(), 100000);

    // verify systemConfig's BooleanConf
    Boolean booleanConfValue =
        systemConfig.getValue(
            "celeborn.test.timeoutMs.only", null, Boolean.TYPE, ConfigType.STRING);
    Assert.assertFalse(booleanConfValue);

    // verify systemConfig's intConf
    Integer intConfValue =
        systemConfig.getValue("celeborn.test.int.only", null, Integer.TYPE, ConfigType.STRING);
    Assert.assertEquals(intConfValue.intValue(), 10);
  }

  public void verifyTenantConfig(ConfigService configService) {
    // ------------- Verify TenantConfig ----------------- //
    DynamicConfig tenantConfig = configService.getTenantConfigFromCache("tenant_id");
    // verify tenantConfig's bytesConf -- use tenantConf
    Long value =
        tenantConfig.getValue(
            CelebornConf.CLIENT_PUSH_BUFFER_INITIAL_SIZE().key(),
            CelebornConf.CLIENT_PUSH_BUFFER_INITIAL_SIZE(),
            Long.TYPE,
            ConfigType.BYTES);
    Assert.assertEquals(value.longValue(), 10240);

    // verify tenantConfig's bytesConf -- defer to systemConf
    value =
        tenantConfig.getValue(
            CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE().key(),
            CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE(),
            Long.TYPE,
            ConfigType.BYTES);
    Assert.assertEquals(value.longValue(), 1024000);

    // verify tenantConfig's bytesConf -- defer to celebornConf
    value =
        tenantConfig.getValue(
            CelebornConf.SHUFFLE_PARTITION_SPLIT_THRESHOLD().key(),
            CelebornConf.SHUFFLE_PARTITION_SPLIT_THRESHOLD(),
            Long.TYPE,
            ConfigType.BYTES);
    Assert.assertEquals(value.longValue(), 1073741824);

    // verify tenantConfig's bytesConf only -- use tenantConf
    value =
        tenantConfig.getValue(
            "celeborn.client.push.buffer.initial.size.only", null, Long.TYPE, ConfigType.BYTES);
    Assert.assertEquals(value.longValue(), 102400);

    // verify tenantConfig's bytesConf with none
    value =
        tenantConfig.getValue(
            "celeborn.client.push.buffer.initial.size.only.none",
            null,
            Long.TYPE,
            ConfigType.BYTES);
    Assert.assertNull(value);

    DynamicConfig tenantConfigNone = configService.getTenantConfigFromCache("tenant_id_none");
    // verify tenantConfig's bytesConf -- defer to systemConf
    value =
        tenantConfigNone.getValue(
            CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE().key(),
            CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE(),
            Long.TYPE,
            ConfigType.BYTES);
    Assert.assertEquals(value.longValue(), 1024000);

    // ------------- Verify with defaultValue ----------------- //
    value =
        tenantConfig.getWithDefaultValue(
            "celeborn.client.push.buffer.initial.size.only", 100L, Long.TYPE, ConfigType.BYTES);
    Assert.assertEquals(value.longValue(), 102400);

    Long withDefaultValue =
        tenantConfigNone.getWithDefaultValue("none", 10L, Long.TYPE, ConfigType.STRING);
    Assert.assertEquals(withDefaultValue.longValue(), 10);
  }

  public void verifyTenantUserConfig(ConfigService configService) {
    // ------------- Verify UserConfig ----------------- //
    DynamicConfig userConfig = configService.getTenantUserConfigFromCache("tenant_id1", "Jerry");
    // verify userConfig's bytesConf -- use userConf
    Long value =
        userConfig.getValue(
            CelebornConf.CLIENT_PUSH_BUFFER_INITIAL_SIZE().key(),
            CelebornConf.CLIENT_PUSH_BUFFER_INITIAL_SIZE(),
            Long.TYPE,
            ConfigType.BYTES);
    Assert.assertEquals(value.longValue(), 1024);

    // verify userConfig's bytesConf -- defer to tenantConf
    value =
        userConfig.getValue(
            CelebornConf.CLIENT_PUSH_QUEUE_CAPACITY().key(),
            CelebornConf.CLIENT_PUSH_QUEUE_CAPACITY(),
            Long.TYPE,
            ConfigType.BYTES);
    Assert.assertEquals(value.longValue(), 1024);

    // verify userConfig's bytesConf -- defer to systemConf
    value =
        userConfig.getValue(
            CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE().key(),
            CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE(),
            Long.TYPE,
            ConfigType.BYTES);
    Assert.assertEquals(value.longValue(), 1024000);

    // verify userConfig's bytesConf -- defer to celebornConf
    value =
        userConfig.getValue(
            CelebornConf.SHUFFLE_PARTITION_SPLIT_THRESHOLD().key(),
            CelebornConf.SHUFFLE_PARTITION_SPLIT_THRESHOLD(),
            Long.TYPE,
            ConfigType.BYTES);
    Assert.assertEquals(value.longValue(), 1073741824);

    // verify userConfig's bytesConf with none
    value =
        userConfig.getValue(
            "celeborn.client.push.buffer.initial.size.only.none",
            null,
            Long.TYPE,
            ConfigType.BYTES);
    Assert.assertNull(value);

    DynamicConfig userConfigNone =
        configService.getTenantUserConfigFromCache("tenant_id", "non_exist");
    // verify userConfig's bytesConf -- defer to tenantConf
    value =
        userConfigNone.getValue(
            CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE().key(),
            CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE(),
            Long.TYPE,
            ConfigType.BYTES);
    Assert.assertEquals(value.longValue(), 1024000);

    Long withDefaultValue =
        userConfigNone.getWithDefaultValue("none", 10L, Long.TYPE, ConfigType.STRING);
    Assert.assertEquals(withDefaultValue.longValue(), 10);
  }

  public void verifyConfigChanged(ConfigService configService) {

    SystemConfig systemConfig = configService.getSystemConfigFromCache();
    // verify systemConfig's intConf
    Integer intConfValue =
        systemConfig.getValue("celeborn.test.int.only", null, Integer.TYPE, ConfigType.STRING);
    Assert.assertEquals(intConfValue.intValue(), 100);

    // verify systemConfig's bytesConf -- defer to celebornConf
    Long value =
        systemConfig.getValue(
            CelebornConf.SHUFFLE_PARTITION_SPLIT_THRESHOLD().key(),
            CelebornConf.SHUFFLE_PARTITION_SPLIT_THRESHOLD(),
            Long.TYPE,
            ConfigType.BYTES);
    Assert.assertEquals(value.longValue(), 1073741824);
  }

  public void verifyServiceManager(
      IServiceManager serviceManager, CelebornConf celebornConf, Instant gmtTime) {
    ClusterInfo clusterInfo = serviceManager.getClusterInfo(celebornConf.clusterName());
    Assert.assertEquals(gmtTime, clusterInfo.getGmtCreate());
    Assert.assertEquals(gmtTime, clusterInfo.getGmtModify());
  }

  private void verifyTags(ConfigService configService) {
    SystemConfig systemConfig = configService.getSystemConfigFromCache();
    Map<String, Set<String>> tags = systemConfig.getTags();

    Set<String> tag1 = tags.getOrDefault("tag1", new HashSet<>());
    Assert.assertEquals(tag1.size(), 2);
    Assert.assertTrue(tag1.contains("host1:1111"));
    Assert.assertTrue(tag1.contains("host2:2222"));

    Set<String> tag2 = tags.getOrDefault("tag2", new HashSet<>());
    Assert.assertEquals(tag2.size(), 2);
    Assert.assertTrue(tag2.contains("host3:3333"));
    Assert.assertTrue(tag2.contains("host4:4444"));

    Set<String> tag3 = tags.getOrDefault("tag3", new HashSet<>());
    Assert.assertEquals(tag3.size(), 0);

    verifyTenantAndUserTagsAsNull(configService);
  }

  private void verifyTagsChanged(ConfigService configService) {
    SystemConfig systemConfig = configService.getSystemConfigFromCache();
    Map<String, Set<String>> tags = systemConfig.getTags();

    Set<String> tag1 = tags.getOrDefault("tag1", new HashSet<>());
    Assert.assertEquals(tag1.size(), 2);
    Assert.assertTrue(tag1.contains("host1:1111"));
    Assert.assertTrue(tag1.contains("host2:2222"));

    Set<String> tag2 = tags.getOrDefault("tag2", new HashSet<>());
    Assert.assertEquals(tag2.size(), 0);

    Set<String> tag3 = tags.getOrDefault("tag3", new HashSet<>());
    Assert.assertEquals(tag3.size(), 2);
    Assert.assertTrue(tag3.contains("host3:3333"));
    Assert.assertTrue(tag3.contains("host4:4444"));
  }

  private void verifyTenantAndUserTagsAsNull(ConfigService configService) {
    TenantConfig tenantConfig = configService.getRawTenantConfigFromCache("tenant_id1");
    Assert.assertNull(tenantConfig.getTags());

    DynamicConfig userConfig = configService.getTenantUserConfigFromCache("tenant_id1", "Jerry");
    Assert.assertNull(userConfig.getTags());
  }
}
