package org.skywalking.apm.agent.core.plugin;

import net.bytebuddy.agent.ByteBuddyAgent;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.dynamic.loading.ByteArrayClassLoader;
import net.bytebuddy.dynamic.loading.PackageDefinitionStrategy;
import net.bytebuddy.matcher.ElementMatchers;
import org.hamcrest.CoreMatchers;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.powermock.modules.junit4.PowerMockRunner;
import org.skywalking.apm.agent.core.plugin.utility.ClassFileExtraction;

import java.lang.instrument.ClassFileTransformer;

import static net.bytebuddy.matcher.ElementMatchers.none;
import static org.hamcrest.MatcherAssert.assertThat;

@RunWith(PowerMockRunner.class)
public class AbstractClassEnhancePluginDefineTest {
    static final String WEAVE_CLASS = "org.skywalking.apm.agent.core.plugin.pluginTargetObject";
    static final String INTERCEPTOR_CLASS = "org.skywalking.apm.agent.core.plugin.MockPluginInterceptor";
    static final String WEAVE_INSTANCE_METHOD_NAME = "instanceMethod";
    static final String WEAVE_INSTANCE_WITH_EXCEPTION_METHOD_NAME = "instanceMethodWithException";
    static final String WEAVE_STATIC_METHOD_NAME = "staticMethod";
    private ClassLoader classLoader;

    @Before
    public void setUp() throws Exception {
        classLoader = new ByteArrayClassLoader.ChildFirst(getClass().getClassLoader(),
            ClassFileExtraction.of(TargetObject.class),
            null,
            ByteArrayClassLoader.PersistenceHandler.MANIFEST,
            PackageDefinitionStrategy.NoOp.INSTANCE);
    }

    @Test
    public void weaveInstanceMethod() throws Exception {
        ByteBuddyAgent.install();
        ClassFileTransformer classFileTransformer = new AgentBuilder.Default()
            .with(AgentBuilder.PoolStrategy.Default.FAST)
            .ignore(none())
            .type(ElementMatchers.is(TargetObject.class), ElementMatchers.is(classLoader)).transform(new MockTargetObjectTransformer())
            .installOnByteBuddyAgent();

        try {
            Class<?> type = classLoader.loadClass(TargetObject.class.getName());
            assertThat(type.getDeclaredMethod(WEAVE_INSTANCE_METHOD_NAME).invoke(type.getDeclaredConstructor(String.class).newInstance("a"))
                , CoreMatchers.<Object>is(WEAVE_INSTANCE_METHOD_NAME + "a"));
        } finally {
            ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
        }
    }

    @Test(expected = RuntimeException.class)
    public void weaveInstanceMethodWITEXCEPTION() throws Exception {
        ByteBuddyAgent.install();
        ClassFileTransformer classFileTransformer = new AgentBuilder.Default()
            .with(AgentBuilder.PoolStrategy.Default.FAST)
            .ignore(none())
            .type(ElementMatchers.is(TargetObject.class), ElementMatchers.is(classLoader)).transform(new MockTargetObjectTransformer())
            .installOnByteBuddyAgent();

        try {
            Class<?> type = classLoader.loadClass(TargetObject.class.getName());
            type.getDeclaredMethod(WEAVE_INSTANCE_WITH_EXCEPTION_METHOD_NAME).invoke(type.getDeclaredConstructor(String.class).newInstance("a"));
        } finally {
            ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
        }
    }

    @Test
    public void weaveStaticMethod() throws Exception {
        ByteBuddyAgent.install();
        ClassFileTransformer classFileTransformer = new AgentBuilder.Default()
            .with(AgentBuilder.PoolStrategy.Default.FAST)
            .ignore(none())
            .type(ElementMatchers.is(TargetObject.class), ElementMatchers.is(classLoader)).transform(new MockTargetObjectTransformer())
            .installOnByteBuddyAgent();

        try {
            Class<?> type = classLoader.loadClass(TargetObject.class.getName());
            assertThat(type.getDeclaredMethod(WEAVE_STATIC_METHOD_NAME).invoke(type), CoreMatchers.<Object>is(WEAVE_STATIC_METHOD_NAME + "_STATIC"));
        } finally {
            ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
        }
    }

    public static class MockTargetObjectTransformer implements AgentBuilder.Transformer {

        @Override
        public DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription typeDescription,
                                                ClassLoader classLoader) {
            try {
                DynamicType.Builder newBuilder = transformInstanceMethod(builder);
                return transformStaticMethod(newBuilder);
            } catch (Exception exception) {
                throw new AssertionError(exception);
            }
        }

        private DynamicType.Builder<?> transformStaticMethod(DynamicType.Builder newBuilder) {
            MockPluginStaticMethodInstrumentation staticMethodInstrumentation = new MockPluginStaticMethodInstrumentation();
            return staticMethodInstrumentation.define(WEAVE_CLASS, newBuilder, AbstractClassEnhancePluginDefineTest.class.getClassLoader());
        }

        private DynamicType.Builder transformInstanceMethod(DynamicType.Builder<?> builder) {
            MockPluginInstanceMethodInstrumentation instrumentation = new MockPluginInstanceMethodInstrumentation();
            return instrumentation.define(WEAVE_CLASS, builder, AbstractClassEnhancePluginDefineTest.class.getClassLoader());
        }
    }

}