package com.atlassian.bitbucket.mesh.coordination;

import com.atlassian.bitbucket.mesh.git.exception.AbortedException;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.repository.Repository;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Throwables;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.util.Collection;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicBoolean;

import static java.util.Objects.requireNonNull;

class DefaultWriteLock implements WriteLock {

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

    private final AtomicBoolean closed;
    private final CompletableFuture<Void> future;
    private final WriteToken.Promise localPromise;
    private final Object lock;
    private final Repository repository;
    private final long timestamp;
    private final Set<String> writeTokensHeld;
    private final Set<String> writeTokensRequired;

    DefaultWriteLock(Repository repository, long timestamp, WriteToken.Promise localPromise,
                     MeshNode localNode, Collection<MeshNode> nodes) {
        this.localPromise = localPromise;
        this.repository = repository;
        this.timestamp = timestamp;

        closed = new AtomicBoolean();
        future = new CompletableFuture<>();
        lock = new Object();
        writeTokensHeld = new HashSet<>();
        writeTokensRequired = new HashSet<>();
        writeTokensRequired.add(localNode.getId());
        nodes.forEach(node -> writeTokensRequired.add(node.getId()));
    }

    @Nonnull
    @Override
    public CompletionStage<Void> asCompletionStage() {
        return future;
    }

    @Override
    public void await() throws InterruptedException {
        try {
            future.get();
        } catch (CancellationException e) {
            close(); //If a caller cancels the CompletionStage, ensure we close the localPromise
            throw new AbortedException("WriteLock canceled", e);
        } catch (ExecutionException e) {
            Throwable cause = e.getCause();
            Throwables.throwIfUnchecked(cause);
            throw new AbortedException("WriteLock failed", cause);
        }
    }

    @Override
    public void close() {
        if (closed.compareAndSet(false, true)) {
            boolean held = future.isDone();
            localPromise.close();
            log.debug("[{}] Write lock closed (released = {})", repository, held);
            future.completeExceptionally(new AbortedException("WriteLock canceled"));
        }
    }

    @Override
    public long getLogicalTimestamp() {
        return timestamp;
    }

    @Override
    public boolean onTokenGranted(@Nonnull MeshNode node) {
        synchronized (lock) {
            String nodeId = requireNonNull(node, "node").getId();
            if (!closed.get() && writeTokensRequired.contains(nodeId)) {
                writeTokensHeld.add(nodeId);
                log.debug("[{}] {} granted write token (granted = {}, needed = {})", repository, node,
                        writeTokensHeld.size(), writeTokensRequired.size());
                if (isAcquirable()) {
                    log.debug("[{}] Write lock acquired (tokens held = {})", repository, writeTokensHeld.size());
                    future.complete(null);
                }
                return true;
            }
            return false;
        }
    }

    @Override
    public boolean remove(@Nonnull MeshNode node) {
        String nodeId = requireNonNull(node, "node").getId();

        synchronized (lock) {
            //If we hold the specified node's token, we always release it
            boolean released = writeTokensHeld.remove(nodeId);
            //Whether we held the node's token or not, ensure we're not requiring it to acquire the lock. If
            //we were, then removing it may mean we acquired the lock
            if (writeTokensRequired.remove(nodeId) && isAcquirable()) {
                if (!future.isDone()) {
                    log.debug("[{}] Write lock acquired (tokens held = {})", repository, writeTokensHeld.size());
                    future.complete(null);
                }
            }

            return released;
        }
    }

    @Override
    public boolean yieldToken(@Nonnull MeshNode node) {
        synchronized (lock) {
            if (isAcquired()) {
                log.debug("[{}] Not yielding write token from {} until the lock has been released", repository, node);
                return false;
            }
            if (writeTokensHeld.remove(requireNonNull(node, "node").getId())) {
                log.debug("[{}] Yielded write token from {}", repository, node);
                return true;
            }
            log.info("[{}] Asked to yield write token from {} which is not or is no longer held", repository, node);
            return false;
        }
    }

    @VisibleForTesting
    int getTokenCount() {
        return writeTokensHeld.size();
    }

    /**
     * @return {@code true} if {@link #future} has completed <i>non-exceptionally</i>, meaning the lock was
     *         {@link #isAcquirable acquirable} <i>and it was acquired</i>; otherwise, {@code false}
     */
    @VisibleForTesting
    boolean isAcquired() {
        return future.isDone() && !future.isCompletedExceptionally();
    }

    /**
     * @return {@code true} if all of the required tokens are held; otherwise, {@code false}
     */
    private boolean isAcquirable() {
        return writeTokensHeld.size() == writeTokensRequired.size();
    }
}
