/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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.facebook.buck.jvm.kotlin;

import static com.google.common.collect.Iterables.transform;

import com.facebook.buck.core.build.execution.context.ExecutionContext;
import com.facebook.buck.core.exceptions.HumanReadableException;
import com.facebook.buck.core.model.BuildTarget;
import com.facebook.buck.core.rulekey.AddToRuleKey;
import com.facebook.buck.core.sourcepath.PathSourcePath;
import com.facebook.buck.core.sourcepath.SourcePath;
import com.facebook.buck.core.sourcepath.resolver.SourcePathResolverAdapter;
import com.facebook.buck.io.filesystem.ProjectFilesystem;
import com.facebook.buck.jvm.java.javax.SynchronizedToolProvider;
import com.facebook.buck.util.ClassLoaderCache;
import com.google.common.base.Joiner;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.ImmutableSortedSet;
import java.io.File;
import java.io.PrintStream;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.net.MalformedURLException;
import java.net.URL;
import java.nio.file.Path;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
import java.util.stream.Collectors;

public class JarBackedReflectedKotlinc implements Kotlinc {

  private static final String COMPILER_CLASS = "org.jetbrains.kotlin.cli.jvm.K2JVMCompiler";
  private static final String EXIT_CODE_CLASS = "org.jetbrains.kotlin.cli.common.ExitCode";
  private static final KotlincVersion VERSION = ImmutableKotlincVersion.ofImpl("in memory");

  private static final Function<Path, URL> PATH_TO_URL =
      p -> {
        try {
          return p.toUri().toURL();
        } catch (MalformedURLException e) {
          throw new RuntimeException(e);
        }
      };

  // Used to hang onto the KotlinDaemonShim for the life of the buckd process
  private static final Map<Set<String>, Object> kotlinShims = new ConcurrentHashMap<>();

  @AddToRuleKey private final ImmutableSet<SourcePath> compilerClassPath;
  private final Path annotationProcessingClassPath;
  private final Path standardLibraryClasspath;

  JarBackedReflectedKotlinc(
      ImmutableSet<SourcePath> compilerClassPath,
      Path annotationProcessingClassPath,
      Path standardLibraryClasspath) {
    this.compilerClassPath = compilerClassPath;
    this.annotationProcessingClassPath = annotationProcessingClassPath;
    this.standardLibraryClasspath = standardLibraryClasspath;
  }

  @Override
  public KotlincVersion getVersion() {
    return VERSION;
  }

  @Override
  public String getDescription(
      ImmutableList<String> options,
      ImmutableSortedSet<Path> javaSourceFilePaths,
      Path pathToSrcsList) {
    StringBuilder builder = new StringBuilder("kotlinc ");
    Joiner.on(" ").appendTo(builder, options);
    builder.append(" ");
    builder.append("@").append(pathToSrcsList);

    return builder.toString();
  }

  @Override
  public String getShortName() {
    return "kotlinc";
  }

  @Override
  public Path getAnnotationProcessorPath(SourcePathResolverAdapter sourcePathResolverAdapter) {
    return annotationProcessingClassPath;
  }

  @Override
  public Path getStdlibPath(SourcePathResolverAdapter sourcePathResolverAdapter) {
    return standardLibraryClasspath;
  }

  @Override
  public ImmutableList<Path> getAdditionalClasspathEntries(
      SourcePathResolverAdapter sourcePathResolverAdapter) {
    return ImmutableList.of(getStdlibPath(sourcePathResolverAdapter));
  }

  @Override
  public ImmutableList<String> getCommandPrefix(SourcePathResolverAdapter resolver) {
    throw new UnsupportedOperationException("In memory kotlinc may not be used externally");
  }

  @Override
  public int buildWithClasspath(
      ExecutionContext context,
      BuildTarget invokingRule,
      ImmutableList<String> options,
      ImmutableSortedSet<Path> kotlinSourceFilePaths,
      Path pathToSrcsList,
      Optional<Path> workingDirectory,
      ProjectFilesystem projectFilesystem) {

    ImmutableList<Path> expandedSources;
    try {
      expandedSources =
          getExpandedSourcePaths(
              projectFilesystem,
              context.getProjectFilesystemFactory(),
              kotlinSourceFilePaths,
              workingDirectory);
    } catch (Throwable throwable) {
      throwable.printStackTrace();
      throw new HumanReadableException(
          "Unable to expand sources for %s into %s", invokingRule, workingDirectory);
    }

    ImmutableList<String> args =
        ImmutableList.<String>builder()
            .addAll(options)
            .addAll(
                transform(
                    expandedSources,
                    path -> projectFilesystem.resolve(path).toAbsolutePath().toString()))
            .build();

    Set<File> compilerIdPaths =
        compilerClassPath.stream()
            .map(p -> ((PathSourcePath) p).getRelativePath())
            .map(Path::toFile)
            .collect(Collectors.toSet());

    try {
      Object compilerShim =
          kotlinShims.computeIfAbsent(
              compilerIdPaths.stream().map(File::getAbsolutePath).collect(Collectors.toSet()),
              k -> loadCompilerShim(context));

      Method compile = compilerShim.getClass().getMethod("exec", PrintStream.class, String[].class);

      Class<?> exitCodeClass = compilerShim.getClass().getClassLoader().loadClass(EXIT_CODE_CLASS);

      Method getCode = exitCodeClass.getMethod("getCode");

      try (UncloseablePrintStream stdErr = new UncloseablePrintStream(context.getStdErr())) {
        Object exitCode = compile.invoke(compilerShim, stdErr, args.toArray(new String[0]));

        return (Integer) getCode.invoke(exitCode);
      }

    } catch (IllegalAccessException
        | InvocationTargetException
        | NoSuchMethodException
        | ClassNotFoundException ex) {
      throw new RuntimeException(ex);
    }
  }

  private Object loadCompilerShim(ExecutionContext context) {
    try {
      ClassLoaderCache classLoaderCache = context.getClassLoaderCache();
      classLoaderCache.addRef();

      ClassLoader classLoader =
          classLoaderCache.getClassLoaderForClassPath(
              SynchronizedToolProvider.getSystemToolClassLoader(),
              ImmutableList.copyOf(
                  compilerClassPath.stream()
                      .map(p -> ((PathSourcePath) p).getRelativePath())
                      .map(PATH_TO_URL)
                      .iterator()));

      return classLoader.loadClass(COMPILER_CLASS).newInstance();
    } catch (Exception ex) {
      throw new RuntimeException(ex);
    }
  }

  @Override
  public ImmutableMap<String, String> getEnvironment(SourcePathResolverAdapter resolver) {
    throw new UnsupportedOperationException("In memory kotlinc may not be used externally");
  }

  private static class UncloseablePrintStream extends PrintStream {
    UncloseablePrintStream(PrintStream delegate) {
      super(delegate);
    }

    @Override
    public void close() {
      // ignore
    }
  }
}

@@ -1546,29 +1546,8 @@
 et;%0A
-import java.io.File;%0A
 impo
@@ -1746,30 +1746,8 @@
 th;%0A
-import java.util.Map;%0A
 impo
@@ -1790,130 +1790,25 @@
 til.
-Set;%0Aimport java.util.concurrent.ConcurrentHashMap;%0Aimport java.util.function.Function;%0Aimport java.util.stream.Collectors
+function.Function
 ;%0A%0Ap
@@ -2375,176 +2375,8 @@
 %7D;%0A%0A
-  // Used to hang onto the KotlinDaemonShim for the life of the buckd process%0A  private static final Map%3CSet%3CString%3E, Object%3E kotlinShims = new ConcurrentHashMap%3C%3E();%0A%0A
   @A
@@ -5204,397 +5204,41 @@
     
-Set%3CFile%3E compilerIdPaths =%0A        compilerClassPath.stream()%0A            .map(p -%3E ((PathSourcePath) p).getRelativePath())%0A            .map(Path::toFile)%0A            .collect(Collectors.toSet());%0A%0A    try %7B%0A      Object compilerShim =%0A          kotlinShims.computeIfAbsent(%0A              compilerIdPaths.stream().map(File::getAbsolutePath).collect(Collectors.toSet()),%0A              k -%3E
+try %7B%0A      Object compilerShim =
  loa
@@ -5259,17 +5259,16 @@
 context)
-)
 ;%0A%0A     

Remove caching of kotlinc