package com.atlassian.bitbucket.mesh.repository;

import com.atlassian.bitbucket.mesh.coordination.LogicalClock;
import com.atlassian.bitbucket.mesh.event.EventListener;
import com.atlassian.bitbucket.mesh.event.EventPublisher;
import com.atlassian.bitbucket.mesh.management.NodeProvisionedEvent;
import com.atlassian.bitbucket.mesh.management.SidebandManager;
import com.atlassian.bitbucket.mesh.management.SupportInfoProvider;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcRepositoryReplicaObservation;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcSidebandFollowerFragment;
import com.fasterxml.jackson.core.JsonGenerator;
import com.google.common.primitives.Ints;
import org.apache.commons.lang3.mutable.MutableObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.BiConsumer;

import static com.atlassian.bitbucket.mesh.repository.ReplicaState.*;
import static java.util.Objects.requireNonNull;

public class DefaultReplicaStateRegistry implements ReplicaStateRegistry, SupportInfoProvider {

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

    private static final RepositoryVisitor<Long> GET_VERSION = new RepositoryVisitor<Long>() {

        @Override
        public Long visit(@Nonnull ManagedRepository repository) {
            return repository.getVersion();
        }

        @Override
        public Long defaultValue(Repository repository) {
            return 0L;
        }
    };

    private final EventPublisher eventPublisher;
    private final LogicalClock logicalClock;
    private final Map<String, ReplicaStateObservation> replicaStates;
    private final SidebandManager sidebandManager;

    public DefaultReplicaStateRegistry(EventPublisher eventPublisher, LogicalClock logicalClock,
                                       SidebandManager sidebandManager) {
        this.eventPublisher = eventPublisher;
        this.logicalClock = logicalClock;
        this.sidebandManager = sidebandManager;

        replicaStates = new ConcurrentHashMap<>();
    }

    @Override
    public int countByState(@Nonnull ReplicaState state) {
        requireNonNull(state, "state");
        return Ints.saturatedCast(replicaStates.values().stream()
                .filter(observation -> observation.getState() == state)
                .count());
    }

    @Nonnull
    @Override
    public RepositoryReplicaDetails getDetails(@Nonnull Repository repository) {
        requireNonNull(repository, "repository");

        ReplicaStateObservation latestObservation = replicaStates.get(repository.getId());
        long version = repository.accept(GET_VERSION);
        return latestObservation == null ?
                new RepositoryReplicaDetails.Builder(CONSISTENT)
                        .observedVersion(version)
                        .version(version)
                        .build() :
                new RepositoryReplicaDetails.Builder(latestObservation.getState())
                        .observedVersion(latestObservation.getVersion())
                        .version(version)
                        .build();
    }

    @Nonnull
    @Override
    public ReplicaState getState(@Nonnull Repository repository) {
        return getState(requireNonNull(repository, "repository").getId());
    }

    @Nonnull
    @Override
    public ReplicaState getState(@Nonnull String repositoryId) {
        requireNonNull(repositoryId, "repositoryId");

        ReplicaStateObservation latestObservation = replicaStates.get(Repository.canonicalId(repositoryId));
        return latestObservation == null ? CONSISTENT : latestObservation.getState();
    }

    @Override
    public void markInconsistent(@Nonnull MarkInconsistentParameters parameters) {
        requireNonNull(parameters, "parameters");

        long timestamp = logicalClock.tick();
        boolean fromControlPlane = parameters.isFromControlPlane();
        parameters.getInconsistentReplicas().forEach((repositoryId, version) ->
                markInconsistent(repositoryId, version, timestamp, !fromControlPlane));

        if (fromControlPlane && parameters.isComplete()) {
            Map<String, Long> provided = parameters.getInconsistentReplicas();
            replicaStates.forEach((repositoryId, observation) -> {
                if (!provided.containsKey(repositoryId)) {
                    notifyControlPlane(repositoryId, observation.getState(), observation.getVersion());
                }
            });
        }
    }

    @EventListener
    public void onNodeProvisioned(@Nonnull NodeProvisionedEvent ignored) {
        log.debug("Clearing replica states");
        replicaStates.clear();
    }

    @EventListener
    public void onReplicaStateObserved(@Nonnull ReplicaStateObservedEvent event) {
        requireNonNull(event, "event");

        ReplicaStateObservation observation = event.getObservation();
        MutableObject<ReplicaStateUpdatedEvent> toPublish = new MutableObject<>();
        String repositoryId = Repository.canonicalId(event.getRepositoryId());
        replicaStates.compute(repositoryId, (id, current) -> {
            if (log.isTraceEnabled()) {
                log.trace("[{}] Our current observation is {}", repositoryId, current);
                log.trace("[{}] Our observed observation is {}", repositoryId, observation);
            }
            if (current == null || observation.compareTo(current) > 0 || observation.isForce()) {
                //observation is newer than the current registration (ordered by version, timestamp), or a force update
                ReplicaState prevState = current == null ? CONSISTENT : current.getState();
                ReplicaState nextState = observation.getState();
                if (log.isTraceEnabled()) {
                    log.trace("[{}] Our prevState is {} and our nextState before modification is {}", repositoryId,
                            prevState, nextState);
                }
                if (prevState == INCONSISTENT && nextState == UNKNOWN) {
                    //don't move from INCONSISTENT -> UNKNOWN
                    nextState = INCONSISTENT;
                }
                if (log.isTraceEnabled()) {
                    log.trace("[{}] Our nextState is now (or still) {}", repositoryId, nextState);
                }

                if (prevState != nextState) {
                    log.debug("[{}] Updating replica state {} -> {}", repositoryId, prevState, nextState);
                    toPublish.setValue(new ReplicaStateUpdatedEvent(repositoryId, nextState));
                }

                return nextState == CONSISTENT ? null : new ReplicaStateObservation(
                        nextState, observation.getLogicalTimestamp(), observation.getVersion());
            }

            log.debug("[{}] Ignoring stale {} observation for {}", repositoryId, observation.getState(), id);
            //current is newer than event or refers to a newer repository version;
            //ignore the event because it is stale
            return current;
        });

        ReplicaStateUpdatedEvent updatedEvent = toPublish.getValue();
        if (updatedEvent != null) {
            if (log.isTraceEnabled()) {
                log.trace("[{}] We do have an updated event {}", repositoryId, updatedEvent);
            }
            if (!event.isReported()) {
                if (log.isTraceEnabled()) {
                    log.trace("[{}] Our updated event had not already been reported", repositoryId);
                }
                //Before we publish the event (which may trigger repair), we need to inform the control plane
                //of the observed state. If the replica is inconsistent, sending this here will guarantee the
                //control plane receives the message before repair can declare the replica consistent again
                notifyControlPlane(repositoryId, updatedEvent.getState(), observation.getVersion());
            }
            eventPublisher.publish(updatedEvent);
        }
    }

    @EventListener
    public void onRepositoryDeleted(@Nonnull RepositoryDeletedEvent event) {
        requireNonNull(event, "event");

        replicaStates.remove(event.getRepository().getId());
    }

    @Override
    public void remove(@Nonnull String repositoryId) {
        ReplicaStateObservation observation = replicaStates.remove(repositoryId);
        if (observation == null) {
            log.debug("No replica state was present in the registry for {}", repositoryId);
        } else {
            log.debug("Removed {} replica state from the registry for {}", observation.getState(), repositoryId);
        }
    }

    @Override
    public void setState(@Nonnull String repositoryId, @Nonnull ReplicaState state, long version, boolean force) {
        onReplicaStateObserved(new ReplicaStateObservedEvent(repositoryId,
                new ReplicaStateObservation(state, logicalClock.tick(), version, force), true));
    }

    @Override
    public void streamInconsistentReplicas(@Nonnull BiConsumer<String, Long> callback) {
        replicaStates.forEach((repositoryId, observation) -> {
            if (observation.getState() != CONSISTENT) {
                callback.accept(repositoryId, observation.getVersion());
            }
        });
    }

    @Override
    public void writeSupportInfo(@Nonnull JsonGenerator json) throws IOException {
        json.writeObjectFieldStart("inconsistent-replicas");
        try {
            for (Map.Entry<String, ReplicaStateObservation> entry : replicaStates.entrySet()) {
                ReplicaStateObservation observation = entry.getValue();
                json.writeStringField(entry.getKey(), observation.getState().name());
            }
        } finally {
            json.writeEndObject();
        }
    }

    private void markInconsistent(String repositoryId, long version, long timestamp,
                                  boolean notifyControlPlane) {
        log.debug("[{}] - Marking inconsistent for version {} with timestamp {} and notifying the control-plane {}",
                repositoryId, version, timestamp, notifyControlPlane);
        MutableObject<ReplicaStateUpdatedEvent> toPublish = new MutableObject<>();
        replicaStates.compute(repositoryId, (id, observation) -> {
            if (observation == null || observation.getVersion() <= version) {
                ReplicaState prevState = observation == null ? CONSISTENT : observation.getState();
                if (prevState != INCONSISTENT) {
                    toPublish.setValue(new ReplicaStateUpdatedEvent(repositoryId, INCONSISTENT));
                }

                return new ReplicaStateObservation(INCONSISTENT, timestamp, version);
            }

            //Nothing to update or publish
            return observation;
        });

        //If updated, publish the event to trigger / schedule repair
        ReplicaStateUpdatedEvent updatedEvent = toPublish.getValue();
        if (updatedEvent != null) {
            if (notifyControlPlane) {
                notifyControlPlane(repositoryId, updatedEvent.getState(), version);
            }
            eventPublisher.publish(updatedEvent);
        }
    }

    private void notifyControlPlane(String repositoryId, ReplicaState state, long version) {
        sidebandManager.send(RpcSidebandFollowerFragment.newBuilder()
                .setRepositoryReplicaObservation(RpcRepositoryReplicaObservation.newBuilder()
                        .setRepository(repositoryId)
                        .setState(state.toRpc())
                        .setVersion(version))
                .build());
    }
}
