/*
 * Copyright DataStax, Inc.
 *
 * Licensed 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 com.datastax.oss.driver.internal.mapper;

import com.datastax.oss.driver.api.core.ConsistencyLevel;
import com.datastax.oss.driver.api.core.CqlIdentifier;
import com.datastax.oss.driver.api.core.MappedAsyncPagingIterable;
import com.datastax.oss.driver.api.core.PagingIterable;
import com.datastax.oss.driver.api.core.cql.AsyncResultSet;
import com.datastax.oss.driver.api.core.cql.BoundStatementBuilder;
import com.datastax.oss.driver.api.core.cql.PreparedStatement;
import com.datastax.oss.driver.api.core.cql.ResultSet;
import com.datastax.oss.driver.api.core.cql.Row;
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
import com.datastax.oss.driver.api.core.cql.Statement;
import com.datastax.oss.driver.api.core.type.DataTypes;
import com.datastax.oss.driver.api.mapper.MapperContext;
import com.datastax.oss.driver.api.mapper.MapperException;
import com.datastax.oss.driver.api.mapper.annotations.Dao;
import com.datastax.oss.driver.api.mapper.annotations.Query;
import com.datastax.oss.driver.api.mapper.entity.EntityHelper;
import com.datastax.oss.driver.api.mapper.entity.saving.NullSavingStrategy;
import com.datastax.oss.driver.internal.core.ConsistencyLevelRegistry;
import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
import com.datastax.oss.driver.internal.core.cql.ResultSets;
import com.datastax.oss.protocol.internal.ProtocolConstants;
import java.time.Duration;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

/** Base class for generated implementations of {@link Dao}-annotated interfaces. */
public class DaoBase {

  /** The keyspace id placeholder in {@link Query#value()}. */
  public static final String KEYSPACE_ID_PLACEHOLDER = "${keyspaceId}";

  /** The table id placeholder in {@link Query#value()}. */
  public static final String TABLE_ID_PLACEHOLDER = "${tableId}";

  /** The qualified table id placeholder in {@link Query#value()}. */
  public static final String QUALIFIED_TABLE_ID_PLACEHOLDER = "${qualifiedTableId}";

  private static final CqlIdentifier APPLIED = CqlIdentifier.fromInternal("[applied]");

  protected static CompletionStage<PreparedStatement> prepare(
      SimpleStatement statement, MapperContext context) {
    if (context.getExecutionProfileName() != null) {
      statement = statement.setExecutionProfileName(context.getExecutionProfileName());
    } else if (context.getExecutionProfile() != null) {
      statement = statement.setExecutionProfile(context.getExecutionProfile());
    }
    return context.getSession().prepareAsync(statement);
  }

  /**
   * Replaces {@link #KEYSPACE_ID_PLACEHOLDER}, {@link #TABLE_ID_PLACEHOLDER} and/or {@link
   * #QUALIFIED_TABLE_ID_PLACEHOLDER} in a query string, and turns it into a statement.
   *
   * <p>This is used for {@link Query} methods.
   *
   * @param queryStringTemplate the query string to process.
   * @param context the context that contains the keyspace and table that the DAO was created with
   *     (if any).
   * @param entityHelper the helper the entity that is returned from this query, or {@code null} if
   *     the query does not return entities.
   */
  protected static SimpleStatement replaceKeyspaceAndTablePlaceholders(
      String queryStringTemplate, MapperContext context, EntityHelper<?> entityHelper) {

    CqlIdentifier keyspaceId =
        (entityHelper != null) ? entityHelper.getKeyspaceId() : context.getKeyspaceId();
    CqlIdentifier tableId =
        (entityHelper != null) ? entityHelper.getTableId() : context.getTableId();

    String queryString = queryStringTemplate;
    if (queryString.contains(KEYSPACE_ID_PLACEHOLDER)) {
      if (keyspaceId == null) {
        throw new MapperException(
            String.format(
                "Cannot substitute %s in query '%s': the DAO wasn't built with a keyspace%s",
                KEYSPACE_ID_PLACEHOLDER,
                queryStringTemplate,
                (entityHelper == null)
                    ? ""
                    : " and entity "
                        + entityHelper.getEntityClass().getSimpleName()
                        + " does not define a default keyspace"));
      } else {
        queryString = queryString.replace(KEYSPACE_ID_PLACEHOLDER, keyspaceId.asCql(false));
      }
    }

    if (queryString.contains(TABLE_ID_PLACEHOLDER)) {
      if (tableId == null) {
        throw new MapperException(
            String.format(
                "Cannot substitute %s in query '%s': the DAO wasn't built with a table",
                TABLE_ID_PLACEHOLDER, queryStringTemplate));
      } else {
        queryString = queryString.replace(TABLE_ID_PLACEHOLDER, tableId.asCql(false));
      }
    }

    if (queryString.contains(QUALIFIED_TABLE_ID_PLACEHOLDER)) {
      if (tableId == null) {
        throw new MapperException(
            String.format(
                "Cannot substitute %s in query '%s': the DAO wasn't built with a table",
                QUALIFIED_TABLE_ID_PLACEHOLDER, queryStringTemplate));
      } else {
        String qualifiedId =
            (keyspaceId == null)
                ? tableId.asCql(false)
                : keyspaceId.asCql(false) + '.' + tableId.asCql(false);
        queryString = queryString.replace(QUALIFIED_TABLE_ID_PLACEHOLDER, qualifiedId);
      }
    }

    return SimpleStatement.newInstance(queryString);
  }

  public BoundStatementBuilder populateBoundStatementWithStatementAttributes(
      BoundStatementBuilder builder,
      String profileName,
      String consistencyLevel,
      String serialConsistencyLevel,
      Boolean idempotent,
      int pageSize,
      String timeout,
      String keyspace) {

    if (!profileName.isEmpty()) {
      builder = builder.setExecutionProfileName(profileName);
    }
    if (!consistencyLevel.isEmpty()) {
      builder = builder.setConsistencyLevel(getConsistencyLevelFromName(consistencyLevel));
    }
    if (!serialConsistencyLevel.isEmpty()) {
      builder =
          builder.setSerialConsistencyLevel(getConsistencyLevelFromName(serialConsistencyLevel));
    }
    if (idempotent != null) {
      builder = builder.setIdempotence(idempotent);
    }
    if (pageSize > 0) {
      builder = builder.setPageSize(pageSize);
    }
    if (!timeout.isEmpty()) {
      builder = builder.setTimeout(Duration.parse(timeout));
    }
    if (!keyspace.isEmpty()) {
      builder = builder.setRoutingKeyspace(keyspace);
    }
    return builder;
  }

  private ConsistencyLevel getConsistencyLevelFromName(String name) {
    InternalDriverContext idContext = (InternalDriverContext) context.getSession().getContext();
    ConsistencyLevelRegistry registry = idContext.getConsistencyLevelRegistry();
    return registry.codeToLevel(registry.nameToCode(name));
  }

  protected final MapperContext context;
  protected final boolean isProtocolVersionV3;

  protected DaoBase(MapperContext context) {
    this.context = context;
    this.isProtocolVersionV3 = isProtocolVersionV3(context);
  }

  protected ResultSet execute(Statement<?> statement) {
    return context.getSession().execute(statement);
  }

  protected boolean executeAndMapWasAppliedToBoolean(Statement<?> statement) {
    ResultSet rs = execute(statement);
    return rs.wasApplied();
  }

  protected long executeAndMapFirstColumnToLong(Statement<?> statement) {
    Row row = executeAndExtractFirstRow(statement);
    return extractCount(row);
  }

  private long extractCount(Row row) {
    if (row == null) {
      throw new MapperException(
          "Expected the query to return at least one row "
              + "(return type long is intended for COUNT queries)");
    }
    if (row.getColumnDefinitions().size() == 0
        || !row.getColumnDefinitions().get(0).getType().equals(DataTypes.BIGINT)) {
      throw new MapperException(
          "Expected the query to return a column with CQL type BIGINT in first position "
              + "(return type long is intended for COUNT queries)");
    }
    return row.getLong(0);
  }

  protected Row executeAndExtractFirstRow(Statement<?> statement) {
    return execute(statement).one();
  }

  protected <EntityT> EntityT executeAndMapToSingleEntity(
      Statement<?> statement, EntityHelper<EntityT> entityHelper) {
    ResultSet rs = execute(statement);
    return asEntity(rs.one(), entityHelper);
  }

  private <EntityT> EntityT asEntity(Row row, EntityHelper<EntityT> entityHelper) {
    return (row == null
            // Special case for INSERT IF NOT EXISTS. If the row did not exist, the query returns
            // only [applied], we want to return null to indicate there was no previous entity
            || (row.getColumnDefinitions().size() == 1
                && row.getColumnDefinitions().get(0).getName().equals(APPLIED)))
        ? null
        : entityHelper.get(row, false);
  }

  protected <EntityT> Optional<EntityT> executeAndMapToOptionalEntity(
      Statement<?> statement, EntityHelper<EntityT> entityHelper) {
    return Optional.ofNullable(executeAndMapToSingleEntity(statement, entityHelper));
  }

  protected <EntityT> PagingIterable<EntityT> executeAndMapToEntityIterable(
      Statement<?> statement, EntityHelper<EntityT> entityHelper) {
    return execute(statement).map(row -> entityHelper.get(row, false));
  }

  protected <EntityT> Stream<EntityT> executeAndMapToEntityStream(
      Statement<?> statement, EntityHelper<EntityT> entityHelper) {
    return StreamSupport.stream(
        execute(statement).map(row -> entityHelper.get(row, false)).spliterator(), false);
  }

  protected CompletableFuture<AsyncResultSet> executeAsync(Statement<?> statement) {
    CompletionStage<AsyncResultSet> stage = context.getSession().executeAsync(statement);
    // We allow DAO interfaces to return CompletableFuture instead of CompletionStage. This method
    // returns CompletableFuture, which makes the implementation code a bit simpler to generate.
    // In practice this has no performance impact, because the default implementation of
    // toCompletableFuture in the JDK is `return this`.
    return stage.toCompletableFuture();
  }

  protected CompletableFuture<Void> executeAsyncAndMapToVoid(Statement<?> statement) {
    return executeAsync(statement).thenApply(rs -> null);
  }

  protected CompletableFuture<Boolean> executeAsyncAndMapWasAppliedToBoolean(
      Statement<?> statement) {
    return executeAsync(statement).thenApply(AsyncResultSet::wasApplied);
  }

  protected CompletableFuture<Long> executeAsyncAndMapFirstColumnToLong(Statement<?> statement) {
    return executeAsyncAndExtractFirstRow(statement).thenApply(this::extractCount);
  }

  protected CompletableFuture<Row> executeAsyncAndExtractFirstRow(Statement<?> statement) {
    return executeAsync(statement).thenApply(AsyncResultSet::one);
  }

  protected <EntityT> CompletableFuture<EntityT> executeAsyncAndMapToSingleEntity(
      Statement<?> statement, EntityHelper<EntityT> entityHelper) {
    return executeAsync(statement).thenApply(rs -> asEntity(rs.one(), entityHelper));
  }

  protected <EntityT> CompletableFuture<Optional<EntityT>> executeAsyncAndMapToOptionalEntity(
      Statement<?> statement, EntityHelper<EntityT> entityHelper) {
    return executeAsync(statement)
        .thenApply(rs -> Optional.ofNullable(asEntity(rs.one(), entityHelper)));
  }

  protected <EntityT>
      CompletableFuture<MappedAsyncPagingIterable<EntityT>> executeAsyncAndMapToEntityIterable(
          Statement<?> statement, EntityHelper<EntityT> entityHelper) {
    return executeAsync(statement).thenApply(rs -> rs.map(row -> entityHelper.get(row, false)));
  }

  protected <EntityT> CompletableFuture<Stream<EntityT>> executeAsyncAndMapToEntityStream(
      Statement<?> statement, EntityHelper<EntityT> entityHelper) {
    return executeAsync(statement)
        .thenApply(ResultSets::newInstance)
        .thenApply(rs -> StreamSupport.stream(rs.map(entityHelper::get).spliterator(), false));
  }

  protected static void throwIfProtocolVersionV3(MapperContext context) {
    if (isProtocolVersionV3(context)) {
      throw new MapperException(
          String.format(
              "You cannot use %s.%s for protocol version V3.",
              NullSavingStrategy.class.getSimpleName(), NullSavingStrategy.DO_NOT_SET.name()));
    }
  }

  protected static boolean isProtocolVersionV3(MapperContext context) {
    return context.getSession().getContext().getProtocolVersion().getCode()
        <= ProtocolConstants.Version.V3;
  }
}
