package com.atlassian.bitbucket.mesh.network;

import com.atlassian.bitbucket.mesh.annotation.Profiled;
import com.atlassian.bitbucket.mesh.config.NodeConfig;
import com.atlassian.bitbucket.mesh.event.EventPublisher;
import com.atlassian.bitbucket.mesh.grpc.MetadataProvider;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import io.grpc.Metadata;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.util.*;

import static com.atlassian.bitbucket.mesh.rpc.MetadataKeys.KEY_REMOTE_CLIENT;
import static com.atlassian.bitbucket.mesh.util.MoreCollectors.toImmutableList;
import static java.util.Objects.requireNonNull;
import static org.apache.commons.lang3.StringUtils.isNotBlank;

public class DefaultMeshManager implements MeshManager, MetadataProvider {

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

    private final EventPublisher eventPublisher;
    private final Object lock;
    private final NodeConfig nodeConfig;

    private volatile Map<String, Set<PartitionReplicaMigration>> migrationsByPartition;
    private volatile Topology topology;

    public DefaultMeshManager(EventPublisher eventPublisher, NodeConfig nodeConfig) {
        this.eventPublisher = eventPublisher;
        this.nodeConfig = nodeConfig;

        lock = new Object();
        migrationsByPartition = Collections.emptyMap();
        topology = new Topology.Builder()
                .nodes(createLocalNode())
                .build();
    }

    @Override
    public void apply(@Nonnull Metadata metadata) {
        //If we've received the topology and know our human-readable name, include it on the request
        String name = getLocalNode().getName();
        if (isNotBlank(name)) {
            metadata.put(KEY_REMOTE_CLIENT, name);
        }
    }

    @Override
    public MeshNode findNodeById(@Nonnull String nodeId) {
        return topology.findNodeById(nodeId);
    }

    @Nonnull
    @Override
    public MeshNode getLocalNode() {
        return topology.getLocalNode();
    }

    @Nonnull
    @Override
    public MeshNode getNodeById(@Nonnull String nodeId) {
        MeshNode node = findNodeById(nodeId);
        if (node == null) {
            throw new NodeNotFoundException(nodeId, "Node with ID " + nodeId + " not found");
        }
        return node;
    }

    @Nonnull
    @Override
    public Collection<MeshNode> getNodes() {
        return topology.getNodes();
    }

    @Override
    public Partition getPartition(@Nonnull String id) {
        return topology.getPartition(id);
    }

    @Nonnull
    @Override
    public Collection<PartitionReplicaMigration> getPartitionMigrations() {
        return migrationsByPartition.values().stream()
                .flatMap(Collection::stream)
                .collect(toImmutableList());
    }

    @Nonnull
    @Override
    public Collection<PartitionReplicaMigration> getPartitionMigrations(@Nonnull String partitionId) {
        return migrationsByPartition.getOrDefault(requireNonNull(partitionId, "partitionId"), Collections.emptySet());
    }

    @Nonnull
    @Override
    public Collection<Partition> getPartitions() {
        return topology.getPartitions();
    }

    @Override
    @Profiled
    public void updateTopology(@Nonnull UpdateTopologyParameters parameters) {
        requireNonNull(parameters, "parameters");

        Collection<MeshNode> oldNodes;
        Map<String, MeshNode> currentNodes;
        synchronized (lock) {
            oldNodes = topology.getNodes();

            Topology newTopology = new Topology.Builder()
                    //Add the local node to ensure that the local node exists in the topology, even when no
                    //partition replicas are mapped to it
                    .nodes(createLocalNode())
                    .nodes(parameters.getNodes())
                    .partitions(parameters.getPartitions())
                    .build();
            log.debug("Updated topology to {}", topology);
            //Avoid partially constructed objects being assigned to the 'topology' field by assigning it to
            //a local variable first.
            topology = newTopology;
            eventPublisher.publish(new TopologyUpdatedEvent());
            currentNodes = new HashMap<>(topology.getNodesMap());

            //Map the migrations by partition ID
            Map<String, ImmutableSet.Builder<PartitionReplicaMigration>> nextMigrationsByPartition = new HashMap<>();
            for (PartitionReplicaMigration migration : parameters.getMigrations()) {
                nextMigrationsByPartition.computeIfAbsent(migration.getPartitionId(), ignored -> ImmutableSet.builder())
                        .add(migration);
            }
            //Replace the migrationsByPartitions with the migrations provided by parameters
            ImmutableMap.Builder<String, Set<PartitionReplicaMigration>> builder = ImmutableMap.builder();
            nextMigrationsByPartition.forEach((partition, migrationsBuilder) ->
                    builder.put(partition, migrationsBuilder.build()));
            migrationsByPartition = builder.build();
        }

        //Check which nodes have been added/removed/updated and raise corresponding events
        for (MeshNode oldNode : oldNodes) {
            MeshNode currentNode = currentNodes.remove(oldNode.getId());
            if (currentNode == null) {
                eventPublisher.publish(oldNode.isLocal() ?
                        new MeshNodeUpdatedEvent(oldNode, getLocalNode()) :
                        new MeshNodeRemovedEvent(oldNode));
            } else if (!(Objects.equals(oldNode.getName(), currentNode.getName()) &&
                    oldNode.getRpcUrl().equals(currentNode.getRpcUrl()) &&
                    oldNode.isLocal() == currentNode.isLocal())) {
                eventPublisher.publish(new MeshNodeUpdatedEvent(oldNode, currentNode));
            }
        }
        //Any remaining nodes in newNodes have been added
        for (MeshNode node : currentNodes.values()) {
            if (!node.isLocal()) {
                //Only raise events for remote nodes
                eventPublisher.publish(new MeshNodeAddedEvent(node));
            }
        }
    }

    private MeshNode createLocalNode() {
        return new MeshNode.Builder(nodeConfig.getId())
                .local(true)
                .name(nodeConfig.getName())
                .rpcUrl(nodeConfig.getRpcUrl())
                .build();
    }
}
