package com.atlassian.bitbucket.mesh.repair;

import com.atlassian.bitbucket.mesh.AbstractStatusException;
import com.atlassian.bitbucket.mesh.client.FetchObjectsParameters;
import com.atlassian.bitbucket.mesh.config.GrpcConfig;
import com.atlassian.bitbucket.mesh.coordination.LogicalClock;
import com.atlassian.bitbucket.mesh.event.EventPublisher;
import com.atlassian.bitbucket.mesh.git.DarkRefMapper;
import com.atlassian.bitbucket.mesh.git.Reflog;
import com.atlassian.bitbucket.mesh.git.exception.RepositoryRepairFailedException;
import com.atlassian.bitbucket.mesh.git.ref.Branch;
import com.atlassian.bitbucket.mesh.grpc.AbstractFutureResponseObserver;
import com.atlassian.bitbucket.mesh.repository.ReplicaState;
import com.atlassian.bitbucket.mesh.repository.ReplicaStateObservation;
import com.atlassian.bitbucket.mesh.repository.ReplicaStateObservedEvent;
import com.atlassian.bitbucket.mesh.repository.Repository;
import com.atlassian.bitbucket.mesh.rpc.util.ByteStringUtils;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.google.common.collect.Collections2;
import com.google.common.collect.Iterables;
import com.google.common.io.Closer;
import com.google.protobuf.ByteString;
import io.grpc.StatusRuntimeException;
import io.grpc.stub.ClientCallStreamObserver;
import io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.time.Instant;
import java.util.Set;
import java.util.function.Function;

import static com.atlassian.bitbucket.mesh.gprc.GrpcSupport.getDeadline;
import static com.atlassian.bitbucket.mesh.rpc.util.ByteStringUtils.emptyToNull;
import static com.google.common.base.Preconditions.checkState;
import static java.util.Collections.singleton;
import static java.util.Objects.requireNonNull;

/**
 * Implements the {@code RepositoryService/Repair} and {@code RepositoryService/RepairMe} protocols from the side that
 * serves as repository replica to repair <em>from</em>. For {@code Repair} that is the client side and for
 * {@code RepairMe} it's the server side.
 * <p>
 * The {@code Repair} and {@code RepairMe} flows are nearly identical and differ only in how the flow is initiated and
 * the capabilities are negotiated between the client and server:
 * <ul>
 *     <li>{@code RepairMe}: the client drives the repair and initiates the flow by requesting the repository state
 *     on the server. The client advertises its capabilities as part of the {@link RpcGetRepositoryStateRequest}.
 *     The server calculates the common capabilities and provides them on the {@link RpcGetRepositoryDetailsResponse}.
 *     From this point on, the client requests the information from the server and repairs its local repository.</li>
 *     <li>{@code Repair}: the client <em>initiates</em> the repair process by sending a {@link RpcRepairRequest}.
 *     The server then takes over and drives the repair using the same flow as {@code RepairMe}.</li>
 * </ul>
 *
 * @see RepairSource
 */
public class RepairTarget
        extends AbstractFutureResponseObserver<RpcRepairTargetFragment, RpcRepairSourceFragment, RepairSummary> {

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

    private final Instant deadline;
    private final EventPublisher eventPublisher;
    private final GrpcConfig grpcConfig;
    private final LogicalClock logicalClock;
    private final InteractiveRepairHelper repairHelper;

    private volatile Set<String> currentReflogBatch;
    private volatile InteractiveRepair.ReflogRepair currentReflogRepair;
    private volatile long logicalTimestamp;
    private volatile InteractiveRepair.PackProcessor packProcessor;
    private volatile InteractiveRepair.RefComparer refComparer;
    private volatile InteractiveRepair.ReflogComparer reflogComparer;
    private volatile InteractiveRepair repair;
    private volatile StreamObserver<RpcRepairTargetFragment> repairSource;
    private volatile boolean repairsApplied;
    private volatile String sourceRepositoryId;
    private volatile ReplicaStateObservation sourceStateObservation;
    private volatile Instant sourceStateTimestamp;
    private volatile String targetRepositoryId;

    /**
     * Constructor used by {@code Repair} by the RPC service implementation.
     */
    public RepairTarget(EventPublisher eventPublisher, GrpcConfig grpcConfig, LogicalClock logicalClock,
                        InteractiveRepairHelper repairHelper, StreamObserver<RpcRepairTargetFragment> repairSource,
                        Duration timeout) {
        this.eventPublisher = eventPublisher;
        this.grpcConfig = grpcConfig;
        this.logicalClock = logicalClock;
        this.repairHelper = repairHelper;
        this.repairSource = repairSource;

        deadline = getDeadline(timeout);
    }

    /**
     * Constructor used by {@code RepairMe} by the RPC client implementation.
     */
    public RepairTarget(EventPublisher eventPublisher, GrpcConfig grpcConfig, LogicalClock logicalClock,
                        InteractiveRepairHelper repairHelper, Duration timeout) {
        this(eventPublisher, grpcConfig, logicalClock, repairHelper, null, timeout);
    }

    @Override
    public void beforeStart(ClientCallStreamObserver<RpcRepairTargetFragment> requestStream) {
        repairSource = requestStream;
        super.beforeStart(requestStream);
    }

    @Override
    public void onError(Throwable t) {
        try {
            super.onError(t);
        } finally {
            releaseRepairResources();
        }
    }

    @Override
    public void onNext(RpcRepairSourceFragment fragment) {
        try {
            //This only updates the logical clock on repair request to initialize logicalTimestamp at the start of
            //the repair flow and for repository state messages because these can trigger timestamped replica state
            //observations.
            switch (fragment.getCommandOneofCase()) {
                case REPAIR_REQUEST:
                    //Step 0: initiate repair when called RepairTarget is on the server side.
                    //The REPAIR_REQUEST fragment is only sent by the client in RepositoryService/Repair to let the
                    //server (us) know which repository should be repaired. This step is not necessary in the
                    //RepositoryService/RepairMe flow because in that case, the client (us) drives the repair
                    //and calls startRepair directly from MeshClient.repairMe.
                    logicalTimestamp = logicalClock.onReceived(fragment.getLogicalTimestamp());
                    RpcRepairRequest repairRequest = fragment.getRepairRequest();
                    String targetRepositoryId = repairRequest.getRepository();
                    String sourceRepositoryId = repairRequest.hasSourceRepository() ?
                            repairRequest.getSourceRepository() : targetRepositoryId;
                    startRepair(sourceRepositoryId, targetRepositoryId);
                    break;
                case REPOSITORY_STATE:
                    //Step 1: compare hashes to determine whether there is something to repair
                    //Step 8: if any repairs have been applied, compare hashes again to determine whether we're now
                    //        back in sync.
                    logicalTimestamp = logicalClock.onReceived(fragment.getLogicalTimestamp());
                    onRepositoryState(fragment.getRepositoryState());
                    break;
                case REPOSITORY_DETAILS:
                    //Step 2: (on metadata hash mismatch) update metadata/create missing repo
                    onRepositoryDetails(fragment.getRepositoryDetails());
                    break;
                case LIST_REFS_CHUNK:
                    //Step 3: (on content hash mismatch) ref advertisement to compare local and source refs
                    onListRefsChunk(fragment.getListRefsChunk());
                    break;
                case LIST_REFLOGS_CHUNK:
                    //Step 4: (on content hash mismatch) reflog advertisement to compare local and source
                    //reflogs. For managed repositories, both the reflogs under shared-objects/logs/r/<repo-id> and
                    //<repo-dir>/logs are compared.
                    onListReflogsChunk(fragment.getListReflogsChunk());
                    break;
                case REFLOG_START:
                    //Step 5: (if reflogs need to be synced) to signal that the source is going to send a reflog across
                    onReflogStart(fragment.getReflogStart());
                    break;
                case REFLOG_CHUNK:
                    //Step 6: (if reflogs need to be synced) provides a chunk of the 'current' reflog being streamed
                    onReflogChunk(fragment.getReflogChunk());
                    break;
                case PACK_CHUNK:
                    //Step 7: (if missing objects have been identified)
                    onPackObjectsChunk(fragment.getPackChunk());
                    break;
                default:
                    log.warn("Unsupported fragment of type {}", fragment.getCommandOneofCase());
            }
        } catch (RuntimeException e) {
            handleError(e);
        }
    }

    /**
     * Starts the repair flow.
     *
     * @param targetRepositoryId the repository to repair
     */
    public void startRepair(@Nonnull String sourceRepositoryId, @Nonnull String targetRepositoryId) {
        this.sourceRepositoryId = requireNonNull(sourceRepositoryId, "sourceRepositoryId");
        this.targetRepositoryId = requireNonNull(targetRepositoryId, "targetRepositoryId");

        if (log.isDebugEnabled()) {
            log.debug("[{}] Starting repair{}", targetRepositoryId,
                    targetRepositoryId.equals(sourceRepositoryId) ? "" : " from " + sourceRepositoryId);
        }
        try {
            repair = repairHelper.startRepair(targetRepositoryId, getTimeout());
            //If we've started repairs, make sure we release all resources when the RPC finishes
            addListener((summary, thrown) -> releaseRepairResources());

            //Request the repair source's repository state so we can decide which aspects need repair
            sourceStateTimestamp = Instant.now();
            sendFragmentToSource(fragment -> fragment.setGetRepositoryState(RpcGetRepositoryStateRequest.newBuilder()
                    .setCapabilities(RepairCapability.ADVERTISEMENT) //advertise all capabilities
                    .setRepository(sourceRepositoryId)));
        } catch (AbstractStatusException e) {
            //There are two call paths into startRepair. One triggered from a gRPC message, in which case the normal
            //exception handling is in place. The other call path is from MeshClient.repairMe. There, the normal
            //exception translation is _not_ in place, so we need to propagate the error to the repairSource explicitly
            //here.
            repairSource.onError(e.toStatusException());
            throw e;
        }
    }

    @Nonnull
    private Duration getTimeout() {
        return Duration.between(Instant.now(), deadline);
    }

    private void handleError(RuntimeException e) {
        if (e instanceof AbstractStatusException) {
            e = ((AbstractStatusException) e).toStatusException();
        } else if (!(e instanceof StatusRuntimeException)) {
            e = new RepositoryRepairFailedException(e).toStatusException();
        }

        log.warn("[{}] Repair failed", targetRepositoryId, e);
        repairSource.onError(e);
    }

    private boolean isMigration() {
        return sourceRepositoryId != null && !sourceRepositoryId.equals(targetRepositoryId);
    }

    private void maybeFetchReflogs() {
        currentReflogBatch = repair.getNextReflogsToRepairBatch(5000);
        if (currentReflogBatch.isEmpty()) {
            //Done repairing shared _and_ normal reflogs, fetch any missing objects
            onMissingObjectsIdentified();
        } else {
            log.debug("[{}] Detected {} reflogs to repair", targetRepositoryId, currentReflogBatch.size());
            sendFragmentToSource(fragment -> fragment.setGetReflogs(RpcGetReflogsRequest.newBuilder()
                    .addAllPaths(ByteStringUtils.toByteStringsLazily(currentReflogBatch))
                    .setRepository(sourceRepositoryId)));
        }
    }

    private void onListReflogsChunk(RpcReflogsChunk chunk) {
        checkState(reflogComparer != null, "Reflog comparison not expected!");

        for (RpcOsFile file : chunk.getFilesList()) {
            reflogComparer.onSourceReflog(Reflog.fromOsFile(file));
        }
        if (!chunk.getMoreResponses()) {
            reflogComparer.close();
            reflogComparer = null;

            maybeFetchReflogs();
        }
    }

    private void onListRefsChunk(RpcRefsChunk chunk) {
        checkState(refComparer != null, "Refs not expected!");
        Repository repository = repair.getRepository();
        Function<String, String> darkRefMapper = repository == null
                ? Function.identity()
                : repository.accept(DarkRefMapper.INSTANCE);

        for (RpcRef rpcRef : chunk.getRefsList()) {
            refComparer.onSourceRef(new Branch.Builder()
                    .id(darkRefMapper.apply(rpcRef.getId().toStringUtf8()))
                    .latestCommit(rpcRef.getObjectId())
                    .build());
        }
        if (!chunk.getMoreResponses()) {
            refComparer.close();
            refComparer = null;

            if (repair.isOutOfSyncRefDetected()) {
                //Send out the request immediately to allow the source to acquire the GC lock concurrently
                log.debug("[{}] Starting reflog repair", targetRepositoryId);
                sendFragmentToSource(fragment -> fragment.setListReflogs(RpcListReflogsRequest.newBuilder()
                        .setRepository(sourceRepositoryId)));

                //Take out the GC lock before comparing reflogs. This prevents vacuuming from updating the shared
                //reflogs and archiving reflogs for closed PRS while we're streaming reflogs.
                log.debug("[{}] Starting reflog repair. Acquiring GC lock.", targetRepositoryId);
                repair.lockGc();

                reflogComparer = repair.compareReflogs();
            } else {
                log.debug("[{}] Ref comparison did not detect any out of sync refs", targetRepositoryId);
                repairsApplied = true;
                onRepairCompleted(false);
            }
        }
    }

    private void onMetadataSynced() {
        if (repair.isContentOutdated()) {
            log.debug("[{}] Content hash mismatch. Requesting refs from source.", repair.getRepository());
            refComparer = repair.compareRefs();
            sendFragmentToSource(fragment -> fragment.setListRefs(RpcListRefsRequest.newBuilder()
                    .setRepository(sourceRepositoryId)));
        } else {
            log.debug("[{}] Content hash is up-to-date. Completing repair.", repair.getRepository());
            onRepairCompleted(true);
        }
    }

    private void onMissingObjectsFetched() {
        log.debug("[{}] Updating refs and publishing staged reflogs", repair.getRepository());

        if (repair.maybeUpdateRefsAndReflogs()) {
            repairsApplied = true;
        }

        onRepairCompleted(true);
    }

    private void onMissingObjectsIdentified() {
        FetchObjectsParameters parameters = repair.getFetchParameters();
        if (parameters == null || parameters.getIncludes().isEmpty()) {
            //All objects are present locally
            onMissingObjectsFetched();
        } else {
            //Request a pack with missing objects and prepare a processor
            packProcessor = repair.newPackProcessor();

            //Send the request
            sendFragmentToSource(fragment -> fragment.setGetPack(RpcPackObjectsStdinFragment.newBuilder()
                    .setRequest(RpcPackObjectsRequest.newBuilder()
                            .setRepository(sourceRepositoryId)
                            .setThin(true))));

            //Send the includes and excludes as stdin bytes
            ByteString stdin = ByteString.EMPTY;
            long size = 0L;
            long maxSize = grpcConfig.getMaxFragmentSize();
            Iterable<String> revs = Collections2.transform(parameters.getIncludes(), rev -> rev + "\n");
            if (!parameters.getExcludes().isEmpty()) {
                revs = Iterables.concat(revs, Iterables.concat(singleton("--not\n"),
                        Collections2.transform(parameters.getExcludes(), rev -> rev + "\n")));
            }
            for (String rev : revs) {
                ByteString rpcRev = ByteString.copyFromUtf8(rev);
                if (size + rpcRev.size() > maxSize) {
                    RpcStreamChunk.Builder chunk = RpcStreamChunk.newBuilder().setData(stdin);
                    sendFragmentToSource(fragment -> fragment.setGetPack(RpcPackObjectsStdinFragment.newBuilder()
                            .setStdin(chunk)));
                    stdin = ByteString.EMPTY;
                    size = 0L;
                }
                stdin = stdin.concat(rpcRev);
                size += rpcRev.size();
            }

            log.debug("[{}] Requesting pack with missing objects ({} includes, {} excludes)", repair.getRepository(),
                    parameters.getIncludes().size(), parameters.getExcludes().size());
            RpcStreamChunk.Builder chunk = RpcStreamChunk.newBuilder().setData(stdin).setClosed(true);
            sendFragmentToSource(fragment -> fragment.setGetPack(RpcPackObjectsStdinFragment.newBuilder()
                    .setStdin(chunk)));
        }
    }

    private void onPackObjectsChunk(RpcStreamChunk chunk) {
        checkState(packProcessor != null, "Not expecting pack data");

        packProcessor.write(chunk.getData(), chunk.getClosed());

        if (chunk.getClosed()) {
            packProcessor.close();
            packProcessor = null;
            onMissingObjectsFetched();
        }
    }

    private void onReflogChunk(RpcStreamChunk chunk) {
        if (currentReflogRepair != null) {
            try {
                for (ByteBuffer buffer : chunk.getData().asReadOnlyByteBufferList()) {
                    currentReflogRepair.write(buffer);
                }
                if (chunk.getClosed()) {
                    currentReflogRepair.close();
                    currentReflogRepair = null;

                    if (currentReflogBatch.isEmpty()) {
                        //Request the next batch of reflogs to repair, if there are any
                        maybeFetchReflogs();
                    }
                }
            } catch (IOException e) {
                throw new RepositoryRepairFailedException(
                        "Failed to write reflog " + currentReflogRepair.getRefId(), e);
            }
        }
    }

    private void onReflogStart(RpcReflogStart request) {
        Repository repository = repair.getRepository();
        Function<String, String> darkRefMapper = repository == null
                ? Function.identity()
                : repository.accept(DarkRefMapper.INSTANCE);
        String refId = request.getPath().toStringUtf8();
        if (currentReflogBatch.remove(refId)) {
            String mappedRefId = darkRefMapper.apply(refId);
            try {
                currentReflogRepair = repair.startReflogRepair(mappedRefId);
            } catch (IOException e) {
                log.warn("[{}] Failed to open reflog for {} for repair", targetRepositoryId, mappedRefId, e);
            }
        } else {
            log.warn("[{}] Ignoring unexpected reflog content for {}", targetRepositoryId, refId);
        }
    }

    private void onRepairCompleted(boolean success) {
        if (repairsApplied) {
            log.debug("[{}] Repairs have been applied. Checking whether the replica is now consistent",
                    targetRepositoryId);
            sourceStateTimestamp = Instant.now();
            //Force a refresh of the hashes if repair was attempted, but failed and no out of sync refs were detected.
            boolean forceRefresh = !repair.isOutOfSyncRefDetected() && !success;
            if (forceRefresh) {
                repair.resetRepositoryHashes();
            }
            sendFragmentToSource(fragment -> fragment.setGetRepositoryState(RpcGetRepositoryStateRequest.newBuilder()
                    .setNoCache(forceRefresh)
                    .setRepository(sourceRepositoryId)));
        } else {
            wrapUp(success);
        }
    }

    private void onRepositoryDetails(RpcGetRepositoryDetailsResponse response) {
        repair.updateMetadata(emptyToNull(response.getDefaultBranch()), response.getMetadataMap(),
                response.getTranscodeEnabled());
        repairsApplied = true;
        onMetadataSynced();
    }

    private void onRepositoryState(RpcGetRepositoryStateResponse response) {
        ReplicaState state = ReplicaState.fromRpc(response.getState());
        log.debug("[{}] Received source state: {} (Metadata: {}, Content: {})",
                targetRepositoryId, state, response.getMetadataHash(), response.getContentHash());
        sourceStateObservation = new ReplicaStateObservation(state, logicalTimestamp, response.getVersion());

        if (repairsApplied) {
            //We've already applied repairs and are now checking whether our replica has caught up completely
            repair.withLock(() -> {
                repair.setSourceState(response.getMetadataHash(), response.getContentHash());
                wrapUp(repair.isConsistent());
            });
            return;
        }

        if (response.getState() == RpcReplicaState.REPLICA_STATE_CONSISTENT) {
            if (response.hasCapabilities()) {
                repair.setCapabilities(RepairCapability.getSharedCapabilities(response.getCapabilities()));
            }

            repair.withLock(() -> {
                //Repair source is up-to-date
                repair.setSourceState(response.getMetadataHash(), response.getContentHash());

                if (repair.isMetadataOutdated()) {
                    log.debug("[{}] Requesting repository details because {}", targetRepositoryId,
                            (repair.getRepository() == null ? "repository is missing" : "metadata is outdated"));

                    //Retrieve the repository details so we can create/update the repository itself
                    sendFragmentToSource(fragment ->
                            fragment.setGetRepositoryDetails(RpcGetRepositoryDetailsRequest.newBuilder()
                                    .setRepository(sourceRepositoryId)));
                } else {
                    log.debug("[{}] Metadata is in sync", targetRepositoryId);
                    onMetadataSynced();
                }
            });
        } else {
            log.warn("[{}] Cannot repair because the source is not up to date ({})",
                    targetRepositoryId, response.getState());
            onRepairCompleted(false);
        }
    }

    private void releaseRepairResources() {
        try (Closer closer = Closer.create()) {
            closer.register(repair);
            if (currentReflogRepair != null) {
                currentReflogRepair.abort();
            }
        } catch (IOException | RuntimeException e) {
            log.warn("[{}] Error while releasing repair resources", targetRepositoryId, e);
        }
    }

    private void sendFragmentToSource(Function<RpcRepairTargetFragment.Builder, RpcRepairTargetFragment.Builder> builder) {
        repairSource.onNext(builder.apply(RpcRepairTargetFragment.newBuilder()
                .setLogicalTimestamp(logicalClock.tick()))
                .build());
    }

    private void wrapUp(boolean consistent) {
        try {
            ReplicaState sourceState = sourceStateObservation.getState();
            ReplicaState targetState = consistent ? sourceState :
                    sourceState == ReplicaState.CONSISTENT ? ReplicaState.INCONSISTENT : ReplicaState.UNKNOWN;
            long sourceVersion = sourceStateObservation.getVersion();

            if (isMigration()) {
                //In the case of a migration (where the source and target IDs do not match), we do not raise
                //a ReplicaStateObservedEvent because the repository was seen to be CONSISTENT/INCONSISTENT with
                //another repository that this is not a replica of. Instead, we treat the update to the repository
                //as a write operation and increment it's version. Note that if this replica is INCONSISTENT after
                //the repair, the repair source is still notified of it in the summary (and can retry repair or
                //replay recorded writes).
                long version = repair.incrementAndGetVersion();
                if (consistent) {
                    log.debug("[{}] Migration from {} completed successfully and version was updated to {}",
                            repair.getRepository(), sourceRepositoryId, version);
                } else {
                    log.debug("[{}] Migration from {} completed, but a concurrent write that has not been synced " +
                            "was detected. The version was updated to {}.", repair.getRepository(),
                            sourceRepositoryId, version);
                }
            } else {
                if (consistent) {
                    if (sourceVersion == 0L) {
                        log.warn("[{}] Not updating version because the repair source did not provide it",
                                repair.getRepository());
                    } else {
                        repair.updateVersion(sourceVersion);
                        log.debug("[{}] Updated version to {}", repair.getRepository(), sourceVersion);
                    }
                }
                eventPublisher.publish(new ReplicaStateObservedEvent(targetRepositoryId,
                        new ReplicaStateObservation(targetState, logicalTimestamp, sourceVersion), false));
            }

            RepairSummary summary = new RepairSummary(sourceStateTimestamp, sourceState, targetState, sourceVersion);
            setResult(summary);

            sendFragmentToSource(fragment -> fragment.setSummary(summary.toRepairSummary()));
            repairSource.onCompleted();
        } finally {
            releaseRepairResources();
        }
    }
}
