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

import com.google.protobuf.InvalidProtocolBufferException;
import io.grpc.stub.StreamObserver;
import io.grpc.testing.GrpcServerRule;
import java.util.ArrayList;
import java.util.List;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.powermock.reflect.Whitebox;
import org.skywalking.apm.agent.core.boot.ServiceManager;
import org.skywalking.apm.agent.core.conf.RemoteDownstreamConfig;
import org.skywalking.apm.agent.core.context.ContextManager;
import org.skywalking.apm.agent.core.context.tag.Tags;
import org.skywalking.apm.agent.core.context.trace.AbstractSpan;
import org.skywalking.apm.agent.core.context.trace.SpanLayer;
import org.skywalking.apm.agent.core.test.tools.AgentServiceRule;
import org.skywalking.apm.agent.core.test.tools.SegmentStorage;
import org.skywalking.apm.agent.core.test.tools.SegmentStoragePoint;
import org.skywalking.apm.agent.core.test.tools.TracingSegmentRunner;
import org.skywalking.apm.network.proto.Downstream;
import org.skywalking.apm.network.proto.SpanObject;
import org.skywalking.apm.network.proto.SpanType;
import org.skywalking.apm.network.proto.TraceSegmentObject;
import org.skywalking.apm.network.proto.TraceSegmentServiceGrpc;
import org.skywalking.apm.network.proto.UpstreamSegment;
import org.skywalking.apm.network.trace.component.ComponentsDefine;

import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.mockito.Mockito.spy;

@RunWith(TracingSegmentRunner.class)
public class TraceSegmentServiceClientTest {

    @Rule
    public AgentServiceRule agentServiceRule = new AgentServiceRule();

    @Rule
    public GrpcServerRule grpcServerRule = new GrpcServerRule().directExecutor();

    @SegmentStoragePoint
    private SegmentStorage storage;

    private TraceSegmentServiceClient serviceClient = new TraceSegmentServiceClient();
    private List<UpstreamSegment> upstreamSegments;

    private TraceSegmentServiceGrpc.TraceSegmentServiceImplBase serviceImplBase = new TraceSegmentServiceGrpc.TraceSegmentServiceImplBase() {
        @Override
        public StreamObserver<UpstreamSegment> collect(final StreamObserver<Downstream> responseObserver) {
            return new StreamObserver<UpstreamSegment>() {
                @Override
                public void onNext(UpstreamSegment value) {
                    upstreamSegments.add(value);
                }

                @Override
                public void onError(Throwable t) {
                }

                @Override
                public void onCompleted() {
                    responseObserver.onNext(Downstream.getDefaultInstance());
                    responseObserver.onCompleted();
                }
            };
        }
    };

    @BeforeClass
    public static void setUpBeforeClass() {
        RemoteDownstreamConfig.Agent.APPLICATION_ID = 1;
        RemoteDownstreamConfig.Agent.APPLICATION_INSTANCE_ID = 1;
    }

    @Before
    public void setUp() throws Throwable {
        Whitebox.setInternalState(ServiceManager.INSTANCE.findService(GRPCChannelManager.class), "reconnect", false);
        spy(serviceClient);

        Whitebox.setInternalState(serviceClient, "serviceStub",
            TraceSegmentServiceGrpc.newStub(grpcServerRule.getChannel()));
        Whitebox.setInternalState(serviceClient, "status", GRPCChannelStatus.CONNECTED);

        upstreamSegments = new ArrayList<UpstreamSegment>();
    }

    @Test
    public void testSendTraceSegmentWithoutException() throws InvalidProtocolBufferException {
        grpcServerRule.getServiceRegistry().addService(serviceImplBase);

        AbstractSpan firstEntrySpan = ContextManager.createEntrySpan("/testFirstEntry", null);
        firstEntrySpan.setComponent(ComponentsDefine.TOMCAT);
        Tags.HTTP.METHOD.set(firstEntrySpan, "GET");
        Tags.URL.set(firstEntrySpan, "127.0.0.1:8080");
        SpanLayer.asHttp(firstEntrySpan);
        ContextManager.stopSpan();

        serviceClient.consume(storage.getTraceSegments());

        assertThat(upstreamSegments.size(), is(1));
        UpstreamSegment upstreamSegment = upstreamSegments.get(0);
        assertThat(upstreamSegment.getGlobalTraceIdsCount(), is(1));
        TraceSegmentObject traceSegmentObject = TraceSegmentObject.parseFrom(upstreamSegment.getSegment());
        assertThat(traceSegmentObject.getRefsCount(), is(0));
        assertThat(traceSegmentObject.getSpansCount(), is(1));

        SpanObject spanObject = traceSegmentObject.getSpans(0);
        assertThat(spanObject.getSpanType(), is(SpanType.Entry));
        assertThat(spanObject.getSpanId(), is(0));
        assertThat(spanObject.getParentSpanId(), is(-1));
    }

    @Test
    public void testSendTraceSegmentWithException() throws InvalidProtocolBufferException {
        grpcServerRule.getServiceRegistry().addService(serviceImplBase);

        AbstractSpan firstEntrySpan = ContextManager.createEntrySpan("/testFirstEntry", null);
        firstEntrySpan.setComponent(ComponentsDefine.TOMCAT);
        Tags.HTTP.METHOD.set(firstEntrySpan, "GET");
        Tags.URL.set(firstEntrySpan, "127.0.0.1:8080");
        SpanLayer.asHttp(firstEntrySpan);
        ContextManager.stopSpan();
        grpcServerRule.getServer().shutdownNow();
        serviceClient.consume(storage.getTraceSegments());

        assertThat(upstreamSegments.size(), is(0));

        boolean reconnect = Whitebox.getInternalState(ServiceManager.INSTANCE.findService(GRPCChannelManager.class), "reconnect");
        assertThat(reconnect, is(true));

    }
}
