package com.atlassian.bitbucket.mesh.test.cluster;

import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import io.grpc.ManagedChannel;
import io.grpc.netty.NegotiationType;
import io.grpc.netty.NettyChannelBuilder;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.net.InetSocketAddress;
import java.net.URI;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import static java.util.Objects.requireNonNull;

public class StandaloneMeshControlPlane implements MeshControlPlane {

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

    private final Map<String, ManagedChannel> channelByNodeId;
    private final ExecutorService executorService;
    private final AtomicLong nextNodeId;
    private final Map<String, RpcNode> nodeByUrl;
    private final Set<String> partitions;

    public StandaloneMeshControlPlane(int partitionCount) {
        partitions = IntStream.range(0, partitionCount)
                .mapToObj(StandaloneMeshControlPlane::hexEncode)
                .collect(Collectors.toCollection(CopyOnWriteArraySet::new));

        channelByNodeId = new ConcurrentHashMap<>();
        executorService = Executors.newCachedThreadPool(new ThreadFactoryBuilder()
                .setDaemon(true)
                .setNameFormat("mesh-cluster:thread-%d")
                .build());
        nextNodeId = new AtomicLong(1000L);
        nodeByUrl = new ConcurrentHashMap<>();
    }

    public boolean addPartition(String partition) {
        boolean added = partitions.add(partition);
        if (added) {
            onTopologyUpdated();
        }
        return added;
    }

    public void destroy() {
        channelByNodeId.values().forEach(ManagedChannel::shutdown);
        channelByNodeId.clear();
        executorService.shutdown();
    }

    @Nonnull
    @Override
    public ManagedChannel openChannel(@Nonnull MeshClusterNode node) {
        return buildChannel(node.getUrl());
    }

    @Override
    public String registerNode(@Nonnull MeshClusterNode node) {
        RpcNode meshNode = addNode(node);
        onTopologyUpdated();
        return meshNode.getId();
    }

    public boolean removePartition(String partition) {
        boolean removed = partitions.remove(partition);
        if (removed) {
            onTopologyUpdated();
        }
        return removed;
    }

    @Override
    public boolean unregisterNode(@Nonnull MeshClusterNode node) {
        RpcNode removed = nodeByUrl.remove(node.getUrl());
        if (removed == null) {
            return false;
        }
        try {
            setConfiguration(removed, RpcSetConfigurationRequest.newBuilder()
                    .setTopology(RpcTopology.getDefaultInstance()));
        } catch (RuntimeException e) {
            log.warn("Failed to set configuration on {} ({})", removed.getId(), removed.getRpcUrl(), e);
        }
        ManagedChannel channel = channelByNodeId.remove(removed.getId());
        if (channel != null) {
            channel.shutdown();
        }
        onTopologyUpdated();
        return true;
    }

    protected ManagedChannel buildChannel(String url) {
        URI uri = URI.create(url);

        String host = uri.getHost();
        int port = uri.getPort();

        log.trace("Creating gRPC channel for {}:{}", host, port);
        return NettyChannelBuilder.forAddress(new InetSocketAddress(host, port))
                .executor(executorService)
                .negotiationType(NegotiationType.PLAINTEXT)
                //This executor handles certain blocking operations performed by gRPC itself (like DNS lookups,
                //just as an example). It may never get used, but if it does we don't want to use the default
                .offloadExecutor(executorService)
                .build();
    }

    protected RpcNodeConfiguration getConfiguration(RpcNode node) {
        return RpcNodeConfiguration.newBuilder()
                .putProperties("node.id", node.getId())
                .putProperties("node.name", node.getName())
                .build();
    }

    protected ManagementServiceGrpc.ManagementServiceBlockingStub getStub(RpcNode node) {
        ManagedChannel channel = channelByNodeId.computeIfAbsent(node.getId(), id -> buildChannel(node.getRpcUrl()));
        return ManagementServiceGrpc.newBlockingStub(channel);
    }

    protected RpcTopology getTopology() {
        RpcTopology.Builder builder = RpcTopology.newBuilder();
        List<RpcNode> nodes = ImmutableList.copyOf(nodeByUrl.values());

        partitions.forEach(partition -> {
            RpcPartition.Builder partitionBuilder = RpcPartition.newBuilder()
                    .setId(partition);
            for (RpcNode node : nodes) {
                partitionBuilder.addReplicas(RpcPartitionReplica.newBuilder()
                        .setId(partition + "-" + node.getId())
                        .setNodeId(node.getId()));
            }
            builder.addPartitions(partitionBuilder);
        });
        builder.addAllNodes(nodes);
        return builder.build();
    }

    protected void setConfiguration(RpcNode node, RpcSetConfigurationRequest.Builder builder) {
        RpcSetConfigurationRequest request = builder.setNodeConfiguration(getConfiguration(node))
                .build();
        getStub(node).setConfiguration(request);
    }

    private static String hexEncode(int id) {
        return StringUtils.leftPad(Integer.toString(id, 16), 4, '0');
    }

    @Nonnull
    private RpcNode addNode(@Nonnull MeshClusterNode node) {
        requireNonNull(node, "node");

        return nodeByUrl.computeIfAbsent(node.getUrl(), ignored -> {
            String id = node.getId() == null ? Long.toString(nextNodeId.getAndIncrement()) : node.getId();
            return RpcNode.newBuilder()
                    .setId(id)
                    .setName(node.getName())
                    .setRpcUrl(node.getUrl())
                    .build();
        });
    }

    private void onTopologyUpdated() {
        RpcTopology topology = getTopology();
        RpcSetConfigurationRequest.Builder builder = RpcSetConfigurationRequest.newBuilder()
                .setTopology(topology);

        for (RpcNode node : topology.getNodesList()) {
            try {
                setConfiguration(node, builder);
            } catch (RuntimeException e) {
                log.error("Failed to update topology on node {} ({})", node.getName(), node.getRpcUrl(), e);
            }
        }
    }
}
