package com.criteo.hadoop.garmadon.agent.tracers.hadoop.resourcemanager;

import com.criteo.hadoop.garmadon.agent.tracers.MethodTracer;
import com.criteo.hadoop.garmadon.agent.tracers.Tracer;
import com.criteo.hadoop.garmadon.agent.utils.AgentAttachmentRule;
import com.criteo.hadoop.garmadon.agent.utils.ClassFileExtraction;
import com.criteo.hadoop.garmadon.schema.events.Header;
import net.bytebuddy.agent.ByteBuddyAgent;
import net.bytebuddy.dynamic.loading.ByteArrayClassLoader;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.yarn.api.records.*;
import org.apache.hadoop.yarn.event.Dispatcher;
import org.apache.hadoop.yarn.event.EventHandler;
import org.apache.hadoop.yarn.server.resourcemanager.RMContext;
import org.apache.hadoop.yarn.server.resourcemanager.ahs.RMApplicationHistoryWriter;
import org.apache.hadoop.yarn.server.resourcemanager.metrics.SystemMetricsPublisher;
import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainerEvent;
import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainerEventType;
import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainerImpl;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.MethodRule;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Constructor;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;

import static org.hamcrest.CoreMatchers.instanceOf;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.Assert.assertNotNull;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

public class RMContainerTracerTest {
    @Rule
    public MethodRule agentAttachmentRule = new AgentAttachmentRule();

    private ClassLoader classLoader;

    @Before
    public void setUp() throws IOException {
        classLoader = new ByteArrayClassLoader.ChildFirst(getClass().getClassLoader(),
            ClassFileExtraction.of(
                Tracer.class,
                MethodTracer.class,
                RMContainerTracer.class,
                RMContainerTracer.RMContainerImplTracer.class
            ),
            ByteArrayClassLoader.PersistenceHandler.MANIFEST);
    }

    @Test
    @AgentAttachmentRule.Enforce
    public void ContainerResourceMonitoringModule_should_attach_to_recordCpuUsage() throws ClassNotFoundException, NoSuchMethodException, IllegalAccessException, InstantiationException, InvocationTargetException, NoSuchFieldException {
        assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));

        final Header[] header = new Header[1];
        final Object[] event = new Object[1];
        RMContainerTracer.initEventHandler((t, h, o) -> {
            header[0] = h;
            event[0] = o;
        });

        ClassFileTransformer classFileTransformer = new RMContainerTracer.RMContainerImplTracer().installOnByteBuddyAgent();

        try {
            Class<?> clazz = classLoader.loadClass(RMContainerImpl.class.getName());
            Constructor<?> constructor = clazz.getDeclaredConstructor(Container.class,
                ApplicationAttemptId.class, NodeId.class, String.class,
                RMContext.class);
            constructor.setAccessible(true);

            RMContext rmContext = mock(RMContext.class);
            Configuration yarnConf = new Configuration();
            when(rmContext.getYarnConfiguration())
                .thenReturn(yarnConf);
            Dispatcher dispatcher = mock(Dispatcher.class);
            EventHandler eventHandler = mock(EventHandler.class);
            when(dispatcher.getEventHandler())
                .thenReturn(eventHandler);
            when(rmContext.getDispatcher())
                .thenReturn(dispatcher);

            RMApplicationHistoryWriter rmApplicationHistoryWriter = mock(RMApplicationHistoryWriter.class);
            when(rmContext.getRMApplicationHistoryWriter())
                .thenReturn(rmApplicationHistoryWriter);

            SystemMetricsPublisher systemMetricsPublisher = mock(SystemMetricsPublisher.class);
            when(rmContext.getSystemMetricsPublisher())
                .thenReturn(systemMetricsPublisher);

            Object inFormat = constructor.newInstance(mock(Container.class), mock(ApplicationAttemptId.class), mock(NodeId.class), "user", rmContext);

            Method m = clazz.getDeclaredMethod("handle", RMContainerEvent.class);

            ContainerId cid = mock(ContainerId.class);
            ApplicationAttemptId applicationAttemptId = mock(ApplicationAttemptId.class);
            when(cid.toString())
                .thenReturn("cid");
            when(cid.getApplicationAttemptId())
                .thenReturn(applicationAttemptId);
            when(applicationAttemptId.toString())
                .thenReturn("appattempt_id");
            ApplicationId applicationId = mock(ApplicationId.class);
            when(applicationAttemptId.getApplicationId())
                .thenReturn(applicationId);
            when(applicationId.toString())
                .thenReturn("app_id");
            m.invoke(inFormat, new RMContainerEvent(cid, RMContainerEventType.START));

            assertNotNull(header[0]);
            assertNotNull(event[0]);

        } finally {
            ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
        }
    }

}