package com.atlassian.bitbucket.mesh.coordination;

import com.atlassian.bitbucket.mesh.network.MeshManager;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.repository.Repository;
import com.google.common.annotations.VisibleForTesting;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.TaskScheduler;

import javax.annotation.Nonnull;
import java.time.Duration;
import java.time.Instant;
import java.util.Collection;
import java.util.Collections;
import java.util.NavigableSet;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ConcurrentSkipListSet;
import java.util.concurrent.atomic.AtomicInteger;

import static java.util.Objects.requireNonNull;

public class DefaultWriteTokenRegistry implements WriteTokenRegistry {

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

    private final LogicalClock logicalClock;
    private final MeshManager meshManager;
    private final Duration retryYieldRequestInterval;
    private final ConcurrentMap<String, RepositoryToken> tokenByRepositorySlug;

    public DefaultWriteTokenRegistry(LogicalClock logicalClock, MeshManager meshManager, TaskScheduler taskScheduler,
                                     Duration retryYieldRequestInterval) {
        this.logicalClock = logicalClock;
        this.meshManager = meshManager;
        this.retryYieldRequestInterval = retryYieldRequestInterval;

        tokenByRepositorySlug = new ConcurrentHashMap<>();
        taskScheduler.scheduleAtFixedRate(this::maybeRetryYieldRequests, retryYieldRequestInterval);
    }

    @Nonnull
    @Override
    public WriteToken.Promise acquireToken(@Nonnull Repository repository, @Nonnull String transactionId,
                                           long logicalTimestamp, @Nonnull WriteTokenHandler handler) {
        requireNonNull(repository, "repository");
        requireNonNull(handler, "handler");
        requireNonNull(transactionId, "transactionId");

        return getToken(repository).acquire(transactionId, logicalTimestamp, handler);
    }

    @Nonnull
    @Override
    public WriteLock startLocalLock(@Nonnull Repository repository, @Nonnull String transactionId) {
        return startLock(repository, transactionId, Collections.emptySet());
    }

    @Nonnull
    @Override
    public DefaultWriteLock startLock(@Nonnull Repository repository, @Nonnull String transactionId,
                                      @Nonnull Collection<MeshNode> nodes) {
        requireNonNull(repository, "repository");
        requireNonNull(transactionId, "transactionId");
        requireNonNull(nodes, "nodes");

        //Create the timestamp for the token request to be used with the write lock
        long timestamp = logicalClock.tick();
        MeshNode localNode = meshManager.getLocalNode();

        //Create a handler that integrates the local node's token with the WriteLock
        WriteLockTokenHandler handler = new WriteLockTokenHandler(localNode);
        //Request the local token for the repository
        WriteToken.Promise promise = acquireToken(repository, transactionId, timestamp, handler);
        DefaultWriteLock writeLock = new DefaultWriteLock(repository, timestamp, promise, localNode, nodes);
        handler.setWriteLock(writeLock);
        return writeLock;
    }

    @VisibleForTesting
    String getHolder(@Nonnull Repository repository) {
        RepositoryToken token = tokenByRepositorySlug.get(requireNonNull(repository, "repository").getSlug());
        return token == null ? null : token.getHolder();
    }

    private RepositoryToken getToken(Repository repository) {
        return tokenByRepositorySlug.compute(repository.getSlug(), (slug, token) -> {
            if (token == null || token.referenceCount.incrementAndGet() == 1) {
                //If referenceCount == 1, we've hit a race condition where the last token request for the
                //repository has just completed and has been removed. When this happens referenceCount reaches 0
                //and the RepositoryToken is invalidated.
                token = new RepositoryToken(slug);
            }
            return token;
        });
    }

    private void maybeRetryYieldRequests() {
        tokenByRepositorySlug.values().forEach(token -> token.maybeRetryYieldRequest(retryYieldRequestInterval));
    }

    private static class TokenPromise implements Comparable<TokenPromise>, WriteToken.Promise {

        private final WriteTokenHandler handler;
        private final long timestamp;
        private final RepositoryToken token;
        private final String transactionId;

        private boolean closed;
        private volatile Instant lastYieldRequestedTimestamp;
        private volatile Instant yieldRequestedTimestamp;
        private volatile WriteTokenStub stub;

        TokenPromise(RepositoryToken token, String transactionId, long timestamp, WriteTokenHandler handler) {
            this.handler = handler;
            this.timestamp = timestamp;
            this.token = token;
            this.transactionId = transactionId;
        }

        @Override
        public void close() {
            synchronized (token.lock) {
                //Close only once
                if (!closed) {
                    closed = true;

                    token.release(this);
                    stub = null;
                }
            }
        }

        @Override
        public int compareTo(@Nonnull TokenPromise other) {
            //Sort by timestamp and then transactionId to achieve consistent ordering on all Mesh nodes
            int result = Long.compare(timestamp, other.timestamp);
            return result == 0 ? transactionId.compareTo(other.transactionId) : result;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) {
                return true;
            }
            if (o == null || getClass() != o.getClass()) {
                return false;
            }
            TokenPromise that = (TokenPromise) o;
            return timestamp == that.timestamp &&
                    token.equals(that.token) &&
                    transactionId.equals(that.transactionId);
        }

        @Override
        public int hashCode() {
            return Objects.hash(timestamp, token, transactionId);
        }

        @Override
        public String toString() {
            return "{key=" + token.slug + ", t=" + timestamp + ", tx=" + transactionId + "}";
        }

        boolean grantToken() {
            synchronized (token.lock) {
                if (closed) {
                    return false;
                }
                if (stub == null && handler != null) {
                    stub = new WriteTokenStub(this);
                    handler.onGranted(stub);
                }
                return true;
            }
        }

        void maybeRetryYieldRequest(Duration retryAfterDuration) {
            Instant instant = lastYieldRequestedTimestamp;
            log.debug("Attempting to retry yielding the token {}", this);
            if (stub != null && instant != null && Instant.now().isAfter(instant.plus(retryAfterDuration))) {
                //Token yield was requested more than retryDuration ago but the token hasn't been yielded yet. Try again
                synchronized (token.lock) {
                    //Double checked locking
                    instant = lastYieldRequestedTimestamp;
                    if (stub != null && instant != null && Instant.now().isAfter(instant.plus(retryAfterDuration))) {
                        log.info("Release of write token for {} was requested at {} but hasn't been yielded yet. " +
                                 "Retrying..", token.slug, instant);
                        requestRelease();
                    }
                }
            }
        }

        void onYielded(WriteTokenStub stub) {
            synchronized (token.lock) {
                if (this.stub == stub) {
                    this.stub = null;
                    token.onYielded(this, false);
                    lastYieldRequestedTimestamp = null;
                    yieldRequestedTimestamp = null;
                } else {
                    log.warn("Write token for {} was yielded after it had already been invalidated. Was the " +
                            "same token yielded multiple times?", token.slug);
                }
            }
        }

        void requestRelease() {
            synchronized (token.lock) {
                if (stub != null) {
                    //This requests the handler to yield the token. If the handler complies and closes the
                    //WriteTokenStub, onYielded is called.
                    lastYieldRequestedTimestamp = Instant.now();
                    if (yieldRequestedTimestamp == null) {
                        yieldRequestedTimestamp = lastYieldRequestedTimestamp;
                    }
                    handler.yieldToken();
                }
            }
        }
    }

    //A note on concurrency: both the onGranted and yieldToken methods are called from the local token's TokenPromise,
    //while synchronized on the RepositoryToken.lock to ensure granting the token and requests to yield the token
    //cannot be processed out of order.
    private static class WriteLockTokenHandler implements WriteTokenHandler {

        private final MeshNode node;

        private volatile WriteToken token;
        private volatile WriteLock writeLock;

        WriteLockTokenHandler(MeshNode node) {
            this.node = node;
        }

        @Override
        public void onGranted(@Nonnull WriteToken token) {
            this.token = requireNonNull(token, "token");
            WriteLock writeLock = this.writeLock;
            if (writeLock != null) {
                writeLock.onTokenGranted(node);
            }
        }

        @Override
        public void yieldToken() {
            WriteLock writeLock = this.writeLock;
            WriteToken token = this.token;
            if (writeLock != null && token != null && writeLock.yieldToken(node)) {
                token.yieldToken();
                this.token = null;
            }
        }

        void setWriteLock(WriteLock writeLock) {
            this.writeLock = writeLock;
            WriteToken token = this.token;
            if (token != null) {
                writeLock.onTokenGranted(node);
            }
        }
    }

    /**
     * Token that is provided to the registered {@link WriteTokenHandler handler} when the token is granted.
     * When this token is closed, it returns the token to the registry and re-activates the {@link TokenPromise}.
     * <p>
     * Note that a new WriteTokenStub instance is created every time the token is granted to the same promise.
     */
    private static class WriteTokenStub implements WriteToken {

        private final TokenPromise promise;

        private WriteTokenStub(TokenPromise promise) {
            this.promise = promise;
        }

        @Override
        public String toString() {
            return "WriteTokenStub{holder=" + promise.transactionId + ", repository=" + promise.token.slug + "}";
        }

        @Override
        public void yieldToken() {
            promise.onYielded(this);
        }
    }

    private class RepositoryToken {

        final Object lock;

        private final NavigableSet<TokenPromise> promises;
        private final AtomicInteger referenceCount;
        private final String slug;

        //The promise that is currently fulfilled (and whose handler holds the token)
        private volatile TokenPromise current;

        RepositoryToken(String slug) {
            this.slug = slug;

            lock = new Object();
            promises = new ConcurrentSkipListSet<>();
            referenceCount = new AtomicInteger(1);
        }

        @Override
        public String toString() {
            return slug;
        }

        @Nonnull
        TokenPromise acquire(String transactionId, long timestamp, WriteTokenHandler handler) {
            TokenPromise promise = new TokenPromise(this, transactionId, timestamp, handler);
            performInitialCallback(promise);
            return promise;
        }

        String getHolder() {
            TokenPromise current = this.current;
            return current == null ? null : current.transactionId;
        }

        void maybeRetryYieldRequest(Duration retryAfterDuration) {
            TokenPromise current = this.current;
            if (current != null) {
                current.maybeRetryYieldRequest(retryAfterDuration);
            }
        }

        void release(TokenPromise promise) {
            if (promises.remove(promise)) {
                if (referenceCount.decrementAndGet() == 0) {
                    log.debug("Token {} released (queue={})", promise, promises);
                    //Nothing is using this token anymore, remove the token from the registry
                    //noinspection ObjectEquality
                    tokenByRepositorySlug.compute(slug, (k, current) -> current == this ? null : current);
                } else {
                    onYielded(promise, true);
                }
            }
        }

        @SuppressWarnings("ObjectEquality")
        private void onYielded(TokenPromise request, boolean released) {
            synchronized (lock) {
                if (current == request) {
                    //Grant the token to the highest priority item in the queue
                    log.debug("Token {} {} (queue={})", request, released ? "released" : "yielded", promises);
                    current = promises.isEmpty() ? null : promises.first();
                    while (current != null) {
                        try {
                            if (current.grantToken()) {
                                log.debug("Token {} granted (queue={})", current, promises);
                                return;
                            }
                            //This can happen if current is closed concurrently (by the token requester)
                            promises.remove(current);
                            current = promises.isEmpty() ? null : promises.first();
                        } catch (RuntimeException e) {
                            log.error("Failed to grant token {}; releasing token", current, e);
                            //Closing the request will mark the token as released from the holder, remove holder from
                            //the queue and attempt to grant the token to the next request in the queue.
                            current.close();
                        }
                    }
                }
            }
        }

        private void performInitialCallback(TokenPromise promise) {
            synchronized (lock) {
                promises.add(promise);
                TokenPromise first = promises.first();
                try {
                    //noinspection ObjectEquality
                    if (first == promise) {
                        //Request is the highest priority request
                        if (current == null) {
                            //No other process holds the token at the moment, so grant it
                            if (promise.grantToken()) {
                                current = promise;
                                log.debug("Token {} granted (queue={})", promise, promises);
                            } else {
                                promises.remove(promise);
                                log.debug("Request for token {} was canceled before it could granted", promise);
                            }
                        } else {
                            //Another process with a lower priority current holds the request.
                            log.debug("Queued token request {} and asking {} to yield {} (queue={})",
                                    promise, current.transactionId, current.token, promises);
                            //Since this request has a higher priority, ask the current holder to yield immediately
                            current.requestRelease();
                        }
                    } else {
                        //There's another request with a higher priority. This request will have to wait for now
                        log.debug("Queued token request {} (queue={})", promise, promises);
                    }
                } catch (RuntimeException e) {
                    log.error("Handler for token request {} failed; releasing token", promise, e);
                    promise.close();
                    throw e;
                }
            }
        }
    }
}
