package com.atlassian.bitbucket.mesh.management;

import com.atlassian.bitbucket.mesh.event.EventListener;
import com.atlassian.bitbucket.mesh.grpc.GrpcServerShutdownEvent;
import com.atlassian.bitbucket.mesh.request.RequestContext;
import com.atlassian.bitbucket.mesh.request.RequestManager;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcHeartbeat;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcSidebandFollowerFragment;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcSidebandLeaderFragment;
import com.google.common.annotations.VisibleForTesting;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.util.ArrayDeque;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicBoolean;

import static java.util.Objects.requireNonNull;
import static org.apache.commons.lang3.StringUtils.defaultString;

public class DefaultSidebandManager implements SidebandManager {

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

    //This queue (effectively) orders sideband channels by when they're opened. The first channel in
    //the queue has been open the longest, and will be the first channel we try to send to. This way
    //Mesh at least _tries_ to send messages to a consistent control plane node
    private final Queue<SidebandChannel> channels;
    private final Queue<RpcSidebandFollowerFragment> queue;
    private final int queueLimit;
    private final RequestManager requestManager;

    public DefaultSidebandManager(RequestManager requestManager, int queueLimit) {
        this.queueLimit = queueLimit;
        this.requestManager = requestManager;

        channels = new ConcurrentLinkedQueue<>();
        queue = new ArrayDeque<>();
    }

    @EventListener
    public void onGrpcServerShutdown(GrpcServerShutdownEvent event) {
        closeChannels();
    }

    @Nonnull
    @Override
    public SidebandReceiver open(@Nonnull OpenSidebandParameters parameters) {
        requireNonNull(parameters, "parameters");

        //Since sideband channels are expected to be long-running RPC requests, it's required that an
        //active request exist when one is opened
        RequestContext requestContext = requestManager.getContextOrThrow();
        //While both getAgent and getRemoteAddress are @Nullable, in practice they won't both be null.
        //We use whichever is present as an identifier for the connected control plane node
        String id = defaultString(requestContext.getAgent(), requestContext.getRemoteAddress());

        SidebandChannel channel = new SidebandChannel(id, parameters.getSender());
        synchronized (channels) {
            //When a new channel is connected, check the queue. This is done while holding a lock to ensure
            //send isn't queuing new fragments while we're trying to send them all
            int queueSize = queue.size();
            if (queueSize > 0) {
                log.info("{}: Sending {} queued fragments", id, queueSize);
                RpcSidebandFollowerFragment fragment;
                while ((fragment = queue.poll()) != null) {
                    channel.send(fragment);
                }
            }

            //Add the channel _after_ draining the queue. This ensures there's no concurrency between trying
            //to drain the queue here and send trying to send another fragment
            channels.add(channel);
            log.info("{}: Sideband connected", id);
        }

        return channel;
    }

    @Override
    public boolean send(@Nonnull RpcSidebandFollowerFragment fragment) {
        int attempted = 0;
        for (SidebandChannel channel : channels) {
            if (channel.send(fragment)) {
                return true;
            }
            ++attempted;
        }

        if (attempted == 0) {
            //If no sideband channels are open, try to queue the fragment and send it when one is opened.
            //The primary use case for this is when the control plane sends inconsistent replicas during
            //initial configuration, triggering repair. If repair completes before the sideband is opened
            //we need to queue the consistent observation and send it when we can
            synchronized (channels) {
                SidebandChannel channel = channels.peek();
                if (channel == null) {
                    //The queue size is limited to prevent excessive memory usage, so if we can't queue
                    //the fragment we need to log that we're dropping it
                    if (queue.size() < queueLimit) {
                        queue.add(fragment);

                        log.info("Queued {} fragment; no sideband channels are connected",
                                fragment.getSidebandOneofCase());
                        return true; //Consider queued the same as sent
                    }
                    log.warn("Discarding {} fragment; no sideband channels are connected and the queue is full",
                            fragment.getSidebandOneofCase());
                } else {
                    //This implies we raced with open, and it acquired the lock before us. In that case,
                    //we want to send the fragment over the now-open channel rather than queuing
                    return channel.send(fragment);
                }
            }
        } else {
            //If all the connected sideband channels fail, we don't queue (for now)
            log.warn("Failed to send {} fragment after trying {} sideband channel(s)",
                    fragment.getSidebandOneofCase(), attempted);
        }

        return false;
    }

    @VisibleForTesting
    boolean isOpen() {
        return channels.stream().anyMatch(SidebandChannel::isOpen);
    }

    private void closeChannels() {
        //Loop over the registered sideband channels and close each while also clearing their registration
        SidebandChannel channel;
        while ((channel = channels.poll()) != null) {
            channel.close();
        }
    }

    private class SidebandChannel implements SidebandReceiver {

        private final AtomicBoolean closed;
        private final String id;
        private final SidebandSender sender;

        SidebandChannel(String id, SidebandSender sender) {
            this.sender = sender;
            this.id = id;

            closed = new AtomicBoolean();
        }

        @Override
        public void close() {
            if (closed.compareAndSet(false, true)) {
                //If the channel is still registered, it means the control plane disconnected. Otherwise,
                //if it's already been removed, it means we're closing connections for shutdown
                boolean disconnect = channels.remove(this);

                try {
                    sender.close();
                } catch (RuntimeException e) {
                    if (!isClosed(e)) { //If the sideband is already closed, ignore the exception
                        log.debug("{}: Ignoring exception while closing sideband", id, e);
                    }
                } finally {
                    log.info("{}: Sideband {}", id, disconnect ? "disconnected" : "closed");
                }
            }
        }

        @Override
        public void onReceive(@Nonnull RpcSidebandLeaderFragment fragment) {
            switch (fragment.getSidebandOneofCase()) {
                case HEARTBEAT:
                    log.debug("{}: Received heartbeat", id);
                    send(RpcSidebandFollowerFragment.newBuilder()
                            .setHeartbeat(RpcHeartbeat.getDefaultInstance())
                            .build());
                    break;
            }
        }

        @Override
        public String toString() {
            return id;
        }

        boolean isOpen() {
            return !closed.get();
        }

        boolean send(@Nonnull RpcSidebandFollowerFragment fragment) {
            if (closed.get()) {
                log.debug("{}: Not sending {} fragment; the channel is closed", id, fragment.getSidebandOneofCase());
                return false;
            }

            try {
                sender.send(fragment);

                log.debug("{}: Sent {} fragment", id, fragment.getSidebandOneofCase());
                return true;
            } catch (RuntimeException e) {
                //If we fail trying to send a fragment, close the sideband and wait for the control
                //plane to reopen it. If this was the last sideband channel still open, we'll queue
                //fragments until a new one opens
                log.warn("{}: Closing sideband after failure to send {} fragment",
                        id, fragment.getSidebandOneofCase(), e);
                close();

                return false;
            }
        }

        private boolean isClosed(RuntimeException e) {
            if (e instanceof IllegalStateException) {
                String message = e.getMessage();

                return "call already closed".equals(message) || "call is closed".equals(message);
            }

            return false;
        }
    }
}
