package com.atlassian.bitbucket.mesh.network;

import com.atlassian.bitbucket.mesh.config.NodeConfig;
import com.atlassian.bitbucket.mesh.event.EventPublisher;
import com.google.common.collect.ImmutableSet;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Captor;
import org.mockito.InOrder;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

import java.util.Collection;

import static com.atlassian.bitbucket.mesh.util.MoreCollectors.toImmutableSet;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.*;
import static org.junit.Assert.*;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class DefaultMeshManagerTest {

    private static final String NODE_ID = "99";
    private static final String NODE_NAME = "node-99";
    private static final String NODE_URL = "https://node-99:7777";
    private static final String NODE1_ID = "1";
    private static final String NODE2_ID = "2";
    private static final MeshNode LOCAL_NODE = new MeshNode.Builder(NODE_ID)
            .local(true)
            .name(NODE_NAME)
            .rpcUrl(NODE_URL)
            .build();
    private static final MeshNode NODE1 = new MeshNode.Builder(NODE1_ID)
            .name("node-1")
            .rpcUrl("https://node-1:7777")
            .build();
    private static final MeshNode NODE2 = new MeshNode.Builder(NODE2_ID)
            .name("node-2")
            .rpcUrl("https://node-2:7777")
            .build();

    @Captor
    private ArgumentCaptor<MeshNodeEvent> eventCaptor;
    @Mock
    private EventPublisher eventPublisher;
    @Mock
    private NodeConfig nodeConfig;

    @Before
    public void setUpConfig() {
        when(nodeConfig.getId()).thenReturn(NODE_ID);
        when(nodeConfig.getName()).thenReturn(NODE_NAME);
        when(nodeConfig.getRpcUrl()).thenReturn(NODE_URL);
    }

    @Test
    public void testGetNodes() {
        DefaultMeshManager manager = createManager();
        manager.updateTopology(new UpdateTopologyParameters.Builder()
                .nodes(LOCAL_NODE, NODE1, NODE2)
                .build());

        MeshNode node1 = manager.findNodeById(NODE1_ID);
        MeshNode node2 = manager.findNodeById(NODE2_ID);
        assertNotNull(node1);
        assertNotNull(node2);
        assertThat(manager.getNodes(), containsInAnyOrder(manager.getLocalNode(), node1, node2));
    }

    @Test
    public void testGetNodeById() {
        DefaultMeshManager manager = createManager();
        manager.updateTopology(new UpdateTopologyParameters.Builder()
                .nodes(NODE1, LOCAL_NODE)
                .build());

        MeshNode node = manager.getNodeById(NODE1_ID);
        assertNotNull(node);
        assertEquals(NODE1_ID, node.getId());
    }

    @Test
    public void testGetNodeByIdReturnsLocalNode() {
        DefaultMeshManager manager = createManager();

        MeshNode node = manager.getNodeById(NODE_ID);
        assertSame(manager.getLocalNode(), node);
        assertEquals(NODE_ID, node.getId());
        assertEquals(NODE_URL, node.getRpcUrl());
        assertTrue(node.isLocal());
    }

    @Test
    public void testGetNodeByIdThrowsIfNotFound() {
        DefaultMeshManager manager = createManager();

        NodeNotFoundException exception = assertThrows(NodeNotFoundException.class,
                () -> manager.getNodeById("missing-id"));

        assertEquals("missing-id", exception.getNodeId());
    }

    @Test
    public void testUpdateTopology() {
        DefaultMeshManager manager = createManager();

        manager.updateTopology(new UpdateTopologyParameters.Builder()
                .partitions(new Partition.Builder("0001")
                        .replicas(new PartitionReplica.Builder("replica1", LOCAL_NODE).build())
                        .replicas(new PartitionReplica.Builder("replica2", NODE1).build())
                        .build())
                .partitions(new Partition.Builder("0017")
                        .replicas(new PartitionReplica.Builder("replica3", LOCAL_NODE).build())
                        .build())
                .build());

        //Verify the local node
        MeshNode localNode = manager.getLocalNode();
        assertEquals(LOCAL_NODE.getId(), localNode.getId());
        assertEquals(LOCAL_NODE.getName(), localNode.getName());
        assertEquals(LOCAL_NODE.getRpcUrl(), localNode.getRpcUrl());
        assertTrue(localNode.isLocal());

        //Verify the partitions
        assertEquals(2, manager.getPartitions().size());

        Partition partition1 = manager.getPartition("0001");
        assertNotNull(partition1);
        assertEquals(2, partition1.getReplicas().size());
        assertEquals("0001", partition1.getId());
        assertEquals(ImmutableSet.of(LOCAL_NODE.getId(), NODE1.getId()),
                partition1.getReplicas().stream().map(replica -> replica.getNode().getId()).collect(toImmutableSet()));

        Partition partition17 = manager.getPartition("0017");
        assertNotNull(partition17);
        assertEquals("0017", partition17.getId());
        assertEquals(1, partition17.getReplicas().size());
        assertEquals(ImmutableSet.of(LOCAL_NODE.getId()),
                partition17.getReplicas().stream().map(replica -> replica.getNode().getId()).collect(toImmutableSet()));

        //This validation _must_ use inOrder because ArgumentCaptor _is not type-aware_. Even though it's
        //generic, and typed as MeshNodeEvent here, it will match the publish(TopologyUpdatedEvent) too.
        //Using inOrder "consumes" the TopologyUpdatedEvent so the capture only matches one invocation.
        //See https://github.com/mockito/mockito/issues/565
        InOrder ordered = inOrder(eventPublisher);
        ordered.verify(eventPublisher).publish(isA(TopologyUpdatedEvent.class));
        ordered.verify(eventPublisher).publish(eventCaptor.capture());

        MeshNodeEvent event = eventCaptor.getValue();
        assertThat(event, instanceOf(MeshNodeAddedEvent.class));
        assertEquals(NODE1, event.getNode());
    }

    @Test
    public void testUpdateTopologyAddsLocalNodeIfMissing() {
        DefaultMeshManager manager = createManager();

        manager.updateTopology(new UpdateTopologyParameters.Builder()
                .nodes(NODE1)
                .build());
        assertEquals(LOCAL_NODE.getId(), manager.getLocalNode().getId());

        verify(eventPublisher).publish(isA(TopologyUpdatedEvent.class));
    }

    @Test
    public void testUpdateTopologyThrowsWhenMultipleLocalNodesProvided() {
        DefaultMeshManager manager = createManager();

        assertThrows(IllegalArgumentException.class, () ->
                manager.updateTopology(new UpdateTopologyParameters.Builder()
                        .nodes(LOCAL_NODE, new MeshNode.Builder(NODE1)
                                .local(true)
                                .build())
                        .build()));

        verifyNoInteractions(eventPublisher);
    }

    @Test
    public void testUpdateTopologyWithChangeToLocalNodeId() {
        DefaultMeshManager manager = createManager();

        MeshNode updatedLocal = new MeshNode.Builder(LOCAL_NODE).rpcUrl("changed").build();
        when(nodeConfig.getRpcUrl()).thenReturn("changed");
        manager.updateTopology(new UpdateTopologyParameters.Builder()
                .nodes(updatedLocal)
                .build());

        //See testUpdateTopology for why this uses inOrder
        InOrder ordered = inOrder(eventPublisher);
        ordered.verify(eventPublisher).publish(isA(TopologyUpdatedEvent.class));
        ordered.verify(eventPublisher).publish(eventCaptor.capture());

        MeshNodeEvent event = eventCaptor.getValue();
        assertThat(event, instanceOf(MeshNodeUpdatedEvent.class));
        assertEquals("changed", event.getNode().getRpcUrl());
        assertEquals(LOCAL_NODE.getRpcUrl(), ((MeshNodeUpdatedEvent) event).getOldValue().getRpcUrl());
    }

    @Test
    public void testUpdateTopologyWithChangeToLocalNodeRpcUrl() {
        DefaultMeshManager manager = createManager();

        MeshNode updatedLocal = new MeshNode.Builder("changed")
                .local(true)
                .name(LOCAL_NODE.getName())
                .rpcUrl(LOCAL_NODE.getRpcUrl())
                .build();
        when(nodeConfig.getId()).thenReturn("changed");
        manager.updateTopology(new UpdateTopologyParameters.Builder()
                .nodes(updatedLocal)
                .build());

        //See testUpdateTopology for why this uses inOrder
        InOrder ordered = inOrder(eventPublisher);
        ordered.verify(eventPublisher).publish(isA(TopologyUpdatedEvent.class));
        ordered.verify(eventPublisher).publish(eventCaptor.capture());

        MeshNodeEvent event = eventCaptor.getValue();
        assertThat(event, instanceOf(MeshNodeUpdatedEvent.class));
        assertEquals("changed", event.getNode().getId());
        assertEquals(LOCAL_NODE.getId(), ((MeshNodeUpdatedEvent) event).getOldValue().getId());
    }

    @Test
    public void testUpdateTopologyWithPartitionMigrations() {
        DefaultMeshManager manager = createManager();

        manager.updateTopology(new UpdateTopologyParameters.Builder()
                .partitions(new Partition.Builder("0001")
                        .replicas(new PartitionReplica.Builder("replica1", LOCAL_NODE).build())
                        .replicas(new PartitionReplica.Builder("replica2", NODE1).build())
                        .build())
                .partitions(new Partition.Builder("0017")
                        .replicas(new PartitionReplica.Builder("replica3", LOCAL_NODE).build())
                        .build())
                .migrations(new PartitionReplicaMigration.Builder()
                        .partition("0001")
                        .sourceNode(LOCAL_NODE)
                        .targetNode(NODE1)
                        .build())
                .build());

        //Verify the local node
        MeshNode localNode = manager.getLocalNode();
        assertEquals(LOCAL_NODE.getId(), localNode.getId());
        assertEquals(LOCAL_NODE.getName(), localNode.getName());
        assertEquals(LOCAL_NODE.getRpcUrl(), localNode.getRpcUrl());
        assertTrue(localNode.isLocal());

        //Verify the partitions
        assertEquals(2, manager.getPartitions().size());

        Partition partition1 = manager.getPartition("0001");
        assertNotNull(partition1);
        assertEquals(2, partition1.getReplicas().size());
        assertEquals("0001", partition1.getId());
        assertEquals(ImmutableSet.of(LOCAL_NODE.getId(), NODE1.getId()),
                partition1.getReplicas().stream().map(replica -> replica.getNode().getId()).collect(toImmutableSet()));

        Partition partition17 = manager.getPartition("0017");
        assertNotNull(partition17);
        assertEquals("0017", partition17.getId());
        assertEquals(1, partition17.getReplicas().size());
        assertEquals(ImmutableSet.of(LOCAL_NODE.getId()),
                partition17.getReplicas().stream().map(replica -> replica.getNode().getId()).collect(toImmutableSet()));

        //Verify the migrations
        Collection<PartitionReplicaMigration> migrations = manager.getPartitionMigrations();
        assertEquals(1, migrations.size());
        PartitionReplicaMigration migration = migrations.iterator().next();
        assertEquals("0001", migration.getPartitionId());
        assertEquals(LOCAL_NODE, migration.getSourceNode());
        assertEquals(NODE1, migration.getTargetNode());

        //Check getPartitionMigrations by partition
        migrations = manager.getPartitionMigrations("0001");
        assertEquals(1, migrations.size());
        assertEquals(migration, migrations.iterator().next());

        assertThat(manager.getPartitionMigrations("0017"), empty());
        assertThat(manager.getPartitionMigrations("9999"), empty());

        //Update the topology again, clearing the migrations
        manager.updateTopology(new UpdateTopologyParameters.Builder()
                .partitions(new Partition.Builder("0001")
                        .replicas(new PartitionReplica.Builder("replica1", LOCAL_NODE).build())
                        .replicas(new PartitionReplica.Builder("replica2", NODE1).build())
                        .build())
                .partitions(new Partition.Builder("0017")
                        .replicas(new PartitionReplica.Builder("replica3", LOCAL_NODE).build())
                        .build())
                .build());

        assertEquals(partition1, manager.getPartition("0001"));
        assertEquals(partition17, manager.getPartition("0017"));
        assertThat(manager.getPartitionMigrations(), empty());
        assertThat(manager.getPartitionMigrations("0001"), empty());
    }

    private DefaultMeshManager createManager() {
        return new DefaultMeshManager(eventPublisher, nodeConfig);
    }
}