/*
 * Copyright 2019 NAVER Corp.
 *
 * 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.navercorp.pinpoint.collector.receiver.grpc;

import com.google.protobuf.Empty;
import com.google.protobuf.GeneratedMessageV3;
import com.navercorp.pinpoint.collector.cluster.GrpcAgentConnection;
import com.navercorp.pinpoint.collector.cluster.ProfilerClusterManager;
import com.navercorp.pinpoint.collector.util.RequestManager;
import com.navercorp.pinpoint.common.server.cluster.ClusterKey;
import com.navercorp.pinpoint.grpc.MessageFormatUtils;
import com.navercorp.pinpoint.grpc.trace.PCmdActiveThreadCount;
import com.navercorp.pinpoint.grpc.trace.PCmdActiveThreadDump;
import com.navercorp.pinpoint.grpc.trace.PCmdActiveThreadLightDump;
import com.navercorp.pinpoint.grpc.trace.PCmdEcho;
import com.navercorp.pinpoint.grpc.trace.PCmdRequest;
import com.navercorp.pinpoint.grpc.trace.PCmdResponse;
import com.navercorp.pinpoint.io.ResponseMessage;
import com.navercorp.pinpoint.rpc.PinpointSocketException;
import com.navercorp.pinpoint.rpc.common.SocketState;
import com.navercorp.pinpoint.rpc.common.SocketStateChangeResult;
import com.navercorp.pinpoint.rpc.common.SocketStateCode;
import com.navercorp.pinpoint.rpc.packet.stream.StreamCode;
import com.navercorp.pinpoint.rpc.packet.stream.StreamResponsePacket;
import com.navercorp.pinpoint.rpc.stream.ClientStreamChannel;
import com.navercorp.pinpoint.rpc.stream.ClientStreamChannelEventHandler;
import com.navercorp.pinpoint.rpc.stream.StreamChannel;
import com.navercorp.pinpoint.rpc.stream.StreamChannelRepository;
import com.navercorp.pinpoint.rpc.stream.StreamException;
import com.navercorp.pinpoint.thrift.dto.command.TRouteResult;
import com.navercorp.pinpoint.thrift.io.CommandHeaderTBaseSerializerFactory;
import com.navercorp.pinpoint.thrift.sender.message.CommandGrpcToThriftMessageConverter;
import com.navercorp.pinpoint.thrift.util.SerializationUtils;
import io.grpc.Status;
import io.grpc.StatusException;
import io.grpc.stub.StreamObserver;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.apache.thrift.TBase;
import org.apache.thrift.TException;

import java.net.InetSocketAddress;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @author Taejin Koo
 */
public class PinpointGrpcServer {

    private final Logger logger = LogManager.getLogger(this.getClass());
    private final boolean isInfo = logger.isInfoEnabled();

    private final SocketState state = new SocketState();
    private final AtomicReference<List<Integer>> supportCommandServiceList = new AtomicReference<>();

    private final CommandGrpcToThriftMessageConverter messageConverter = new CommandGrpcToThriftMessageConverter();
    private final CommandHeaderTBaseSerializerFactory commandHeaderTBaseSerializerFactory = CommandHeaderTBaseSerializerFactory.getDefaultInstance();

    private final StreamChannelRepository streamChannelRepository = new StreamChannelRepository();

    private final InetSocketAddress remoteAddress;
    private final ClusterKey clusterKey;
    private final RequestManager<ResponseMessage> requestManager;
    private final ProfilerClusterManager profilerClusterManager;
    private final StreamObserver<PCmdRequest> requestObserver;

    private volatile Runnable onCloseHandler;

    public PinpointGrpcServer(InetSocketAddress remoteAddress, ClusterKey clusterKey, RequestManager<ResponseMessage> requestManager, ProfilerClusterManager profilerClusterManager, StreamObserver<PCmdRequest> requestObserver) {
        this.remoteAddress = Objects.requireNonNull(remoteAddress, "remoteAddress");
        this.clusterKey = Objects.requireNonNull(clusterKey, "clusterKey");
        this.requestManager = Objects.requireNonNull(requestManager, "requestManager");
        this.profilerClusterManager = Objects.requireNonNull(profilerClusterManager, "profilerClusterManager");
        this.requestObserver = Objects.requireNonNull(requestObserver, "requestObserver");
    }

    public void setOnCloseHandler(Runnable onCloseHandler) {
        this.onCloseHandler = onCloseHandler;
    }

    public void connected() {
        toState(SocketStateCode.CONNECTED);
    }

    public boolean handleHandshake(List<Integer> supportCommandServiceList) {
        if (isInfo) {
            logger.info("{} handleHandshake() started. data:{}", clusterKey, supportCommandServiceList);
        }

        boolean isFirst = this.supportCommandServiceList.compareAndSet(null, supportCommandServiceList);
        if (isFirst) {
            toState(SocketStateCode.RUN_WITHOUT_HANDSHAKE);
            SocketStateChangeResult socketStateChangeResult = toState(SocketStateCode.RUN_DUPLEX);
            return socketStateChangeResult.isChange();
        }
        return false;
    }

    private SocketStateChangeResult toState(SocketStateCode socketStateCode) {
        SocketStateChangeResult result = state.to(socketStateCode);
        if (result.isChange()) {
            if (SocketStateCode.RUN_DUPLEX == socketStateCode) {
                GrpcAgentConnection grpcAgentConnection = new GrpcAgentConnection(this, supportCommandServiceList.get());
                profilerClusterManager.register(grpcAgentConnection);
            } else if (SocketStateCode.isClosed(socketStateCode)) {
                GrpcAgentConnection grpcAgentConnection = new GrpcAgentConnection(this, Collections.emptyList());
                profilerClusterManager.unregister(grpcAgentConnection);
            }
        } else {
            logger.warn("Failed to change state. agent:{}, result:{}", clusterKey, result);
        }

        logger.info("transition grpc server of {}: {}", clusterKey, result);
        return result;
    }

    public CompletableFuture<ResponseMessage> request(GeneratedMessageV3 message) {
        if (!state.checkState(SocketStateCode.RUN_DUPLEX)) {
            return createFailedFuture(new IllegalStateException("failed to request message. caused:illegal State"));
        }

        PCmdRequest request = createRequest(message);
        if (request == null) {
            return createFailedFuture(new PinpointSocketException(TRouteResult.NOT_SUPPORTED_REQUEST.name()));
        }

        CompletableFuture<ResponseMessage> future = requestManager.register(request.getRequestId());
        requestObserver.onNext(request);
        return future;
    }

    // 1st message : client(collector) -> server(agent)
    public ClientStreamChannel openStream(GeneratedMessageV3 message, ClientStreamChannelEventHandler streamChannelEventHandler) throws StreamException {
        if (!state.checkState(SocketStateCode.RUN_DUPLEX)) {
            throw new StreamException(StreamCode.STATE_NOT_CONNECTED);
        }

        PCmdRequest request = createRequest(message);
        if (request == null) {
            throw new StreamException(StreamCode.TYPE_UNSUPPORT);
        }

        GrpcClientStreamChannel grpcClientStreamChannel = new GrpcClientStreamChannel(remoteAddress, request.getRequestId(), streamChannelRepository, streamChannelEventHandler);
        try {
            grpcClientStreamChannel.init();
            grpcClientStreamChannel.connect(new Runnable() {
                @Override
                public void run() {
                    requestObserver.onNext(request);
                }
            }, 1000);
        } catch (StreamException e) {
            grpcClientStreamChannel.close(e.getStreamCode());
            throw e;
        }
        return grpcClientStreamChannel;
    }

    private PCmdRequest createRequest(GeneratedMessageV3 message) {
        PCmdRequest.Builder requestBuilder = PCmdRequest.newBuilder();

        final int requestId = requestManager.nextRequestId();
        requestBuilder.setRequestId(requestId);

        if (message instanceof PCmdEcho) {
            requestBuilder.setCommandEcho((PCmdEcho) message);
            return requestBuilder.build();
        } else if (message instanceof PCmdActiveThreadCount) {
            requestBuilder.setCommandActiveThreadCount((PCmdActiveThreadCount) message);
            return requestBuilder.build();
        } else if (message instanceof PCmdActiveThreadDump) {
            requestBuilder.setCommandActiveThreadDump((PCmdActiveThreadDump) message);
            return requestBuilder.build();
        } else if (message instanceof PCmdActiveThreadLightDump) {
            requestBuilder.setCommandActiveThreadLightDump((PCmdActiveThreadLightDump) message);
            return requestBuilder.build();
        } else {
            return null;
        }
    }

    public void handleMessage(int responseId, GeneratedMessageV3 message) {
        if (!state.checkState(SocketStateCode.RUN_DUPLEX)) {
            logger.warn("failed to handle message. caused:illegal State");
            return;
        }

        if (isInfo) {
            logger.info("{} handleMessage:{}", clusterKey, MessageFormatUtils.debugLog(message));
        }

        final CompletableFuture<ResponseMessage> responseFuture = requestManager.messageReceived(responseId, clusterKey::format);
        if (responseFuture == null) {
            logger.warn("Response future is timeout responseId:{}", responseId);
            return;
        }
        try {
            TBase<?, ?> tMessage = messageConverter.toMessage(message);
            byte[] serialize = SerializationUtils.serialize(tMessage, commandHeaderTBaseSerializerFactory);

            ResponseMessage responseMessage = ResponseMessage.wrap(serialize);
            responseFuture.complete(responseMessage);
        } catch (TException e) {
            responseFuture.completeExceptionally(new PinpointSocketException(e.getMessage()));
        }
    }

    public void handleFail(PCmdResponse failMessage) {
        String message = failMessage.getMessage().getValue();
        if (message != null) {
            setFailMessageToFuture(failMessage.getResponseId(), message);
        } else {
            setFailMessageToFuture(failMessage.getResponseId(), "failed to route message");
        }
    }

    // 2nd message : server(agent) -> client(collector)
    public boolean handleStreamCreateMessage(int streamId, StreamObserver<Empty> connectionObserver) {
        if (isInfo) {
            logger.info("handleStreamCreateMessage. streamId:{}", streamId);
        }

        StreamChannel streamChannel = streamChannelRepository.getStreamChannel(streamId);
        if (streamChannel == null) {
            logger.warn("Can't find suitable streamChannel.(streamId:{})", streamId);
            return false;
        }

        GrpcClientStreamChannel grpcClientStreamChannel = (GrpcClientStreamChannel) streamChannel;
        grpcClientStreamChannel.setConnectionObserver(connectionObserver);
        return grpcClientStreamChannel.changeStateConnected();
    }

    public void handleStreamMessage(int streamId, GeneratedMessageV3 message) throws StreamException {
        if (logger.isDebugEnabled()) {
            logger.debug("handleStreamMessage() streamId:{}, message:{}", streamId, MessageFormatUtils.debugLog(message));
        }

        GrpcClientStreamChannel streamChannel = (GrpcClientStreamChannel) streamChannelRepository.getStreamChannel(streamId);
        if (streamChannel == null) {
            throw new StreamException(StreamCode.ID_NOT_FOUND, "Can't find suitable streamChannel.(streamId:" + streamId + ")");
        }

        TBase<?, ?> tBase = messageConverter.toMessage(message);
        if (tBase == null) {
            throw new StreamException(StreamCode.TYPE_ERROR, "Failed to convert message.(message:" + MessageFormatUtils.debugLog(message).toString() + ")");
        }

        if (logger.isDebugEnabled()) {
            logger.debug("handleStreamMessage() streamId:{}, message:{}", streamId, tBase);
        }
        try {
            byte[] serialize = SerializationUtils.serialize(tBase, commandHeaderTBaseSerializerFactory);
            streamChannel.handleStreamResponsePacket(new StreamResponsePacket(streamId, serialize));
        } catch (TException t) {
            throw new StreamException(StreamCode.TYPE_UNKNOWN, "Failed to serialize message.(tBase:" + tBase + ")");
        } catch (StreamException e) {
            throw e;
        }
    }

    public void handleStreamDisconnected(int streamId) {
        handleStreamDisconnected(streamId, null);
    }

    public void handleStreamDisconnected(int streamId, Throwable t) {
        GrpcClientStreamChannel streamChannel = (GrpcClientStreamChannel) streamChannelRepository.getStreamChannel(streamId);
        if (streamChannel == null) {
            logger.warn("Can't find suitable streamChannel.(streamId:{})", streamId);
            return;
        }

        if (t == null) {
            streamChannel.disconnect(StreamCode.STATE_CLOSED);
        } else {
            streamChannel.disconnect(StreamCode.CONNECTION_ERRROR);
        }
    }

    public void disconnected() {
        close(SocketStateCode.BEING_CLOSE_BY_CLIENT);
    }

    public void close() {
        close(SocketStateCode.BEING_CLOSE_BY_SERVER);
    }

    public void close(SocketStateCode toState) {
        logger.info("close() will be started. ( remoteAddress:{}, clusterKey:{}, closeState:{}", remoteAddress, clusterKey, toState);

        if (onCloseHandler != null) {
            onCloseHandler.run();
        }

        synchronized (this) {
            try {
                SocketStateCode currentState = getState();
                if (SocketStateCode.isRun(currentState)) {
                    if (toState == SocketStateCode.BEING_CLOSE_BY_SERVER || toState == SocketStateCode.BEING_CLOSE_BY_CLIENT) {
                        toState(toState);
                        requestObserver.onCompleted();
                    } else {
                        toState(toState);
                        requestObserver.onError(new StatusException(Status.UNKNOWN));
                    }
                }

            } catch (Exception e) {
                // It could throw exception when requestObserver is already completed.
                logger.warn("Exception occurred while requestObserver invokes onCompleted. message:{}", e.getMessage(), e);
            }

            try {
                SocketStateCode currentStateCode = getState();
                if (SocketStateCode.BEING_CLOSE_BY_SERVER == currentStateCode) {
                    toState(SocketStateCode.CLOSED_BY_SERVER);
                } else if (SocketStateCode.BEING_CLOSE_BY_CLIENT == currentStateCode) {
                    toState(SocketStateCode.CLOSED_BY_CLIENT);
                } else if (SocketStateCode.isClosed(currentStateCode)) {
                    logger.warn("stop(). Socket has closed state({}).", currentStateCode);
                } else {
                    toState(SocketStateCode.ERROR_ILLEGAL_STATE_CHANGE);
                    logger.warn("stop(). Socket has unexpected state({})", currentStateCode);
                }
            } finally {
                logger.info("{} <=> local all streamChannels will be close.", clusterKey);
                streamChannelRepository.close(StreamCode.STATE_CLOSED);
            }
        }

    }

    private void setFailMessageToFuture(int responseId, String message) {
        CompletableFuture<ResponseMessage> future = requestManager.removeMessageFuture(responseId);
        if (future != null) {
            future.completeExceptionally(new PinpointSocketException(message));
        }
    }

    public InetSocketAddress getRemoteAddress() {
        return remoteAddress;
    }

    public SocketStateCode getState() {
        return state.getCurrentState();
    }

    public ClusterKey getClusterKey() {
        return clusterKey;
    }

    public CompletableFuture<ResponseMessage> createFailedFuture(Exception failException) {
        CompletableFuture<ResponseMessage> failedFuture = new CompletableFuture<>();
        failedFuture.completeExceptionally(failException);
        return failedFuture;
    }
}