/*
 * Copyright (c) 2018 Uber Technologies, 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.uber.profiling;

import java.lang.instrument.Instrumentation;

public final class Agent {

    private static AgentImpl agentImpl = new AgentImpl();

    private Agent() {
    }

    public static void agentmain(final String args, final Instrumentation instrumentation) {
        premain(args, instrumentation);
    }

    public static void premain(final String args, final Instrumentation instrumentation) {
        System.out.println("Java Agent " + AgentImpl.VERSION + " premain args: " + args);

        Arguments arguments = Arguments.parseArgs(args);
        arguments.runConfigProvider();
        agentImpl.run(arguments, instrumentation, null);
    }
}

static {
    ClassLoader.registerAsParallelCapable();
}

/**
 * Registers the caller as parallel capable.
 * The registration succeeds if and only if all of the following
 * conditions are met:
 * <ol>
 * <li> no instance of the caller has been created</li>
 * <li> all of the super classes (except class Object) of the caller are
 * registered as parallel capable</li>
 * </ol>
 * <p>Note that once a class loader is registered as parallel capable, there
 * is no way to change it back.</p>
 *
 * @return  true if the caller is successfully registered as
 *          parallel capable and false if otherwise.
 *
 * @since   1.7
 */
@CallerSensitive
protected static boolean registerAsParallelCapable() {
    Class<? extends ClassLoader> callerClass =
        Reflection.getCallerClass().asSubclass(ClassLoader.class);
    return ParallelLoaders.register(callerClass);
}

/**
 * Encapsulates the set of parallel capable loader types.
 */
private static class ParallelLoaders {
    private ParallelLoaders() {}

    // the set of parallel capable loader types
    private static final Set<Class<? extends ClassLoader>> loaderTypes =
        Collections.newSetFromMap(
            new WeakHashMap<Class<? extends ClassLoader>, Boolean>());
    static {
        synchronized (loaderTypes) { loaderTypes.add(ClassLoader.class); }
    }

    /**
     * Registers the given class loader type as parallel capabale.
     * Returns {@code true} is successfully registered; {@code false} if
     * loader's super class is not registered.
     */
    static boolean register(Class<? extends ClassLoader> c) {
        synchronized (loaderTypes) {
            if (loaderTypes.contains(c.getSuperclass())) {
                // register the class loader as parallel capable
                // if and only if all of its super classes are.
                // Note: given current classloading sequence, if
                // the immediate super class is parallel capable,
                // all the super classes higher up must be too.
                loaderTypes.add(c);
                return true;
            } else {
                return false;
            }
        }
    }

    /**
     * Returns {@code true} if the given class loader type is
     * registered as parallel capable.
     */
    static boolean isRegistered(Class<? extends ClassLoader> c) {
        synchronized (loaderTypes) {
            return loaderTypes.contains(c);
        }
    }
}

// Maps class name to the corresponding lock object when the current
// class loader is parallel capable.
// Note: VM also uses this field to decide if the current class loader
// is parallel capable and the appropriate lock object for class loading.
private final ConcurrentHashMap<String, Object> parallelLockMap;

private ClassLoader(Void unused, ClassLoader parent) {
    this.parent = parent;
    if (ParallelLoaders.isRegistered(this.getClass())) {
        parallelLockMap = new ConcurrentHashMap<>();
        package2certs = new ConcurrentHashMap<>();
        assertionLock = new Object();
    } else {
        // no finer-grained lock; lock on the classloader instance
        parallelLockMap = null;
        package2certs = new Hashtable<>();
        assertionLock = this;
    }
}

/**
 * Returns the lock object for class loading operations.
 * For backward compatibility, the default implementation of this method
 * behaves as follows. If this ClassLoader object is registered as
 * parallel capable, the method returns a dedicated object associated
 * with the specified class name. Otherwise, the method returns this
 * ClassLoader object.
 *
 * @param  className
 *         The name of the to-be-loaded class
 *
 * @return the lock for class loading operations
 *
 * @throws NullPointerException
 *         If registered as parallel capable and <tt>className</tt> is null
 *
 * @see #loadClass(String, boolean)
 *
 * @since  1.7
 */
protected Object getClassLoadingLock(String className) {
    Object lock = this;
    if (parallelLockMap != null) {
        Object newLock = new Object();
        lock = parallelLockMap.putIfAbsent(className, newLock);
        if (lock == null) {
            lock = newLock;
        }
    }
    return lock;
}

/**
 * Loads the class with the specified <a href="#name">binary name</a>.  The
 * default implementation of this method searches for classes in the
 * following order:
 *
 * <ol>
 *
 *   <li><p> Invoke {@link #findLoadedClass(String)} to check if the class
 *   has already been loaded.  </p></li>
 *
 *   <li><p> Invoke the {@link #loadClass(String) <tt>loadClass</tt>} method
 *   on the parent class loader.  If the parent is <tt>null</tt> the class
 *   loader built-in to the virtual machine is used, instead.  </p></li>
 *
 *   <li><p> Invoke the {@link #findClass(String)} method to find the
 *   class.  </p></li>
 *
 * </ol>
 *
 * <p> If the class was found using the above steps, and the
 * <tt>resolve</tt> flag is true, this method will then invoke the {@link
 * #resolveClass(Class)} method on the resulting <tt>Class</tt> object.
 *
 * <p> Subclasses of <tt>ClassLoader</tt> are encouraged to override {@link
 * #findClass(String)}, rather than this method.  </p>
 *
 * <p> Unless overridden, this method synchronizes on the result of
 * {@link #getClassLoadingLock <tt>getClassLoadingLock</tt>} method
 * during the entire class loading process.
 *
 * @param  name
 *         The <a href="#name">binary name</a> of the class
 *
 * @param  resolve
 *         If <tt>true</tt> then resolve the class
 *
 * @return  The resulting <tt>Class</tt> object
 *
 * @throws  ClassNotFoundException
 *          If the class could not be found
 */
protected Class<?> loadClass(String name, boolean resolve)
    throws ClassNotFoundException
{
    synchronized (getClassLoadingLock(name)) {
        // First, check if the class has already been loaded
        Class<?> c = findLoadedClass(name);
        if (c == null) {
            long t0 = System.nanoTime();
            try {
                if (parent != null) {
                    c = parent.loadClass(name, false);
                } else {
                    c = findBootstrapClassOrNull(name);
                }
            } catch (ClassNotFoundException e) {
                // ClassNotFoundException thrown if class not found
                // from the non-null parent class loader
            }

            if (c == null) {
                // If still not found, then invoke findClass in order
                // to find the class.
                long t1 = System.nanoTime();
                c = findClass(name);

                // this is the defining class loader; record the stats
                sun.misc.PerfCounter.getParentDelegationTime().addTime(t1 - t0);
                sun.misc.PerfCounter.getFindClassTime().addElapsedTimeFrom(t1);
                sun.misc.PerfCounter.getFindClasses().increment();
            }
        }
        if (resolve) {
            resolveClass(c);
        }
        return c;
    }
}