/*
 * Copyright 2022 Conductor Authors.
 * <p>
 * Licensed 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
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * 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 com.netflix.conductor.core.execution.tasks;

import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import com.netflix.conductor.annotations.VisibleForTesting;
import com.netflix.conductor.common.metadata.tasks.TaskDef;
import com.netflix.conductor.common.metadata.workflow.WorkflowTask;
import com.netflix.conductor.common.utils.TaskUtils;
import com.netflix.conductor.core.dal.ExecutionDAOFacade;
import com.netflix.conductor.core.events.ScriptEvaluator;
import com.netflix.conductor.core.execution.WorkflowExecutor;
import com.netflix.conductor.core.utils.ParametersUtils;
import com.netflix.conductor.model.TaskModel;
import com.netflix.conductor.model.WorkflowModel;

import static com.netflix.conductor.common.metadata.tasks.TaskType.TASK_TYPE_DO_WHILE;

@Component(TASK_TYPE_DO_WHILE)
public class DoWhile extends WorkflowSystemTask {

    private static final Logger LOGGER = LoggerFactory.getLogger(DoWhile.class);

    private final ParametersUtils parametersUtils;
    private final ExecutionDAOFacade executionDAOFacade;

    public DoWhile(ParametersUtils parametersUtils, ExecutionDAOFacade executionDAOFacade) {
        super(TASK_TYPE_DO_WHILE);
        this.parametersUtils = parametersUtils;
        this.executionDAOFacade = executionDAOFacade;
    }

    @Override
    public void cancel(WorkflowModel workflow, TaskModel task, WorkflowExecutor executor) {
        task.setStatus(TaskModel.Status.CANCELED);
    }

    @Override
    public boolean execute(
            WorkflowModel workflow, TaskModel doWhileTaskModel, WorkflowExecutor workflowExecutor) {

        boolean hasFailures = false;
        StringBuilder failureReason = new StringBuilder();
        Map<String, Object> output = new HashMap<>();

        /*
         * Get the latest set of tasks (the ones that have the highest retry count). We don't want to evaluate any tasks
         * that have already failed if there is a more current one (a later retry count).
         */
        Map<String, TaskModel> relevantTasks = new LinkedHashMap<>();
        TaskModel relevantTask;
        for (TaskModel t : workflow.getTasks()) {
            if (doWhileTaskModel
                            .getWorkflowTask()
                            .has(TaskUtils.removeIterationFromTaskRefName(t.getReferenceTaskName()))
                    && !doWhileTaskModel.getReferenceTaskName().equals(t.getReferenceTaskName())
                    && doWhileTaskModel.getIteration() == t.getIteration()) {
                relevantTask = relevantTasks.get(t.getReferenceTaskName());
                if (relevantTask == null || t.getRetryCount() > relevantTask.getRetryCount()) {
                    relevantTasks.put(t.getReferenceTaskName(), t);
                }
            }
        }
        Collection<TaskModel> loopOverTasks = relevantTasks.values();

        if (LOGGER.isDebugEnabled()) {
            LOGGER.debug(
                    "Workflow {} waiting for tasks {} to complete iteration {}",
                    workflow.getWorkflowId(),
                    loopOverTasks.stream()
                            .map(TaskModel::getReferenceTaskName)
                            .collect(Collectors.toList()),
                    doWhileTaskModel.getIteration());
        }

        // if the loopOverTasks collection is empty, no tasks inside the loop have been scheduled.
        // so schedule it and exit the method.
        if (loopOverTasks.isEmpty()) {
            // For list iteration, check if the items list is empty before scheduling iteration 1.
            // An empty items list means there is nothing to iterate over, so complete immediately.
            if (isListIteration(doWhileTaskModel)) {
                List<Object> itemsList = evaluateItemsList(workflow, doWhileTaskModel);
                if (itemsList.isEmpty()) {
                    LOGGER.debug(
                            "Task {} has an empty items list, completing without executing loop tasks",
                            doWhileTaskModel.getTaskId());
                    doWhileTaskModel.addOutput("iteration", 0);
                    return markTaskSuccess(doWhileTaskModel);
                }
            }

            doWhileTaskModel.setIteration(1);
            doWhileTaskModel.addOutput("iteration", doWhileTaskModel.getIteration());

            // For list iteration, inject loopItem and loopIndex
            injectLoopVariables(workflow, doWhileTaskModel);

            return scheduleNextIteration(doWhileTaskModel, workflow, workflowExecutor);
        }

        for (TaskModel loopOverTask : loopOverTasks) {
            TaskModel.Status taskStatus = loopOverTask.getStatus();
            hasFailures = !taskStatus.isSuccessful();
            if (hasFailures) {
                failureReason.append(loopOverTask.getReasonForIncompletion()).append(" ");
            }
            output.put(
                    TaskUtils.removeIterationFromTaskRefName(loopOverTask.getReferenceTaskName()),
                    loopOverTask.getOutputData());
            if (hasFailures) {
                break;
            }
        }
        doWhileTaskModel.addOutput(String.valueOf(doWhileTaskModel.getIteration()), output);

        Optional<Integer> keepLastN =
                Optional.ofNullable(doWhileTaskModel.getWorkflowTask().getInputParameters())
                        .map(parameters -> parameters.get("keepLastN"))
                        .map(value -> (Integer) value);
        if (keepLastN.isPresent() && doWhileTaskModel.getIteration() > keepLastN.get()) {
            Integer iteration = doWhileTaskModel.getIteration();
            IntStream.rangeClosed(1, iteration - keepLastN.get())
                    .mapToObj(Integer::toString)
                    .forEach(doWhileTaskModel::removeOutput);

            // Remove old iteration tasks from the database
            removeIterations(workflow, doWhileTaskModel, keepLastN.get());
        }

        if (hasFailures) {
            LOGGER.debug(
                    "Task {} failed in {} iteration",
                    doWhileTaskModel.getTaskId(),
                    doWhileTaskModel.getIteration() + 1);
            return markTaskFailure(
                    doWhileTaskModel, TaskModel.Status.FAILED, failureReason.toString());
        }

        if (!isIterationComplete(doWhileTaskModel, relevantTasks)) {
            // current iteration is not complete (all tasks inside the loop are not terminal)
            return false;
        }

        // if we are here, the iteration is complete, and we need to check if there is a next
        // iteration by evaluating the loopCondition
        boolean shouldContinue;
        try {
            shouldContinue = evaluateCondition(workflow, doWhileTaskModel);
            LOGGER.debug(
                    "Task {} condition evaluated to {}",
                    doWhileTaskModel.getTaskId(),
                    shouldContinue);
            if (shouldContinue) {
                doWhileTaskModel.setIteration(doWhileTaskModel.getIteration() + 1);
                doWhileTaskModel.addOutput("iteration", doWhileTaskModel.getIteration());

                // For list iteration, inject loopItem and loopIndex for next iteration
                injectLoopVariables(workflow, doWhileTaskModel);

                return scheduleNextIteration(doWhileTaskModel, workflow, workflowExecutor);
            } else {
                LOGGER.debug(
                        "Task {} took {} iterations to complete",
                        doWhileTaskModel.getTaskId(),
                        doWhileTaskModel.getIteration() + 1);
                return markTaskSuccess(doWhileTaskModel);
            }
        } catch (Exception e) {
            String message =
                    String.format(
                            "Unable to evaluate condition %s, exception %s",
                            doWhileTaskModel.getWorkflowTask().getLoopCondition(), e.getMessage());
            LOGGER.error(message);
            return markTaskFailure(
                    doWhileTaskModel, TaskModel.Status.FAILED_WITH_TERMINAL_ERROR, message);
        }
    }

    /**
     * Removes old iterations from the workflow to prevent database bloat. This method identifies
     * and deletes tasks from iterations that exceed the keepLastN retention policy.
     *
     * @param workflow The workflow model containing all tasks
     * @param doWhileTaskModel The DO_WHILE task model
     * @param keepLastN Number of most recent iterations to keep
     */
    @VisibleForTesting
    void removeIterations(WorkflowModel workflow, TaskModel doWhileTaskModel, int keepLastN) {
        int currentIteration = doWhileTaskModel.getIteration();

        // Calculate which iterations should be removed (all iterations before currentIteration -
        // keepLastN)
        int iterationsToRemove = currentIteration - keepLastN;

        if (iterationsToRemove <= 0) {
            // Nothing to remove yet
            return;
        }

        LOGGER.debug(
                "Removing iterations 1 to {} for DO_WHILE task {} (keeping last {} iterations)",
                iterationsToRemove,
                doWhileTaskModel.getReferenceTaskName(),
                keepLastN);

        // Find and remove tasks from old iterations
        List<TaskModel> tasksToRemove =
                workflow.getTasks().stream()
                        .filter(
                                task -> {
                                    // Check if this task belongs to the DO_WHILE loop
                                    String taskRefWithoutIteration =
                                            TaskUtils.removeIterationFromTaskRefName(
                                                    task.getReferenceTaskName());
                                    boolean belongsToLoop =
                                            doWhileTaskModel
                                                            .getWorkflowTask()
                                                            .has(taskRefWithoutIteration)
                                                    && !doWhileTaskModel
                                                            .getReferenceTaskName()
                                                            .equals(task.getReferenceTaskName());

                                    // Check if this task is from an old iteration that should be
                                    // removed
                                    boolean isOldIteration =
                                            task.getIteration() <= iterationsToRemove;

                                    return belongsToLoop && isOldIteration;
                                })
                        .collect(Collectors.toList());

        // Remove each task from the database
        for (TaskModel taskToRemove : tasksToRemove) {
            try {
                LOGGER.debug(
                        "Removing task {} (iteration {}) from workflow {}",
                        taskToRemove.getReferenceTaskName(),
                        taskToRemove.getIteration(),
                        workflow.getWorkflowId());
                executionDAOFacade.removeTask(taskToRemove.getTaskId());
            } catch (Exception e) {
                LOGGER.error(
                        "Failed to remove task {} (iteration {}) from workflow {}",
                        taskToRemove.getReferenceTaskName(),
                        taskToRemove.getIteration(),
                        workflow.getWorkflowId(),
                        e);
                // Continue with other tasks even if one fails
            }
        }

        LOGGER.info(
                "Removed {} tasks from {} old iterations for DO_WHILE task {} in workflow {}",
                tasksToRemove.size(),
                iterationsToRemove,
                doWhileTaskModel.getReferenceTaskName(),
                workflow.getWorkflowId());
    }

    /**
     * Check if all tasks in the current iteration have reached terminal state.
     *
     * @param doWhileTaskModel The {@link TaskModel} of DO_WHILE.
     * @param referenceNameToModel Map of taskReferenceName to {@link TaskModel}.
     * @return true if all tasks in DO_WHILE.loopOver are in <code>referenceNameToModel</code> and
     *     reached terminal state.
     */
    private boolean isIterationComplete(
            TaskModel doWhileTaskModel, Map<String, TaskModel> referenceNameToModel) {
        List<WorkflowTask> workflowTasksInsideDoWhile =
                doWhileTaskModel.getWorkflowTask().getLoopOver();
        int iteration = doWhileTaskModel.getIteration();
        boolean allTasksTerminal = true;
        for (WorkflowTask workflowTaskInsideDoWhile : workflowTasksInsideDoWhile) {
            String taskReferenceName =
                    TaskUtils.appendIteration(
                            workflowTaskInsideDoWhile.getTaskReferenceName(), iteration);
            if (referenceNameToModel.containsKey(taskReferenceName)) {
                TaskModel taskModel = referenceNameToModel.get(taskReferenceName);
                if (!taskModel.getStatus().isTerminal()) {
                    allTasksTerminal = false;
                    break;
                }
            } else {
                allTasksTerminal = false;
                break;
            }
        }

        if (!allTasksTerminal) {
            // Cases where tasks directly inside loop over are not completed.
            // loopOver -> [task1 -> COMPLETED, task2 -> IN_PROGRESS]
            return false;
        }

        // Check all the tasks in referenceNameToModel are completed or not. These are set of tasks
        // which are not directly inside loopOver tasks, but they are under hierarchy
        // loopOver -> [decisionTask -> COMPLETED [ task1 -> COMPLETED, task2 -> IN_PROGRESS]]
        if (referenceNameToModel.values().stream()
                .anyMatch(taskModel -> !taskModel.getStatus().isTerminal())) {
            return false;
        }

        // Check that every terminal task's successor within the DO_WHILE hierarchy has been
        // scheduled. This guards against premature iteration advancement caused by intra-loop
        // ordering in decide(): INLINE (and other sync system tasks) share the
        // tasksToBeScheduled loop with DO_WHILE. If the sync task executes first it becomes
        // terminal in memory, but the decider hasn't yet run to schedule its successor. Without
        // this check DO_WHILE would declare the iteration complete and advance.
        // loopOver -> [SWITCH -> COMPLETED [ task1 -> COMPLETED, task2 -> NOT_YET_SCHEDULED ]]
        String doWhileRef = doWhileTaskModel.getWorkflowTask().getTaskReferenceName();
        for (TaskModel task : referenceNameToModel.values()) {
            if (task.getStatus().isTerminal()) {
                String refNameWithoutIteration =
                        TaskUtils.removeIterationFromTaskRefName(task.getReferenceTaskName());
                WorkflowTask nextWorkflowTask =
                        doWhileTaskModel.getWorkflowTask().next(refNameWithoutIteration, null);
                // A non-null next task that is still within the DO_WHILE hierarchy (i.e. not the
                // DO_WHILE task itself, which is returned for the last task in the sequence) means
                // there is a successor that must be scheduled before the iteration is complete.
                if (nextWorkflowTask != null
                        && !doWhileRef.equals(nextWorkflowTask.getTaskReferenceName())
                        && doWhileTaskModel
                                .getWorkflowTask()
                                .has(nextWorkflowTask.getTaskReferenceName())) {
                    String nextTaskRef =
                            TaskUtils.appendIteration(
                                    nextWorkflowTask.getTaskReferenceName(), iteration);
                    if (!referenceNameToModel.containsKey(nextTaskRef)) {
                        // Successor task not yet scheduled — iteration is not complete.
                        return false;
                    }
                }
            }
        }

        return true;
    }

    boolean scheduleNextIteration(
            TaskModel doWhileTaskModel, WorkflowModel workflow, WorkflowExecutor workflowExecutor) {
        LOGGER.debug(
                "Scheduling loop tasks for task {} as condition {} evaluated to true",
                doWhileTaskModel.getTaskId(),
                doWhileTaskModel.getWorkflowTask().getLoopCondition());
        workflowExecutor.scheduleNextIteration(doWhileTaskModel, workflow);
        return true; // Return true even though status not changed. Iteration has to be updated in
        // execution DAO.
    }

    boolean markTaskFailure(TaskModel taskModel, TaskModel.Status status, String failureReason) {
        LOGGER.error("Marking task {} failed with error.", taskModel.getTaskId());
        taskModel.setReasonForIncompletion(failureReason);
        taskModel.setStatus(status);
        return true;
    }

    boolean markTaskSuccess(TaskModel taskModel) {
        LOGGER.debug(
                "Task {} took {} iterations to complete",
                taskModel.getTaskId(),
                taskModel.getIteration() + 1);
        taskModel.setStatus(TaskModel.Status.COMPLETED);
        return true;
    }

    /**
     * Inject loopItem and loopIndex variables into the DO_WHILE task output for list iteration.
     * Tasks inside the loop can access these via workflow expressions.
     *
     * @param workflow The workflow model
     * @param doWhileTask The DO_WHILE task model
     */
    @VisibleForTesting
    void injectLoopVariables(WorkflowModel workflow, TaskModel doWhileTask) {
        if (!isListIteration(doWhileTask)) {
            return;
        }

        List<Object> itemsList = evaluateItemsList(workflow, doWhileTask);
        int currentIteration = doWhileTask.getIteration();
        int loopIndex = currentIteration - 1; // 0-based index

        // Add loopIndex to output
        doWhileTask.addOutput("loopIndex", loopIndex);

        // Add loopItem to output if within bounds
        if (loopIndex >= 0 && loopIndex < itemsList.size()) {
            Object loopItem = itemsList.get(loopIndex);
            doWhileTask.addOutput("loopItem", loopItem);
            LOGGER.debug(
                    "Injected loop variables for task {}: loopIndex={}, loopItem={}",
                    doWhileTask.getTaskId(),
                    loopIndex,
                    loopItem);
        } else {
            LOGGER.warn(
                    "loopIndex {} is out of bounds for items list of size {} in task {}",
                    loopIndex,
                    itemsList.size(),
                    doWhileTask.getTaskId());
        }
    }

    /**
     * Check if this DO_WHILE task is using list iteration mode (has 'items' parameter or '_items'
     * in inputParameters for Orkes compatibility)
     *
     * @param task The DO_WHILE task model
     * @return true if the task has an 'items' parameter set or '_items' in inputParameters
     */
    @VisibleForTesting
    boolean isListIteration(TaskModel task) {
        // Check new OSS approach: items field on WorkflowTask
        String items = task.getWorkflowTask().getItems();
        if (items != null && !items.trim().isEmpty()) {
            return true;
        }

        // Check Orkes compatibility: _items in inputParameters
        Map<String, Object> inputParams = task.getWorkflowTask().getInputParameters();
        if (inputParams != null && inputParams.containsKey("_items")) {
            Object itemsValue = inputParams.get("_items");
            return itemsValue != null
                    && (itemsValue instanceof String && !((String) itemsValue).trim().isEmpty()
                            || itemsValue instanceof Collection
                            || itemsValue instanceof Object[]);
        }

        return false;
    }

    /**
     * Evaluate the 'items' parameter to get the list of items to iterate over. Supports both new
     * OSS approach (items field on WorkflowTask) and Orkes compatibility (_items in
     * inputParameters).
     *
     * @param workflow The workflow model
     * @param task The DO_WHILE task model
     * @return List of items to iterate over, or empty list if items cannot be evaluated
     */
    @VisibleForTesting
    List<Object> evaluateItemsList(WorkflowModel workflow, TaskModel task) {
        TaskDef taskDefinition = task.getTaskDefinition().orElse(null);
        Object itemsValue = null;

        // Priority 1: Check new OSS approach - items field on WorkflowTask
        String itemsParam = task.getWorkflowTask().getItems();
        if (itemsParam != null && !itemsParam.trim().isEmpty()) {
            // Create a temporary input parameters map with the items parameter
            Map<String, Object> tempInputParams = new HashMap<>();
            tempInputParams.put("items", itemsParam);

            // Use ParametersUtils to evaluate the expression
            Map<String, Object> evaluatedParams =
                    parametersUtils.getTaskInputV2(
                            tempInputParams, workflow, task.getTaskId(), taskDefinition);

            itemsValue = evaluatedParams.get("items");
        }

        // Priority 2: Check Orkes compatibility - _items in inputParameters
        if (itemsValue == null) {
            Map<String, Object> evaluatedInputParams =
                    parametersUtils.getTaskInputV2(
                            task.getWorkflowTask().getInputParameters(),
                            workflow,
                            task.getTaskId(),
                            taskDefinition);

            if (evaluatedInputParams.containsKey("_items")) {
                itemsValue = evaluatedInputParams.get("_items");
            }
        }

        // Convert itemsValue to List<Object>
        if (itemsValue instanceof List) {
            return (List<Object>) itemsValue;
        } else if (itemsValue instanceof Collection) {
            return new ArrayList<>((Collection<?>) itemsValue);
        } else if (itemsValue instanceof Object[]) {
            return Arrays.asList((Object[]) itemsValue);
        } else if (itemsValue != null) {
            // If it's a single value, wrap it in a list
            return Collections.singletonList(itemsValue);
        }

        return Collections.emptyList();
    }

    @VisibleForTesting
    boolean evaluateCondition(WorkflowModel workflow, TaskModel task) {
        TaskDef taskDefinition = task.getTaskDefinition().orElse(null);
        // Use paramUtils to compute the task input
        Map<String, Object> conditionInput =
                parametersUtils.getTaskInputV2(
                        task.getWorkflowTask().getInputParameters(),
                        workflow,
                        task.getTaskId(),
                        taskDefinition);
        conditionInput.put(task.getReferenceTaskName(), task.getOutputData());
        List<TaskModel> loopOver =
                workflow.getTasks().stream()
                        .filter(
                                t ->
                                        (task.getWorkflowTask()
                                                        .has(
                                                                TaskUtils
                                                                        .removeIterationFromTaskRefName(
                                                                                t
                                                                                        .getReferenceTaskName()))
                                                && !task.getReferenceTaskName()
                                                        .equals(t.getReferenceTaskName())))
                        .collect(Collectors.toList());

        for (TaskModel loopOverTask : loopOver) {
            conditionInput.put(
                    TaskUtils.removeIterationFromTaskRefName(loopOverTask.getReferenceTaskName()),
                    loopOverTask.getOutputData());
        }

        // Check if we're in list iteration mode
        if (isListIteration(task)) {
            List<Object> itemsList = evaluateItemsList(workflow, task);
            int currentIteration = task.getIteration();

            // Inject loopIndex and loopItem into condition input
            // loopIndex is 0-based (currentIteration - 1 because iteration starts at 1)
            int loopIndex = currentIteration - 1;
            conditionInput.put("loopIndex", loopIndex);

            // Inject loopItem if we're within bounds
            if (loopIndex >= 0 && loopIndex < itemsList.size()) {
                conditionInput.put("loopItem", itemsList.get(loopIndex));
            }

            // For list iteration, continue if we haven't reached the end of the list
            // The condition is: loopIndex < itemsList.size() - 1 (there's another item after
            // current)
            boolean hasMoreItems = loopIndex < itemsList.size() - 1;

            // If there's a loopCondition, evaluate it AND combine with hasMoreItems
            // Otherwise, just use hasMoreItems
            String condition = task.getWorkflowTask().getLoopCondition();
            if (condition != null && !condition.trim().isEmpty()) {
                LOGGER.debug(
                        "List iteration: Evaluating condition: {} with loopIndex={}, loopItem={}",
                        condition,
                        loopIndex,
                        conditionInput.get("loopItem"));
                boolean conditionResult = ScriptEvaluator.evalBool(condition, conditionInput);
                // Continue only if BOTH condition is true AND there are more items
                return conditionResult && hasMoreItems;
            } else {
                LOGGER.debug(
                        "List iteration: loopIndex={}, items.size={}, hasMoreItems={}",
                        loopIndex,
                        itemsList.size(),
                        hasMoreItems);
                return hasMoreItems;
            }
        }

        // Counter-based iteration (backward compatibility)
        String condition = task.getWorkflowTask().getLoopCondition();
        boolean result = false;
        if (condition != null) {
            LOGGER.debug("Condition: {} is being evaluated", condition);
            // Evaluate the expression by using the Nashorn based script evaluator
            result = ScriptEvaluator.evalBool(condition, conditionInput);
        }
        return result;
    }
}
