/*
 * 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.UNSUPPORTED;
import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.values;

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.Query;
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.util.generation.GeneratedCodePatterns;
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 DaoQueryMethodGenerator extends DaoMethodGenerator {

  private final String queryString;
  private final NullSavingStrategyValidation nullSavingStrategyValidation;

  public DaoQueryMethodGenerator(
      ExecutableElement methodElement,
      Map<Name, TypeElement> typeParameters,
      TypeElement processedType,
      DaoImplementationSharedCode enclosingClass,
      ProcessorContext context) {
    super(methodElement, typeParameters, processedType, enclosingClass, context);
    this.queryString = methodElement.getAnnotation(Query.class).value();
    nullSavingStrategyValidation = new NullSavingStrategyValidation(context);
  }

  protected Set<DaoReturnTypeKind> getSupportedReturnTypes() {
    ImmutableSet.Builder<DaoReturnTypeKind> builder = ImmutableSet.builder();
    for (DefaultDaoReturnTypeKind value : values()) {
      if (value != UNSUPPORTED) {
        builder.add(value);
      }
    }
    return builder.build();
  }

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

  @Override
  public Optional<MethodSpec> generate() {

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

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

    CodeBlock.Builder createStatementBlock = CodeBlock.builder();

    List<? extends VariableElement> parameters = methodElement.getParameters();

    VariableElement boundStatementFunction = findBoundStatementFunction(methodElement);
    if (boundStatementFunction != null) {
      parameters = parameters.subList(0, methodElement.getParameters().size() - 1);
    }

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

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

    createStatementBlock.addStatement(
        "$1T nullSavingStrategy = $1T.$2L", NullSavingStrategy.class, nullSavingStrategy);

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

    if (validateCqlNamesPresent(parameters)) {
      GeneratedCodePatterns.bindParameters(
          parameters, createStatementBlock, enclosingClass, context, true);

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

      return crudMethod(createStatementBlock, returnType, helperFieldName);
    } else {
      return Optional.empty();
    }
  }

  private void generatePrepareRequest(
      MethodSpec.Builder methodBuilder, String requestName, String helperFieldName) {
    methodBuilder.addStatement(
        "$T $L = replaceKeyspaceAndTablePlaceholders($S, context, $L)",
        SimpleStatement.class,
        requestName,
        queryString,
        helperFieldName);
  }
}
