# SPDX-FileCopyrightText: 2025-2026 Global Computing Lab
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

from __future__ import annotations

from pprint import pprint
from typing import TYPE_CHECKING, Mapping

import networkx as nx

from a4x.orchestration.annotations import AnnotationType
from a4x.orchestration.site import Site
from a4x.orchestration.task import Task

if TYPE_CHECKING:
    from a4x.orchestration.file import File
    from a4x.orchestration.plugin import Plugin
    from a4x.orchestration.utils import StrCompatPathLike


class Workflow(AnnotationType):
    """A class representing an entire workflow."""

    def __init__(self, name: str, description: str = "") -> None:
        """Initialize the workflow.

        :param name: the name of the workflow
        :param description: a description of the workflow
        """
        super().__init__()
        self.name = name
        self.description = description
        self.graph = nx.DiGraph()
        self.task_inputs = set()
        self.task_outputs = set()
        self.sites = set()
        self.global_environ = {}

    def add_site(self, site: Site) -> None:
        """Add a single site to the workflow's set of sites.

        :param site: the site to add
        """
        if not isinstance(site, Site):
            raise TypeError("Sites must be of type 'a4x.orchestration.Site'")
        if site not in self.sites:
            self.sites.add(site)

    def add_sites(self, *sites) -> None:
        """Add multiple sites to the workflow's set of sites.

        :param sites: the sites to add
        """
        for s in sites:
            self.add_site(s)

    def add_task(self, task: Task) -> None:
        """Add a single task to the workflow.

        This method essentially adds a vertex to the NetworkX DAG underlying the
        Workflow class.

        :param task: the task to add
        """
        if not isinstance(task, Task):
            raise TypeError("Tasks provided to 'add_task' must be of type 'Task'")
        self.task_inputs.update(task.inputs)
        self.task_outputs.update(task.outputs)
        self.graph.add_node(task.task_name, task=task)
        if task.site is not None:
            self.add_site(task.site)

    def add_tasks(self, *tasks) -> None:
        """Add multiple tasks to the workflow.

        This method essentially adds multiple vertex to the NetworkX DAG underlying
        the Workflow class.

        :param tasks: the tasks to add
        """
        for t in tasks:
            self.add_task(t)

    def add_edge(self, src_task: Task, dest_task: Task) -> None:
        """Add a dependency between two tasks in the workflow.

        This method essentially adds an edge between two vertices in the NetworkX
        DAG underlying the Workflow class. If either task is not in the workflow already,
        it will be added into the workflow before creating the dependency.

        :param src_task: the parent task (i.e., the edge's source vertex)
        :param dest_task: the child task (i.e., the edge's destination vertex)
        """
        if src_task.task_name not in self.graph:
            self.add_task(src_task)
        if dest_task.task_name not in self.graph:
            self.add_task(dest_task)
        self.graph.add_edge(src_task.task_name, dest_task.task_name)

    def add_dependency(self, task: Task, parents: list[Task], children: list[Task]) -> None:
        """Add dependencies between the specified task and its parents and children.

        This method is similar to :code:`Workflow.add_edge`, but it adds multiple dependencies/edges
        to the workflow in one call. It creates edges from each task in :code:`parents` to the main
        task (i.e., the :code:`task` parameter) and from the main task to each task in :code:`children`.

        :param task: the main task for adding dependencies
        :param parents: the tasks that should become parents of the main task
        :param children: the tasks that should become the children of the main task
        """
        if not isinstance(task, Task):
            raise TypeError("Tasks provided to 'add_dependency' must be of type 'Task'")
        if task.task_name not in self.graph:
            self.add_task(task)
        for p in parents:
            if not isinstance(p, Task):
                raise TypeError("Tasks provided to 'add_dependency' must be of type 'Task'")
            if p.task_name not in self.graph:
                self.add_task(p)
            self.add_edge(p, task)
        for c in children:
            if not isinstance(c, Task):
                raise TypeError("Tasks provided to 'add_dependency' must be of type 'Task'")
            if c.task_name not in self.graph:
                self.add_task(c)
            self.add_edge(task, c)

    def generate_dependencies_from_task_inputs_outputs(self, override: bool = True) -> None:
        """Generate dependencies between tasks based on each task's :code:`inputs` and :code:`outputs` properties.

        This method links tasks such that one or more of task A's outputs are also one or more of task B's inputs.
        After determining these links, this method uses :code:`Workflow.add_edge` to create edges in the
        NetworkX DAG underlying the Workflow class that correspond with the links built from inputs and outputs.

        :param override: if true, delete all existing dependencies in the workflow before generating new dependencies. Otherwise, this method will raise an error if any dependencies already exist.
        """  # noqa: E501
        if not override and self.num_edges != 0:
            raise RuntimeError(
                "Cannot generate task dependencies when dependencies already exist and 'override' is False"
            )
        if override and self.num_edges != 0:
            self.graph.clear_edges()
        input_taskmap = {}
        for _, task in self.graph.nodes(data="task"):  # type: ignore
            for inp in task.inputs:
                if inp in input_taskmap.keys():
                    input_taskmap[inp].append(task)
                else:
                    input_taskmap[inp] = [task]
        for _, task in self.graph.nodes(data="task"):  # type: ignore
            for outp in task.outputs:
                if outp in input_taskmap.keys():
                    for child_task in input_taskmap[outp]:
                        self.add_edge(task, child_task)

    def get_task_by_name(self, name: str) -> Task:
        """Get a Task object from the workflow based on its task name.

        :return: the requested Task
        """
        return dict(self.graph.nodes(data="task"))[name]

    def get_tasks_in_topological_order(self) -> list[Task]:
        """Get a list of Task objects in the workflow in topoological order.

        This method simply uses NetworkX to sort the task names topologically, and then it converts
        those sorted task names to Task objects.

        :return: a topologically sorted list of Tasks in the workflow
        """
        return [self.graph.nodes[tn]["task"] for tn in nx.lexicographical_topological_sort(self.graph)]

    def convert(self, plugin_key: str, ignore_builtin_plugins: bool = False, **kwargs) -> Plugin:
        """Convert the workflow to a WMS/RJMS-specific plugin.

        :param plugin_key: the identifier of the plugin to generate. Must be an element of the output of :code:`a4x.orchestration.get_plugin_names`
        :param ignore_builtin_plugins: if true, do not consider plugins built into A4X-Orchestration directly
        :param kwargs: remaining keyword arguments that are passed as-is to the plugin
        """  # noqa: E501
        from a4x.orchestration import convert_to_plugin

        return convert_to_plugin(self, plugin_key, ignore_builtin_plugins, **kwargs)

    def resolve(self) -> None:
        """Resolve all files and tasks in the workflow.

        This method does two main things. First, it resolves all File objects in the workflow,
        which involves determining whether the file is virtual or not and, if not, using pathlib to get a resolved path.
        Second, this method will add any task execution sites to the workflow's set of sites.
        """
        for _, task in self.graph.nodes(data="task"):
            task._resolve_site()  # noqa: SLF001
            if task.site is not None:
                self.sites.add(task.site)
        for site in self.sites:
            site.resolve_site_files()

    @property
    def num_nodes(self) -> int:
        """The number of tasks (i.e., graph nodes) in the workflow."""
        return self.graph.number_of_nodes()

    @property
    def num_edges(self) -> int:
        """The number of dependencies (i.e., graph edges) in the workflow."""
        return self.graph.number_of_edges()

    @property
    def num_task_dependencies(self) -> int:
        """The number of dependencies (i.e., graph edges) in the workflow."""
        return self.num_edges

    @property
    def max_parallel_tasks(self) -> int:
        """Maximum number of tasks that could possibly be run in parallel.

        This property uses Dilworth's Theory to determine the size of the maximum
        antichain in the DAG.
        """
        # This approach requires the workflow to be a DAG
        if not nx.is_directed_acyclic_graph(self.graph):
            raise ValueError("Cannot get max_parallel_tasks if the graph is not a DAG")
        # Get the list of vertex labels so we don't have to repeatedly invoked "nodes()"
        nodes = list(self.graph.nodes())
        # Exit early if there are no vertices
        if len(nodes) == 0:
            return 0
        # Dilworth's theory says that the size of the maximum antichain is equal to the
        # minimum number of chains needed to cover all the nodes. This means the problem
        # translates to a minimum path covering of the DAG. The size of the minimum
        # path covering is known to equal:
        #   Number of vertices - Maximum matching on the reachability bipartite graph
        # When building this reachability bipartite graph, we need to have an edge between
        # any two vertices connected by a path. The best way to collect the set of these
        # edges is to create a transitive closure of the original graph.
        trans_closure = nx.transitive_closure_dag(self.graph)
        # Create a bipartite graph containing copies of the labels of each vertex
        # in the workflow graph in each partition
        bipartite_graph = nx.Graph()
        part0_nodes = []
        part1_nodes = []
        for n in nodes:
            # We add "L" and "R" to the tuples so that networkx doesn't try to merge vertices
            # in the two partitions
            part0_nodes.append((n, "L"))
            part1_nodes.append((n, "R"))
        bipartite_graph.add_nodes_from(part0_nodes, bipartite=0)
        bipartite_graph.add_nodes_from(part1_nodes, bipartite=1)

        # Add edges to the bipartite graph between any two vertices in the original
        # graph that are connected by a path. We loop over the transitive closure of the original
        # graph here to ensure we connect vertices based on path connection, not just direct (u, v)
        # edges
        for u, v in trans_closure.edges():
            if u != v:
                bipartite_graph.add_edge((u, "L"), (v, "R"))

        # Calculate the maximum matching of the bipartite graph
        # Note: this returns a dictionary where keys are vertex labels in some matching and values
        #       are the corresponding vertex the key is matched to
        max_matching = nx.bipartite.maximum_matching(bipartite_graph, top_nodes=part0_nodes)
        # NetworkX will return a dictionary with 2 entries per matching. Since we only
        # want to count each matching edge once, we divide the number of entries by 2.
        # We use floor division to ensure the result is an integer.
        matching_size = len(max_matching) // 2
        # As explained above, the maximum antichain equals the number of nodes minus the
        # size of the maximum matching of the reachability bipartite graph
        return len(nodes) - matching_size

    @property
    def width(self) -> int:
        """Maximum number of tasks that could possibly be run in parallel.

        This property uses Dilworth's Theory to determine the size of the maximum
        antichain in the DAG.
        """
        return self.max_parallel_tasks

    def __len__(self) -> int:
        return self.num_nodes

    @property
    def task_inputs_and_outputs(self) -> set[File]:
        """All File objects in the :code:`inputs` and :code:`outputs` fields of the workflow tasks."""
        return self.all_task_inputs | self.all_task_outputs

    @property
    def all_task_inputs(self) -> set[File]:
        """All File objects in the :code:`inputs` field of the workflow tasks."""
        files = set()
        for task in self.get_tasks_in_topological_order():
            files |= set(task.inputs)
        return files

    @property
    def all_task_outputs(self) -> set[File]:
        """All File objects in the :code:`outputs` field of the workflow tasks."""
        files = set()
        for task in self.get_tasks_in_topological_order():
            files |= set(task.outputs)
        return files

    @property
    def task_inputs_from_graph(self) -> dict[str, list[File]]:
        """A mapping of task names to input File objects."""
        return {
            node_tup[0]: node_tup[1].inputs
            for node_tup in self.graph.nodes(data="task")  # type: ignore
        }

    @property
    def task_outputs_from_graph(self) -> dict[str, list[File]]:
        """A mapping of task names to output File objects."""
        return {
            node_tup[0]: node_tup[1].outputs
            for node_tup in self.graph.nodes(data="task")  # type: ignore
        }

    @property
    def root_tasks(self) -> list[Task]:
        """The root tasks in the workflow (i.e., tasks with no parent tasks)."""
        root_task_names = [task_name for task_name, in_degree in self.graph.in_degree() if in_degree == 0]
        return [self.graph.nodes[name]["task"] for name in root_task_names]

    @property
    def leaf_tasks(self) -> list[Task]:
        """The leaf tasks in the workflow (i.e., tasks with no children tasks)."""
        root_task_names = [task_name for task_name, out_degree in self.graph.out_degree() if out_degree == 0]
        return [self.graph.nodes[name]["task"] for name in root_task_names]

    @property
    def input_files(self) -> set[File]:
        """The set of input files for the root tasks in the graph."""
        input_set = set()
        for task in self.root_tasks:
            input_set |= set(task.inputs)
        return input_set

    @property
    def output_files(self) -> set[File]:
        """The set of output files for the leaf tasks in the graph."""
        output_set = set()
        for task in self.leaf_tasks:
            output_set |= set(task.outputs)
        return output_set

    @property
    def intermediate_files(self) -> set[File]:
        """The set of intermediate files in the graph."""
        intermediate_set = set()
        root_tasks = self.root_tasks
        leaf_tasks = self.leaf_tasks
        for task in self.get_tasks_in_topological_order():
            if task in root_tasks and task in leaf_tasks:
                continue
            if task not in root_tasks:
                intermediate_set |= set(task.inputs)
            if task not in leaf_tasks:
                intermediate_set |= set(task.outputs)
        return intermediate_set

    @property
    def environment(self) -> Mapping:
        """Extra environment variables to specify in the workflow."""
        return self.global_environ

    @environment.setter
    def environment(self, environ: Mapping) -> None:
        if not isinstance(environ, Mapping):
            raise TypeError("The 'environment' property must be a mapping type")
        self.global_environ = environ

    def __repr__(self) -> str:
        return f"Workflow({self.name})"

    def to_config(self, fpath: StrCompatPathLike) -> None:
        """Create a YAML config file for the workflow.

        :param fpath: the path of the YAML file to create
        """
        from a4x.orchestration.config import write_config

        write_config(self, fpath)

    @staticmethod
    def from_config(fpath: StrCompatPathLike) -> Workflow:
        """Create an A4X Workflow object from a YAML config file.

        :param fpath: the path of the YAML file to load
        :return: the created A4X Workflow object
        """
        from a4x.orchestration.config import read_config

        return read_config(fpath)


def print_workflow_info(wflow: Workflow) -> None:
    """Print the contents of a workflow to STDOUT.

    :param wflow: the workflow to print
    """
    print("Workflow:", wflow.name)
    print("Description:", wflow.description)
    if len(wflow.annotations) > 0:
        print("Annotations:")
        pprint(wflow.annotations, indent=4)
    print()

    wflow_tasks = wflow.get_tasks_in_topological_order()
    wflow_deps = {t.task_name: list(wflow.graph.successors(t.task_name)) for t in wflow_tasks}

    print("Tasks:")
    print("======")
    for task in wflow_tasks:
        print("  - Name:", task.task_name)
        print(
            "    Site:",
            task.site if task.site is not None else None,
        )
        print("    Inputs:")
        for inp in task.inputs:
            print("    * Path/Name:", inp)
            print("      Directory:", inp.directory)
            if inp.directory is not None:
                print("      Storage Type:", inp.directory.storage_type)
                print("      Persistency:", inp.directory.persistency)
                if inp.directory.site is not None:
                    print("      Site:", inp.directory.site.name)
                    print("      Scheduler:", inp.directory.site.scheduler)
        print("    Outputs:")
        for inp in task.outputs:
            print("    * Path/Name:", inp.path_attr)
            print("      Directory:", str(inp.directory))
            if inp.directory is not None:
                print("      Storage Type:", inp.directory.storage_type)
                print("      Persistency:", inp.directory.persistency)
                if inp.directory.site is not None:
                    print("      Site:", inp.directory.site.name)
                    print("      Scheduler:", inp.directory.site.scheduler)
        print("    Commands:")
        for cmd in task.commands:
            print("    * Command:", cmd.command_or_exe)
            print("      Args:", cmd.args)
        if len(task.annotations) > 0:
            print("    Annotations:", task.annotations)
    print()

    print("Task Dependencies:")
    print("==================")
    for parent, children in wflow_deps.items():
        print(f"  - {parent} -> {children}")
