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

from __future__ import annotations

from typing import Any, MutableSequence

from typing_extensions import Self

from a4x.orchestration.command import Command
from a4x.orchestration.file import File
from a4x.orchestration.resources import SchedulableWork
from a4x.orchestration.site import Site, get_default_site


class Task(SchedulableWork, MutableSequence):
    """
    A class representing a single Task in the workflow.

    This class behaves like a Python :code:`list` containing :code:`Command`
    objects. When adding commands into a :code:`Task`, users can add either
    strings or :code:`Command` objects. When provided strings, the :code:`Task`
    class will implicitly convert the strings into :code:`Command` objects with
    no scheduling settings.

    This class also supports all the same properties and methods for
    scheduling settings as :code:`Command`. Regarding scheduling, users
    can think of the :code:`Command` class as representing parallel launches
    (with e.g., 'mpiexec', 'srun', 'flux run'), while the :code:`Task` class
    represents batch submissions (with e.g., 'sbatch', 'flux batch').
    """

    def __init__(self, name: str, description: str = "") -> None:
        """Create a Task.

        :param name: the name (i.e., unique identifier) of the Task
        :param description: a description of the Task
        """
        super().__init__()
        self.task_name = name
        self.description = description
        self.commands = []
        self.inputs = []
        self.outputs = []
        self.site = None
        self.add_input_extra_kwargs = {}
        self.add_output_extra_kwargs = {}

    def set_site(self, site: Site) -> None:
        """Set the execution site for this task.

        :param site: the site
        """
        if not isinstance(site, Site):
            raise TypeError("The site for a Task must be of type 'Site'")
        self.site = site

    def __getitem__(self, idx: int | slice) -> Command:
        """Get a command from the Task.

        :param idx: the index into the list of :code:`Command` objects
        :return: one or more requested :code:`Command` object(s)
        :raises IndexError: if the index in out-of-bounds
        """
        return self.commands[idx]

    def __setitem__(self, idx: int | slice, val: str | Command) -> None:
        """Set a command in the Task.

        :param idx: the index into the list of :code:`Command` objects
        :param val: the command to add
        :raises TypeError: if :code:`val` is not a string or :code:`Command`
        :raises IndexError: if the index in out-of-bounds
        """
        if not isinstance(val, (str, Command)):
            raise TypeError("Tasks can only store Commands or strings (which get implicitly converted into Commands)")
        cmd = val
        if isinstance(val, str):
            cmd = Command(val)
        self.commands[idx] = cmd

    def __delitem__(self, idx: int | slice) -> None:
        """Remove a command from the Task by index.

        :param idx: the index into the list of :code:`Command` objects
        :raises IndexError: if the index in out-of-bounds
        """
        del self.commands[idx]

    def __len__(self) -> int:
        """Get the number of commands in the Task.

        :return: the number of commands in the Task
        :rtype: int
        """
        return len(self.commands)

    def insert(self, idx: int | slice, val: str | Command) -> None:
        """Insert a command at a specific index in the Task's command list.

        :param idx: the index into the list of :code:`Command` objects
        :param val: the command to add
        :raises TypeError: if :code:`val` is not a string or :code:`Command`
        """
        if not isinstance(val, (str, Command)):
            raise TypeError("Tasks can only store Commands or strings (which get implicitly converted into Commands)")
        cmd = val
        if isinstance(val, str):
            cmd = Command(val)
        self.commands.insert(idx, cmd)

    def add_inputs(self, *inputs, **extra_kwargs) -> Self:
        """Add one or more inputs to the Task.

        :param *inputs: the inputs to add to the Task
        :param **extra_kwargs: extra keyword arguments to associate with input addition. WMS plugins can choose to use or not to use these arguments
        """  # noqa: E501
        self.add_input_extra_kwargs = extra_kwargs
        self.inputs = list(inputs)
        if any(not isinstance(i, File) for i in self.inputs):
            raise TypeError("All positional arguments to 'add_inputs' should be of type 'a4x.orchestration.File'")
        return self

    def add_outputs(self, *outputs, **extra_kwargs) -> Self:
        """Add one or more outputs to the Task.

        :param *outputs: the outputs to add to the Task
        :param **extra_kwargs: extra keyword arguments to associate with output addition. WMS plugins can choose to use or not to use these arguments
        """  # noqa: E501
        self.add_output_extra_kwargs = extra_kwargs
        self.outputs = list(outputs)
        if any(not isinstance(o, File) for o in self.outputs):
            raise TypeError("All positional arguments to 'add_outputs' should be of type 'a4x.orchestration.File'")
        return self

    def get_inputs(self) -> list[File]:
        """Get all inputs to the Task.

        :return: all inputs to the Task
        """
        return self.inputs

    def get_outputs(self) -> list[File]:
        """Get all outputs to the Task.

        :return: all outputs to the Task
        """
        return self.outputs

    def __hash__(self) -> int:
        return hash(self.task_name)

    def __repr__(self) -> str:
        return f"Task(name={self.task_name},site={self.site},commands={self.commands})"

    def __eq__(self, other: Any) -> bool:
        return (
            isinstance(other, Task)
            and self.task_name == other.task_name
            and self.commands == other.commands
            and self.inputs == other.inputs
            and self.outputs == other.outputs
            and (
                (self.site is None and other.site is None)
                or (self.site is not None and other.site is not None and self.site == other.site)
            )
        )

    def __lt__(self, other: Task) -> bool:
        return self.task_name < other.task_name

    def _resolve_site(self) -> None:
        """Resolve all input and output File objects in the Task and populate the execution site, if not already populated."""  # noqa: E501
        default_site = get_default_site()
        if self.site is None:
            file_sites = set()
            for in_file in self.inputs:
                if in_file.site is not None:
                    file_sites.add(in_file.site)
            for out_file in self.outputs:
                if out_file.site is not None:
                    file_sites.add(out_file.site)
            if len(file_sites) > 1:
                raise RuntimeError(
                    f"Task '{self.task_name}' does not have a site set and input/output files are in multiple sites. Either explicitly set the site for this task with 'Task.set_site' or adjust your inputs/outputs"  # noqa: E501
                )
            if len(file_sites) == 1:
                self.site = next(iter(file_sites))
            else:
                # Else use the default site as the site for this task
                self.site = None
            if self.site is None:
                self.site = default_site
