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

import static com.datastax.oss.driver.api.mapper.annotations.SchemaHint.TargetElement;

import com.datastax.oss.driver.api.core.CqlIdentifier;
import com.datastax.oss.driver.api.core.metadata.schema.KeyspaceMetadata;
import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata;
import com.datastax.oss.driver.api.core.type.DataType;
import com.datastax.oss.driver.api.core.type.UserDefinedType;
import com.datastax.oss.driver.api.core.type.reflect.GenericType;
import com.datastax.oss.driver.api.mapper.annotations.SchemaHint;
import com.datastax.oss.driver.internal.mapper.processor.MethodGenerator;
import com.datastax.oss.driver.internal.mapper.processor.dao.LoggingGenerator;
import com.squareup.javapoet.CodeBlock;
import com.squareup.javapoet.MethodSpec;
import com.squareup.javapoet.TypeName;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;
import javax.lang.model.element.Modifier;
import javax.lang.model.element.TypeElement;

public class EntityHelperSchemaValidationMethodGenerator implements MethodGenerator {

  private final EntityDefinition entityDefinition;
  private TypeElement entityTypeElement;
  private LoggingGenerator loggingGenerator;
  private EntityHelperGenerator entityHelperGenerator;

  public EntityHelperSchemaValidationMethodGenerator(
      EntityDefinition entityDefinition,
      TypeElement entityTypeElement,
      LoggingGenerator loggingGenerator,
      EntityHelperGenerator entityHelperGenerator) {
    this.entityDefinition = entityDefinition;
    this.entityTypeElement = entityTypeElement;
    this.loggingGenerator = loggingGenerator;
    this.entityHelperGenerator = entityHelperGenerator;
  }

  @Override
  public Optional<MethodSpec> generate() {
    MethodSpec.Builder methodBuilder =
        MethodSpec.methodBuilder("validateEntityFields")
            .addAnnotation(Override.class)
            .addModifiers(Modifier.PUBLIC)
            .returns(TypeName.VOID);

    Optional<TargetElement> targetElement =
        Optional.ofNullable(entityTypeElement.getAnnotation(SchemaHint.class))
            .map(SchemaHint::targetElement);

    if (targetElement.isPresent() && targetElement.get() == TargetElement.NONE) {
      methodBuilder.addComment(
          "Nothing to do, validation was disabled with @SchemaHint(targetElement = NONE)");
    } else {
      // get keyspaceId from context, and if not present fallback to keyspace set on session
      methodBuilder.addStatement(
          "$1T keyspaceId = this.keyspaceId != null ? this.keyspaceId : context.getSession().getKeyspace().orElse(null)",
          CqlIdentifier.class);

      methodBuilder.addStatement("String entityClassName = $S", entityDefinition.getClassName());
      generateKeyspaceNull(methodBuilder);

      generateKeyspaceNameWrong(methodBuilder);

      methodBuilder.addStatement(
          "$1T<$2T> keyspace = context.getSession().getMetadata().getKeyspace(keyspaceId)",
          Optional.class,
          KeyspaceMetadata.class);

      // Generates expected names to be present in cql (table or udt)
      List<CodeBlock> expectedCqlNames =
          entityDefinition.getAllColumns().stream()
              .map(PropertyDefinition::getCqlName)
              .collect(Collectors.toList());
      methodBuilder.addStatement(
          "$1T<$2T> expectedCqlNames = new $3T<>()",
          List.class,
          CqlIdentifier.class,
          ArrayList.class);
      for (CodeBlock expectedCqlName : expectedCqlNames) {
        methodBuilder.addStatement(
            "expectedCqlNames.add($1T.fromCql($2L))", CqlIdentifier.class, expectedCqlName);
      }

      methodBuilder.addStatement(
          "$1T<$2T> tableMetadata = keyspace.flatMap(v -> v.getTable(tableId))",
          Optional.class,
          TableMetadata.class);

      // Generated UserDefineTypes metadata
      methodBuilder.addStatement(
          "$1T<$2T> userDefinedType = keyspace.flatMap(v -> v.getUserDefinedType(tableId))",
          Optional.class,
          UserDefinedType.class);

      generateValidationChecks(methodBuilder, targetElement);

      logMissingMetadata(methodBuilder);
    }
    return Optional.of(methodBuilder.build());
  }

  private void logMissingMetadata(MethodSpec.Builder methodBuilder) {
    methodBuilder.addComment(
        "warn if there is not keyspace.table for defined entity - it means that table is missing, or schema it out of date.");
    methodBuilder.beginControlFlow("else");
    loggingGenerator.warn(
        methodBuilder,
        "[{}] There is no ks.table or UDT: {}.{} for the entity class: {}, or metadata is out of date.",
        CodeBlock.of("context.getSession().getName()"),
        CodeBlock.of("keyspaceId"),
        CodeBlock.of("tableId"),
        CodeBlock.of("entityClassName"));
    methodBuilder.endControlFlow();
  }

  // handle case where keyspace name is not present in metadata keyspaces
  private void generateKeyspaceNameWrong(MethodSpec.Builder methodBuilder) {
    methodBuilder.beginControlFlow(
        "if(!keyspaceNamePresent(context.getSession().getMetadata().getKeyspaces(), keyspaceId))");
    loggingGenerator.warn(
        methodBuilder,
        "[{}] Unable to validate table: {} for the entity class: {} "
            + "because the session metadata has no information about the keyspace: {}.",
        CodeBlock.of("context.getSession().getName()"),
        CodeBlock.of("tableId"),
        CodeBlock.of("entityClassName"),
        CodeBlock.of("keyspaceId"));
    methodBuilder.addStatement("return");
    methodBuilder.endControlFlow();
  }

  // Handle case where keyspaceId = null.
  // In such case we cannot infer and validate schema for table or udt
  private void generateKeyspaceNull(MethodSpec.Builder methodBuilder) {
    methodBuilder.beginControlFlow("if (keyspaceId == null)");
    loggingGenerator.warn(
        methodBuilder,
        "[{}] Unable to validate table: {} for the entity class: {} because the keyspace "
            + "is unknown (the entity does not declare a default keyspace, and neither the "
            + "session nor the DAO were created with a keyspace). The DAO will only work if it "
            + "uses fully-qualified queries with @Query or @QueryProvider.",
        CodeBlock.of("context.getSession().getName()"),
        CodeBlock.of("tableId"),
        CodeBlock.of("entityClassName"));
    methodBuilder.addStatement("return");
    methodBuilder.endControlFlow();
  }

  private void generateValidationChecks(
      MethodSpec.Builder methodBuilder, Optional<TargetElement> targetElement) {
    // if SchemaHint was not provided explicitly try to match TABLE, then fallback to UDT
    if (!targetElement.isPresent()) {
      validateColumnsInTable(methodBuilder);
      validateColumnsInUdt(methodBuilder, true);
    }
    // if explicitly provided SchemaHint is TABLE, then generate only TABLE check
    else if (targetElement.get().equals(TargetElement.TABLE)) {
      validateColumnsInTable(methodBuilder);
    }
    // if explicitly provided SchemaHint is UDT, then generate only UDT check
    else if (targetElement.get().equals(TargetElement.UDT)) {
      validateColumnsInUdt(methodBuilder, false);
    }
  }

  private void validateColumnsInTable(MethodSpec.Builder methodBuilder) {
    methodBuilder.beginControlFlow("if (tableMetadata.isPresent())");

    generateMissingClusteringColumnsCheck(methodBuilder);

    generateMissingPKsCheck(methodBuilder);

    generateMissingColumnsCheck(methodBuilder);

    generateColumnsTypeCheck(methodBuilder);

    methodBuilder.endControlFlow();
  }

  private void generateColumnsTypeCheck(MethodSpec.Builder methodBuilder) {
    methodBuilder.addComment("validation of types");
    generateExpectedTypesPerColumn(methodBuilder);

    methodBuilder.addStatement(
        "$1T<$2T> missingTableTypes = findTypeMismatches(expectedTypesPerColumn, tableMetadata.get().getColumns(), context.getSession().getContext().getCodecRegistry())",
        List.class,
        String.class);
    methodBuilder.addStatement(
        "throwMissingTableTypesIfNotEmpty(missingTableTypes, keyspaceId, tableId, entityClassName)");
  }

  private void generateMissingColumnsCheck(MethodSpec.Builder methodBuilder) {
    methodBuilder.addComment("validation of all columns");

    methodBuilder.addStatement(
        "$1T<$2T> missingTableCqlNames = findMissingCqlIdentifiers(expectedCqlNames, tableMetadata.get().getColumns().keySet())",
        List.class,
        CqlIdentifier.class);

    // Throw if there are any missingTableCqlNames
    CodeBlock missingCqlColumnExceptionMessage =
        CodeBlock.of(
            "String.format(\"The CQL ks.table: %s.%s has missing columns: %s that are defined in the entity class: %s\", "
                + "keyspaceId, tableId, missingTableCqlNames, entityClassName)");
    methodBuilder.beginControlFlow("if (!missingTableCqlNames.isEmpty())");
    methodBuilder.addStatement(
        "throw new $1T($2L)", IllegalArgumentException.class, missingCqlColumnExceptionMessage);
    methodBuilder.endControlFlow();
  }

  private void generateMissingPKsCheck(MethodSpec.Builder methodBuilder) {
    methodBuilder.addComment("validation of missing PKs");
    List<CodeBlock> expectedCqlPKs =
        entityDefinition.getPartitionKey().stream()
            .map(PropertyDefinition::getCqlName)
            .collect(Collectors.toList());

    methodBuilder.addStatement(
        "$1T<$2T> expectedCqlPKs = new $3T<>()", List.class, CqlIdentifier.class, ArrayList.class);
    for (CodeBlock expectedCqlName : expectedCqlPKs) {
      methodBuilder.addStatement(
          "expectedCqlPKs.add($1T.fromCql($2L))", CqlIdentifier.class, expectedCqlName);
    }
    methodBuilder.addStatement(
        "$1T<$2T> missingTablePksNames = findMissingColumns(expectedCqlPKs, tableMetadata.get().getPartitionKey())",
        List.class,
        CqlIdentifier.class);

    // throw if there are any missing PK columns
    CodeBlock missingCqlColumnExceptionMessage =
        CodeBlock.of(
            "String.format(\"The CQL ks.table: %s.%s has missing Primary Key columns: %s that are defined in the entity class: %s\", "
                + "keyspaceId, tableId, missingTablePksNames, entityClassName)");
    methodBuilder.beginControlFlow("if (!missingTablePksNames.isEmpty())");
    methodBuilder.addStatement(
        "throw new $1T($2L)", IllegalArgumentException.class, missingCqlColumnExceptionMessage);
    methodBuilder.endControlFlow();
  }

  private void generateMissingClusteringColumnsCheck(MethodSpec.Builder methodBuilder) {
    List<CodeBlock> expectedCqlClusteringColumns =
        entityDefinition.getClusteringColumns().stream()
            .map(PropertyDefinition::getCqlName)
            .collect(Collectors.toList());

    if (!expectedCqlClusteringColumns.isEmpty()) {
      methodBuilder.addComment("validation of missing Clustering Columns");
      methodBuilder.addStatement(
          "$1T<$2T> expectedCqlClusteringColumns = new $3T<>()",
          List.class,
          CqlIdentifier.class,
          ArrayList.class);
      for (CodeBlock expectedCqlName : expectedCqlClusteringColumns) {
        methodBuilder.addStatement(
            "expectedCqlClusteringColumns.add($1T.fromCql($2L))",
            CqlIdentifier.class,
            expectedCqlName);
      }

      methodBuilder.addStatement(
          "$1T<$2T> missingTableClusteringColumnNames = findMissingColumns(expectedCqlClusteringColumns, tableMetadata.get().getClusteringColumns().keySet())",
          List.class,
          CqlIdentifier.class);

      // throw if there are any missing Clustering Columns columns
      CodeBlock missingCqlColumnExceptionMessage =
          CodeBlock.of(
              "String.format(\"The CQL ks.table: %s.%s has missing Clustering columns: %s that are defined in the entity class: %s\", "
                  + "keyspaceId, tableId, missingTableClusteringColumnNames, entityClassName)");
      methodBuilder.beginControlFlow("if (!missingTableClusteringColumnNames.isEmpty())");
      methodBuilder.addStatement(
          "throw new $1T($2L)", IllegalArgumentException.class, missingCqlColumnExceptionMessage);
      methodBuilder.endControlFlow();
    }
  }

  // Finds out missingTableCqlNames - columns that are present in Entity Mapping but NOT present in
  // UDT table
  private void validateColumnsInUdt(MethodSpec.Builder methodBuilder, boolean generateElse) {
    if (generateElse) {
      methodBuilder.beginControlFlow("else if (userDefinedType.isPresent())");
    } else {
      methodBuilder.beginControlFlow("if (userDefinedType.isPresent())");
    }

    generateUdtMissingColumnsCheck(methodBuilder);

    generateUdtColumnsTypeCheck(methodBuilder);

    methodBuilder.endControlFlow();
  }

  private void generateUdtColumnsTypeCheck(MethodSpec.Builder methodBuilder) {
    methodBuilder.addComment("validation of UDT types");
    generateExpectedTypesPerColumn(methodBuilder);

    methodBuilder.addStatement(
        "$1T<$2T> expectedColumns = userDefinedType.get().getFieldNames()",
        List.class,
        CqlIdentifier.class);
    methodBuilder.addStatement(
        "$1T<$2T> expectedTypes = userDefinedType.get().getFieldTypes()",
        List.class,
        DataType.class);

    methodBuilder.addStatement(
        "$1T<$2T> missingTableTypes = findTypeMismatches(expectedTypesPerColumn, expectedColumns, expectedTypes, context.getSession().getContext().getCodecRegistry())",
        List.class,
        String.class);
    methodBuilder.addStatement(
        "throwMissingUdtTypesIfNotEmpty(missingTableTypes, keyspaceId, tableId, entityClassName)");
  }

  private void generateUdtMissingColumnsCheck(MethodSpec.Builder methodBuilder) {
    methodBuilder.addComment("validation of UDT columns");
    methodBuilder.addStatement(
        "$1T<$2T> columns = userDefinedType.get().getFieldNames()",
        List.class,
        CqlIdentifier.class);

    methodBuilder.addStatement(
        "$1T<$2T> missingTableCqlNames = findMissingCqlIdentifiers(expectedCqlNames, columns)",
        List.class,
        CqlIdentifier.class);

    // Throw if there are any missingTableCqlNames
    CodeBlock missingCqlUdtExceptionMessage =
        CodeBlock.of(
            "String.format(\"The CQL ks.udt: %s.%s has missing columns: %s that are defined in the entity class: %s\", "
                + "keyspaceId, tableId, missingTableCqlNames, entityClassName)");
    methodBuilder.beginControlFlow("if (!missingTableCqlNames.isEmpty())");
    methodBuilder.addStatement(
        "throw new $1T($2L)", IllegalArgumentException.class, missingCqlUdtExceptionMessage);
    methodBuilder.endControlFlow();
  }

  private void generateExpectedTypesPerColumn(MethodSpec.Builder methodBuilder) {
    methodBuilder.addStatement(
        "$1T<$2T, $3T<?>> expectedTypesPerColumn = new $4T<>()",
        Map.class,
        CqlIdentifier.class,
        GenericType.class,
        LinkedHashMap.class);

    Map<CodeBlock, TypeName> expectedTypesPerColumn =
        entityDefinition.getAllColumns().stream()
            .collect(
                Collectors.toMap(PropertyDefinition::getCqlName, v -> v.getType().asRawTypeName()));

    for (Map.Entry<CodeBlock, TypeName> expected : expectedTypesPerColumn.entrySet()) {
      methodBuilder.addStatement(
          "expectedTypesPerColumn.put($1T.fromCql($2L), $3L)",
          CqlIdentifier.class,
          expected.getKey(),
          entityHelperGenerator.addGenericTypeConstant(expected.getValue().box()));
    }
  }
}
