package com.atlassian.bitbucket.mesh.grpc;

import com.atlassian.bitbucket.mesh.health.rpc.v1.HealthGrpc;
import com.atlassian.bitbucket.mesh.management.NodeManager;
import com.atlassian.bitbucket.mesh.rpc.CapabilityServiceGrpc;
import com.atlassian.bitbucket.mesh.rpc.v1.ManagementServiceGrpc;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcError;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcErrorUtils;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcNodeUnconfigured;
import com.google.common.collect.ImmutableSet;
import io.grpc.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.lang.management.ManagementFactory;
import java.time.Duration;
import java.util.Set;

/**
 * Intercepts requests and ensures the node has been configured by the control plane before allowing requests to
 * proceed. Requests received before the node has been configured are blocked briefly, prior to being rejected,
 * to give the control plane time to send configuration. If, after that timeout period, configuration still has
 * not been applied, the request is rejected.
 */
public class RequireConfiguredServerInterceptor implements ServerInterceptor {

    static final Set<String> ALWAYS_ALLOWED = ImmutableSet.of(
            CapabilityServiceGrpc.getGetCapabilitiesMethod().getFullMethodName(),
            HealthGrpc.getCheckMethod().getFullMethodName(),
            HealthGrpc.getWatchMethod().getFullMethodName(),
            ManagementServiceGrpc.getHeartbeatMethod().getFullMethodName(),
            ManagementServiceGrpc.getProvisionNodeMethod().getFullMethodName(),
            ManagementServiceGrpc.getSidebandMethod().getFullMethodName(),
            ManagementServiceGrpc.getSetConfigurationMethod().getFullMethodName()
    );

    private static final Logger log = LoggerFactory.getLogger(RequireConfiguredServerInterceptor.class);

    private final NodeManager nodeManager;
    private final Duration timeout;

    public RequireConfiguredServerInterceptor(NodeManager nodeManager, Duration timeout) {
        this.nodeManager = nodeManager;
        this.timeout = timeout;
    }

    @Override
    public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call, Metadata headers,
                                                                 ServerCallHandler<ReqT, RespT> next) {
        if (nodeManager.isConfigured() || ALWAYS_ALLOWED.contains(call.getMethodDescriptor().getFullMethodName())) {
            return next.startCall(call, headers);
        }

        Status status = checkConfigured();
        if (status == null) {
            return next.startCall(call, headers);
        }

        //If the node has not been configured, reject the request
        call.close(status, RpcErrorUtils.toMetadata(RpcError.newBuilder()
                .setNodeUnconfigured(RpcNodeUnconfigured.newBuilder()
                        .setUptime(getUptime()))
                .build()));

        //Nothing to do; we've already rejected the request
        return new ServerCall.Listener<ReqT>() { };
    }

    private static long getUptime() {
        try {
            double uptime = ManagementFactory.getRuntimeMXBean().getUptime();

            return Math.round(uptime);
        } catch (Throwable ignored) {
            //If anything prevents getting the RuntimeMXBean or the uptime, ignore it and return
            //a marker value
            return -1L;
        }
    }

    private Status checkConfigured() {
        try {
            if (nodeManager.awaitConfiguration(timeout)) {
                return null;
            }

            //If either message is logged, here, request details will be included automatically; they
            //don't need to be manually formatted in
            log.debug("Blocked request; the node has not been configured");
            return Status.UNAVAILABLE
                    .withDescription("The node has not been configured");
        } catch (InterruptedException e) {
            log.warn("Blocked request; interrupted while awaiting configuration");
            return Status.CANCELLED
                    .withDescription("Interrupted while awaiting configuration");
        }
    }
}
