package com.atlassian.bitbucket.mesh.repair;

import com.atlassian.bitbucket.mesh.annotation.Profiled;
import com.atlassian.bitbucket.mesh.client.MeshClient;
import com.atlassian.bitbucket.mesh.client.RepairRepositoryParameters;
import com.atlassian.bitbucket.mesh.config.Backoff;
import com.atlassian.bitbucket.mesh.config.RepairConfig;
import com.atlassian.bitbucket.mesh.event.EventListener;
import com.atlassian.bitbucket.mesh.execution.ExecutionManager;
import com.atlassian.bitbucket.mesh.git.exception.RepositoryRepairFailedException;
import com.atlassian.bitbucket.mesh.network.MeshManager;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.network.Partition;
import com.atlassian.bitbucket.mesh.network.PartitionReplica;
import com.atlassian.bitbucket.mesh.repository.*;
import com.atlassian.bitbucket.mesh.util.ExecutorUtils;
import com.google.common.annotations.VisibleForTesting;
import io.grpc.Status;
import org.apache.commons.lang3.mutable.MutableBoolean;
import org.apache.commons.lang3.mutable.MutableObject;
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.List;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.stream.Collectors;

import static com.atlassian.bitbucket.mesh.util.MoreCollectors.toImmutableList;
import static java.util.Objects.requireNonNull;

public class DefaultRepairManager implements RepairManager {

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

    private final Set<String> abandoned;
    private final RepairConfig config;
    private final ExecutionManager executionManager;
    private final ExecutorService executorService;
    private final MeshClient meshClient;
    private final MeshManager meshManager;
    private final ReplicaStateRegistry replicaStateRegistry;
    private final TaskScheduler taskScheduler;

    public DefaultRepairManager(RepairConfig config, ExecutionManager executionManager, ExecutorService executorService,
                                MeshClient meshClient, MeshManager meshManager,
                                ReplicaStateRegistry replicaStateRegistry, TaskScheduler taskScheduler) {
        this.config = config;
        this.executionManager = executionManager;
        this.executorService = executorService;
        this.meshClient = meshClient;
        this.meshManager = meshManager;
        this.replicaStateRegistry = replicaStateRegistry;
        this.taskScheduler = taskScheduler;

        abandoned = ConcurrentHashMap.newKeySet();
    }

    public void destroy() {
        ExecutorUtils.shutdown(executorService, log);
    }

    /**
     * State observations imply repository activity. If repairs have been abandoned for an active repository, and
     * the observation indicates the repository is still in need of repair, repair attempts are restarted.
     * <p>
     * Repositories that are having a difficult time repairing (for example, a very busy repository where repairs
     * end up chasing consistency but never quite catching up) may go through multiple cycles of repairing, being
     * abandoned, restarting and repairing more, being abandoned, etc. So long as the repository is active, repairs
     * will never be <i>fully</i> abandoned.
     *
     * @param event provides the observed replica state
     */
    @EventListener
    public void onReplicaStateObserved(ReplicaStateObservedEvent event) {
        //The order here is intentional. We always want to purge abandoned status for any active repository,
        //but we only want to schedule repair if the observation implies a repairable state. In practice the
        //state should always be either INCONSISTENT or UNKNOWN, if repairs have previously been abandoned,
        //since it's only possible for such a replica to become CONSISTENT again via repair, but this handles
        //the case where that invariant is somehow violated
        if (abandoned.remove(event.getRepositoryId()) && shouldRepair(event.getObservation().getState())) {
            scheduleRepair(event.getRepositoryId());
        }
    }

    /**
     * Triggers repair when a replica transitions to an {@link ReplicaState#INCONSISTENT inconsistent} or
     * {@link ReplicaState#UNKNOWN unknown} state.
     * <p>
     * If the scheduled repairs run out of attempts, per the configured limit, they will be abandoned. That
     * means the replica state will not be updated again (since only repair can fix an inconsistent replica),
     * repairs will only be rescheduled if a new {@link ReplicaStateObservedEvent state observation} is raised,
     * implying the repository is still active.
     *
     * @param event provides the updated state
     */
    @EventListener
    public void onReplicaStateUpdated(ReplicaStateUpdatedEvent event) {
        if (shouldRepair(event.getState())) {
            scheduleRepair(event.getRepositoryId());
        }
    }

    /**
     * Ensures repairs are not marked abandoned for deleted repositories.
     *
     * @param event provides the repository that was deleted
     */
    @EventListener
    public void onRepositoryDeleted(RepositoryDeletedEvent event) {
        //Just in case we'd previously abandoned repairs for the repository, purge that status
        abandoned.remove(event.getRepository().getId());
    }

    @Nonnull
    @Override
    @Profiled
    public ReplicaState repair(@Nonnull String repositoryId, @Nonnull Collection<MeshNode> sourceNodes,
                               @Nonnull MeshNode targetNode) {
        return repair(repositoryId, sourceNodes, repositoryId, targetNode);
    }

    @Nonnull
    @Override
    public ReplicaState repair(@Nonnull String sourceRepositoryId, @Nonnull Collection<MeshNode> sourceNodes,
                               @Nonnull String targetRepositoryId, @Nonnull MeshNode targetNode) {
        requireNonNull(sourceRepositoryId, "sourceRepositoryId");
        requireNonNull(sourceNodes, "sourceNodes");
        requireNonNull(targetRepositoryId, "targetRepositoryId");
        requireNonNull(targetNode, "targetNode");

        RepairRepositoryParameters parameters = new RepairRepositoryParameters.Builder(targetRepositoryId)
                .force(true)
                .retryOnFailure(false)
                .sourceRepository(sourceRepositoryId)
                .timeout(config.getTimeout().getSeconds())
                .build();
        if (targetNode.isLocal()) {
            sourceNodes = sourceNodes.stream().filter(MeshNode::isRemote).collect(Collectors.toList());
            if (sourceNodes.isEmpty()) {
                throw new IllegalArgumentException("Cannot repair local replica of " + targetRepositoryId +
                        "; no (remote) source replicas have been provided");
            }

            log.debug("[{}] Repairing local replica from {}", targetRepositoryId, sourceNodes);
            return new RepairTask(parameters, sourceNodes).repair();
        }

        if (!sourceNodes.isEmpty() && sourceNodes.iterator().next().isLocal()) {
            ReplicaState state = replicaStateRegistry.getState(sourceRepositoryId);
            if (state == ReplicaState.CONSISTENT) {
                log.debug("[{}] Repairing replica on {}", sourceRepositoryId, targetNode);
                return meshClient.repair(targetNode, parameters).getTargetState();
            } else {
                log.debug("[{}] Cannot repair replica on {} because the local replica is {}",
                        sourceRepositoryId, targetNode, state);
                return ReplicaState.UNKNOWN;
            }
        }
        throw new IllegalArgumentException("Cannot repair " + targetRepositoryId +
                "; neither the source nor the target node correspond to the local node");
    }

    @VisibleForTesting
    boolean isAbandoned(@Nonnull String repositoryId) {
        return abandoned.contains(repositoryId);
    }

    private static boolean shouldRepair(ReplicaState state) {
        return state == ReplicaState.INCONSISTENT || state == ReplicaState.UNKNOWN;
    }

    private List<MeshNode> getRemoteReplicas(String repositoryId) {
        return ManagedRepository.Id.parseId(repositoryId)
                .map(id -> {
                    Partition partition = meshManager.getPartition(id.getPartition());
                    if (partition == null) {
                        return null;
                    }
                    return partition.getReplicas().stream()
                            .map(PartitionReplica::getNode)
                            .filter(MeshNode::isRemote)
                            .collect(toImmutableList());
                }).orElseGet(Collections::emptyList);
    }

    private void scheduleRepair(String repositoryId) {
        RepairRepositoryParameters parameters = new RepairRepositoryParameters.Builder(repositoryId)
                .retryOnFailure(true)
                .timeout(config.getTimeout().getSeconds())
                .build();
        executorService.submit(new RepairTask(parameters, getRemoteReplicas(repositoryId)));
    }

    private class RepairTask implements Runnable {

        private final RepairRepositoryParameters parameters;
        private final Collection<MeshNode> sources;

        private int attempts;
        private Backoff.Execution backoff;

        RepairTask(RepairRepositoryParameters parameters, Collection<MeshNode> sources) {
            this.parameters = parameters;
            this.sources = sources;
        }

        @Override
        public void run() {
            try {
                repair();
            } catch (RuntimeException e) {
                //This is logged to debug because repair() should have already logged why repairs failed
                log.debug("[{}] Repair failed", parameters.getRepositoryId(), e);
            }
        }

        ReplicaState repair() {
            String repositoryId = parameters.getRepositoryId();
            if (!parameters.isForce() && replicaStateRegistry.getState(repositoryId) == ReplicaState.CONSISTENT) {
                //Already consistent, nothing left to do
                return ReplicaState.CONSISTENT;
            }

            MutableObject<RuntimeException> exception = new MutableObject<>();
            long start = System.currentTimeMillis();
            int replicasAttempted = 0;
            for (MeshNode node : sources) {
                MutableBoolean alreadyRunning = new MutableBoolean();
                RepairSummary summary = executionManager.withContext(ctx -> {
                    try {
                        log.debug("[{}] Starting repair from {}", repositoryId, node);
                        return meshClient.repairMe(node, parameters);
                    } catch (RepositoryRepairAlreadyRunningException e) {
                        alreadyRunning.setTrue();
                    } catch (RuntimeException repairException) {
                        exception.setValue(repairException);
                        Status status = Status.fromThrowable(repairException);
                        if (status.getCode() == Status.Code.UNAVAILABLE) {
                            log.info("[{}] Cannot repair from {} because it is unavailable. Trying next replica..",
                                    repositoryId, node);
                        } else {
                            log.info("[{}] Repair from {} failed. Trying next replica..",
                                    repositoryId, node, repairException);
                        }
                    }
                    return null;
                });
                if (alreadyRunning.booleanValue()) {
                    log.info("[{}] Aborting repair because it is already running", repositoryId);
                    return ReplicaState.REPAIRING;
                }

                ++replicasAttempted;
                if (summary == null) {
                    //Repair failed, try the next replica
                    continue;
                }
                if (summary.getTargetState() == ReplicaState.CONSISTENT) {
                    log.info("[{}] Repair from {} completed in {} ms{}",
                            repositoryId, node, System.currentTimeMillis() - start,
                            replicasAttempted > 1 ? " (tried " + replicasAttempted + " nodes)" : "");
                    //Ensure the repository isn't marked abandoned if it's consistent
                    abandoned.remove(repositoryId);
                    return ReplicaState.CONSISTENT;
                }
                if (summary.getSourceState() == ReplicaState.CONSISTENT) {
                    if (++attempts < config.getMaxAttempts()) {
                        log.info("[{}] Still not consistent after repair from {}. Rescheduling repair",
                                repositoryId, node);
                        executorService.execute(this);
                    } else {
                        log.warn("[{}] Abandoning repair after {} attempts", repositoryId, attempts);
                        abandoned.add(repositoryId);
                    }
                    return summary.getTargetState();
                }
                //Source is not consistent itself
                log.info("Node {} did not have a consistent replica of {}. Trying next replica..",
                        node, repositoryId);
            }

            if (replicasAttempted > 0) {
                if (++attempts < config.getMaxAttempts() && parameters.isRetryOnFailure()) {
                    //This may seem odd, but we want to enforce a maximum thread count in the executorService
                    //and allow its core threads to time out to reduce the amount of threads in use when nothing
                    //needs to be repaired. This is not possible with ScheduledExecutorService so we use
                    //taskScheduler to delay the execution, but still execute on the repair thread pool.
                    Duration delay = nextRepairDelay();
                    log.warn("[{}] No up-to-date replica is available for repair (tried {} nodes). " +
                            "Retrying in {}s (attempt {}/{})", repositoryId, replicasAttempted, delay.getSeconds(),
                            attempts + 1, config.getMaxAttempts());
                    taskScheduler.schedule(() -> executorService.execute(this), Instant.now().plus(delay));
                } else {
                    log.warn("[{}] Abandoning repair! No node with an up-to-date replica is available (tried {} nodes)",
                            repositoryId, replicasAttempted);
                    abandoned.add(repositoryId);
                }
            }

            throw new RepositoryRepairFailedException("Could not repair " + repositoryId, exception.getValue());
        }

        private Duration nextRepairDelay() {
            if (backoff == null) {
                backoff = config.getBackoff().start();
            }
            return backoff.nextDelay();
        }
    }
}
