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

import com.atlassian.bitbucket.mesh.management.NodeManager;
import com.atlassian.bitbucket.mesh.management.UpdateConfigurationParameters;
import com.atlassian.bitbucket.mesh.network.*;
import org.apache.commons.lang3.StringUtils;
import org.springframework.test.context.TestContext;
import org.springframework.test.context.TestExecutionListener;

import javax.annotation.Nonnull;

/**
 * {@link TestExecutionListener} that ensures that the topology is configured for a single node Mesh cluster
 * with 64 partitions.
 */
public class TopologyTestExecutionListener implements TestExecutionListener {

    private final MeshManager meshManager;
    private final NodeManager nodeManager;

    private UpdateConfigurationParameters singleNodeTopologyParameters;

    public TopologyTestExecutionListener(MeshManager meshManager, NodeManager nodeManager) {
        this.meshManager = meshManager;
        this.nodeManager = nodeManager;
    }

    @Override
    public void beforeTestMethod(@Nonnull TestContext testContext) {
        setupSingleNodeTopology();
    }

    private void setupSingleNodeTopology() {
        if (meshManager.getNodes().size() != 1 || meshManager.getPartitions().size() != 64) {
            if (singleNodeTopologyParameters == null) {
                //In the tests, we create managed repositories in partition 0-64. Some tests require us
                //to have these partitions registered in the topology.
                singleNodeTopologyParameters = new UpdateConfigurationParameters.Builder()
                        .topology(createTestTopologyParameters())
                        .build();
            }

            nodeManager.updateConfiguration(singleNodeTopologyParameters);
        }
    }

    private UpdateTopologyParameters createTestTopologyParameters() {
        MeshNode localNode = meshManager.getLocalNode();
        UpdateTopologyParameters.Builder builder = new UpdateTopologyParameters.Builder()
                .nodes(localNode);
        for (int partition = 0; partition < 64; ++partition) {
            String partitionId = StringUtils.leftPad(Integer.toHexString(partition), 4, '0');
            builder.partitions(new Partition.Builder(partitionId)
                    .replicas(new PartitionReplica.Builder(partitionId + "-1", localNode).build())
                    .build());
        }
        return builder.build();
    }
}
