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

import org.springframework.test.context.TestContext;
import org.springframework.test.context.support.AbstractTestExecutionListener;

import javax.annotation.Nonnull;
import java.lang.reflect.Method;

public class MeshClusterTestExecutionListener extends AbstractTestExecutionListener {

    private static final String MESH_CLUSTER_CLASS_CONFIG = "bbs.mesh.cluster.class";
    private static final String MESH_CLUSTER_METHOD_CONFIG = "bbs.mesh.cluster.method";

    private final MeshClusterManager clusterManager;

    public MeshClusterTestExecutionListener(MeshClusterManager clusterManager) {
        this.clusterManager = clusterManager;
    }

    @Override
    public void afterTestClass(TestContext testContext) {
        if (testContext.removeAttribute(MESH_CLUSTER_CLASS_CONFIG) != null) {
            tearDownCluster();
        }
    }

    @Override
    public void afterTestMethod(TestContext testContext) {
        if (testContext.removeAttribute(MESH_CLUSTER_METHOD_CONFIG) != null) {
            //A MeshCluster was configured at the method level
            MeshCluster annotation = (MeshCluster) testContext.getAttribute(MESH_CLUSTER_CLASS_CONFIG);
            if (annotation == null) {
                tearDownCluster();
            } else {
                //Restore the class level configuration
                setUpCluster(annotation);
            }
        }
    }

    @Override
    public void beforeTestMethod(@Nonnull TestContext testContext) {
        MeshCluster annotation = getAnnotation(testContext.getTestMethod());
        if (annotation == null) {
            //Fall back to the class level annotation if there is one
            annotation = (MeshCluster) testContext.getAttribute(MESH_CLUSTER_CLASS_CONFIG);
        } else {
            testContext.setAttribute(MESH_CLUSTER_METHOD_CONFIG, annotation);
        }
        if (annotation != null) {
            setUpCluster(annotation);
        }
    }

    @Override
    public void prepareTestInstance(TestContext testContext) {
        MeshCluster annotation = getAnnotation(testContext.getTestClass());
        if (annotation != null) {
            testContext.setAttribute(MESH_CLUSTER_CLASS_CONFIG, annotation);
            setUpCluster(annotation);
        }
    }

    private static MeshCluster getAnnotation(Method method) {
        return method.getAnnotation(MeshCluster.class);
    }

    private static MeshCluster getAnnotation(Class<?> clazz) {
        MeshCluster annotation = clazz.getAnnotation(MeshCluster.class);
        while (annotation == null) {
            clazz = clazz.getSuperclass();
            if (clazz == null) {
                //Annotation not found in the class hierarchy
                return null;
            }
            annotation = clazz.getAnnotation(MeshCluster.class);
        }
        return annotation;
    }

    private void setUpCluster(MeshCluster annotation) {
        clusterManager.initCluster(MeshClusterConfig.forAnnotation(annotation));
    }

    private void tearDownCluster() {
        clusterManager.getAllNodes().forEach(node -> {
            node.reset();
            if (node.isPartOfCluster()) {
                node.leaveCluster(false);
            }
        });
    }
}
