/*
 * 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.wayang.api.sql.calcite.jdbc;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.primitives.Ints;
import org.apache.calcite.avatica.ColumnMetaData;
import org.apache.calcite.avatica.util.DateTimeUtils;
import org.apache.calcite.linq4j.function.Function0;
import org.apache.calcite.linq4j.function.Function1;
import org.apache.calcite.sql.SqlDialect;
import org.apache.calcite.sql.SqlDialectFactory;
import org.apache.calcite.util.ImmutableNullableList;
import org.apache.calcite.util.Pair;
import org.apache.calcite.util.Util;
import org.apache.commons.dbcp2.BasicDataSource;
import org.checkerframework.checker.nullness.qual.Nullable;

import javax.sql.DataSource;
import java.sql.*;
import java.util.List;
import java.util.TimeZone;

/**
 * Utilities for the JDBC provider.
 */
final class JdbcUtils {
  private JdbcUtils() {
    throw new AssertionError("no instances!");
  }

  /** Returns a function that, given a {@link ResultSet}, returns a function
   * that will yield successive rows from that result set. */
  static Function1<ResultSet, Function0<@Nullable Object[]>> rowBuilderFactory(
      final List<Pair<ColumnMetaData.Rep, Integer>> list) {
    ColumnMetaData.Rep[] reps =
        Pair.left(list).toArray(new ColumnMetaData.Rep[0]);
    int[] types = Ints.toArray(Pair.right(list));
    return resultSet -> new ObjectArrayRowBuilder1(resultSet, reps, types);
  }

  /** Returns a function that, given a {@link ResultSet}, returns a function
   * that will yield successive rows from that result set;
   * as {@link #rowBuilderFactory(List)} except that values are in Calcite's
   * internal format (e.g. DATE represented as int). */
  static Function1<ResultSet, Function0<@Nullable Object[]>> rowBuilderFactory2(
      final List<Pair<ColumnMetaData.Rep, Integer>> list) {
    ColumnMetaData.Rep[] reps =
        Pair.left(list).toArray(new ColumnMetaData.Rep[0]);
    int[] types = Ints.toArray(Pair.right(list));
    return resultSet -> new ObjectArrayRowBuilder2(resultSet, reps, types);
  }

  /** Pool of dialects. */
  static class DialectPool {
    public static final DialectPool INSTANCE = new DialectPool();

    private final LoadingCache<Pair<SqlDialectFactory, DataSource>, SqlDialect> cache =
        CacheBuilder.newBuilder().softValues()
            .build(CacheLoader.from(DialectPool::dialect));

    private static SqlDialect dialect(
        Pair<SqlDialectFactory, DataSource> key) {
      SqlDialectFactory dialectFactory = key.left;
      DataSource dataSource = key.right;
      Connection connection = null;
      try {
        connection = dataSource.getConnection();
        DatabaseMetaData metaData = connection.getMetaData();
        SqlDialect dialect = dialectFactory.create(metaData);
        connection.close();
        connection = null;
        return dialect;
      } catch (SQLException e) {
        throw new RuntimeException(e);
      } finally {
        if (connection != null) {
          try {
            connection.close();
          } catch (SQLException e) {
            // ignore
          }
        }
      }
    }

    public SqlDialect get(SqlDialectFactory dialectFactory, DataSource dataSource) {
      final Pair<SqlDialectFactory, DataSource> key =
          Pair.of(dialectFactory, dataSource);
      return cache.getUnchecked(key);
    }
  }

  /** Builder that calls {@link ResultSet#getObject(int)} for every column,
   * or {@code getXxx} if the result type is a primitive {@code xxx},
   * and returns an array of objects for each row. */
  abstract static class ObjectArrayRowBuilder
      implements Function0<@Nullable Object[]> {
    protected final ResultSet resultSet;
    protected final int columnCount;
    protected final ColumnMetaData.Rep[] reps;
    protected final int[] types;

    ObjectArrayRowBuilder(ResultSet resultSet, ColumnMetaData.Rep[] reps,
        int[] types) {
      this.resultSet = resultSet;
      this.reps = reps;
      this.types = types;
      try {
        this.columnCount = resultSet.getMetaData().getColumnCount();
      } catch (SQLException e) {
        throw Util.throwAsRuntime(e);
      }
    }

    @Override public @Nullable Object[] apply() {
      try {
        final @Nullable Object[] values = new Object[columnCount];
        for (int i = 0; i < columnCount; i++) {
          values[i] = value(i);
        }
        return values;
      } catch (SQLException e) {
        throw new RuntimeException(e);
      }
    }

    /**
     * Gets a value from a given column in a JDBC result set.
     *
     * @param i Ordinal of column (1-based, per JDBC)
     */
    protected abstract @Nullable Object value(int i) throws SQLException;

    long timestampToLong(Timestamp v) {
      return v.getTime();
    }

    long timeToLong(Time v) {
      return v.getTime();
    }

    long dateToLong(Date v) {
      return v.getTime();
    }
  }

  /** Row builder that shifts DATE, TIME, TIMESTAMP values into local time
   * zone. */
  static class ObjectArrayRowBuilder1 extends ObjectArrayRowBuilder {
    final TimeZone timeZone = TimeZone.getDefault();

    ObjectArrayRowBuilder1(ResultSet resultSet, ColumnMetaData.Rep[] reps,
        int[] types) {
      super(resultSet, reps, types);
    }

    @Override protected @Nullable Object value(int i) throws SQLException {
      // MySQL returns timestamps shifted into local time. Using
      // getTimestamp(int, Calendar) with a UTC calendar should prevent this,
      // but does not. So we shift explicitly.
      switch (types[i]) {
      case Types.TIMESTAMP:
        final Timestamp timestamp = resultSet.getTimestamp(i + 1);
        return timestamp == null ? null : new Timestamp(timestampToLong(timestamp));
      case Types.TIME:
        final Time time = resultSet.getTime(i + 1);
        return time == null ? null : new Time(timeToLong(time));
      case Types.DATE:
        final Date date = resultSet.getDate(i + 1);
        return date == null ? null : new Date(dateToLong(date));
      default:
        break;
      }
      return reps[i].jdbcGet(resultSet, i + 1);
    }

    @Override long timestampToLong(Timestamp v) {
      long time = v.getTime();
      int offset = timeZone.getOffset(time);
      return time + offset;
    }

    @Override long timeToLong(Time v) {
      long time = v.getTime();
      int offset = timeZone.getOffset(time);
      return (time + offset) % DateTimeUtils.MILLIS_PER_DAY;
    }

    @Override long dateToLong(Date v) {
      long time = v.getTime();
      int offset = timeZone.getOffset(time);
      return time + offset;
    }
  }

  /** Row builder that converts JDBC values into internal values. */
  static class ObjectArrayRowBuilder2 extends ObjectArrayRowBuilder1 {
    ObjectArrayRowBuilder2(ResultSet resultSet, ColumnMetaData.Rep[] reps,
        int[] types) {
      super(resultSet, reps, types);
    }

    @Override protected @Nullable Object value(int i) throws SQLException {
      switch (types[i]) {
      case Types.TIMESTAMP:
        final Timestamp timestamp = resultSet.getTimestamp(i + 1);
        return timestamp == null ? null : timestampToLong(timestamp);
      case Types.TIME:
        final Time time = resultSet.getTime(i + 1);
        return time == null ? null : (int) timeToLong(time);
      case Types.DATE:
        final Date date = resultSet.getDate(i + 1);
        return date == null ? null
            : (int) (dateToLong(date) / DateTimeUtils.MILLIS_PER_DAY);
      default:
        return reps[i].jdbcGet(resultSet, i + 1);
      }
    }
  }

  /** Ensures that if two data sources have the same definition, they will use
   * the same object.
   *
   * <p>This in turn makes it easier to cache
   * {@link SqlDialect} objects. Otherwise, each time we
   * see a new data source, we have to open a connection to find out what
   * database product and version it is. */
  static class DataSourcePool {
    public static final DataSourcePool INSTANCE = new DataSourcePool();

    private final LoadingCache<List<@Nullable String>, BasicDataSource> cache =
        CacheBuilder.newBuilder().softValues()
            .build(CacheLoader.from(DataSourcePool::dataSource));

    private static BasicDataSource dataSource(
          List<? extends @Nullable String> key) {
      BasicDataSource dataSource = new BasicDataSource();
      dataSource.setUrl(key.get(0));
      dataSource.setUsername(key.get(1));
      dataSource.setPassword(key.get(2));
      dataSource.setDriverClassName(key.get(3));
      return dataSource;
    }

    public DataSource get(String url, @Nullable String driverClassName,
                          @Nullable String username, @Nullable String password) {
      // Get data source objects from a cache, so that we don't have to sniff
      // out what kind of database they are quite as often.
      final List<@Nullable String> key =
          ImmutableNullableList.of(url, username, password, driverClassName);
      return cache.getUnchecked(key);
    }
  }
}
