#  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
#
#  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:
#
#       https://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.
"""DAG (Directed Acyclic Graph) Runners."""

import threading
import time
from collections import defaultdict
from enum import Enum
from typing import Any, Callable, Dict, List, Optional

from zenml.logger import get_logger

logger = get_logger(__name__)


def reverse_dag(dag: Dict[str, List[str]]) -> Dict[str, List[str]]:
    """Reverse a DAG.

    Args:
        dag: Adjacency list representation of a DAG.

    Returns:
        Adjacency list representation of the reversed DAG.
    """
    reversed_dag = defaultdict(list)

    # Reverse all edges in the graph.
    for node, upstream_nodes in dag.items():
        for upstream_node in upstream_nodes:
            reversed_dag[upstream_node].append(node)

    # Add nodes without incoming edges back in.
    for node in dag:
        if node not in reversed_dag:
            reversed_dag[node] = []

    return reversed_dag


class NodeStatus(Enum):
    """Status of the execution of a node."""

    NOT_STARTED = "not_started"
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class ThreadedDagRunner:
    """Multi-threaded DAG Runner.

    This class expects a DAG of strings in adjacency list representation, as
    well as a custom `run_fn` as input, then calls `run_fn(node)` for each
    string node in the DAG.

    Steps that can be executed in parallel will be started in separate threads.
    """

    def __init__(
        self,
        dag: Dict[str, List[str]],
        run_fn: Callable[[str], Any],
        preparation_fn: Optional[Callable[[str], bool]] = None,
        finalize_fn: Optional[Callable[[Dict[str, NodeStatus]], None]] = None,
        parallel_node_startup_waiting_period: float = 0.0,
        max_parallelism: Optional[int] = None,
        continue_fn: Optional[Callable[[], bool]] = None,
    ) -> None:
        """Define attributes and initialize all nodes in waiting state.

        Args:
            dag: Adjacency list representation of a DAG.
                E.g.: [(1->2), (1->3), (2->4), (3->4)] should be represented as
                `dag={2: [1], 3: [1], 4: [2, 3]}`
            run_fn: A function `run_fn(node)` that runs a single node
            preparation_fn: A function that is called before the node is run.
                If provided, the function return value determines whether the
                node should be run or can be skipped.
            finalize_fn: A function `finalize_fn(node_states)` that is called
                when all nodes have completed.
            parallel_node_startup_waiting_period: Delay in seconds to wait in
                between starting parallel nodes.
            max_parallelism: Maximum number of nodes to run in parallel
            continue_fn: A function that returns True if the run should continue
                after each step execution, False if it should stop (e.g., due
                to cancellation). If None, execution continues normally.

        Raises:
            ValueError: If max_parallelism is not greater than 0.
        """
        if max_parallelism is not None and max_parallelism <= 0:
            raise ValueError("max_parallelism must be greater than 0")

        self.parallel_node_startup_waiting_period = (
            parallel_node_startup_waiting_period
        )
        self.max_parallelism = max_parallelism
        self.dag = dag
        self.reversed_dag = reverse_dag(dag)
        self.run_fn = run_fn
        self.preparation_fn = preparation_fn
        self.finalize_fn = finalize_fn
        self.continue_fn = continue_fn
        self.nodes = dag.keys()
        self.node_states = {
            node: NodeStatus.NOT_STARTED for node in self.nodes
        }
        self._lock = threading.Lock()

        self._stop_requested = False

    def _can_run(self, node: str) -> bool:
        """Determine whether a node is ready to be run.

        This is the case if the node has not run yet and all of its upstream
        node have already completed.

        Args:
            node: The node.

        Returns:
            True if the node can run else False.
        """
        if not self.node_states[node] == NodeStatus.NOT_STARTED:
            return False

        # Check that all upstream nodes of this node have already completed.
        for upstream_node in self.dag[node]:
            if not self.node_states[upstream_node] == NodeStatus.COMPLETED:
                return False

        return True

    def _prepare_node_run(self, node: str) -> None:
        """Prepare a node run.

        Args:
            node: The node.
        """
        if self.max_parallelism is None:
            with self._lock:
                self.node_states[node] = NodeStatus.RUNNING
        else:
            while True:
                with self._lock:
                    logger.debug(f"Checking if {node} can run.")
                    running_nodes = len(
                        [
                            state
                            for state in self.node_states.values()
                            if state == NodeStatus.RUNNING
                        ]
                    )
                    if running_nodes < self.max_parallelism:
                        self.node_states[node] = NodeStatus.RUNNING
                        break

                logger.debug(f"Waiting for {running_nodes} nodes to finish.")
                time.sleep(1)

    def _run_node(self, node: str) -> None:
        """Run a single node.

        Calls the user-defined run_fn, then calls `self._finish_node`.

        Args:
            node: The node.
        """
        self._prepare_node_run(node)

        # Check if execution should continue (e.g., check for cancellation)
        if self.continue_fn:
            self._stop_requested = (
                self._stop_requested or not self.continue_fn()
            )
            if self._stop_requested:
                self._finish_node(node, cancelled=True)
                return

        if self.preparation_fn:
            run_required = self.preparation_fn(node)
            if not run_required:
                self._finish_node(node)
                return

        try:
            self.run_fn(node)
            self._finish_node(node)
        except Exception as e:
            self._finish_node(node, failed=True)
            logger.exception(f"Node `{node}` failed: {e}")

    def _run_node_in_thread(self, node: str) -> threading.Thread:
        """Run a single node in a separate thread.

        Args:
            node: The node.

        Returns:
            The thread in which the node was run.
        """
        assert self.node_states[node] == NodeStatus.NOT_STARTED
        with self._lock:
            self.node_states[node] = NodeStatus.PENDING

        # Run node in new thread.
        thread = threading.Thread(
            name=node, target=self._run_node, args=(node,)
        )
        thread.start()
        return thread

    def _finish_node(
        self, node: str, failed: bool = False, cancelled: bool = False
    ) -> None:
        """Mark a node as finished and potentially start new nodes.

        Args:
            node: The node to mark as finished.
            failed: Whether the node failed.
            cancelled: Whether the node was cancelled.
        """
        with self._lock:
            if failed:
                self.node_states[node] = NodeStatus.FAILED
            elif cancelled:
                self.node_states[node] = NodeStatus.CANCELLED
            else:
                self.node_states[node] = NodeStatus.COMPLETED

        if failed or cancelled:
            # If the node failed or was cancelled, we don't need to run any downstream nodes.
            return

        # Run downstream nodes.
        threads: List[threading.Thread] = []
        for downstream_node in self.reversed_dag[node]:
            if self._can_run(downstream_node):
                if threads and self.parallel_node_startup_waiting_period > 0:
                    time.sleep(self.parallel_node_startup_waiting_period)

                thread = self._run_node_in_thread(downstream_node)
                threads.append(thread)

        # Wait for all downstream nodes to complete.
        for thread in threads:
            thread.join()

    def run(self) -> None:
        """Call `self.run_fn` on all nodes in `self.dag`.

        The order of execution is determined using topological sort.
        Each node is run in a separate thread to enable parallelism.
        """
        # Run all nodes that can be started immediately.
        # These will, in turn, start other nodes once all of their respective
        # upstream nodes have completed.
        threads: List[threading.Thread] = []
        for node in self.nodes:
            if self._can_run(node):
                if threads and self.parallel_node_startup_waiting_period > 0:
                    time.sleep(self.parallel_node_startup_waiting_period)

                thread = self._run_node_in_thread(node)
                threads.append(thread)

        # Wait till all nodes have completed.
        for thread in threads:
            thread.join()

        # Call the finalize function.
        if self.finalize_fn:
            self.finalize_fn(self.node_states)

        # Print a status report.
        failed_nodes = []
        skipped_nodes = []
        for node in self.nodes:
            if self.node_states[node] == NodeStatus.FAILED:
                failed_nodes.append(node)
            elif self.node_states[node] == NodeStatus.NOT_STARTED:
                skipped_nodes.append(node)

        if failed_nodes:
            logger.error(
                "The following nodes failed: " + ", ".join(failed_nodes)
            )
        if skipped_nodes:
            logger.warning(
                "The following nodes were not run because they depend on other "
                "nodes that didn't complete: " + ", ".join(skipped_nodes)
            )
        if not failed_nodes and not skipped_nodes:
            logger.info("All nodes completed successfully.")
