package com.atlassian.bitbucket.mesh.repository;

import com.atlassian.bitbucket.mesh.execution.ExecutionLocal;
import com.atlassian.bitbucket.mesh.execution.ExecutionManager;
import com.atlassian.bitbucket.mesh.request.RequestContext;
import com.atlassian.bitbucket.mesh.request.RequestManager;
import com.atlassian.bitbucket.mesh.rpc.v1.RepositoryServiceGrpc;
import com.atlassian.bitbucket.mesh.rpc.v1.TransactionServiceGrpc;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Sets;
import io.grpc.MethodDescriptor;
import io.grpc.ServerCall;
import org.apache.commons.lang3.mutable.MutableObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.time.Duration;
import java.time.Instant;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.ConcurrentHashMap;

import static java.util.Objects.requireNonNull;

public class DefaultRpcRepositoryAccessTracker implements RpcRepositoryAccessTracker {

    /**
     * Methods that are only called by peer nodes, and should not be counted as a repository access.
     * These methods are excluded because we only want to count client RPC calls and not prevent draining
     * from completing if we only repair / replication calls from peer nodes are active.
     */
    private static final Set<MethodDescriptor<?, ?>> INTERNAL_RPC_METHODS = ImmutableSet.of(
            RepositoryServiceGrpc.getRepairMeMethod(),
            RepositoryServiceGrpc.getRepairMethod(),
            TransactionServiceGrpc.getTransactMethod());
    private static final Logger log = LoggerFactory.getLogger(DefaultRpcRepositoryAccessTracker.class);

    private final ExecutionLocal<Set<Key>> accessedScopes;
    private final RepositoryVisitor<Void> repositoryAccessVisitor;
    private final RequestManager requestManager;
    private final Map<Key, DrainableRequestTracker> trackers;

    public DefaultRpcRepositoryAccessTracker(ExecutionManager executionManager, RequestManager requestManager) {
        this.requestManager = requestManager;

        accessedScopes = executionManager.newExecutionLocal(this::onRequestEnd);

        trackers = new ConcurrentHashMap<>();
        repositoryAccessVisitor = new RepositoryVisitor<Void>() {

            @Override
            public Void visit(@Nonnull ManagedRepository repository) {
                onManagedRepositoryAccessed(repository);
                return null;
            }

            @Override
            public Void defaultValue(Repository repository) {
                onNonManagedRepositoryAccessed(repository);
                return null;
            }
        };
    }

    @Nonnull
    @Override
    public CompletionStage<Void> drainAll() {
        return drain(Key.all());
    }

    @Nonnull
    @Override
    public CompletionStage<Void> drainHierarchy(@Nonnull String partitionId, @Nonnull String hierarchyId) {
        return drain(Key.hierarchy(partitionId, hierarchyId));
    }

    @Nonnull
    @Override
    public CompletionStage<Void> drainPartition(@Nonnull String partitionId) {
        return drain(Key.partition(partitionId));
    }

    @Override
    public int getActiveRequestCount() {
        return getActiveRequestCount(Key.all());
    }

    @Override
    public int getActiveRequestCountByHierarchy(@Nonnull String partitionId, @Nonnull String hierarchyId) {
        return getActiveRequestCount(Key.hierarchy(partitionId, hierarchyId));
    }

    @Override
    public int getActiveRequestCountByPartition(@Nonnull String partitionId) {
        return getActiveRequestCount(Key.partition(partitionId));
    }

    @Override
    public void onRepositoryAccessed(@Nonnull Repository repository) {
        requireNonNull(repository, "repository");

        if (!isCalledFromInternalRpc()) {
            //Don't track repository accesses from "internal" RPCs such as write replication
            repository.accept(repositoryAccessVisitor);
        }
    }

    private static boolean isInternalRpc(RequestContext context) {
        Object rawRequest = context.getRawRequest();
        if (rawRequest instanceof ServerCall) {
            return INTERNAL_RPC_METHODS.contains(((ServerCall<?, ?>) rawRequest).getMethodDescriptor());
        }
        return false;
    }

    private CompletionStage<Void> drain(Key key) {
        MutableObject<Instant> start = new MutableObject<>();
        DrainableRequestTracker tracker = trackers.computeIfPresent(key, (k, v) -> {
            log.info("Draining {}...", key);
            start.setValue(Instant.now());
            v.startDraining();
            return v;
        });

        if (tracker == null) {
            //If there is no tracker for the key, there are no active requests so we're already done draining
            log.info("Drained {} in 0ms", key);
            return CompletableFuture.completedFuture(null);
        }

        return tracker.asCompletionStage().thenApply(ignored -> {
            log.info("Drained {} in {}ms", key, Duration.between(start.getValue(), Instant.now()).toMillis());
            return null;
        });
    }

    private int getActiveRequestCount(Key key) {
        DrainableRequestTracker tracker = trackers.get(key);
        return tracker == null ? 0 : tracker.getActiveRequestCount();
    }

    private boolean isCalledFromInternalRpc() {
        return requestManager.getContext()
                .map(DefaultRpcRepositoryAccessTracker::isInternalRpc)
                .orElse(false); //no RPC active
    }

    private void onManagedRepositoryAccessed(@Nonnull ManagedRepository repository) {
        Key allKey = Key.all();
        Key partitionKey = Key.partition(repository.getPartition());
        Key hierarchyKey = Key.hierarchy(repository.getPartition(), repository.getHierarchyId());

        Set<Key> scopes = accessedScopes.computeIfAbsent(() -> {
            Set<Key> keys = Sets.newConcurrentHashSet();
            log.trace("Adding new Set [{}] in the ExecutionLocal", System.identityHashCode(keys));
            return keys;
        });
        if (scopes != null) {
            //scopes will be null if no execution context is active (so no active RPC)
            if (scopes.add(allKey)) {
                onRequestStarted(allKey);
            }
            if (scopes.add(partitionKey)) {
                onRequestStarted(partitionKey);
            }
            if (scopes.add(hierarchyKey)) {
                onRequestStarted(hierarchyKey);
            }
        }
    }

    private void onNonManagedRepositoryAccessed(@Nonnull Repository ignored) {
        Key allKey = Key.all();

        Set<Key> scopes = accessedScopes.computeIfAbsent(() -> {
            Set<Key> keys = Sets.newConcurrentHashSet();
            log.trace("Adding new Set [{}] in the ExecutionLocal", System.identityHashCode(keys));
            return keys;
        });
        if (scopes != null) {
            //scopes will be null if no execution context is active (so no active RPC)
            if (scopes.add(allKey)) {
                onRequestStarted(allKey);
            }
        }
    }

    private void onRequestStarted(Key key) {
        trackers.compute(key, (k, v) -> {
            DrainableRequestTracker tracker = v == null ? new DrainableRequestTracker() : v;
            int count = tracker.onRequestStarted();
            log.debug("{} has {} active requests", key, count);
            return tracker;
        });
    }

    private void onRequestEnd(Set<Key> accessedScopes) {
        log.trace("Request with ExecutionLocal containing Set [{}] has ended", System.identityHashCode(accessedScopes));
        for (Key key : accessedScopes) {
            trackers.compute(key, (k, tracker) -> {
                if (tracker == null || tracker.onRequestEnded() == 0) {
                    //Last request ended for the scope in question, remove the tracker to save memory
                    //and clear trackers that have completed draining
                    return null;
                }
                return tracker;
            });
        }
    }

    private enum TrackerType {
        ALL,
        HIERARCHY,
        PARTITION
    }

    private static class Key {

        private static final Key ALL = new Key(TrackerType.ALL, null);

        private final TrackerType type;
        private final String id;

        Key(@Nonnull TrackerType type, String id) {
            this.type = type;
            this.id = id;
        }

        static Key all() {
            return ALL;
        }

        static Key hierarchy(@Nonnull String partitionId, @Nonnull String hierarchyId) {
            return new Key(TrackerType.HIERARCHY, "p/" + requireNonNull(partitionId, "partitionId") +
                    "/h/" + requireNonNull(hierarchyId, "hierarchyId"));
        }

        static Key partition(@Nonnull String partitionId) {
            return new Key(TrackerType.PARTITION, requireNonNull(partitionId, "partitionId"));
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) {
                return true;
            }
            if (o == null || getClass() != o.getClass()) {
                return false;
            }
            Key key = (Key) o;
            return type == key.type &&
                    Objects.equals(id, key.id);
        }

        @Override
        public int hashCode() {
            return Objects.hash(type, id);
        }

        @Override
        public String toString() {
            return type.name().toLowerCase() + (id == null ? "" : " " + id);
        }
    }
}
