/*
 * 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.gravitino.iceberg.common.utils;

import static org.apache.hadoop.fs.CommonConfigurationKeysPublic.HADOOP_SECURITY_AUTHENTICATION;
import static org.apache.hadoop.fs.CommonConfigurationKeysPublic.HADOOP_SECURITY_AUTHORIZATION;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.Maps;
import java.sql.SQLException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Locale;
import java.util.Map;
import org.apache.gravitino.catalog.lakehouse.iceberg.IcebergCatalogBackend;
import org.apache.gravitino.catalog.lakehouse.iceberg.IcebergConstants;
import org.apache.gravitino.exceptions.ConnectionFailedException;
import org.apache.gravitino.iceberg.common.ClosableHiveCatalog;
import org.apache.gravitino.iceberg.common.IcebergConfig;
import org.apache.gravitino.iceberg.common.authentication.AuthenticationConfig;
import org.apache.gravitino.iceberg.common.rest.auth.UserPrincipalForwardingAuthManager;
import org.apache.hadoop.hdfs.HdfsConfiguration;
import org.apache.iceberg.CatalogProperties;
import org.apache.iceberg.CatalogUtil;
import org.apache.iceberg.catalog.Catalog;
import org.apache.iceberg.hive.HiveCatalog;
import org.apache.iceberg.hive.HiveCatalogWithMetadataLocationSupport;
import org.apache.iceberg.inmemory.InMemoryCatalog;
import org.apache.iceberg.io.ResolvingFileIO;
import org.apache.iceberg.jdbc.JdbcCatalog;
import org.apache.iceberg.jdbc.JdbcCatalogWithMetadataLocationSupport;
import org.apache.iceberg.jdbc.UncheckedSQLException;
import org.apache.iceberg.memory.MemoryCatalogWithMetadataLocationSupport;
import org.apache.iceberg.rest.RESTCatalog;
import org.apache.iceberg.rest.auth.AuthProperties;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class IcebergCatalogUtil {

  private static final Logger LOG = LoggerFactory.getLogger(IcebergCatalogUtil.class);

  private static InMemoryCatalog loadMemoryCatalog(IcebergConfig icebergConfig) {
    String icebergCatalogName = icebergConfig.getCatalogBackendName();
    InMemoryCatalog memoryCatalog = new MemoryCatalogWithMetadataLocationSupport();
    Map<String, String> resultProperties = icebergConfig.getIcebergCatalogProperties();
    if (!resultProperties.containsKey(IcebergConstants.WAREHOUSE)) {
      resultProperties.put(IcebergConstants.WAREHOUSE, "/tmp");
    }
    applyDefaultResolvingFileIO(resultProperties);
    memoryCatalog.initialize(icebergCatalogName, resultProperties);
    return memoryCatalog;
  }

  private static HiveCatalog loadHiveCatalog(IcebergConfig icebergConfig) {
    ClosableHiveCatalog hiveCatalog = new HiveCatalogWithMetadataLocationSupport();
    HdfsConfiguration hdfsConfiguration = new HdfsConfiguration();
    String icebergCatalogName = icebergConfig.getCatalogBackendName();

    Map<String, String> properties = icebergConfig.getIcebergCatalogProperties();
    applyDefaultResolvingFileIO(properties);
    properties.forEach(hdfsConfiguration::set);
    AuthenticationConfig authenticationConfig = new AuthenticationConfig(properties);
    if (authenticationConfig.isSimpleAuth()) {
      hiveCatalog.setConf(hdfsConfiguration);
      hiveCatalog.initialize(icebergCatalogName, properties);
      return hiveCatalog;
    } else if (authenticationConfig.isKerberosAuth()) {
      Map<String, String> resultProperties = new HashMap<>(properties);
      resultProperties.put(CatalogProperties.CLIENT_POOL_CACHE_KEYS, "USER_NAME");
      hdfsConfiguration.set(HADOOP_SECURITY_AUTHORIZATION, "true");
      hdfsConfiguration.set(HADOOP_SECURITY_AUTHENTICATION, "kerberos");
      hiveCatalog.setConf(hdfsConfiguration);
      hiveCatalog.initialize(icebergCatalogName, resultProperties);
      return hiveCatalog;
    } else {
      throw new UnsupportedOperationException(
          "Unsupported authentication method: " + authenticationConfig.getAuthType());
    }
  }

  @SuppressWarnings("FormatStringAnnotation")
  private static JdbcCatalog loadJdbcCatalog(IcebergConfig icebergConfig) {
    String driverClassName = icebergConfig.getJdbcDriver();
    String icebergCatalogName = icebergConfig.getCatalogBackendName();

    Map<String, String> properties = icebergConfig.getIcebergCatalogProperties();
    applyDefaultResolvingFileIO(properties);
    try {
      // Load the jdbc driver
      Class.forName(driverClassName);
    } catch (ClassNotFoundException e) {
      throw new IllegalArgumentException("Couldn't load jdbc driver " + driverClassName);
    }
    JdbcCatalog jdbcCatalog =
        new JdbcCatalogWithMetadataLocationSupport(
            icebergConfig.get(IcebergConfig.JDBC_INIT_TABLES));

    // Default to V1 schema to support view operations; can be overridden by explicit config.
    properties.putIfAbsent(IcebergConstants.ICEBERG_JDBC_SCHEMA_VERSION, "V1");

    // Default to strict mode so that creating a table or view in a non-existent namespace fails
    // with NoSuchNamespaceException (HTTP 404) instead of implicitly creating the namespace,
    // matching the Iceberg REST spec and the memory backend behavior. Can be overridden by
    // explicit config.
    properties.putIfAbsent(IcebergConstants.ICEBERG_JDBC_STRICT_MODE, "true");

    HdfsConfiguration hdfsConfiguration = new HdfsConfiguration();
    properties.forEach(hdfsConfiguration::set);
    jdbcCatalog.setConf(hdfsConfiguration);
    try {
      jdbcCatalog.initialize(icebergCatalogName, properties);
    } catch (UncheckedSQLException e) {
      if (e.getCause() instanceof SQLException
          && e.getCause().getMessage().contains("Access denied")) {
        throw new ConnectionFailedException(e, e.getMessage());
      }
      throw e;
    }
    return jdbcCatalog;
  }

  private static Catalog loadRestCatalog(IcebergConfig icebergConfig) {
    String icebergCatalogName = icebergConfig.getCatalogBackendName();
    RESTCatalog restCatalog = new RESTCatalog();
    HdfsConfiguration hdfsConfiguration = new HdfsConfiguration();
    Map<String, String> properties = Maps.newHashMap(icebergConfig.getIcebergCatalogProperties());
    applyDefaultResolvingFileIO(properties);

    // REST catalog must use forward access token from the user request
    properties.put(AuthProperties.AUTH_TYPE, UserPrincipalForwardingAuthManager.class.getName());
    applyRestCatalogHttpTimeoutProperties(icebergConfig, properties);

    properties.forEach(hdfsConfiguration::set);
    restCatalog.setConf(hdfsConfiguration);
    restCatalog.initialize(icebergCatalogName, properties);
    return restCatalog;
  }

  private static Catalog loadCustomCatalog(IcebergConfig icebergConfig) {
    String customCatalogName = icebergConfig.getCatalogBackendName();
    String className = icebergConfig.get(IcebergConfig.CATALOG_BACKEND_IMPL);
    Map<String, String> properties = icebergConfig.getIcebergCatalogProperties();
    applyDefaultResolvingFileIO(properties);
    return CatalogUtil.loadCatalog(
        className, customCatalogName, properties, new HdfsConfiguration());
  }

  @VisibleForTesting
  public static void applyDefaultResolvingFileIO(Map<String, String> properties) {
    properties.putIfAbsent(IcebergConstants.IO_IMPL, ResolvingFileIO.class.getName());
  }

  @VisibleForTesting
  static void applyRestCatalogHttpTimeoutProperties(
      IcebergConfig icebergConfig, Map<String, String> properties) {
    properties.put(
        IcebergConstants.ICEBERG_REST_CLIENT_CONNECTION_TIMEOUT_MS,
        String.valueOf(
            icebergConfig.get(IcebergConfig.REST_CATALOG_BACKEND_CLIENT_CONNECTION_TIMEOUT_MS)));
    properties.put(
        IcebergConstants.ICEBERG_REST_CLIENT_SOCKET_TIMEOUT_MS,
        String.valueOf(
            icebergConfig.get(IcebergConfig.REST_CATALOG_BACKEND_CLIENT_SOCKET_TIMEOUT_MS)));
  }

  @VisibleForTesting
  static Catalog loadCatalogBackend(String catalogType) {
    return loadCatalogBackend(
        IcebergCatalogBackend.valueOf(catalogType.toUpperCase(Locale.ROOT)),
        new IcebergConfig(Collections.emptyMap()));
  }

  public static Catalog loadCatalogBackend(
      IcebergCatalogBackend catalogBackend, IcebergConfig icebergConfig) {
    LOG.info("Load catalog backend of {}", catalogBackend);
    switch (catalogBackend) {
      case MEMORY:
        return loadMemoryCatalog(icebergConfig);
      case HIVE:
        return loadHiveCatalog(icebergConfig);
      case JDBC:
        return loadJdbcCatalog(icebergConfig);
      case REST:
        return loadRestCatalog(icebergConfig);
      case CUSTOM:
        return loadCustomCatalog(icebergConfig);
      default:
        throw new RuntimeException(
            catalogBackend
                + " catalog is not supported yet, supported catalogs: [memory]"
                + catalogBackend);
    }
  }

  private IcebergCatalogUtil() {}
}
