/*
 * 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.processor.dao;

import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.BOOLEAN;
import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.BOUND_STATEMENT;
import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.CUSTOM;
import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.FUTURE_OF_ASYNC_RESULT_SET;
import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.FUTURE_OF_BOOLEAN;
import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.FUTURE_OF_VOID;
import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.REACTIVE_RESULT_SET;
import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.RESULT_SET;
import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.VOID;

import com.datastax.oss.driver.api.core.cql.BoundStatement;
import com.datastax.oss.driver.api.core.cql.BoundStatementBuilder;
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
import com.datastax.oss.driver.api.mapper.annotations.Update;
import com.datastax.oss.driver.api.mapper.entity.saving.NullSavingStrategy;
import com.datastax.oss.driver.internal.mapper.processor.ProcessorContext;
import com.datastax.oss.driver.internal.mapper.processor.entity.EntityDefinition;
import com.datastax.oss.driver.internal.mapper.processor.entity.PropertyDefinition;
import com.datastax.oss.driver.internal.mapper.processor.util.generation.GeneratedCodePatterns;
import com.datastax.oss.driver.internal.querybuilder.update.DefaultUpdate;
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableSet;
import com.squareup.javapoet.ClassName;
import com.squareup.javapoet.CodeBlock;
import com.squareup.javapoet.MethodSpec;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import javax.lang.model.element.ExecutableElement;
import javax.lang.model.element.Name;
import javax.lang.model.element.TypeElement;
import javax.lang.model.element.VariableElement;

public class DaoUpdateMethodGenerator extends DaoMethodGenerator {

  private final NullSavingStrategyValidation nullSavingStrategyValidation;

  public DaoUpdateMethodGenerator(
      ExecutableElement methodElement,
      Map<Name, TypeElement> typeParameters,
      TypeElement processedType,
      DaoImplementationSharedCode enclosingClass,
      ProcessorContext context) {
    super(methodElement, typeParameters, processedType, enclosingClass, context);
    nullSavingStrategyValidation = new NullSavingStrategyValidation(context);
  }

  protected Set<DaoReturnTypeKind> getSupportedReturnTypes() {
    return ImmutableSet.of(
        VOID,
        FUTURE_OF_VOID,
        RESULT_SET,
        BOUND_STATEMENT,
        FUTURE_OF_ASYNC_RESULT_SET,
        BOOLEAN,
        FUTURE_OF_BOOLEAN,
        REACTIVE_RESULT_SET,
        CUSTOM);
  }

  @Override
  public boolean requiresReactive() {
    // Validate the return type:
    DaoReturnType returnType =
        parseAndValidateReturnType(getSupportedReturnTypes(), Update.class.getSimpleName());
    if (returnType == null) {
      return false;
    }
    return returnType.requiresReactive();
  }

  @Override
  public Optional<MethodSpec> generate() {

    // Validate the parameters:
    // - the first one must be the entity.
    // - the others are completely free-form (they'll be used as additional bind variables)
    // A Function<BoundStatementBuilder, BoundStatementBuilder> can be added in last position.
    List<? extends VariableElement> parameters = methodElement.getParameters();
    VariableElement boundStatementFunction = findBoundStatementFunction(methodElement);
    if (boundStatementFunction != null) {
      parameters = parameters.subList(0, parameters.size() - 1);
    }
    TypeElement entityElement =
        parameters.isEmpty()
            ? null
            : EntityUtils.asEntityElement(parameters.get(0), typeParameters);
    if (entityElement == null) {
      context
          .getMessager()
          .error(
              methodElement,
              "%s methods must take the entity to update as the first parameter",
              Update.class.getSimpleName());
      return Optional.empty();
    }
    warnIfCqlNamePresent(parameters.subList(0, 1));
    EntityDefinition entityDefinition = context.getEntityFactory().getDefinition(entityElement);

    // Validate the return type:
    DaoReturnType returnType =
        parseAndValidateReturnType(getSupportedReturnTypes(), Update.class.getSimpleName());
    if (returnType == null) {
      return Optional.empty();
    }

    // Generate the method:
    String helperFieldName = enclosingClass.addEntityHelperField(ClassName.get(entityElement));
    String statementName =
        enclosingClass.addPreparedStatement(
            methodElement,
            (methodBuilder, requestName) ->
                generatePrepareRequest(methodBuilder, requestName, helperFieldName));

    CodeBlock.Builder createStatementBlock = CodeBlock.builder();

    createStatementBlock.addStatement(
        "$T boundStatementBuilder = $L.boundStatementBuilder()",
        BoundStatementBuilder.class,
        statementName);

    populateBuilderWithStatementAttributes(createStatementBlock, methodElement);
    populateBuilderWithFunction(createStatementBlock, boundStatementFunction);

    String entityParameterName = parameters.get(0).getSimpleName().toString();

    Update annotation = methodElement.getAnnotation(Update.class);
    String customWhereClause = annotation.customWhereClause();

    NullSavingStrategy nullSavingStrategy =
        nullSavingStrategyValidation.getNullSavingStrategy(
            Update.class, Update::nullSavingStrategy, methodElement, enclosingClass);

    if (customWhereClause.isEmpty()) {
      // We generated an update by primary key (see maybeAddWhereClause), all entity properties are
      // present as placeholders.
      createStatementBlock.addStatement(
          "$1L.set($2L, boundStatementBuilder, $3T.$4L, false)",
          helperFieldName,
          entityParameterName,
          NullSavingStrategy.class,
          nullSavingStrategy);
    } else {
      createStatementBlock.addStatement(
          "$1T nullSavingStrategy = $1T.$2L", NullSavingStrategy.class, nullSavingStrategy);

      // Only non-PK properties are present in SET ... clauses.
      // (if the custom clause has custom placeholders, this will be addressed below)
      for (PropertyDefinition property : entityDefinition.getRegularColumns()) {
        GeneratedCodePatterns.setValue(
            property.getCqlName(),
            property.getType(),
            CodeBlock.of("$L.$L()", entityParameterName, property.getGetterName()),
            "boundStatementBuilder",
            createStatementBlock,
            enclosingClass,
            true,
            false);
      }
    }

    // Handle all remaining parameters as additional bound values in customWhereClause or
    // customIfClause
    if (parameters.size() > 1) {
      List<? extends VariableElement> bindMarkers = parameters.subList(1, parameters.size());
      if (validateCqlNamesPresent(bindMarkers)) {
        GeneratedCodePatterns.bindParameters(
            bindMarkers, createStatementBlock, enclosingClass, context, false);
      } else {
        return Optional.empty();
      }
    }

    createStatementBlock.addStatement(
        "$T boundStatement = boundStatementBuilder.build()", BoundStatement.class);

    return crudMethod(createStatementBlock, returnType, helperFieldName);
  }

  private void generatePrepareRequest(
      MethodSpec.Builder methodBuilder, String requestName, String helperFieldName) {
    Update annotation = methodElement.getAnnotation(Update.class);

    maybeAddWhereClause(
        methodBuilder, requestName, helperFieldName, annotation.customWhereClause());
    maybeAddTtl(annotation.ttl(), methodBuilder);
    maybeAddTimestamp(annotation.timestamp(), methodBuilder);
    methodBuilder.addCode(")");
    maybeAddIfClause(methodBuilder, annotation);

    methodBuilder.addCode(".asCql()");
    methodBuilder.addCode(")$];\n");
  }

  private void maybeAddWhereClause(
      MethodSpec.Builder methodBuilder,
      String requestName,
      String helperFieldName,
      String customWhereClause) {

    if (customWhereClause.isEmpty()) {
      methodBuilder.addCode(
          "$[$1T $2L = $1T.newInstance((($4T)$3L.updateByPrimaryKey()",
          SimpleStatement.class,
          requestName,
          helperFieldName,
          DefaultUpdate.class);
    } else {
      methodBuilder.addCode(
          "$[$1T $2L = $1T.newInstance((($5T)$3L.updateStart().whereRaw($4S)",
          SimpleStatement.class,
          requestName,
          helperFieldName,
          customWhereClause,
          DefaultUpdate.class);
    }
  }

  private void maybeAddIfClause(MethodSpec.Builder methodBuilder, Update annotation) {
    if (annotation.ifExists() && !annotation.customIfClause().isEmpty()) {
      context
          .getMessager()
          .error(
              methodElement,
              "Invalid annotation parameters: %s cannot have both ifExists and customIfClause",
              Update.class.getSimpleName());
    }

    if (annotation.ifExists()) {
      methodBuilder.addCode(".ifExists()");
    }

    if (!annotation.customIfClause().isEmpty()) {
      methodBuilder.addCode(".ifRaw($S)", annotation.customIfClause());
    }
  }
}
