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

from __future__ import annotations

from abc import abstractmethod
from pathlib import Path
from typing import Any

from pydantic import (
    BaseModel,
    NonNegativeInt,
    PositiveInt,
    ValidationError,
    model_validator,
)
from ruamel.yaml import YAML
from typing_extensions import Literal, Self

from a4x.orchestration.annotations import AnnotationType
from a4x.orchestration.command import Command
from a4x.orchestration.file import File
from a4x.orchestration.resources import Resources, SchedulableWork, Slot
from a4x.orchestration.site import (
    Directory,
    PersistencyType,
    Scheduler,
    Site,
    StorageType,
)
from a4x.orchestration.task import Task
from a4x.orchestration.utils import StrCompatPathLike, StrCompatPathLikeForIsInstance
from a4x.orchestration.workflow import Workflow


class A4XOrchestrationConfigBase(BaseModel):
    """Base class for building the A4X-Orchestration config spec"""

    def _validate_sites_directories_and_files(
        self,
        valid_sites: set[str],
        valid_directories: set[str],
        valid_files: set[str],
    ) -> None:
        """Validate sites, directories, or files represented in the current component of the config.

        :param valid_sites: a set of valid site names
        :param valid_directories: a set of valid directory names
        :param valid_files: a set of valid file names
        """
        raise NotImplementedError(
            f"INTERNAL ERROR: _validate_sites_directories_and_files is not implemented for {type(self).__name__}",
        )

    @abstractmethod
    def _to_a4x(
        self,
        curr_a4x_object: Any | None = None,
        site_map: dict[str, Site] | None = None,
        file_map: dict[str, File] | None = None,
    ) -> Any:
        """Convert this component of the config to an A4X object or update an existing A4X object.

        :param curr_a4x_object: the A4X object to update, if provided
        :param site_map: a mapping of site name to A4X Site objects
        :param file_map: a mapping of file name to A4X File objects
        :return: the created A4X object or :code:`curr_a4x_object` if it was updated instead
        """

    @classmethod
    @abstractmethod
    def _from_a4x(
        cls,
        curr_a4x_object: Any,
        curr_config_object: A4XOrchestrationConfigBase | None = None,
    ) -> Self:
        """Convert the provided A4X object into this component of the config or update an existing config component.

        :param curr_a4x_object: the A4X object to convert
        :param curr_config_object: the config component to update, if provided
        :return: the created config component or :code:`curr_config_object` if it was updated instead
        """


class AnnotationConfig(A4XOrchestrationConfigBase):
    """The base class for any config component that can specify annotations."""

    annotations: dict[str, Any] = {}
    """User-specified annotations on the config component"""

    def _to_a4x(
        self,
        curr_a4x_object: AnnotationType | None = None,
        site_map: dict[str, Site] | None = None,
        file_map: dict[str, File] | None = None,
    ) -> AnnotationType:
        # This function requires an A4X object to update
        if curr_a4x_object is None or not isinstance(curr_a4x_object, AnnotationType):
            raise RuntimeError("INTERNAL ERROR: Cannot call _to_a4x without an A4X object for AnnotationConfig")
        # Set the A4X object's annotations field with this model
        curr_a4x_object.annotations = self.annotations.copy()
        return curr_a4x_object

    @classmethod
    def _from_a4x(
        cls,
        curr_a4x_object: AnnotationType,
        curr_config_object: AnnotationConfig | None = None,
    ) -> Self:
        # This function requires a config component to update
        if curr_config_object is None or not isinstance(curr_config_object, AnnotationConfig):
            raise RuntimeError(
                "INTERNAL ERROR: Cannot call _from_a4x without an existing config object for AnnotationConfig",
            )
        # Set the config component's annotations field with the A4X object
        curr_config_object.annotations = curr_a4x_object.annotations.copy()
        return curr_config_object


class SiteAndDirectoryConfig(A4XOrchestrationConfigBase):
    """A specification of a directory on a specific site."""

    site: str
    """The name of the site where the directory exists"""
    directory: str
    """The name of the directory"""

    def _validate_sites_directories_and_files(
        self,
        valid_sites: set[str],
        valid_directories: set[str],
        valid_files: set[str],
    ) -> None:
        # Check if the site and directory identifiers are specified in the 'sites' section of the config
        if self.site not in valid_sites:
            raise ValueError(f"Unrecognized site ({self.site}) in site/directory entry")
        if self.directory not in valid_directories:
            raise ValueError(f"Unrecognized directory ({self.directory}) in site/directory entry")

    def _to_a4x(
        self,
        curr_a4x_object: Any | None = None,
        site_map: dict[str, Site] | None = None,
        file_map: dict[str, File] | None = None,
    ) -> Directory:
        # Get the actual A4X Site object from site_map
        try:
            site = site_map[self.site]
        except KeyError as e:
            raise KeyError(f"Invalid site name: '{self.site}'") from e
        # Get the actual A4X Directory object from the Site object
        try:
            return site[self.directory]
        except KeyError as e:
            raise KeyError(f"Invalid directory identifier: '{self.directory}'") from e

    @classmethod
    def _from_a4x(
        cls,
        curr_a4x_object: Directory,
        curr_config_object: A4XOrchestrationConfigBase | None = None,
    ) -> Self:
        # Create the config component from an A4X Directory object
        return cls(
            site=curr_a4x_object.site.name,
            directory=curr_a4x_object.identifier,
        )


class ResourcesConfig(A4XOrchestrationConfigBase):
    """A specification of resources without defining a slot."""

    num_procs: PositiveInt
    """Number of processes/ranks in the resource request"""
    cores_per_proc: PositiveInt | None = None
    """Number of CPU cores per process"""
    gpus_per_proc: NonNegativeInt | None = None
    """Number of GPUs per process"""
    num_nodes: PositiveInt | None = None
    """Number of nodes in the resource request"""
    exclusive: bool = False
    """If true, allocate nodes exclusively"""
    exclusive_node_per_proc: bool = False
    """If true, ensure each process has an entire node-worth of resources"""

    def _to_a4x(
        self,
        curr_a4x_object: SchedulableWork | None = None,
        site_map: dict[str, Site] | None = None,
        file_map: dict[str, File] | None = None,
    ) -> SchedulableWork:
        # This function requires an A4X object to update
        if curr_a4x_object is None or not isinstance(curr_a4x_object, SchedulableWork):
            raise RuntimeError(
                "INTERNAL ERROR: Cannot call _to_a4x for ResourcesConfig without a valid SchedulableWork object"
            )
        # The fields in ResourcesConfig map 1-to-1 to A4X's SchedulableWork.set_resources method
        curr_a4x_object.set_resources(
            num_procs=self.num_procs,
            cores_per_proc=self.cores_per_proc,
            gpus_per_proc=self.gpus_per_proc,
            num_nodes=self.num_nodes,
            allocate_nodes_exclusively=self.exclusive,
            exclusive_node_per_proc=self.exclusive_node_per_proc,
        )
        return curr_a4x_object

    @classmethod
    def _from_a4x(
        cls,
        curr_a4x_object: Any,
        curr_config_object: A4XOrchestrationConfigBase | None = None,
    ) -> Self:
        raise NotImplementedError(
            "ResourcesConfig._from_a4x is not implemented because generated configs will use ResourcesWSlotConfig"
        )


class SlotConfig(A4XOrchestrationConfigBase):
    """A specification of the resources in a process slot."""

    num_nodes: PositiveInt | None = None
    """Number of nodes per process"""
    num_cores: PositiveInt | None = None
    """Number of CPU cores per process"""
    num_gpus: NonNegativeInt | None = None
    """Number of GPUs per process"""

    @model_validator(mode="after")
    def _mandate_at_least_one_field_required(self) -> Self:
        # Raise an error if none of the fields are specified
        if not any(getattr(self, field) is not None for field in ("num_nodes", "num_cores", "num_gpus")):
            raise ValueError("At least one of 'num_nodes', 'num_cores', or 'num_gpus' is required for a slot")
        return self

    def _to_a4x(
        self,
        curr_a4x_object: Any | None = None,
        site_map: dict[str, Site] | None = None,
        file_map: dict[str, File] | None = None,
    ) -> Slot:
        # Create the slot from the fields of this model
        return Slot(num_nodes=self.num_nodes, num_cores=self.num_cores, num_gpus=self.num_gpus)

    @classmethod
    def _from_a4x(
        cls,
        curr_a4x_object: Slot,
        curr_config_object: A4XOrchestrationConfigBase | None = None,
    ) -> Self:
        # Create the SlotConfig from the fields of the Slot
        return cls(
            num_nodes=curr_a4x_object.nodes,
            num_cores=curr_a4x_object.cores,
            num_gpus=curr_a4x_object.gpus,
        )


class ResourcesWSlotConfig(A4XOrchestrationConfigBase):
    """A specification of resources using slots for per-process resources."""

    num_procs: PositiveInt
    """The total number of processes for the resource specification"""
    slot: SlotConfig
    """The slot defining the resources per process"""
    num_nodes: PositiveInt | None = None
    """The total number of nodes for the resource specification"""
    num_slots_per_node: PositiveInt | None = None
    """The number of slots to allocate per node"""
    exclusive: bool = False
    """If true, allocate nodes exclusively"""

    def _to_a4x(
        self,
        curr_a4x_object: Any | None = None,
        site_map: dict[str, Site] | None = None,
        file_map: dict[str, File] | None = None,
    ) -> Resources:
        # Create an A4X Slot object from the slot specification
        slot = self.slot._to_a4x()  # noqa: SLF001
        # Create an A4X Resources object
        return Resources(
            num_procs=self.num_procs,
            per_proc_resources=slot,
            num_nodes=self.num_nodes,
            exclusive=self.exclusive,
        )

    @classmethod
    def _from_a4x(
        cls,
        curr_a4x_object: Resources,
        curr_config_object: A4XOrchestrationConfigBase | None = None,
    ) -> Self:
        # Create the resource specification from the A4X Resources object
        return ResourcesWSlotConfig(
            num_procs=curr_a4x_object.num_procs,
            slot=SlotConfig._from_a4x(curr_a4x_object=curr_a4x_object.resources_per_slot),  # noqa: SLF001
            num_nodes=curr_a4x_object.num_nodes,
            num_slots_per_node=curr_a4x_object.num_slots_per_node,
            exclusive=curr_a4x_object.exclusive,
        )


class SchedulableConfig(AnnotationConfig):
    """A specification of scheduling parameters."""

    duration: int | str | None = None
    """The walltime/duration of the job or task"""
    queue: str | None = None
    """The submission queue of the job or task"""
    project: str | None = None
    """The project (or "account" in Slurm terms) for the job or task"""
    cwd: str | SiteAndDirectoryConfig | None = None
    """The working directory of the job or task"""
    environment: dict[str, Any] = {}
    """Any environment variables for the job or task"""
    stdin: str | None = None
    """If provided, redirect STDIN to the specified file"""
    stdout: str | None = None
    """If provided, redirect STDOUT to the specified file"""
    stderr: str | None = None
    """If provided, redirect STDERR to the specified file"""
    resources: ResourcesConfig | None = None
    """A specification of resources for the job or task without slots"""
    resources_w_slot: ResourcesWSlotConfig | None = None
    """A specification of resources for the job or task with slots"""

    @model_validator(mode="after")
    def validate_resources_and_resources_w_slot(self) -> Self:
        # Ensure that we don't have both 'resources' and 'resources_w_slot' set
        if self.resources is not None and self.resources_w_slot is not None:
            raise ValueError("Either 'resources' or 'resources_w_slot' can be set, but not both")
        return self

    def _validate_sites_directories_and_files(
        self,
        valid_sites: set[str],
        valid_directories: set[str],
        valid_files: set[str],
    ) -> None:
        # If the 'cwd' field is a SiteAndDirectoryConfig object, call its _validate_sites_directories_and_files
        # method to ensure the config is valid
        if isinstance(self.cwd, SiteAndDirectoryConfig):
            self.cwd._validate_sites_directories_and_files(valid_sites, valid_directories, valid_files)  # noqa: SLF001

    def _to_a4x(
        self,
        curr_a4x_object: SchedulableWork | None = None,
        site_map: dict[str, Site] | None = None,
        file_map: dict[str, File] | None = None,
    ) -> SchedulableWork:
        # This method updates an existing A4X SchedulableWork object.
        # So, if 'curr_a4x_object' is not provided, error out
        if curr_a4x_object is None or not isinstance(curr_a4x_object, SchedulableWork):
            raise RuntimeError(
                "INTERNAL ERROR: Cannot call _to_a4x for SchedulableConfig without a valid SchedulableWork object"
            )
        # Call parent class's version of this function to handle annotations
        curr_a4x_object = super()._to_a4x(curr_a4x_object=curr_a4x_object, site_map=site_map, file_map=file_map)
        # Most of the fields are copied as-is into the A4X object
        curr_a4x_object.duration = self.duration
        curr_a4x_object.queue = self.queue
        curr_a4x_object.project = self.project
        curr_a4x_object.environment = self.environment.copy()
        # Copy the 'cwd' field as-is if it's a string. Otherwise, recursively call this method to get an A4X Directory
        if isinstance(self.cwd, SiteAndDirectoryConfig):
            self.cwd = self.cwd._to_a4x(site_map=site_map)  # noqa: SLF001
        else:
            curr_a4x_object.cwd = self.cwd
        # For stdin, stdout, and stderr, check if the field's value is the name/identifier for a file in the config.
        # If so, get the corresponding A4X File object from 'file_map'. Otherwise, copy the value as-is
        if self.stdin is not None and self.stdin in file_map:
            curr_a4x_object.stdin = file_map[self.stdin]
        else:
            curr_a4x_object.stdin = self.stdin
        if self.stdout is not None and self.stdout in file_map:
            curr_a4x_object.stdout = file_map[self.stdout]
        else:
            curr_a4x_object.stdout = self.stdout
        if self.stderr is not None and self.stderr in file_map:
            curr_a4x_object.stderr = file_map[self.stderr]
        else:
            curr_a4x_object.stderr = self.stderr
        # Handle 'resources' and 'resources_w_slot' by calling their respective versions of '_to_a4x'
        if self.resources is not None:
            curr_a4x_object = self.resources._to_a4x(  # noqa: SLF001
                curr_a4x_object=curr_a4x_object, site_map=site_map, file_map=file_map
            )
        if self.resources_w_slot is not None:
            curr_a4x_object.jobspec_settings.resources = self.resources_w_slot._to_a4x(  # noqa: SLF001
                curr_a4x_object=curr_a4x_object, site_map=site_map, file_map=file_map
            )

    @classmethod
    def _from_a4x(
        cls,
        curr_a4x_object: SchedulableWork,
        curr_config_object: SchedulableConfig | None = None,
    ) -> Self:
        # This method updates an existing config component.
        # So, if 'curr_config_object' is not provided, error out
        if curr_config_object is None or not isinstance(curr_config_object, SchedulableConfig):
            raise RuntimeError(
                "INTERNAL ERROR: the curr_config_object parameter to FileConfig._to_a4x must be provided and a SchedulableConfig object"  # noqa: E501
            )
        # If the A4X object's 'cwd' field is a Directory object, invoke SiteAndDirectory._from_a4x to get
        # the correct value of the 'cwd' field. Otherwise, the value is copied as-is from the A4X object
        cwd = curr_a4x_object.cwd
        if cwd is not None and isinstance(cwd, Directory):
            cwd = SiteAndDirectoryConfig._from_a4x(  # noqa: SLF001
                curr_a4x_object=curr_a4x_object.cwd,
            )
        # Copy most fields as-is
        curr_config_object.duration = curr_a4x_object.duration
        curr_config_object.queue = curr_a4x_object.queue
        curr_config_object.project = curr_a4x_object.project
        curr_config_object.cwd = cwd
        curr_config_object.environment = curr_a4x_object.environment.copy()
        # For stdin, stdout, stderr, set the field to either the as-is value from the A4X object
        # or File.path_attr if the A4X object's field has a File value
        curr_config_object.stdin = (
            str(curr_a4x_object.stdin.path_attr) if isinstance(curr_a4x_object.stdin, File) else curr_a4x_object.stdin
        )
        curr_config_object.stdout = (
            str(curr_a4x_object.stdout.path_attr)
            if isinstance(curr_a4x_object.stdout, File)
            else curr_a4x_object.stdout
        )
        curr_config_object.stderr = (
            str(curr_a4x_object.stderr.path_attr)
            if isinstance(curr_a4x_object.stderr, File)
            else curr_a4x_object.stderr
        )
        # We always use 'resources_w_slot' when exporting a workflow to a config.
        # So, invoke ResourcesWSlotConfig._from_a4x to build that field
        resources = curr_a4x_object.get_resources()
        if resources is not None:
            curr_config_object.resources_w_slot = ResourcesWSlotConfig._from_a4x(resources)  # noqa: SLF001
        # Call parent class's version of this function to handle annotations
        super()._from_a4x(curr_a4x_object=curr_a4x_object, curr_config_object=curr_config_object)
        return curr_config_object


class DirectoryConfig(AnnotationConfig):
    """A specification for a directory."""

    name: str
    """Identifier for the directory"""
    path: Path
    """Absolute file path for the directory"""
    storage_type: StorageType = StorageType.UNKNOWN
    """The type of storage (e.g., local, shared) for the directory"""
    persistency: PersistencyType = PersistencyType.UNKNOWN
    """The persistency of storage (e.g., persistent, scratch) for the directory"""

    def _to_a4x(
        self,
        curr_a4x_object: Site | None = None,
        site_map: dict[str, Site] | None = None,
        file_map: dict[str, File] | None = None,
    ) -> Directory:
        # This method both creates a new A4X Directory object and modifies an existing
        # A4X Site object. So, we first confirm that the Site object is passed to 'curr_a4x_object'
        if curr_a4x_object is None or not isinstance(curr_a4x_object, Site):
            raise RuntimeError("INTERNAL ERROR: Cannot call _to_a4x for DirectoryConfig without a valid Site object")
        # Create the Directory object using the Site
        new_dir = curr_a4x_object.add_directory(
            identifier=self.name,
            path=self.path,
            storage_type=self.storage_type,
            persistency=self.persistency,
        )
        # Call the parent class version of this function to handle annotations
        super()._to_a4x(curr_a4x_object=new_dir)
        return new_dir

    @classmethod
    def _from_a4x(
        cls,
        curr_a4x_object: Directory,
        curr_config_object: A4XOrchestrationConfigBase | None = None,
    ) -> Self:
        # Create a DirectoryConfig from an A4X Directory object
        new_dir_config = cls(
            name=curr_a4x_object.identifier,
            path=curr_a4x_object.path,
            storage_type=curr_a4x_object.storage_type,
            persistency=curr_a4x_object.persistency,
        )
        # Call the parent class version of this function to handle annotations
        return super()._from_a4x(curr_a4x_object=curr_a4x_object, curr_config_object=new_dir_config)


class SiteConfig(AnnotationConfig):
    """A specification for execution and data storage sites."""

    name: str
    """Identifier for the site"""
    scheduler: Scheduler = Scheduler.UNKNOWN
    """The scheduler for the site"""
    directories: list[DirectoryConfig] = []
    """Directories that exist on the site"""

    def _to_a4x(
        self,
        curr_a4x_object: Any | None = None,
        site_map: dict[str, Site] | None = None,
        file_map: dict[str, File] | None = None,
    ) -> Site:
        # Create the Site object
        site = Site(self.name, self.scheduler)
        # Call the parent class version of this function to handle annotations
        super()._to_a4x(curr_a4x_object=site)
        # Loop over the 'directories' field to convert DirectoryConfig objects
        # into A4X Directory objects
        for dir_config in self.directories:
            dir_config._to_a4x(curr_a4x_object=site)  # noqa: SLF001
        return site

    @classmethod
    def _from_a4x(
        cls,
        curr_a4x_object: Site,
        curr_config_object: list[DirectoryConfig] | None = None,
    ) -> Self:
        # By this point, directories have already been converted into a list of DirectoryConfig
        # objects. So, we validate that this list has been passed in
        if (
            curr_config_object is None
            or not isinstance(curr_config_object, list)
            or any(not isinstance(dir_config, DirectoryConfig) for dir_config in curr_config_object)
        ):
            raise RuntimeError(
                "INTERNAL ERROR: the 'curr_config_object' parameter to SiteConfig._from_a4x must be a list of DirectoryConfig objects"  # noqa: E501
            )
        # Create the SiteConfig
        new_site_config = cls(
            name=curr_a4x_object.name,
            scheduler=curr_a4x_object.scheduler,
            directories=curr_config_object,
        )
        # Call the parent class version of this method to handle annotations
        return super()._from_a4x(curr_a4x_object=curr_a4x_object, curr_config_object=new_site_config)


class FileConfig(A4XOrchestrationConfigBase):
    """A specification for a file in a workflow."""

    name: str | Path
    """The name (or relative path) of the file"""
    site: str | None = None
    """The site on which the file resides"""
    directory: str | None = None
    """The directory in which the file resides"""

    @model_validator(mode="after")
    def _validate_both_or_neither_site_and_directory(self) -> Self:
        # Ensure that either both the 'site' and 'directory' fields are provided or that neither are provided
        if not (self.site is None and self.directory is None) and not (
            self.site is not None and self.directory is not None
        ):
            raise ValueError(
                "The 'site' and 'directory' fields for a file must either both be provided or both be null"
            )
        return self

    def _to_a4x(
        self,
        curr_a4x_object: Any | None = None,
        site_map: dict[str, Site] | None = None,
        file_map: dict[str, File] | None = None,
    ) -> File:
        directory = None
        # Use the 'site' and 'directory' fields to get the A4X Directory object, if those fields are provided
        if self.site is not None and self.directory is not None:
            # The 'site_map' argument is required if the 'site' and 'directory' fields are provided
            if site_map is None:
                raise RuntimeError(
                    "INTERNAL ERROR: the site_map parameter to FileConfig._to_a4x must be provided when the site and directory are specified"  # noqa: E501
                )
            try:
                site = site_map[self.site]
            except KeyError as e:
                raise KeyError(f"Invalid site name: '{self.site}'") from e
            try:
                directory = site[self.directory]
            except KeyError as e:
                raise KeyError(f"Invalid directory identifier: '{self.directory}'") from e
        # Create the A4X File object
        return File(self.name, directory=directory)

    @classmethod
    def _from_a4x(
        cls,
        curr_a4x_object: File,
        curr_config_object: A4XOrchestrationConfigBase | None = None,
    ) -> Self:
        file_site = None
        file_dir = None
        # If the A4X File object passed as curr_a4x_object has a 'directory' attribute,
        # use that Directory object to set the 'site' and 'directory' fields of the config component
        if curr_a4x_object.directory is not None:
            file_dir = curr_a4x_object.directory.identifier
            file_site = curr_a4x_object.site.name
        # Create the config component
        return cls(
            name=str(curr_a4x_object.path_attr),
            site=file_site,
            directory=file_dir,
        )


class CommandArgObjectConfig(A4XOrchestrationConfigBase):
    """A specification of an A4X File or Directory object in a command's arguments."""

    type: Literal["file", "directory"]
    """The type of A4X object being represented"""
    val: str | SiteAndDirectoryConfig
    """The specification of the A4X object being represented"""

    def _validate_sites_directories_and_files(
        self,
        valid_sites: set[str],
        valid_directories: set[str],
        valid_files: set[str],
    ) -> None:
        # If the argument represents an A4X File object, ensure the 'val' field is a string
        # and ensure that string is in the list of valid files for the workflow
        if self.type == "file":
            if not isinstance(self.val, str):
                raise ValueError("When 'type' is 'file', the 'val' field must be the string identifier of a file")
            if self.val not in valid_files:
                raise ValueError(f"Unrecognized file ({self.val}) in 'val' field of a file/directory command argument")
        # If the argument represents an A4X Directory object, ensure the 'val' field is a 'SiteAndDirectoryConfig'
        # and call that class's '_validate_sites_directories_and_files' method to validate the contents
        else:
            if not isinstance(self.val, SiteAndDirectoryConfig):
                raise ValueError(
                    "When 'type' is 'directory', the 'val' field must be a combination of 'site' and 'directory'"
                )
            self.val._validate_sites_directories_and_files(valid_sites, valid_directories, valid_files)  # noqa: SLF001

    def _to_a4x(
        self,
        curr_a4x_object: Any | None = None,
        site_map: dict[str, Site] | None = None,
        file_map: dict[str, File] | None = None,
    ) -> File | Directory:
        # Both file_map and site_map are required for this method
        if file_map is None or site_map is None:
            raise RuntimeError(
                "INTERNAL ERROR: a file_map and site_map must be provided to CommandArgObjectConfig._to_a4x"
            )
        # If the argument represents an A4X File object, lookup the File object from file_map
        if self.type == "file":
            # NOTE: we should never hit this error. This is here for safety
            if isinstance(self.val, SiteAndDirectoryConfig):
                raise TypeError(
                    "When setting 'type' to 'file' in a Command argument, the 'val' field must be a single string"
                )
            try:
                return file_map[self.val]
            except KeyError:
                # If the file is not in file_map, we just return the file identifier as-is
                return self.val
        # If the argument represents a directory, call SiteAndDirectoryConfig._to_a4x to validate its contents
        elif isinstance(self.val, SiteAndDirectoryConfig):
            return self.val._to_a4x(site_map=site_map)  # noqa: SLF001
        # If something somehow fails, we return the value as-is
        else:
            return self.val

    @classmethod
    def _from_a4x(
        cls,
        curr_a4x_object: Any,
        curr_config_object: A4XOrchestrationConfigBase | None = None,
    ) -> Self | Any:
        # If the argument represents an A4X File object, get the value from the object's 'path_attr' field
        if isinstance(curr_a4x_object, File):
            return cls(type="file", val=str(curr_a4x_object.path_attr))
        # If the argument represents an A4X Directory object, delegate the value creation to SiteAndDirectoryConfig
        if isinstance(curr_a4x_object, Directory):
            return cls(
                type="directory",
                val=SiteAndDirectoryConfig._from_a4x(curr_a4x_object=curr_a4x_object),  # noqa: SLF001
            )
        # If the argument represents neither an A4X File nor Directory object, then we assume that the value
        # should NOT be converted to a CommandArgObjectConfig, and we return the input as-is.
        return curr_a4x_object


class CommandConfig(SchedulableConfig):
    """A specification of a command in a workflow task."""

    cmd: str
    """The executable to run in the command or a block of shell code representing the command"""
    args: list[str | int | float | CommandArgObjectConfig] = []
    """The arguments to the executable. Should only be specified if the 'cmd' field is just a executable path"""
    description: str = ""
    """A description for the command"""

    def _validate_sites_directories_and_files(
        self,
        valid_sites: set[str],
        valid_directories: set[str],
        valid_files: set[str],
    ) -> None:
        # Recursively call this method on any argument of type CommandArgObjectConfig for validation
        for curr_arg in self.args:
            if isinstance(curr_arg, CommandArgObjectConfig):
                curr_arg._validate_sites_directories_and_files(valid_sites, valid_directories, valid_files)  # noqa: SLF001

    def _to_a4x(
        self,
        curr_a4x_object: Any | None = None,
        site_map: dict[str, Site] | None = None,
        file_map: dict[str, File] | None = None,
    ) -> Command:
        # This method requires both file_map and site_map
        if file_map is None or site_map is None:
            raise RuntimeError("INTERNAL ERROR: a file_map and site_map must be provided to CommandConfig._to_a4x")
        # Get the command from the config
        # If the command matches a key in file_map, assume that the File object should take priority
        local_cmd = self.cmd
        if self.cmd in file_map:
            local_cmd = file_map[self.cmd]
        # Generate the list of args using CommandArgObjectConfig
        local_args = []
        for curr_arg in self.args:
            if isinstance(curr_arg, CommandArgObjectConfig):
                local_args.append(curr_arg._to_a4x(site_map=site_map, file_map=file_map))  # noqa: SLF001
            else:
                local_args.append(curr_arg)
        # Create the A4X Command object and populate the description field
        a4x_cmd = Command(local_cmd, *local_args)
        a4x_cmd.description = self.description
        # Call the parent class version of this method to handle scheduling parameters and annotations
        super()._to_a4x(curr_a4x_object=a4x_cmd, site_map=site_map, file_map=file_map)
        return a4x_cmd

    @classmethod
    def _from_a4x(
        cls,
        curr_a4x_object: Command,
        curr_config_object: A4XOrchestrationConfigBase | None = None,
    ) -> Self:
        new_command_config = None
        # How we convert from an A4X Command object to a config component depends on whether we
        # are dealing with a block command or not.
        #
        # When dealing with a block command, we just copy fields (in particular the 'command_or_exe' field)
        # as-is, and we then call AnnotationConfig._from_a4x to handle annotations.
        #
        # When dealing with a non-block command, we use CommandArgObjectConfig to handle args. Then,
        # we create the Command config by copying the remaining information as-is. Finally, we call
        # the parent class's version of this method to handle scheduling parameters
        if curr_a4x_object.is_block_command:
            new_command_config = cls(
                cmd=curr_a4x_object.command_or_exe,
                description=curr_a4x_object.description if curr_a4x_object.description is not None else "",
            )
            new_command_config = AnnotationConfig._from_a4x(  # noqa: SLF001
                curr_a4x_object=curr_a4x_object, curr_config_object=new_command_config
            )
        else:
            command_args = [CommandArgObjectConfig._from_a4x(arg) for arg in curr_a4x_object.args]  # noqa: SLF001
            new_command_config = cls(
                cmd=curr_a4x_object.command_or_exe
                if not isinstance(curr_a4x_object.command_or_exe, File)
                else str(curr_a4x_object.command_or_exe.path_attr),
                args=command_args,
                description=curr_a4x_object.description if curr_a4x_object.description is not None else "",
            )
            new_command_config = super()._from_a4x(
                curr_a4x_object=curr_a4x_object, curr_config_object=new_command_config
            )
        return new_command_config


class TaskConfig(SchedulableConfig):
    """A specification for a task in the workflow."""

    name: str
    """The name of the task"""
    commands: list[CommandConfig]
    """The commands defining the task. Can be thought of as the parts of a shell script defining what the task does"""
    description: str = ""
    """A description for the task"""
    dependencies: list[str] = []
    """The names of other tasks that must complete before this one"""
    site: str | None = None
    """The site on which this task should execute"""
    inputs: list[str] = []
    """The names of input files to this task"""
    outputs: list[str] = []
    """The names of output files to this task"""
    output_extra_kwargs: dict[str, Any] = {}
    """WMS-specific extra key-value pairs to provide when specifying outputs"""
    input_extra_kwargs: dict[str, Any] = {}
    """WMS-specific extra key-value pairs to provide when specifying inputs"""

    def _validate_sites_directories_and_files(
        self,
        valid_sites: set[str],
        valid_directories: set[str],
        valid_files: set[str],
    ) -> None:
        # If the 'site' field is provided, ensure it refers to a valid site in the workflow
        if self.site is not None:
            if self.site not in valid_sites:
                raise ValueError(f"Unrecognized site ({self.site}) in task {self.name}")
        # If the 'inputs' field is provided, ensure each element refers to a valid file in the workflow
        for f in self.inputs:
            if f not in valid_files:
                raise ValueError(f"Unrecognized input file ({f})) in task {self.name}")
        # If the 'outputs' field is provided, ensure each element refers to a valid file in the workflow
        for f in self.outputs:
            if f not in valid_files:
                raise ValueError(f"Unrecognized output file ({f})) in task {self.name}")
        # Validate each command in the task
        for cmd in self.commands:
            cmd._validate_sites_directories_and_files(valid_sites, valid_directories, valid_files)  # noqa: SLF001

    def _to_a4x(
        self,
        curr_a4x_object: Any | None = None,
        site_map: dict[str, Site] | None = None,
        file_map: dict[str, File] | None = None,
    ) -> tuple[Task, list[str]]:
        # Both file_map and site_map are required in this method
        if file_map is None or site_map is None:
            raise RuntimeError("INTERNAL ERROR: a file_map and site_map must be provided to TaskConfig._to_a4x")
        # Create the A4X task object
        task = Task(self.name, self.description)
        # Create the commands by calling CommandConfig._to_a4x
        a4x_cmds = [cmd._to_a4x(site_map=site_map, file_map=file_map) for cmd in self.commands]  # noqa: SLF001
        # Add the commands to the task
        task.extend(a4x_cmds)
        # If the 'site' field is provided, set the site on the A4X Task
        if self.site is not None:
            try:
                task.set_site(site_map[self.site])
            except KeyError as e:
                raise KeyError(f"Unknown site: {self.site}") from e
        # Add inputs and outputs with extra kwargs, if provided
        if len(self.inputs) > 0:
            task.add_inputs(*[file_map[inf] for inf in self.inputs], **self.input_extra_kwargs)
        if len(self.outputs) > 0:
            task.add_outputs(*[file_map[inf] for inf in self.outputs], **self.output_extra_kwargs)
        # Call the parent class version of this method to handle scheduling arguments and annotations
        super()._to_a4x(curr_a4x_object=task, site_map=site_map, file_map=file_map)
        # Here's we return not only the A4X Task object, but also a list of dependencies so that the WorkflowConfig
        # class can rebuild the DAG
        return task, self.dependencies

    @classmethod
    def _from_a4x(
        cls,
        curr_a4x_object: tuple[Task, list[str]],
        curr_config_object: A4XOrchestrationConfigBase | None = None,
    ) -> Self:
        # This method takes both the A4X Task and a list of the names of its predecessors in the workflow DAG
        # So, the first thing we do is unpack them
        task_obj, parent_tasks = curr_a4x_object
        # Call CommandConfig._from_a4x to build the config for the Task's commands
        command_configs = [CommandConfig._from_a4x(curr_a4x_object=cmd) for cmd in task_obj.commands]  # noqa: SLF001
        # Create the config component for the Task
        task_config = cls(
            name=task_obj.task_name,
            commands=command_configs,
            description=task_obj.description,
            dependencies=parent_tasks,
            site=task_obj.site.name if task_obj.site is not None else None,
            inputs=[str(f.path_attr) for f in task_obj.inputs],
            outputs=[str(f.path_attr) for f in task_obj.outputs],
            output_extra_kwargs=task_obj.add_output_extra_kwargs,
            input_extra_kwargs=task_obj.add_input_extra_kwargs,
        )
        # Call the parent class version of this function to handle scheduling arguments and annotations
        super()._from_a4x(curr_a4x_object=task_obj, curr_config_object=task_config)
        return task_config


class WorkflowConfig(AnnotationConfig):
    """A specification for an entire A4X workflow."""

    name: str
    """The name of the workflow"""
    tasks: list[TaskConfig]
    """The tasks comprising the workflow"""
    description: str = ""
    """A description of the workflow"""
    sites: list[SiteConfig] = []
    """The sites involved in the workflow"""
    files: list[FileConfig] = []
    """The files involved in the workflow"""
    environment: dict[str, Any] = {}
    """Workflow-global environment modifications"""

    @model_validator(mode="after")
    def _validate_sites_directories_and_files(self) -> Self:
        """Validate that the workflow specification has valid sites, directories, and files.

        This method runs automatically through Pydantic and launches the validation process.
        """
        valid_sites = set()
        valid_directories = set()
        # Get the set of valid files in the workflow
        valid_files = {f.name for f in self.files}
        # Loop over all sites and populate sets of valid sites and directories
        for site_config in self.sites:
            valid_sites.add(site_config.name)
            for directory in site_config.directories:
                valid_directories.add(directory.name)
        # This loop drives the validation process by calling TaskConfig._validate_sites_directories_and_files
        for task in self.tasks:
            task._validate_sites_directories_and_files(valid_sites, valid_directories, valid_files)  # noqa: SLF001
        return self

    def _to_a4x(
        self,
        curr_a4x_object: Any | None = None,
        site_map: dict[str, Site] | None = None,
        file_map: dict[str, File] | None = None,
    ) -> Workflow:
        # Create the A4X Workflow object
        wflow = Workflow(self.name, self.description)
        # Copy the workflow environment as-is
        wflow.environment = self.environment.copy()
        # Loop over 'sites' field and build the A4X Site and Directory objects
        wflow_site_map = {}
        for site_config in self.sites:
            a4x_site = site_config._to_a4x()  # noqa: SLF001
            wflow_site_map[a4x_site.name] = a4x_site
        # Loop over 'files' field and build the A4X File objects
        wflow_file_map = {}
        for file_config in self.files:
            a4x_file = file_config._to_a4x(site_map=wflow_site_map)  # noqa: SLF001
            wflow_file_map[str(a4x_file.path_attr)] = a4x_file
        # Loop over the 'tasks' field and build the A4X Task objects using sets of Site and File objects made previously
        task_name_deps = {}
        task_map = {}
        for task_config in self.tasks:
            a4x_task, task_deps = task_config._to_a4x(site_map=wflow_site_map, file_map=wflow_file_map)  # noqa: SLF001
            task_map[a4x_task.task_name] = a4x_task
            task_name_deps[a4x_task] = task_deps.copy()
        # Create edges in the DAG to fully build the workflow
        task_deps = {task: [task_map[d] for d in deps] for task, deps in task_name_deps.items()}
        for child_task, parents in task_deps.items():
            for parent_task in parents:
                wflow.add_edge(parent_task, child_task)
        # Call the parent class version of this method to handle annotations
        super()._to_a4x(curr_a4x_object=wflow)
        return wflow

    @classmethod
    def _from_a4x(
        cls,
        curr_a4x_object: Workflow,
        curr_config_object: A4XOrchestrationConfigBase | None = None,
    ) -> Self:
        task_configs = []
        site_configs = []
        file_configs = []
        # Get all tasks, sites, and files in the workflow
        tasks = curr_a4x_object.get_tasks_in_topological_order()
        sites, files = _search_for_sites_and_files_from_tasks(curr_a4x_object.sites, tasks)
        # For each Site, build a new SiteConfig and a DirectoryConfig for each directory in the A4X Site
        for curr_site in sites:
            curr_directories = []
            for directory in curr_site.values():
                curr_directories.append(
                    DirectoryConfig._from_a4x(  # noqa: SLF001
                        curr_a4x_object=directory,
                    )
                )
            site_configs.append(SiteConfig._from_a4x(curr_a4x_object=curr_site, curr_config_object=curr_directories))  # noqa: SLF001
        # Create a new FileConfig for each A4X File object in the workflow
        for file in files:
            file_configs.append(
                FileConfig._from_a4x(  # noqa: SLF001
                    curr_a4x_object=file,
                )
            )
        # Get a record of the parents of each task
        parent_tasks = [list(curr_a4x_object.graph.predecessors(t.task_name)) for t in tasks]
        # Create the TaskConfig objects from each Task and its parents
        task_configs = [TaskConfig._from_a4x((t, parents)) for t, parents in zip(tasks, parent_tasks)]  # noqa: SLF001
        # Create the WorkflowConfig
        wflow_config = cls(
            name=curr_a4x_object.name,
            tasks=task_configs,
            description=curr_a4x_object.description,
            sites=site_configs,
            files=file_configs,
            environment=curr_a4x_object.environment,
        )
        # Call the parent class version of this method to handle annotations
        super()._from_a4x(curr_a4x_object=curr_a4x_object, curr_config_object=wflow_config)
        return wflow_config


class A4XOrchestrationConfig(BaseModel):
    """Top-level model representing the entire config."""

    workflow: WorkflowConfig
    """The workflow"""


def _search_for_sites_and_files_from_tasks(sites: set[Site], tasks: list[Task]) -> tuple[set[Site], set[File]]:  # noqa: PLR0912
    """Collect all Site and File objects from the Tasks in the workflow.

    :param sites: the initial set of Site objects to start with, pulled from Workflow.sites
    :param tasks: the set of Task objects to consider in building the sets of Site and File objects
    :return: a tuple of the set of Sites and the set of Files
    """
    files = set()
    # First find files and sites from the tasks
    for task in tasks:
        # If the task has a site, add it to the set of sites
        if task.site is not None:
            sites.add(task.site)
        # Add all File objects from task inputs to the file set
        for input_file in task.inputs:
            files.add(input_file)
            # If the File is not virtual (i.e., it has a directory and site defined), add the Site of the File
            # to the site set
            if input_file.directory is not None and input_file.directory.site is not None:
                sites.add(input_file.directory.site)
        # Add all File objects from task outputs to the file set
        for output_file in task.outputs:
            files.add(output_file)
            # If the File is not virtual (i.e., it has a directory and site defined), add the Site of the File
            # to the site set
            if output_file.directory is not None and output_file.directory.site is not None:
                sites.add(output_file.directory.site)
        # The Command objects can also store File and Directory objects, so we loop
        # over the commands to collect File and Site objects
        for cmd in task.commands:
            # If the 'command_or_exe' attribute is a File object, add it to the file set.
            # Also, if the File object is non-virtual, add the Site of the File to the site set.
            if isinstance(cmd.command_or_exe, File):
                files.add(cmd.command_or_exe)
                if cmd.command_or_exe.directory is not None and cmd.command_or_exe.directory.site is not None:
                    sites.add(cmd.command_or_exe.directory.site)
            # Loop over the args of the Command to find File or Directory objects
            for arg in cmd.args:
                # If the arg is a File object, add it to the file set.
                # Also, if the File object is non-virtual, add the Site of the File to the site set.
                if isinstance(arg, File):
                    files.add(arg)
                    if arg.directory is not None and arg.directory.site is not None:
                        sites.add(arg.directory.site)
                # If the arg is a Directory object, add its site to the site set.
                elif isinstance(arg, Directory):
                    if arg.site is not None:
                        sites.add(arg.site)
    # Finally, loop over the site set and the directories for each site, and add the File objects associated
    # with each directory to the file set
    for site in sites:
        for directory in site.values():
            files = files | directory.files
    return sites, files


def write_config(wflow: Workflow, fname: StrCompatPathLike) -> None:
    """Write the workflow to a YAML file.

    :param wflow: the A4X Workflow object to write
    :param fname: the path to the YAML file to be generated
    """
    yaml = YAML()
    fname_path = Path(fname).expanduser().resolve()
    # Call _from_a4x to create the config representation of the workflow
    wflow_config = WorkflowConfig._from_a4x(curr_a4x_object=wflow)  # noqa: SLF001
    # Add the workflow config to A4XOrchestrationConfig for forrmatting and then dump to a JSON-compatible format
    final_config = A4XOrchestrationConfig(workflow=wflow_config).model_dump(mode="json", exclude_defaults=True)
    # Write the config to the YAML file
    with fname_path.open("w") as f:
        yaml.dump(final_config, f)


def read_config(fname: StrCompatPathLike) -> Workflow:
    """Read a workflow from a YAML file.

    :param fname: the path to the YAML file to read
    :return: the A4X Workflow object built from the YAML file
    """
    # Ensure the config path is something that can be converted to pathlib.Path
    if not isinstance(fname, StrCompatPathLikeForIsInstance):
        raise TypeError("The 'fname' argument must be a string or a type compliant with os.PathLike")
    config = None
    yaml = YAML()
    # Validate that the config actually exists
    fname_path = Path(fname).expanduser().resolve()
    if not fname_path.is_file():
        raise FileNotFoundError(f"Cannot find YAML config: {fname}")
    # Read the config with ruamel.yaml
    with fname_path.open("r") as f:
        config = yaml.load(f)
    if config is None:
        raise RuntimeError(f"Cannot load YAML config from {fname}")
    # Use the Pydantic model to validate the config's contents
    try:
        model = A4XOrchestrationConfig(**config)
    except ValidationError as e:
        raise ValidationError.from_exception_data(
            title=f"Failed to validate the contents of YAML config file: {fname}", line_errors=e.errors()
        ) from e
    # Create the A4X Workflow object and return
    return model.workflow._to_a4x()  # noqa: SLF001
