package com.atlassian.stash.internal.maintenance.latch;

import com.atlassian.johnson.Johnson;
import com.atlassian.johnson.event.Event;
import com.atlassian.johnson.event.EventLevel;
import com.atlassian.johnson.event.EventType;
import com.atlassian.bitbucket.cluster.ClusterNode;
import com.atlassian.stash.internal.hazelcast.NodeIdMemberSelector;
import com.atlassian.bitbucket.util.Chainable;
import com.google.common.collect.ImmutableSet;
import com.hazelcast.core.Cluster;
import com.hazelcast.core.IExecutorService;
import com.hazelcast.core.Member;
import com.hazelcast.spring.context.SpringAware;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.Serializable;
import java.util.Collection;
import java.util.Objects;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.Callable;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;

/**
 * Base implementation of {@link Latch latches} that support both local and cluster-wide latching and draining.
 * Subclasses need only implement {@link #acquireLocally()}, {@link #drainLocally(long, TimeUnit, boolean)} and
 * {@link #unlatchLocally()}.
 */
public abstract class ClusterableLatch implements Latch {

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

    private final Set<String> drainedMembers;

    private volatile boolean acquired;
    private volatile boolean drained;
    private volatile boolean drainedLocally;
    private volatile String id;
    private volatile boolean unlatched;

    protected final Cluster cluster;
    protected final IExecutorService executor;
    protected final String latchServiceBeanName;
    protected final LatchMode mode;
    protected final Object lock;

    protected ClusterableLatch(LatchMode mode, Cluster cluster, IExecutorService executor, String latchServiceBeanName) {
        this.cluster = cluster;
        this.executor = executor;
        this.latchServiceBeanName = latchServiceBeanName;
        this.mode = mode;

        drainedMembers = new CopyOnWriteArraySet<>();
        lock = new Object();
    }

    // this method is _not_ part of the Latch interface. LatchableServices are expected to call this method straight
    // after construction to acquire the latch either locally or in the cluster
    public void acquire(@Nullable String latchId) {
        if (acquired && id != null && id.equals(latchId)) {
            // already acquired under the requested id
            return;
        }

        checkState(!acquired, "Latch has already been acquired");

        String newId = latchId == null ? UUID.randomUUID().toString() : latchId;
        // acquire the latch locally
        synchronized (lock) {
            checkState(!acquired, "Latch has already been acquired");
            acquireLocally();

            id = newId;
            acquired = true;
        }

        if (mode == LatchMode.CLUSTER && latchId == null) {
            // no latchId was provided, so this is the node initiating a cluster latch. The local node is already latched,
            // now latch the other nodes
            Collection<Member> remoteMembers = getRemoteMembers();
            if (!remoteMembers.isEmpty()) {
                AcquireCallback callback = new AcquireCallback();
                executor.submitToMembers(new AcquireLatchTask(latchServiceBeanName, newId), remoteMembers, callback);

                try {
                    callback.await(2, TimeUnit.MINUTES);
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    log.warn("Interrupted while waiting for latch to be acquired in the cluster");
                }
                if (callback.getLatchedMembers().size() != remoteMembers.size()) {
                    // some nodes failed to latch - unlatch all latched nodes
                    unlatchLocally();
                    unlatchOrPassivateNodes(callback.getLatchedMembers());

                    throw new IllegalStateException("Failed to acquire the latch on all nodes in the cluster");
                }
            }
        }
    }

    @Override
    public boolean drain(long timeout, @Nonnull TimeUnit timeUnit) {
        return drain(timeout, timeUnit, false);
    }

    @Override
    public boolean forceDrain(long timeout, @Nonnull TimeUnit timeUnit) {
        return drain(timeout, timeUnit, true);
    }

    @Nonnull
    @Override
    public LatchMode getMode() {
        return mode;
    }

    public boolean isAcquired() {
        return acquired;
    }

    public void onNodeJoined(ClusterNode node) {
        if (acquired && !unlatched) {
            // latch the newly added node as well
            executor.submitToMembers(new AcquireLatchTask(latchServiceBeanName, id),
                    new NodeIdMemberSelector(node.getId()));
        }
    }

    @Override
    public void unlatch() {
        checkState(!unlatched, "This latch is no longer active");

        if (mode == LatchMode.LOCAL) {
            doUnlatchLocally();
        } else {
            unlatchOrPassivateNodes(null);
        }
    }

    protected abstract void acquireLocally();

    protected abstract boolean drainLocally(long timeout, @Nonnull TimeUnit timeUnit, boolean force);

    protected abstract void unlatchLocally();

    private boolean doDrainLocally(long timeout, TimeUnit timeUnit, boolean force) {
        if (!drainedLocally) {
            drainedLocally |= drainLocally(timeout, timeUnit, force);
        }
        return drainedLocally;
    }

    private void doUnlatchLocally() {
        if (unlatched) {
            return;
        }

        synchronized (lock) {
            if (unlatched) {
                return;
            }

            unlatchLocally();
            unlatched = true;
        }
    }

    private boolean drain(long timeout, @Nonnull TimeUnit timeUnit, boolean force) {
        checkState(acquired, "Latch has not been acquired yet");

        if (unlatched) {
            log.debug("This latch is no longer active");
            return false;
        }
        if (drained) {
            return true;
        }

        if (mode == LatchMode.LOCAL) {
            drained |= doDrainLocally(timeout, timeUnit, force);
        } else {
            // drain *all* the nodes
            Collection<Member> nonDrained = getNonDrainedMembers();
            if (!nonDrained.isEmpty()) {
                DrainCallback callback = new DrainCallback();
                executor.submitToMembers(new DrainTask(latchServiceBeanName, id, timeout, timeUnit, force), nonDrained, callback);

                try {
                    // give the callback a bit of extra time to account for possible latency
                    callback.await(500 + timeUnit.toMillis(timeout), TimeUnit.MILLISECONDS);
                    drained = callback.isDrained();
                } catch (InterruptedException e) {
                    // restore the interrupted flag
                    Thread.currentThread().interrupt();
                }
            }
        }
        return drained;
    }

    private Collection<Member> getNonDrainedMembers() {
        return Chainable.chain(cluster.getMembers()).filter(member -> !drainedMembers.contains(member.getUuid())).toList();
    }

    private Collection<Member> getRemoteMembers() {
        return Chainable.chain(cluster.getMembers()).filter(member -> !member.localMember()).toList();
    }

    private void unlatchOrPassivateNodes(Collection<Member> members) {
        if (members != null && members.isEmpty()) {
            // nothing to do
            return;
        }

        UnlatchCallback callback = new UnlatchCallback();
        UnlatchTask unlatchTask = new UnlatchTask(latchServiceBeanName, id);
        if (members != null) {
            executor.submitToMembers(unlatchTask, members, callback);
        } else {
            executor.submitToAllMembers(unlatchTask, callback);
        }

        try {
            callback.await(1, TimeUnit.MINUTES);

            if (callback.isSuccess()) {
                log.debug("Cluster unlatch completed successfully");
            } else {
                log.warn("Failed to unlatch some cluster members: {}", callback.errorMembers);
                executor.executeOnMembers(new UnlatchFailedJohnsonTask(), callback.errorMembers);
            }
        } catch (InterruptedException e) {
            // restore interrupted flag
            Thread.currentThread().interrupt();
            log.warn("Interrupted while waiting for latch to be released on other cluster nodes. Not all nodes may " +
                    "have unlatched", e);
        }
    }

    private static class AcquireCallback extends ResultCollectingExecutionCallback<Object> {
        private final Set<Member> latchedMembers = new CopyOnWriteArraySet<>();

        public Set<Member> getLatchedMembers() {
            return ImmutableSet.copyOf(latchedMembers);
        }

        @Override
        protected void onSuccess(Member member, Object value) {
            latchedMembers.add(member);
        }
    }

    private static class AcquireLatchTask extends ClusterableLatchTask implements Callable<Void> {

        private final String latchId;

        private AcquireLatchTask(String latchableServiceBeanName, String latchId) {
            super(latchableServiceBeanName, latchId);
            this.latchId = latchId;
        }

        @Override
        public Void call() throws Exception {
            checkState(service != null, "LatchableService %s was not injected", beanName);
            service.acquireLatch(LatchMode.CLUSTER, latchId);
            return null;
        }
    }

    private class DrainCallback extends ResultCollectingExecutionCallback<Boolean> {
        private final AtomicBoolean drained = new AtomicBoolean(true);

        public boolean isDrained() {
            return drained.get();
        }

        @Override
        protected void onError(Member member, Throwable throwable) {
            drained.set(false);
        }

        @Override
        protected void onSuccess(Member member, Boolean value) {
            if (Boolean.TRUE.equals(value)) {
                drainedMembers.add(member.getUuid());
                log.debug("Node {} drained successfully", member.getUuid());
            } else {
                drained.set(false);
                log.debug("Node {} did not drain: {}", member.getUuid(), value != null ? value.toString() : "");
            }
        }
    }

    private static class DrainTask extends ClusterableLatchTask implements Callable<Boolean> {
        private static final long serialVersionUID = 2L;

        private final long timeoutMs;
        private final boolean force;

        public DrainTask(String latchServiceBeanName, String latchId, long timeout, TimeUnit timeUnit, boolean force) {
            super(latchServiceBeanName, latchId);
            this.force = force;
            timeoutMs = timeUnit.toMillis(timeout);
        }

        @Override
        public Boolean call() throws Exception {
            if (latch == null) {
                // This can happen, for example, if we receive a viewMaintenanceStatus request (which calls drain with
                // zero timeout to retrieve the drained state) before we have latched the service.  Just return false.
                return false;
            }
            checkState(latchId.equals(latch.id), "Unexpected latch for %s: expected latch ID %s but got %s", beanName, latchId, latch.id);

            return latch.doDrainLocally(timeoutMs, TimeUnit.MILLISECONDS, force);
        }
    }

    private static class UnlatchCallback extends ResultCollectingExecutionCallback<Void> {
        protected final Set<Member> errorMembers = new CopyOnWriteArraySet<>();

        @Override
        protected void onError(Member member, Throwable throwable) {
            errorMembers.add(member);
        }
    }

    private static class UnlatchTask extends ClusterableLatchTask implements Runnable {
        private static final long serialVersionUID = 1L;

        public UnlatchTask(String latchServiceBeanName, String latchId) {
            super(latchServiceBeanName, latchId);
        }

        @Override
        public void run() {
            checkState(latch != null, "Latch for %s was not injected", beanName);
            checkState(latchId.equals(latch.id), "Unexpected latch for %s: expected latch ID %s but got %s", beanName, latchId, latch.id);

            latch.doUnlatchLocally();
        }
    }

    @SpringAware
    protected abstract static class ClusterableLatchTask implements Serializable {
        protected final String beanName;
        protected final String latchId;

        protected transient LatchableService service;
        protected transient ClusterableLatch latch;

        protected ClusterableLatchTask(String beanName, String latchId) {
            this.beanName = checkNotNull(beanName, "beanName");
            this.latchId = checkNotNull(latchId, "latchId");
        }

        @Autowired
        public void setApplicationContext(ApplicationContext applicationContext) {
            // Look up the LatchableService by it's beanName, then retrieve the latch.
            try {
                service = applicationContext.getBean(beanName, LatchableService.class);
                Latch l = service.getCurrentLatch();
                if (l != null) {
                    checkState(l instanceof ClusterableLatch, "Latch %s is not a ClusterableLatch", l);
                    latch = (ClusterableLatch) l;
                    checkState(Objects.equals(latchId, latch.id),
                            "An unexpected latch was found. Expected %s, found %s", latchId, latch.id);
                }
            } catch (BeansException e) {
                log.error("Latchable service '{}' not found - cannot drain the cluster latch", beanName, e);
                throw new IllegalStateException(e);
            }
        }
    }

    private static class UnlatchFailedJohnsonTask implements Runnable, Serializable {
        private static final long serialVersionUID = 1L;

        private static final String UNLATCH_FAILED_EVENT_TYPE = "unlatch-failed";
        private static final String UNLATCH_FAILED_EVENT_LEVEL = "error";

        @Override
        public void run() {
            Event event = new Event(EventType.get(UNLATCH_FAILED_EVENT_TYPE),
                    "Failed to unlatch this cluster node",
                    EventLevel.get(UNLATCH_FAILED_EVENT_LEVEL));
            Johnson.getEventContainer().addEvent(event);
        }
    }
}
