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

import com.datastax.oss.driver.api.core.CqlIdentifier;
import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
import com.datastax.oss.driver.api.mapper.annotations.Dao;
import com.datastax.oss.driver.api.mapper.annotations.DaoFactory;
import com.datastax.oss.driver.api.mapper.annotations.DaoKeyspace;
import com.datastax.oss.driver.api.mapper.annotations.DaoProfile;
import com.datastax.oss.driver.api.mapper.annotations.DaoTable;
import com.datastax.oss.driver.api.mapper.annotations.Mapper;
import com.datastax.oss.driver.internal.mapper.DaoCacheKey;
import com.datastax.oss.driver.internal.mapper.processor.GeneratedNames;
import com.datastax.oss.driver.internal.mapper.processor.MethodGenerator;
import com.datastax.oss.driver.internal.mapper.processor.ProcessorContext;
import com.datastax.oss.driver.internal.mapper.processor.util.generation.GeneratedCodePatterns;
import com.squareup.javapoet.ClassName;
import com.squareup.javapoet.MethodSpec;
import com.squareup.javapoet.TypeName;
import java.util.Optional;
import javax.lang.model.element.Element;
import javax.lang.model.element.ExecutableElement;
import javax.lang.model.element.TypeElement;
import javax.lang.model.element.VariableElement;
import javax.lang.model.type.DeclaredType;
import javax.lang.model.type.TypeKind;
import javax.lang.model.type.TypeMirror;

/**
 * Generates the implementation of a DAO-producing method in a {@link Mapper}-annotated interface.
 */
public class MapperDaoFactoryMethodGenerator implements MethodGenerator {

  private final ExecutableElement methodElement;
  private final MapperImplementationSharedCode enclosingClass;
  private final ProcessorContext context;

  public MapperDaoFactoryMethodGenerator(
      ExecutableElement methodElement,
      MapperImplementationSharedCode enclosingClass,
      ProcessorContext context) {
    this.methodElement = methodElement;
    this.enclosingClass = enclosingClass;
    this.context = context;
  }

  @Override
  public Optional<MethodSpec> generate() {

    // Validate the return type, which tells us what DAO to build, and whether the method should be
    // async.
    ClassName daoImplementationName = null;
    boolean isAsync = false;
    TypeMirror returnTypeMirror = methodElement.getReturnType();
    if (returnTypeMirror.getKind() == TypeKind.DECLARED) {
      DeclaredType declaredReturnType = (DeclaredType) returnTypeMirror;
      if (declaredReturnType.getTypeArguments().isEmpty()) {
        Element returnTypeElement = declaredReturnType.asElement();
        if (returnTypeElement.getAnnotation(Dao.class) != null) {
          daoImplementationName =
              GeneratedNames.daoImplementation(((TypeElement) returnTypeElement));
        }
      } else if (context.getClassUtils().isFuture(declaredReturnType)) {
        TypeMirror typeArgument = declaredReturnType.getTypeArguments().get(0);
        if (typeArgument.getKind() == TypeKind.DECLARED) {
          Element typeArgumentElement = ((DeclaredType) typeArgument).asElement();
          if (typeArgumentElement.getAnnotation(Dao.class) != null) {
            daoImplementationName =
                GeneratedNames.daoImplementation(((TypeElement) typeArgumentElement));
            isAsync = true;
          }
        }
      }
    }
    if (daoImplementationName == null) {
      context
          .getMessager()
          .error(
              methodElement,
              "Invalid return type: %s methods must return a %s-annotated interface, "
                  + "or future thereof",
              DaoFactory.class.getSimpleName(),
              Dao.class.getSimpleName());
      return Optional.empty();
    }

    // Validate the arguments
    String keyspaceArgumentName = null;
    String tableArgumentName = null;
    String profileArgumentName = null;
    boolean profileIsClass = false;

    for (VariableElement parameterElement : methodElement.getParameters()) {
      if (parameterElement.getAnnotation(DaoKeyspace.class) != null) {
        keyspaceArgumentName =
            validateKeyspaceOrTableParameter(
                parameterElement, keyspaceArgumentName, DaoKeyspace.class, context);
        if (keyspaceArgumentName == null) {
          return Optional.empty();
        }
      } else if (parameterElement.getAnnotation(DaoTable.class) != null) {
        tableArgumentName =
            validateKeyspaceOrTableParameter(
                parameterElement, tableArgumentName, DaoTable.class, context);
        if (tableArgumentName == null) {
          return Optional.empty();
        }
      } else if (parameterElement.getAnnotation(DaoProfile.class) != null) {

        profileArgumentName =
            validateExecutionProfile(parameterElement, profileArgumentName, context);
        profileIsClass =
            context.getClassUtils().isSame(parameterElement.asType(), DriverExecutionProfile.class);
        if (profileArgumentName == null) {
          return Optional.empty();
        }
      } else {
        context
            .getMessager()
            .error(
                methodElement,
                "Invalid parameter annotations: "
                    + "%s method parameters must be annotated with @%s, @%s or @%s",
                DaoFactory.class.getSimpleName(),
                DaoKeyspace.class.getSimpleName(),
                DaoTable.class.getSimpleName(),
                DaoProfile.class.getSimpleName());
        return Optional.empty();
      }
    }
    boolean isCachedByMethodArguments =
        (keyspaceArgumentName != null || tableArgumentName != null || profileArgumentName != null);

    TypeName returnTypeName = ClassName.get(methodElement.getReturnType());
    String suggestedFieldName = methodElement.getSimpleName() + "Cache";
    String fieldName =
        isCachedByMethodArguments
            ? enclosingClass.addDaoMapField(suggestedFieldName, returnTypeName)
            : enclosingClass.addDaoSimpleField(
                suggestedFieldName, returnTypeName, daoImplementationName, isAsync);

    MethodSpec.Builder overridingMethodBuilder = GeneratedCodePatterns.override(methodElement);

    if (isCachedByMethodArguments) {
      // DaoCacheKey key = new DaoCacheKey(<ks>, <table>, <profileName>, <profile>)
      // where <ks>, <table> is either the name of the parameter or "(CqlIdentifier)null"
      overridingMethodBuilder.addCode("$1T key = new $1T(", DaoCacheKey.class);
      if (keyspaceArgumentName == null) {
        overridingMethodBuilder.addCode("($T)null", CqlIdentifier.class);
      } else {
        overridingMethodBuilder.addCode("$L", keyspaceArgumentName);
      }
      overridingMethodBuilder.addCode(", ");
      if (tableArgumentName == null) {
        overridingMethodBuilder.addCode("($T)null", CqlIdentifier.class);
      } else {
        overridingMethodBuilder.addCode("$L", tableArgumentName);
      }
      overridingMethodBuilder.addCode(", ");
      if (profileArgumentName == null) {
        overridingMethodBuilder.addCode("null, null);\n");
      } else {
        if (profileIsClass) {
          overridingMethodBuilder.addCode("null, $L);\n", profileArgumentName);
        } else {
          overridingMethodBuilder.addCode("$L, null);\n", profileArgumentName);
        }
      }

      overridingMethodBuilder.addStatement(
          "return $L.computeIfAbsent(key, "
              + "k -> $T.$L(context.withDaoParameters(k.getKeyspaceId(), k.getTableId(), "
              + "k.getExecutionProfileName(), k.getExecutionProfile())))",
          fieldName,
          daoImplementationName,
          isAsync ? "initAsync" : "init");
    } else {
      overridingMethodBuilder.addStatement("return $L.get()", fieldName);
    }
    return Optional.of(overridingMethodBuilder.build());
  }

  private String validateKeyspaceOrTableParameter(
      VariableElement candidate, String previous, Class<?> annotation, ProcessorContext context) {
    if (!isSingleAnnotation(candidate, previous, annotation, context)) {
      return null;
    }
    TypeMirror type = candidate.asType();
    if (!context.getClassUtils().isSame(type, String.class)
        && !context.getClassUtils().isSame(type, CqlIdentifier.class)) {
      context
          .getMessager()
          .error(
              candidate,
              "Invalid parameter type: @%s-annotated parameter of %s methods must be of type %s or %s",
              annotation.getSimpleName(),
              DaoFactory.class.getSimpleName(),
              String.class.getSimpleName(),
              CqlIdentifier.class.getSimpleName());
      return null;
    }
    return candidate.getSimpleName().toString();
  }

  private String validateExecutionProfile(
      VariableElement candidate, String previous, ProcessorContext context) {
    if (!isSingleAnnotation(candidate, previous, DaoProfile.class, context)) {
      return null;
    }
    TypeMirror type = candidate.asType();
    if (!context.getClassUtils().isSame(type, String.class)
        && !context.getClassUtils().isSame(type, DriverExecutionProfile.class)) {
      context
          .getMessager()
          .error(
              candidate,
              "Invalid parameter type: @%s-annotated parameter of %s methods must be of type %s or %s ",
              DaoProfile.class.getSimpleName(),
              DaoFactory.class.getSimpleName(),
              String.class.getSimpleName(),
              DriverExecutionProfile.class.getSimpleName());
      return null;
    }
    return candidate.getSimpleName().toString();
  }

  private boolean isSingleAnnotation(
      VariableElement candidate, String previous, Class<?> annotation, ProcessorContext context) {
    if (previous != null) {
      context
          .getMessager()
          .error(
              candidate,
              "Invalid parameter annotations: "
                  + "only one %s method parameter can be annotated with @%s",
              DaoFactory.class.getSimpleName(),
              annotation.getSimpleName());
      return false;
    }
    return true;
  }
}
