/*
 * Copyright (C) 2021 The Dagger Authors.
 *
 * 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 dagger.hilt.processor.internal.root;

import static javax.lang.model.element.Modifier.PUBLIC;

import androidx.room.compiler.processing.XFiler.Mode;
import androidx.room.compiler.processing.XProcessingEnv;
import androidx.room.compiler.processing.XTypeElement;
import com.google.common.collect.ImmutableSet;
import com.squareup.javapoet.AnnotationSpec;
import com.squareup.javapoet.ClassName;
import com.squareup.javapoet.JavaFile;
import com.squareup.javapoet.TypeSpec;
import dagger.hilt.processor.internal.AggregatedElements;
import dagger.hilt.processor.internal.ClassNames;
import dagger.hilt.processor.internal.Processors;
import java.io.IOException;
import java.util.HashSet;
import java.util.Optional;
import java.util.Set;

/** Generates an {@link dagger.hilt.internal.componenttreedeps.ComponentTreeDeps}. */
final class ComponentTreeDepsGenerator {
  // Keeps track of already generated proxies. For correctness, this same instance of
  // ComponentTreeDepsGenerator must be used for a given round.
  private final Set<ClassName> generatedProxies = new HashSet<>();
  private final XProcessingEnv env;
  private final Mode mode;

  ComponentTreeDepsGenerator(XProcessingEnv env, Mode mode) {
    this.env = env;
    this.mode = mode;
  }

  void generate(ComponentTreeDepsMetadata metadata) throws IOException {
    ClassName name = metadata.name();
    TypeSpec.Builder builder =
        TypeSpec.classBuilder(name)
            // No originating element since this is generated by the aggregating processor.
            .addAnnotation(componentTreeDepsAnnotation(metadata));

    Processors.addGeneratedAnnotation(builder, env, ClassNames.ROOT_PROCESSOR.toString());

    env.getFiler().write(JavaFile.builder(name.packageName(), builder.build()).build(), mode);
  }

  AnnotationSpec componentTreeDepsAnnotation(ComponentTreeDepsMetadata metadata)
      throws IOException {
    AnnotationSpec.Builder builder = AnnotationSpec.builder(ClassNames.COMPONENT_TREE_DEPS);
    addDeps(builder, metadata.aggregatedRootDeps(), "rootDeps");
    addDeps(builder, metadata.defineComponentDeps(), "defineComponentDeps");
    addDeps(builder, metadata.aliasOfDeps(), "aliasOfDeps");
    addDeps(builder, metadata.aggregatedDeps(), "aggregatedDeps");
    addDeps(builder, metadata.aggregatedUninstallModulesDeps(), "uninstallModulesDeps");
    addDeps(builder, metadata.aggregatedEarlyEntryPointDeps(), "earlyEntryPointDeps");
    return builder.build();
  }

  private void addDeps(AnnotationSpec.Builder builder, ImmutableSet<XTypeElement> deps, String name)
      throws IOException {
    for (XTypeElement dep : deps) {
      builder.addMember(name, "$T.class", maybeWrapInPublicProxy(dep));
    }
  }

  /**
   * This method will return the public proxy for {@code dep} if it is not public, otherwise it will
   * return {@code dep} itself. It will also generate the proxy if it doesn't already exist.
   *
   * <p>Note: These proxies are only used for serialization. The proxy will be unwrapped when
   * converting to {@link ComponentTreeDepsMetadata}.
   *
   * <p>Note: The public proxy is needed because Hilt versions < 2.35 generated package-private
   * aggregating elements, which can't be referenced directly in the {@code @ComponentTreeDeps}.
   */
  private ClassName maybeWrapInPublicProxy(XTypeElement dep) {
    Optional<ClassName> proxyName = AggregatedElements.aggregatedElementProxyName(dep);
    if (proxyName.isPresent()) {
      // Check the set of already generated proxies to ensure we don't regenerate the proxy in
      // this round. Also check that the element doesn't already exist to ensure we don't regenerate
      // a proxy generated in a previous round.
      if (generatedProxies.add(proxyName.get())
          && env.findTypeElement(proxyName.get().canonicalName()) == null) {
        generateProxy(dep, proxyName.get());
      }
      return proxyName.get();
    }
    return dep.getClassName();
  }

  private void generateProxy(XTypeElement dep, ClassName proxyName) {
    TypeSpec.Builder builder =
        TypeSpec.classBuilder(proxyName)
            .addModifiers(PUBLIC)
            // No originating element since this is generated by the aggregating processor.
            .addAnnotation(
                AnnotationSpec.builder(ClassNames.AGGREGATED_ELEMENT_PROXY)
                    .addMember("value", "$T.class", dep.getClassName())
                    .build());

    Processors.addGeneratedAnnotation(builder, env, ClassNames.ROOT_PROCESSOR.toString());
    env.getFiler().write(JavaFile.builder(proxyName.packageName(), builder.build()).build(), mode);
  }
}