/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.apache.bookkeeper.stream.storage.impl.sc;

import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import java.io.IOException;
import java.time.Duration;
import java.util.Collections;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import lombok.AccessLevel;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.clients.utils.NetUtils;
import org.apache.bookkeeper.common.component.AbstractLifecycleComponent;
import org.apache.bookkeeper.common.concurrent.FutureUtils;
import org.apache.bookkeeper.stats.StatsLogger;
import org.apache.bookkeeper.stream.proto.cluster.ClusterAssignmentData;
import org.apache.bookkeeper.stream.proto.cluster.ServerAssignmentData;
import org.apache.bookkeeper.stream.proto.common.Endpoint;
import org.apache.bookkeeper.stream.storage.api.cluster.ClusterMetadataStore;
import org.apache.bookkeeper.stream.storage.api.sc.StorageContainer;
import org.apache.bookkeeper.stream.storage.api.sc.StorageContainerManager;
import org.apache.bookkeeper.stream.storage.api.sc.StorageContainerRegistry;
import org.apache.bookkeeper.stream.storage.conf.StorageConfiguration;
import org.apache.bookkeeper.util.collections.ConcurrentLongHashMap;

/**
 * A zookeeper based implementation of {@link StorageContainerManager}.
 */
@Slf4j
public class ZkStorageContainerManager
    extends AbstractLifecycleComponent<StorageConfiguration>
    implements StorageContainerManager, Consumer<Void> {

    private final Endpoint endpoint;
    private final ClusterMetadataStore metadataStore;
    private final StorageContainerRegistry registry;
    private final ScheduledExecutorService executor;

    // for routing purpose
    private volatile ClusterAssignmentData clusterAssignmentData;
    private volatile Map<Endpoint, ServerAssignmentData> clusterAssignmentMap;
    private volatile ServerAssignmentData myAssignmentData;
    private volatile ConcurrentLongHashMap<Endpoint> containerAssignmentMap;

    // a probe task to probe containers and make sure this manager running containers as assigned
    private ScheduledFuture<?> containerProbeTask;
    private final Duration probeInterval;

    @Getter(AccessLevel.PACKAGE)
    private final Map<Long, StorageContainer> liveContainers;
    @Getter(AccessLevel.PACKAGE)
    private final Set<Long> pendingStartStopContainers;

    public ZkStorageContainerManager(Endpoint myEndpoint,
                                     StorageConfiguration conf,
                                     ClusterMetadataStore clusterMetadataStore,
                                     StorageContainerRegistry registry,
                                     StatsLogger statsLogger) {
        super("zk-storage-container-manager", conf, statsLogger);
        this.endpoint = myEndpoint;
        this.metadataStore = clusterMetadataStore;
        this.registry = registry;
        this.executor = Executors.newSingleThreadScheduledExecutor(
            new ThreadFactoryBuilder().setNameFormat("zk-storage-container-manager").build());
        this.liveContainers = Collections.synchronizedMap(Maps.newConcurrentMap());
        this.pendingStartStopContainers = Collections.synchronizedSet(Sets.newConcurrentHashSet());
        this.containerAssignmentMap = ConcurrentLongHashMap.<Endpoint>newBuilder().build();
        this.clusterAssignmentMap = Maps.newHashMap();
        // probe the containers every 1/2 of controller scheduling interval. this ensures the manager
        // can attempt to start containers before controller reassign them.
        this.probeInterval = Duration.ofMillis(conf.getClusterControllerScheduleIntervalMs() / 2);
    }

    @Override
    protected void doStart() {
        // watch the cluster assignment data
        metadataStore.watchClusterAssignmentData(this, executor);
        log.info("Watched cluster assignment data.");

        // schedule the container probe task
        containerProbeTask = executor.scheduleAtFixedRate(
            this::probeContainers, 0, probeInterval.toMillis(), TimeUnit.MILLISECONDS);
        log.info("Scheduled storage container probe task at every {} ms", probeInterval.toMillis());
    }

    @Override
    protected void doStop() {
        // unwatch the cluster assignment data
        metadataStore.unwatchClusterAssignmentData(this);

        // cancel the probe task
        if (!containerProbeTask.cancel(true)) {
            log.warn("Failed to cancel the container probe task.");
        }

        stopContainers();
    }

    @Override
    protected void doClose() throws IOException {
        // close the registry to shutdown the containers
        registry.close();
        // shutdown the scheduler
        executor.shutdown();
    }

    @Override
    public Endpoint getStorageContainer(long scId) {
        return containerAssignmentMap.get(scId);
    }

    void probeContainers() {
        boolean isMyAssignmentRefreshed = refreshMyAssignment();
        if (!isMyAssignmentRefreshed) {
            // no change to my assignment, quitting
            return;
        }

        if (myAssignmentData == null) {
            // I don't have any containers assigned to me, so stop containers that I am running.
            stopContainers();
        } else {
            processMyAssignment(myAssignmentData);
        }
    }

    private boolean refreshMyAssignment() {
        ClusterAssignmentData clusterAssignmentData = metadataStore.getClusterAssignmentData();

        if (null == clusterAssignmentData) {
            log.info("Cluster assignment data is empty, so skip refreshing");
            return false;
        }

        Map<Endpoint, ServerAssignmentData> newAssignmentMap = clusterAssignmentData.getServersMap().entrySet()
            .stream()
            .collect(Collectors.toMap(
                e -> NetUtils.parseEndpoint(e.getKey()),
                e -> e.getValue()));

        Set<Endpoint> oldAssignedServers = clusterAssignmentMap.keySet();
        Set<Endpoint> newAssignedServers = newAssignmentMap.keySet();

        Set<Endpoint> serversJoined = Sets.difference(newAssignedServers, oldAssignedServers).immutableCopy();
        Set<Endpoint> serversLeft = Sets.difference(oldAssignedServers, newAssignedServers).immutableCopy();
        Set<Endpoint> commonServers = Sets.intersection(newAssignedServers, oldAssignedServers).immutableCopy();

        processServersLeft(serversLeft, clusterAssignmentMap);
        processServersJoined(serversJoined, newAssignmentMap);
        processServersAssignmentChanged(commonServers, clusterAssignmentMap, newAssignmentMap);

        this.clusterAssignmentMap = newAssignmentMap;
        myAssignmentData = newAssignmentMap.get(endpoint);
        return true;
    }

    private void processServersJoined(Set<Endpoint> serversJoined,
                                      Map<Endpoint, ServerAssignmentData> newAssignmentMap) {
        if (!serversJoined.isEmpty()) {
            log.info("Servers joined : {}", serversJoined);
        }
        serversJoined.forEach(ep -> {
            ServerAssignmentData sad = newAssignmentMap.get(ep);
            if (null != sad) {
                sad.getContainersList().forEach(container -> containerAssignmentMap.put(container, ep));
            }
        });
    }

    private void processServersLeft(Set<Endpoint> serversLeft,
                                    Map<Endpoint, ServerAssignmentData> oldAssignmentMap) {
        if (!serversLeft.isEmpty()) {
            log.info("Servers left : {}", serversLeft);
        }
        serversLeft.forEach(ep -> {
            ServerAssignmentData sad = oldAssignmentMap.get(ep);
            if (null != sad) {
                sad.getContainersList().forEach(container -> containerAssignmentMap.remove(container, ep));
            }
        });
    }

    private void processServersAssignmentChanged(Set<Endpoint> commonServers,
                                                 Map<Endpoint, ServerAssignmentData> oldAssignmentMap,
                                                 Map<Endpoint, ServerAssignmentData> newAssignmentMap) {
        commonServers.forEach(ep -> {

            ServerAssignmentData oldSad = oldAssignmentMap.getOrDefault(ep, ServerAssignmentData.getDefaultInstance());
            ServerAssignmentData newSad = newAssignmentMap.getOrDefault(ep, ServerAssignmentData.getDefaultInstance());

            if (oldSad.equals(newSad)) {
                return;
            } else {
                log.info("Server assignment is change for {}:\nold assignment: {}\nnew assignment: {}",
                    NetUtils.endpointToString(ep), oldSad, newSad);
                oldSad.getContainersList().forEach(container -> containerAssignmentMap.remove(container, ep));
                newSad.getContainersList().forEach(container -> containerAssignmentMap.put(container, ep));
            }

        });
    }


    private void stopContainers() {
        Set<Long> liveContainerSet = ImmutableSet.copyOf(liveContainers.keySet());
        liveContainerSet.forEach(this::stopStorageContainer);
    }

    private void processMyAssignment(ServerAssignmentData myAssignment) {
        Set<Long> assignedContainerSet = myAssignment.getContainersList().stream().collect(Collectors.toSet());
        Set<Long> liveContainerSet = Sets.newHashSet(liveContainers.keySet());

        Set<Long> containersToStart =
            Sets.newHashSet(Sets.difference(assignedContainerSet, liveContainerSet).immutableCopy());
        Set<Long> containersToStop =
            Sets.newHashSet(Sets.difference(liveContainerSet, assignedContainerSet).immutableCopy());

        // if the containers are already in the pending start/stop list, we don't touch it until they are completed.

        containersToStart =
            Sets.filter(containersToStart, container -> !pendingStartStopContainers.contains(container));
        containersToStop =
            Sets.filter(containersToStop, container -> !pendingStartStopContainers.contains(container));

        if (!containersToStart.isEmpty() || !containersToStop.isEmpty()) {
            log.info("Process container changes:\n\tIdeal = {}\n\tLive = {}\n\t"
                    + "Pending = {}\n\tToStart = {}\n\tToStop = {}",
                assignedContainerSet,
                liveContainerSet,
                pendingStartStopContainers,
                containersToStart,
                containersToStop);
        }

        containersToStart.forEach(this::startStorageContainer);
        containersToStop.forEach(this::stopStorageContainer);
    }

    private CompletableFuture<StorageContainer> startStorageContainer(long scId) {
        log.info("Starting storage container ({})", scId);
        StorageContainer sc = liveContainers.get(scId);
        if (null != sc) {
            log.warn("Storage container ({}) is already started", scId);
            return FutureUtils.value(sc);
        } else {
            // register the container to pending list
            pendingStartStopContainers.add(scId);
            return registry
                .startStorageContainer(scId)
                .whenComplete((container, cause) -> {
                    try {
                        if (null != cause) {
                            log.warn("Failed to start storage container ({})", scId, cause);
                        } else {
                            log.info("Successfully started storage container ({})", scId);
                            addStorageContainer(scId, container);
                        }
                    } finally {
                        pendingStartStopContainers.remove(scId);
                    }
                });
        }
    }

    private CompletableFuture<Void> stopStorageContainer(long scId) {
        log.info("Stopping storage container ({})", scId);
        StorageContainer sc = liveContainers.get(scId);
        if (null == sc) {
            log.warn("Storage container ({}) is not alive anymore", scId);
            return FutureUtils.Void();
        } else {
            // register the container to pending list
            pendingStartStopContainers.add(scId);
            return registry
                .stopStorageContainer(scId, sc)
                .whenComplete((container, cause) -> {
                    try {
                        if (cause != null) {
                            log.warn("Failed to stop storage container ({})", scId, cause);
                        } else {
                            log.info("Successfully stopped storage container ({})", scId);
                            removeStorageContainer(scId, sc);
                        }

                    } finally {
                        pendingStartStopContainers.remove(scId);
                    }
                });
        }
    }

    private StorageContainer addStorageContainer(long scId, StorageContainer sc) {
        StorageContainer oldSc = liveContainers.putIfAbsent(scId, sc);
        if (null == oldSc) {
            log.info("Storage container ({}) is added to live set.", sc);
            return sc;
        } else {
            log.warn("Storage container ({}) has already been added to live set", sc);
            sc.stop();
            return oldSc;
        }
    }

    private void removeStorageContainer(long scId, StorageContainer sc) {
        if (liveContainers.remove(scId, sc)) {
            log.info("Storage container ({}) is removed from live set.", scId);
        } else {
            log.warn("Storage container ({}) can't be removed from live set.", scId);
        }
    }

    //
    // Callback on cluster assignment data changed.
    //

    @Override
    public void accept(Void aVoid) {
        // any time if the cluster assignment data is changed, schedule a probe task
        this.executor.submit(() -> probeContainers());
    }
}
