/*
 * Copyright The OpenTelemetry Authors
 * SPDX-License-Identifier: Apache-2.0
 */

package io.opentelemetry.javaagent.tooling.instrumentation.indy;

import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
import io.opentelemetry.javaagent.tooling.BytecodeWithUrl;
import io.opentelemetry.javaagent.tooling.muzzle.InstrumentationModuleMuzzle;
import java.io.IOException;
import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodType;
import java.net.URL;
import java.security.PrivilegedAction;
import java.security.ProtectionDomain;
import java.util.Collections;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.matcher.ElementMatcher;
import net.bytebuddy.matcher.StringMatcher;

/**
 * Class loader used to load the helper classes from {@link
 * io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule}s, so that those
 * classes have access to both the agent/extension classes and the instrumented application classes.
 *
 * <p>This class loader implements the following classloading delegation strategy:
 *
 * <ul>
 *   <li>First, injected classes are considered (usually the helper classes from the
 *       InstrumentationModule)
 *   <li>Next, the class loader looks in the agent or extension class loader, depending on where the
 *       InstrumentationModule comes from
 *   <li>Finally, the instrumented application class loader is checked for the class
 * </ul>
 *
 * <p>In addition, this class loader ensures that the lookup of corresponding .class resources
 * follow the same delegation strategy, so that bytecode inspection tools work correctly.
 */
public class InstrumentationModuleClassLoader extends ClassLoader {

  static {
    ClassLoader.registerAsParallelCapable();
  }

  private static final ClassLoader BOOT_LOADER = new ClassLoader() {};

  private static final Map<String, BytecodeWithUrl> ALWAYS_INJECTED_CLASSES =
      Collections.singletonMap(
          LookupExposer.class.getName(), BytecodeWithUrl.create(LookupExposer.class).cached());
  private static final ProtectionDomain PROTECTION_DOMAIN = getProtectionDomain();
  private static final MethodHandle FIND_PACKAGE_METHOD = getFindPackageMethod();

  private final Map<String, BytecodeWithUrl> additionalInjectedClasses;
  private final ClassLoader agentOrExtensionCl;
  private volatile MethodHandles.Lookup cachedLookup;

  @Nullable private final ClassLoader instrumentedCl;

  /**
   * Only class names matching this matcher will be attempted to be loaded from the {@link
   * #agentOrExtensionCl}. If a class is requested and it does not match this matcher, the lookup in
   * {@link #agentOrExtensionCl} will be skipped.
   */
  private final ElementMatcher<String> agentClassNamesMatcher;

  private final Set<InstrumentationModule> installedModules;

  public InstrumentationModuleClassLoader(
      ClassLoader instrumentedCl, ClassLoader agentOrExtensionCl) {
    this(
        instrumentedCl,
        agentOrExtensionCl,
        new StringMatcher("io.opentelemetry.javaagent", StringMatcher.Mode.STARTS_WITH));
  }

  InstrumentationModuleClassLoader(
      @Nullable ClassLoader instrumentedCl,
      ClassLoader agentOrExtensionCl,
      ElementMatcher<String> classesToLoadFromAgentOrExtensionCl) {
    // agent/extension-class loader is "main"-parent, but class lookup is overridden
    super(agentOrExtensionCl);
    additionalInjectedClasses = new ConcurrentHashMap<>();
    installedModules = Collections.newSetFromMap(new ConcurrentHashMap<>());
    this.agentOrExtensionCl = agentOrExtensionCl;
    this.instrumentedCl = instrumentedCl;
    this.agentClassNamesMatcher = classesToLoadFromAgentOrExtensionCl;
  }

  /**
   * Provides a Lookup within this class loader. See {@link LookupExposer} for the details.
   *
   * @return a lookup capable of accessing public types in this class loader
   */
  public MethodHandles.Lookup getLookup() {
    if (cachedLookup == null) {
      // Load the injected copy of LookupExposer and invoke it
      try {
        // we don't mind the race condition causing the initialization to run multiple times here
        Class<?> lookupExposer = loadClass(LookupExposer.class.getName());
        cachedLookup = (MethodHandles.Lookup) lookupExposer.getMethod("getLookup").invoke(null);
      } catch (Exception e) {
        throw new IllegalStateException(e);
      }
    }
    return cachedLookup;
  }

  public synchronized void installModule(InstrumentationModule module) {
    if (module.getClass().getClassLoader() != agentOrExtensionCl) {
      throw new IllegalArgumentException(
          module.getClass().getName() + " is not loaded by " + agentOrExtensionCl);
    }
    if (!installedModules.add(module)) {
      return;
    }
    Map<String, BytecodeWithUrl> classesToInject =
        getClassesToInject(module).stream()
            .collect(
                Collectors.toMap(
                    className -> className,
                    className -> BytecodeWithUrl.create(className, agentOrExtensionCl)));
    installInjectedClasses(classesToInject);
  }

  public synchronized boolean hasModuleInstalled(InstrumentationModule module) {
    return installedModules.contains(module);
  }

  // Visible for testing
  synchronized void installInjectedClasses(Map<String, BytecodeWithUrl> classesToInject) {
    classesToInject.forEach(additionalInjectedClasses::putIfAbsent);
  }

  private static Set<String> getClassesToInject(InstrumentationModule module) {
    Set<String> toInject = new HashSet<>(InstrumentationModuleMuzzle.getHelperClassNames(module));
    // TODO (Jonas): Make muzzle include advice classes as helper classes
    // so that we don't have to include them here
    toInject.addAll(getModuleAdviceNames(module));
    if (module instanceof ExperimentalInstrumentationModule) {
      toInject.removeAll(((ExperimentalInstrumentationModule) module).injectedClassNames());
    }
    return toInject;
  }

  private static Set<String> getModuleAdviceNames(InstrumentationModule module) {
    Set<String> adviceNames = new HashSet<>();
    TypeTransformer nameCollector =
        new TypeTransformer() {
          @Override
          public void applyAdviceToMethod(
              ElementMatcher<? super MethodDescription> methodMatcher, String adviceClassName) {
            adviceNames.add(adviceClassName);
          }

          @Override
          public void applyTransformer(AgentBuilder.Transformer transformer) {}
        };
    for (TypeInstrumentation instr : module.typeInstrumentations()) {
      instr.transform(nameCollector);
    }
    return adviceNames;
  }

  public static final Map<String, byte[]> bytecodeOverride = new ConcurrentHashMap<>();

  @Override
  @SuppressWarnings("removal") // AccessController is deprecated for removal
  protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
    synchronized (getClassLoadingLock(name)) {
      Class<?> result = findLoadedClass(name);

      // This CL is self-first: Injected class are loaded BEFORE a parent lookup
      if (result == null) {
        BytecodeWithUrl injected = getInjectedClass(name);
        if (injected != null) {
          byte[] bytecode =
              bytecodeOverride.get(name) != null
                  ? bytecodeOverride.get(name)
                  : injected.getBytecode();
          if (System.getSecurityManager() == null) {
            result = defineClassWithPackage(name, bytecode);
          } else {
            result =
                java.security.AccessController.doPrivileged(
                    (PrivilegedAction<Class<?>>) () -> defineClassWithPackage(name, bytecode));
          }
        }
      }
      if (result == null && shouldLoadFromAgent(name)) {
        result = tryLoad(agentOrExtensionCl, name);
      }
      if (result == null) {
        result = tryLoad(instrumentedCl, name);
      }

      if (result != null) {
        if (resolve) {
          resolveClass(result);
        }
        return result;
      } else {
        throw new ClassNotFoundException(name);
      }
    }
  }

  private boolean shouldLoadFromAgent(String dotClassName) {
    return agentClassNamesMatcher.matches(dotClassName);
  }

  private static Class<?> tryLoad(@Nullable ClassLoader cl, String name) {
    try {
      return Class.forName(name, false, cl);
    } catch (ClassNotFoundException e) {
      return null;
    }
  }

  @Override
  public URL getResource(String resourceName) {
    String className = resourceToClassName(resourceName);
    if (className == null) {
      // delegate to just the default parent (the agent class loader)
      return super.getResource(resourceName);
    }
    // for classes use the same precedence as in loadClass
    BytecodeWithUrl injected = getInjectedClass(className);
    if (injected != null) {
      return injected.getUrl();
    }
    URL fromAgentCl = agentOrExtensionCl.getResource(resourceName);
    if (fromAgentCl != null) {
      return fromAgentCl;
    }

    if (instrumentedCl != null) {
      return instrumentedCl.getResource(resourceName);
    } else {
      return BOOT_LOADER.getResource(resourceName);
    }
  }

  @Override
  public Enumeration<URL> getResources(String resourceName) throws IOException {
    String className = resourceToClassName(resourceName);
    if (className == null) {
      return super.getResources(resourceName);
    }
    URL resource = getResource(resourceName);
    List<URL> result =
        resource != null ? Collections.singletonList(resource) : Collections.emptyList();
    return Collections.enumeration(result);
  }

  @Nullable
  private static String resourceToClassName(String resourceName) {
    if (!resourceName.endsWith(".class")) {
      return null;
    }
    String className = resourceName;
    if (className.startsWith("/")) {
      className = className.substring(1);
    }
    className = className.replace('/', '.');
    className = className.substring(0, className.length() - ".class".length());
    return className;
  }

  @Nullable
  private BytecodeWithUrl getInjectedClass(String name) {
    BytecodeWithUrl alwaysInjected = ALWAYS_INJECTED_CLASSES.get(name);
    if (alwaysInjected != null) {
      return alwaysInjected;
    }
    return additionalInjectedClasses.get(name);
  }

  private Class<?> defineClassWithPackage(String name, byte[] bytecode) {
    int lastDotIndex = name.lastIndexOf('.');
    if (lastDotIndex != -1) {
      String packageName = name.substring(0, lastDotIndex);
      safeDefinePackage(packageName);
    }
    return defineClass(name, bytecode, 0, bytecode.length, PROTECTION_DOMAIN);
  }

  private void safeDefinePackage(String packageName) {
    if (findPackage(packageName) == null) {
      try {
        definePackage(packageName, null, null, null, null, null, null, null);
      } catch (IllegalArgumentException e) {
        // Can happen if two classes from the same package are loaded concurrently
        if (findPackage(packageName) == null) {
          // package still doesn't exist, the IllegalArgumentException must be for a different
          // reason than a race condition
          throw e;
        }
      }
    }
  }

  /**
   * Invokes {@link #getPackage(String)} for Java 8 and {@link #getDefinedPackage(String)} for Java
   * 9+.
   *
   * <p>Package-private for testing.
   *
   * @param name the name of the package find
   * @return the found package or null if it was not found.
   */
  @SuppressWarnings({"deprecation", "InvalidLink"})
  Package findPackage(String name) {
    try {
      return (Package) FIND_PACKAGE_METHOD.invoke(this, name);
    } catch (Throwable t) {
      throw new IllegalStateException(t);
    }
  }

  @SuppressWarnings("removal") // AccessController is deprecated for removal
  private static ProtectionDomain getProtectionDomain() {
    if (System.getSecurityManager() == null) {
      return InstrumentationModuleClassLoader.class.getProtectionDomain();
    }
    return java.security.AccessController.doPrivileged(
        (PrivilegedAction<ProtectionDomain>)
            ((Class<?>) InstrumentationModuleClassLoader.class)::getProtectionDomain);
  }

  private static MethodHandle getFindPackageMethod() {
    MethodType methodType = MethodType.methodType(Package.class, String.class);
    MethodHandles.Lookup lookup = MethodHandles.lookup();
    try {
      return lookup.findVirtual(ClassLoader.class, "getDefinedPackage", methodType);
    } catch (NoSuchMethodException | IllegalAccessException e) {
      // In Java 8 getDefinedPackage does not exist (HotSpot) or is not accessible (OpenJ9)
      try {
        return lookup.findVirtual(ClassLoader.class, "getPackage", methodType);
      } catch (NoSuchMethodException ex) {
        throw new IllegalStateException("expected method to always exist!", ex);
      } catch (IllegalAccessException ex2) {
        throw new IllegalStateException("Method should be accessible from here", ex2);
      }
    }
  }
}
