package com.atlassian.bitbucket.mesh.grpc;

import com.atlassian.bitbucket.mesh.app.LifecyclePhases;
import com.atlassian.bitbucket.mesh.config.ConfigurationUpdatedEvent;
import com.atlassian.bitbucket.mesh.config.GrpcConfig;
import com.atlassian.bitbucket.mesh.event.EventListener;
import com.atlassian.bitbucket.mesh.event.EventPublisher;
import com.atlassian.bitbucket.mesh.spring.AbstractSmartLifecycle;
import com.atlassian.bitbucket.mesh.util.ConcurrencyUtils;
import com.google.common.collect.ImmutableSet;
import com.google.common.util.concurrent.MoreExecutors;
import io.grpc.*;
import io.grpc.internal.ServerImpl;
import io.grpc.netty.NettyServerBuilder;
import io.netty.util.internal.PlatformDependent;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.io.File;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.net.InetSocketAddress;
import java.nio.file.Path;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import static com.atlassian.bitbucket.mesh.MeshConstants.EXECUTOR_RPC_SERVER;
import static com.atlassian.bitbucket.mesh.rpc.ManagedConfigurationProperties.*;

public class GrpcServer extends AbstractSmartLifecycle {

    private static final Set<String> CONFIG_KEYS_REQUIRING_RESTART = ImmutableSet.of(
            GRPC_MAX_MESSAGE_SIZE, GRPC_MAX_METADATA_SIZE, GRPC_SERVER_KEEP_ALIVE_TIME,
            GRPC_SERVER_KEEP_ALIVE_WITHOUT_CALLS);

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

    private final GrpcConfig config;
    private final EventPublisher eventPublisher;
    private final ExecutorService executor;
    private final ServerInterceptor[] interceptors;
    private final BindableService[] services;

    private volatile NettyServerBuilder lastSuccessfulServerBuilder;
    private volatile Server server;
    private boolean sslEnabled;

    public GrpcServer(GrpcConfig config, EventPublisher eventPublisher,
                      ServerInterceptor[] interceptors, BindableService[] services) {
        this.config = config;
        this.eventPublisher = eventPublisher;
        this.interceptors = interceptors;
        this.services = services;

        executor = Executors.newCachedThreadPool(ConcurrencyUtils.createDaemonThreadFactory(EXECUTOR_RPC_SERVER));
    }

    @Override
    public int getPhase() {
        return LifecyclePhases.GRPC;
    }

    @Nonnull
    public GrpcStatistics getStatistics() {
        InternalChannelz.ServerStats stats;
        try {
            //This is an "internal" API, which in theory we shouldn't use. Unfortunately, gRPC doesn't
            //have any public API that exposes these statistics and we can't provide our own CallTracer
            //because that's package-private internal API (and can't be set on the builder anyway).
            //
            //While this API returns a ListenableFuture, the implementation uses a simple SettableFuture
            //and always returns it complete--so calling get() here will never block.
            stats = ((ServerImpl) server).getStats().get();
        } catch (ExecutionException e) {
            //This shouldn't be possible, with the existing implementation, but the API mandates it
            throw new IllegalStateException("Failed to load statistics", e.getCause());
        } catch (InterruptedException e) {
            //This shouldn't be possible either
            throw new IllegalStateException("Timed out loading statistics", e);
        }

        return new GrpcStatistics() {

            @Override
            public long getFailedRequests() {
                return stats.callsFailed;
            }

            @Override
            public long getRunningRequests() {
                return stats.callsStarted - stats.callsSucceeded - stats.callsFailed;
            }

            @Override
            public long getSuccessfulRequests() {
                return stats.callsSucceeded;
            }

            @Override
            public long getTotalRequests() {
                return stats.callsStarted;
            }
        };
    }

    @EventListener
    public void onConfigurationUpdated(ConfigurationUpdatedEvent event) {
        if (config.isServerRestartOnConfigurationChange() &&
                event.getKeys().stream().anyMatch(CONFIG_KEYS_REQUIRING_RESTART::contains)) {
            executor.execute(this::restart);
        }
    }

    @Override
    public void start() {
        startServer();
        super.start();
    }

    @Override
    public void stop() {
        super.stop();

        try {
            stopServer();
        } finally {
            //Shutdown the executor after shutting down the server
            if (MoreExecutors.shutdownAndAwaitTermination(executor, 10L, TimeUnit.SECONDS)) {
                log.debug("The request executor has shutdown");
            } else {
                log.warn("Failed to shutdown the request executor; there may be active requests");
            }
        }
    }

    private void configureTransportSecurity(ServerBuilder<?> builder) {
        Path certChainPath = config.getServerSslCertChainPath();
        Path privateKeyPath = config.getServerSslPrivateKeyPath();
        if (certChainPath != null && privateKeyPath != null) {
            File certChain = certChainPath.toFile();
            File privateKey = privateKeyPath.toFile();
            if (certChain.isFile() && privateKey.isFile()) {
                builder.useTransportSecurity(certChain, privateKey);
                sslEnabled = true;
                return;
            }
            if (!certChain.isFile()) {
                log.warn("Unable to configure SSL for gRPC server due to problem loading certificate chain from {}",
                        certChainPath);
            }
            if (!privateKey.isFile()) {
                log.warn("Unable to configure SSL for gRPC server due to problem loading private key from {}",
                        privateKeyPath);
            }
        }
        sslEnabled = false;
    }

    private void restart() {
        log.info("Restarting gRPC server");
        stopServer();
        int maxAttempts = 3;
        for (int i = 1; i <= maxAttempts; ++i) {
            try {
                startServer();
                return;
            } catch (RuntimeException e) {
                log.warn("gRPC server restart failed (attempt {}/{})", i, maxAttempts, e);
                if (i != maxAttempts) {
                    try {
                        Thread.sleep(15_000L);
                    } catch (InterruptedException interruptedException) {
                        Thread.currentThread().interrupt();
                        break;
                    }
                } else {
                    if (lastSuccessfulServerBuilder != null) {
                        log.warn("gRPC server failed to start up. Attempting to revert to the previous configuration");

                        try {
                            server = lastSuccessfulServerBuilder.build().start();
                        } catch (IOException innerException) {
                            log.warn("Failed to revert the gRPC server to the previous configuration");
                        }
                    }
                }
            }
        }
    }

    private void startServer() {
        Server current = server;
        if (current != null && !current.isShutdown()) {
            log.info("gRPC server is already running");
            return;
        }

        //We only want the agent to be able to only listen on localhost, so we need to use NettyServerBuilder directly.
        //If we use ServerBuilder.forPort it binds on every local address, not just localhost.
        int port = config.getServerPort();
        NettyServerBuilder builder =
                NettyServerBuilder.forAddress(new InetSocketAddress(config.getServerAddress(), port))
                        .executor(executor)
                        .maxInboundMessageSize(config.getMaxMessageSize())
                        .maxInboundMetadataSize(config.getMaxMetadataSize())
                        //See https://github.com/grpc/proposal/blob/master/A8-client-side-keepalive.md for more details
                        //on what these settings mean and how the client and server settings need to be in tune with
                        //each other
                        .permitKeepAliveTime(config.getKeepAliveTime().toMillis(), TimeUnit.MILLISECONDS)
                        .permitKeepAliveWithoutCalls(config.isKeepAliveWithoutCallsAllowed());
        for (ServerInterceptor interceptor : interceptors) {
            builder.intercept(interceptor);
        }
        for (BindableService service : services) {
            builder.addService(service.bindService());
        }
        configureTransportSecurity(builder);

        try {
            server = builder.build().start();
            log.info("gRPC server started (port: {}, SSL: {}, maxDirectMemory: {} MB)", port, sslEnabled,
                    PlatformDependent.maxDirectMemory() / (1024 * 1024L));
            //Remember the last successful server builder, so we can revert to it if startup after a config change fails
            lastSuccessfulServerBuilder = builder;
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }
    }

    private void stopServer() {
        Server server = this.server;
        if (server != null) {
            try {
                server.shutdown();

                //Signal that the gRPC server is going to shut down to allow long running RPCs to wrap up
                eventPublisher.publish(new GrpcServerShutdownEvent());

                if (server.awaitTermination(config.getServerShutdownTimeout().getSeconds(), TimeUnit.SECONDS)) {
                    log.debug("gRPC server has shutdown gracefully");
                } else {
                    log.warn("gRPC server did not shutdown within the timeout; forcing shutdown");

                    server.shutdownNow();
                    if (server.awaitTermination(5L, TimeUnit.SECONDS)) {
                        log.debug("gRPC server has been forced to shutdown");
                    } else {
                        log.warn("gRPC server did not shutdown; it will be abandoned");
                    }
                }
                this.server = null;
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }
    }
}
