package com.atlassian.bitbucket.mesh.repository;

import com.atlassian.bitbucket.mesh.network.MeshManager;
import com.atlassian.bitbucket.mesh.network.TopologyUpdatedEvent;
import com.atlassian.bitbucket.mesh.partition.Partition;
import com.atlassian.bitbucket.mesh.util.ExecutorUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.event.EventListener;

import java.time.Duration;
import java.time.Instant;
import java.util.concurrent.Callable;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.atomic.AtomicBoolean;

import static java.util.concurrent.CompletableFuture.supplyAsync;
import static java.util.concurrent.TimeUnit.SECONDS;
import static java.util.stream.StreamSupport.stream;

public class ManagedDataStoreCleanupManager {

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

    private final Duration cleanupDelay;
    private final ScheduledExecutorService cleanupExecutor;
    private final ManagedDataStore managedDataStore;
    private final MeshManager meshManager;
    private final AtomicBoolean shutdown;

    public ManagedDataStoreCleanupManager(Duration cleanupDelay, ScheduledExecutorService cleanupExecutor,
                                          ManagedDataStore managedDataStore, MeshManager meshManager) {
        this.cleanupDelay = cleanupDelay;
        this.cleanupExecutor = cleanupExecutor;
        this.managedDataStore = managedDataStore;
        this.meshManager = meshManager;

        this.shutdown = new AtomicBoolean();
    }

    @EventListener(TopologyUpdatedEvent.class)
    public void onTopologyUpdatedEvent() {
        if (!shutdown.get()) {
            try {
                supplyAsync(() -> new MarkDeletedTask(managedDataStore, meshManager).call(), cleanupExecutor)
                        .whenComplete((markersCreated, throwable) -> {
                            if (throwable != null) {
                                log.error("Exception while marking partitions as deleted", throwable);
                                return;
                            }
                            if (!markersCreated) {
                                log.debug("No new markers created, not scheduling cleanup task");
                                return;
                            }
                            try {
                                cleanupExecutor.schedule(() ->
                                                new CleanupTask(cleanupDelay, managedDataStore, meshManager).run(),
                                        cleanupDelay.getSeconds(), SECONDS);
                            } catch (RejectedExecutionException e) {
                                log.warn("Couldn't schedule cleanup task after topology updated", e);
                            }
                        });
            } catch (RejectedExecutionException e) {
                log.warn("Couldn't schedule cleanup task after topology updated", e);
            }
        }
    }

    public void shutdown() {
        if (shutdown.compareAndSet(false, true)) {
            ExecutorUtils.shutdown(cleanupExecutor, log);
        }
    }

    private static class CleanupTask implements Runnable {

        private final Duration cleanupDelay;
        private final ManagedDataStore managedDataStore;
        private final MeshManager meshManager;

        public CleanupTask(Duration cleanupDelay, ManagedDataStore managedDataStore, MeshManager meshManager) {
            this.cleanupDelay = cleanupDelay;
            this.managedDataStore = managedDataStore;
            this.meshManager = meshManager;
        }

        @Override
        public void run() {
            Instant now = Instant.now();
            managedDataStore.streamDeletedPartitions(((partition, deletionTime) -> {
                if (meshManager.getPartition(partition.getHex()) != null) {
                    log.info("Partition marked for deletion still exists in topology: {}. Clearing marker.", partition.getHex());
                    managedDataStore.clearPartitionDeleted(partition);
                } else if (deletionTime.plus(cleanupDelay).isBefore(now)) {
                    log.debug("{} has been marked for deletion at {}, deleting", partition, deletionTime);
                    managedDataStore.deletePartitionFiles(partition);
                } else {
                    log.debug("{} was marked for deletion at {}, which was less than {} ago, not deleting",
                            partition, deletionTime, cleanupDelay);
                }
            }));
        }
    }

    private static class MarkDeletedTask implements Callable<Boolean> {

        private final ManagedDataStore managedDataStore;
        private final MeshManager meshManager;

        private MarkDeletedTask(ManagedDataStore managedDataStore, MeshManager meshManager) {
            this.managedDataStore = managedDataStore;
            this.meshManager = meshManager;
        }

        @Override
        public Boolean call() {
            //Find all partitions on the data store, then compare it with the topology and mark all partitions
            //that don't exist in the topology as deleted.
            return stream(managedDataStore.getAllPartitions().spliterator(), false)
                    .filter(partition -> meshManager.getPartition(partition) == null)
                    .map(Partition::new)
                    .map(partition -> {
                        log.debug("Found {} on disk but it's not in the topology anymore, marking as deleted", partition);
                        boolean marked = managedDataStore.markPartitionDeleted(partition);
                        if (!marked) {
                            log.debug("{} was not marked as deleted", partition);
                        }
                        return marked;
                    })
                    .reduce(false, (a, b) -> a || b);
        }
    }
}
