/*
 * 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.dolphinscheduler.server.master.failover;

import org.apache.dolphinscheduler.common.constants.Constants;
import org.apache.dolphinscheduler.dao.entity.WorkflowInstance;
import org.apache.dolphinscheduler.dao.repository.WorkflowInstanceDao;
import org.apache.dolphinscheduler.plugin.task.api.enums.TaskExecutionStatus;
import org.apache.dolphinscheduler.registry.api.RegistryClient;
import org.apache.dolphinscheduler.registry.api.RegistryLock;
import org.apache.dolphinscheduler.registry.api.enums.RegistryNodeType;
import org.apache.dolphinscheduler.registry.api.utils.RegistryUtils;
import org.apache.dolphinscheduler.server.master.cluster.ClusterManager;
import org.apache.dolphinscheduler.server.master.cluster.MasterServerMetadata;
import org.apache.dolphinscheduler.server.master.cluster.WorkerServerMetadata;
import org.apache.dolphinscheduler.server.master.engine.IWorkflowRepository;
import org.apache.dolphinscheduler.server.master.engine.system.event.GlobalMasterFailoverEvent;
import org.apache.dolphinscheduler.server.master.engine.system.event.MasterFailoverEvent;
import org.apache.dolphinscheduler.server.master.engine.system.event.WorkerFailoverEvent;
import org.apache.dolphinscheduler.server.master.engine.task.execution.ITaskExecution;
import org.apache.dolphinscheduler.server.master.engine.workflow.execution.IWorkflowExecution;

import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.time.StopWatch;

import java.util.Collection;
import java.util.Date;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import lombok.extern.slf4j.Slf4j;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

@Slf4j
@Component
public class FailoverCoordinator implements IFailoverCoordinator {

    @Autowired
    private RegistryClient registryClient;

    @Autowired
    private ClusterManager clusterManager;

    @Autowired
    private IWorkflowRepository workflowRepository;

    @Autowired
    private TaskFailover taskFailover;

    @Autowired
    private WorkflowInstanceDao workflowInstanceDao;

    @Autowired
    private WorkflowFailover workflowFailover;

    @Override
    public void globalMasterFailover(final GlobalMasterFailoverEvent globalMasterFailoverEvent) {
        final StopWatch failoverTimeCost = StopWatch.createStarted();
        log.info("Global master failover starting");
        final List<String> masterAddressWhichContainsUnFinishedWorkflow =
                workflowInstanceDao.queryNeedFailoverMasters();
        for (final String masterAddress : masterAddressWhichContainsUnFinishedWorkflow) {
            final Optional<MasterServerMetadata> aliveMasterOptional =
                    clusterManager.getMasterClusters().getServer(masterAddress);
            if (aliveMasterOptional.isPresent()) {
                // If the master is alive, then we use the alive master's startup time as the failover deadline.
                final MasterServerMetadata aliveMasterServerMetadata = aliveMasterOptional.get();
                log.info("The master[{}] is alive, do global master failover on it", aliveMasterServerMetadata);
                doMasterFailover(
                        masterAddress,
                        aliveMasterServerMetadata.getServerStartupTime(),
                        RegistryUtils.getGlobalMasterFailoverNodePath(
                                masterAddress));
            } else {
                // If the master is not alive, then we use the event time as the failover deadline.
                log.info("The master[{}] is not alive, do global master failover on it", masterAddress);
                doMasterFailover(
                        masterAddress,
                        globalMasterFailoverEvent.getEventTime().getTime(),
                        RegistryUtils.getGlobalMasterFailoverNodePath(masterAddress));
            }
        }

        failoverTimeCost.stop();
        log.info("Global master failover finished, cost: {}/ms", failoverTimeCost.getTime());
    }

    @Override
    public void failoverMaster(final MasterFailoverEvent masterFailoverEvent) {
        final MasterServerMetadata masterServerMetadata = masterFailoverEvent.getMasterServerMetadata();
        log.info("Master[{}] failover starting", masterServerMetadata);
        final String masterAddress = masterServerMetadata.getAddress();

        final Optional<MasterServerMetadata> aliveMasterOptional =
                clusterManager.getMasterClusters().getServer(masterAddress);
        if (aliveMasterOptional.isPresent()) {
            final MasterServerMetadata aliveMasterServerMetadata = aliveMasterOptional.get();
            if (aliveMasterServerMetadata.getServerStartupTime() == masterServerMetadata.getServerStartupTime()) {
                log.info("The master[{}] is alive, maybe it reconnect to registry skip failover", masterServerMetadata);
                return;
            }
        }
        doMasterFailover(
                masterServerMetadata.getAddress(),
                masterFailoverEvent.getEventTime().getTime(),
                RegistryUtils.getFailoveredNodePath(
                        masterServerMetadata.getAddress(),
                        masterServerMetadata.getServerStartupTime(),
                        masterServerMetadata.getProcessId()));
    }

    /**
     * Do master failover.
     * <p> Will failover the workflow which is scheduled by the master and the workflow's fire time is before the maxWorkflowFireTime.
     */
    private void doMasterFailover(final String masterAddress,
                                  final long workflowFailoverDeadline,
                                  final String masterFailoverNodePath) {
        // We use lock to avoid multiple master failover at the same time.
        // Once the workflow has been failovered, then it's state will be changed to FAILOVER
        // Once the FAILOVER workflow has been refired, then it's host will be changed to the new master and have a new
        // start time.
        // So if a master has been failovered multiple times, there is no problem.
        final String masterFailoverLockPath = RegistryUtils.getMasterFailoverLockPath(masterAddress);
        final StopWatch failoverTimeCost = StopWatch.createStarted();
        try (RegistryLock ignored = registryClient.getLock(masterFailoverLockPath)) {
            // If the master has already been failovered, then we skip the failover.
            if (registryClient.exists(masterFailoverNodePath)
                    && String.valueOf(workflowFailoverDeadline).equals(registryClient.get(masterFailoverNodePath))) {
                log.error("The master[{}/{}] is exist at: {}, means it has already been failovered, skip failover",
                        masterAddress,
                        workflowFailoverDeadline,
                        masterFailoverNodePath);
                return;
            }
            final List<WorkflowInstance> needFailoverWorkflows =
                    getFailoverWorkflowsForMaster(masterAddress, new Date(workflowFailoverDeadline));
            needFailoverWorkflows.forEach(workflowFailover::failoverWorkflow);
            registryClient.persist(masterFailoverNodePath, String.valueOf(workflowFailoverDeadline));
            failoverTimeCost.stop();
            log.info("Master[{}] failover {} workflows finished, cost: {}/ms",
                    masterAddress,
                    needFailoverWorkflows.size(),
                    failoverTimeCost.getTime());
        }
    }

    private List<WorkflowInstance> getFailoverWorkflowsForMaster(final String masterAddress,
                                                                 final Date masterCrashTime) {
        // todo: use page query
        final List<WorkflowInstance> workflowInstances =
                workflowInstanceDao.queryNeedFailoverWorkflowInstances(masterAddress);
        return workflowInstances.stream()
                .filter(workflowInstance -> {

                    if (workflowRepository.contains(workflowInstance.getId())) {
                        return false;
                    }

                    // todo: If the first time run workflow have the restartTime, then we can only check this
                    final Date restartTime = workflowInstance.getRestartTime();
                    if (restartTime != null) {
                        return restartTime.before(masterCrashTime);
                    }

                    final Date startTime = workflowInstance.getStartTime();
                    return startTime.before(masterCrashTime);
                })
                .collect(Collectors.toList());
    }

    @Override
    public void failoverWorker(final WorkerFailoverEvent workerFailoverEvent) {
        final WorkerServerMetadata workerServerMetadata = workerFailoverEvent.getWorkerServerMetadata();
        log.info("Worker[{}] failover starting", workerServerMetadata);

        final Optional<WorkerServerMetadata> aliveWorkerOptional =
                clusterManager.getWorkerClusters().getServer(workerServerMetadata.getAddress());
        if (aliveWorkerOptional.isPresent()) {
            final WorkerServerMetadata aliveWorkerServerMetadata = aliveWorkerOptional.get();
            if (aliveWorkerServerMetadata.getServerStartupTime() == workerServerMetadata.getServerStartupTime()) {
                log.info("The worker[{}] is alive, maybe it reconnect to registry skip failover", workerServerMetadata);
                return;
            }
        }
        doWorkerFailover(
                workerServerMetadata.getAddress(),
                System.currentTimeMillis(),
                RegistryUtils.getFailoveredNodePath(
                        workerServerMetadata.getAddress(),
                        workerServerMetadata.getServerStartupTime(),
                        workerServerMetadata.getProcessId()));
    }

    @Override
    public void cleanHistoryFailoverFinishedMarks() {
        // Clean the history failover finished nodes
        // which failover is before the current time minus 1 week
        final Collection<String> failoverFinishedMarkSuffixes =
                registryClient.getChildrenKeys(RegistryNodeType.FAILOVER_FINISH_NODES.getRegistryPath());
        if (CollectionUtils.isEmpty(failoverFinishedMarkSuffixes)) {
            return;
        }
        for (final String failoverFinishedMarkSuffix : failoverFinishedMarkSuffixes) {
            final String failoverFinishedMarkFullPath = RegistryNodeType.FAILOVER_FINISH_NODES.getRegistryPath()
                    + Constants.SINGLE_SLASH + failoverFinishedMarkSuffix;
            try {
                final String failoverFinishTime = registryClient.get(failoverFinishedMarkFullPath);
                if (System.currentTimeMillis() - Long.parseLong(failoverFinishTime) > TimeUnit.DAYS.toMillis(7)) {
                    registryClient.remove(failoverFinishedMarkFullPath);
                    log.info(
                            "Clear the failover finished node: {} which failover time is before the current time minus 1 week",
                            failoverFinishedMarkFullPath);
                }
            } catch (Exception ex) {
                log.error("Failed to clean the failoverFinishedNode: {}", failoverFinishedMarkFullPath, ex);
            }
        }
    }

    private void doWorkerFailover(final String workerAddress,
                                  final long taskFailoverDeadline,
                                  final String workerFailoverNodePath) {
        final StopWatch failoverTimeCost = StopWatch.createStarted();
        // we don't check the workerFailoverNodePath exist, since the worker may be failovered multiple master

        final List<ITaskExecution> needFailoverTasks =
                getFailoverTaskForWorker(workerAddress, new Date(taskFailoverDeadline));
        needFailoverTasks.forEach(taskFailover::failoverTask);

        registryClient.persist(
                workerFailoverNodePath,
                String.valueOf(System.currentTimeMillis()));
        failoverTimeCost.stop();
        log.info("Worker[{}] failover {} tasks finished, cost: {}/ms",
                workerAddress,
                needFailoverTasks.size(),
                failoverTimeCost.getDuration());
    }

    private List<ITaskExecution> getFailoverTaskForWorker(final String workerAddress,
                                                          final Date taskFailoverDeadline) {
        return workflowRepository.getAll()
                .stream()
                .map(IWorkflowExecution::getWorkflowExecutionGraph)
                .flatMap(workflowExecutionGraph -> workflowExecutionGraph.getActiveTaskExecution().stream())
                .filter(ITaskExecution::isTaskInstanceInitialized)
                .filter(taskExecution -> workerAddress
                        .equals(taskExecution.getTaskInstance().getHost()))
                .filter(taskExecution -> {
                    final TaskExecutionStatus state = taskExecution.getTaskInstance().getState();
                    return state == TaskExecutionStatus.DISPATCH || state == TaskExecutionStatus.RUNNING_EXECUTION;
                })
                .filter(taskExecution -> {
                    // The submitTime should not be null.
                    // This is a bad case unless someone manually set the submitTime to null.
                    final Date submitTime = taskExecution.getTaskInstance().getSubmitTime();
                    return submitTime != null && submitTime.before(taskFailoverDeadline);
                })
                .collect(Collectors.toList());
    }

}
