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

from __future__ import annotations

import pathlib
from typing import TYPE_CHECKING, Any, Iterator, MutableMapping

from a4x.orchestration.annotations import AnnotationType
from a4x.orchestration.strenum import StrEnum

if TYPE_CHECKING:
    from a4x.orchestration.utils import StrCompatPathLike


class StorageType(StrEnum):
    """
    An enum representing a type of storage.
    """

    UNKNOWN = "unknown"
    LOCAL = "local"
    SHARED = "shared"


class PersistencyType(StrEnum):
    """
    An enum representing the persistency of storage.
    """

    UNKNOWN = "unknown"
    SCRATCH = "scratch"
    PERSISTENT = "persistent"


class Scheduler(StrEnum):
    """
    An enum representing the resource and job management system (RJMS) of a site.
    """

    UNKNOWN = "unknown"
    CONDOR = "condor"
    LSF = "lsf"
    PBS = "pbs"
    SGE = "sge"
    SLURM = "slurm"
    FLUX = "flux"


class Site(MutableMapping, AnnotationType):
    """A class representing a site (e.g., a system) where tasks can be run and/or files can be stored."""

    def __init__(
        self,
        site_or_host_name: str,
        scheduler_type: Scheduler | str = Scheduler.UNKNOWN,
    ) -> None:
        """Initialize a Site.

        :param site_or_host_name: the identifier/name of the site. Oftentimes the hostname of a system.
        :param scheduler_type: the type of scheduler for the site.
        """
        AnnotationType.__init__(self)
        if not isinstance(site_or_host_name, str):
            raise TypeError("The 'site_or_host_name' argument must be of type 'str'")
        if site_or_host_name == "":
            raise ValueError("The 'site_or_host_name' argument cannot be an empty string")
        self.name = site_or_host_name
        if isinstance(scheduler_type, Scheduler):
            self.scheduler_type = scheduler_type
        elif isinstance(scheduler_type, str):
            self.scheduler_type = Scheduler(scheduler_type)
        else:
            raise TypeError("The 'scheduler_type' argument must be of type 'str' or 'a4x.orchestration.Scheduler'")
        self.directories_attr = {}

    @property
    def scheduler(self) -> Scheduler:
        """The scheduler (e.g., Slurm or Flux) for the site."""
        return self.scheduler_type

    @scheduler.setter
    def scheduler(self, scheduler_type: Scheduler | str) -> None:
        if isinstance(scheduler_type, Scheduler):
            self.scheduler_type = scheduler_type
        elif isinstance(scheduler_type, str):
            self.scheduler_type = Scheduler(scheduler_type)
        else:
            raise TypeError("The 'scheduler_type' argument must be of type 'str' or 'a4x.orchestration.Scheduler'")

    def add_directory(
        self,
        identifier: str,
        path: StrCompatPathLike,
        storage_type: StorageType = StorageType.UNKNOWN,
        persistency: PersistencyType = PersistencyType.UNKNOWN,
    ) -> Directory:
        """Define a directory on the site, and register it with the Site object.

        :param identifier: the identifier for the directory. Identifiers can be used with the Site class's dict-like API to lookup directories
        :param path: the absolute path of the directory on the site
        :param storage_type: the type of storage (e.g., shared or local) on which the directory is located
        :param persistency: the persistency of storage (e.g., scratch or persistent) on which the directory is located
        :return: the A4X Directory object representing the directory
        """  # noqa: E501
        if identifier in self.directories_attr:
            curr_dir = self.directories_attr[identifier]
            if (
                curr_dir.path != pathlib.Path(path)
                or curr_dir.storage_type != storage_type
                or curr_dir.persistency != persistency
            ):
                raise KeyError(
                    f"Cannot add a directory under identifier '{identifier}' because that identifier already exists. If you want to change an existing directory entry, use the Site class's dict-like interface"  # noqa: E501
                )
            return curr_dir
        # The Directory constructor self-registers the Directory in this Site
        return Directory(
            identifier=identifier,
            path=path,
            site=self,
            storage_type=storage_type,
            persistency=persistency,
        )

    def __getitem__(self, key: str) -> Directory:
        return self.directories_attr[key]

    def __setitem__(self, key: str, val: Directory) -> None:
        if not isinstance(val, Directory):
            raise TypeError("Values of 'Site' must be of type 'Directory'")
        self.directories_attr[key] = val
        if val.identifier != key and val.identifier in self.directories_attr:
            del self.directories_attr[val.identifier]
        val.identifier = key

    def __delitem__(self, key: str) -> None:
        del self.directories_attr[key]

    def __len__(self) -> int:
        return len(self.directories_attr)

    def __iter__(self) -> Iterator:
        return iter(self.directories_attr)

    def resolve_site_files(self) -> None:
        """Resolve all files associated with the site."""
        for directory in self.directories_attr.values():
            directory._resolve_files()  # noqa: SLF001

    @property
    def directories(self) -> set[Directory]:
        """The set of all known directories on the site."""
        return set(self.directories_attr.values())

    def __eq__(self, other: Any) -> bool:
        return isinstance(other, Site) and self.name == other.name and self.scheduler_type == other.scheduler_type

    def __hash__(self) -> int:
        return hash((self.name, self.scheduler_type))

    def __repr__(self) -> str:
        return f"Site(name={self.name},sched={self.scheduler_type})"


__DEFAULT_SITE_FOR_TASKS: Site | None = None


def set_default_site(site: Site) -> None:
    """Define a default site that tasks will run on if no task is manually specified by the user.

    :param site: the default site
    """
    global __DEFAULT_SITE_FOR_TASKS
    __DEFAULT_SITE_FOR_TASKS = site


def get_default_site() -> Site:
    """Get the current default site that tasks will run on if no task is manually specified by the user.

    :return: the default site
    """
    return __DEFAULT_SITE_FOR_TASKS


class Directory(AnnotationType):
    """A class representing a directory on a given site."""

    def __init__(
        self,
        identifier: str,
        path: StrCompatPathLike,
        site: Site,
        storage_type: StorageType | str = StorageType.UNKNOWN,
        persistency: PersistencyType | str = PersistencyType.UNKNOWN,
    ) -> None:
        """Initialize the Directory.

        :param identifier: the identifier for the directory. Identifiers can be used with the Site class's dict-like API to lookup directories
        :param path: the absolute path of the directory on the site
        :param site: the site on which the directory is located
        :param storage_type: the type of storage (e.g., shared or local) on which the directory is located
        :param persistency: the persistency of storage (e.g., scratch or persistent) on which the directory is located
        """  # noqa: E501
        super().__init__()
        if not isinstance(storage_type, (StorageType, str)):
            raise TypeError(
                "The 'storage_type' parameter must be of type 'StorageType' or a corresponding string value"
            )
        if not isinstance(persistency, (PersistencyType, str)):
            raise TypeError(
                "The 'persistency' parameter must be of type 'PersistencyType' or a corresponding string value"
            )
        self.identifier = identifier
        self.path = pathlib.Path(path)
        if not self.path.is_absolute():
            raise ValueError(
                "The 'path' argument must represent an absolute path (e.g., a string starting with '/' for Linux)"
            )
        self.storage_type_attr = StorageType(storage_type)
        self.persistency_attr = PersistencyType(persistency)
        self.site = site
        self.files = set()
        if self.identifier not in self.site:
            self.site[self.identifier] = self

    @property
    def storage_type(self) -> StorageType:
        """The storage type of the directory in which the file resides."""
        return self.storage_type_attr

    @property
    def persistency(self) -> PersistencyType:
        """The persistency type of the directory in which the file resides."""
        return self.persistency_attr

    def get_site(self) -> Site:
        """Return the site on which the directory is located."""
        return self.site

    def __truediv__(self, other: StrCompatPathLike) -> Directory:
        identifier = f"{self.identifier}/{other!s}"
        if identifier in self.site:
            return self.site[identifier]
        return Directory(
            identifier=identifier,
            path=self.path / other,
            site=self.site,
            storage_type=self.storage_type_attr,
            persistency=self.persistency_attr,
        )

    def _add_file(self, f) -> None:  # noqa: ANN001
        # Add an import and type check here to avoid cyclic imports
        from a4x.orchestration.file import File

        if not isinstance(f, File):
            raise TypeError("Files added to a directory must be of type a4x.orchestration.File")
        if f not in self.files:
            self.files.add(f)

    def _remove_file(self, f) -> None:  # noqa: ANN001
        # Add an import and type check here to avoid cyclic imports
        from a4x.orchestration.file import File

        if not isinstance(f, File):
            raise TypeError("Files added to a directory must be of type a4x.orchestration.File")
        if f in self.files:
            self.files.remove(f)

    def _resolve_files(self) -> None:
        for f in self.files:
            f.resolve()

    def __eq__(self, other: Any) -> bool:
        # NOTE: getattr logic is to work around partial initialization issues that can
        #       be caused by pickling and copy.deepcopy
        return (
            isinstance(other, Directory)
            and getattr(self, "path", None) == getattr(other, "path", None)
            and getattr(self, "site", None) == getattr(other, "site", None)
            and getattr(self, "storage_type_attr", None) == getattr(other, "storage_type_attr", None)
            and getattr(self, "persistency_attr", None) == getattr(other, "persistency_attr", None)
        )

    def __hash__(self) -> int:
        # NOTE: hasattr logic is to work around partial initialization issues that can
        #       be caused by pickling and copy.deepcopy
        if any(
            not hasattr(self, attr_name)
            for attr_name in (
                "path",
                "site",
                "storage_type_attr",
                "persistency_attr",
            )
        ):
            return id(self)
        return hash((self.path, self.site, self.storage_type_attr, self.persistency_attr))

    def __str__(self) -> str:
        # NOTE: hasattr logic is to work around partial initialization issues that can
        #       be caused by pickling and copy.deepcopy
        if not hasattr(self, "path"):
            return "UNKNOWN"
        return str(self.path)

    def __repr__(self) -> str:
        # NOTE: hasattr logic is to work around partial initialization issues that can
        #       be caused by pickling and copy.deepcopy
        if not hasattr(self, "path"):
            path = "UNKNOWN"
        else:
            path = str(self.path)
        if not hasattr(self, "site"):
            site = "UNKNOWN"
        else:
            site = repr(self.site)
        if not hasattr(self, "storage_type_attr"):
            storage_type = "UNKNOWN"
        else:
            storage_type = str(self.storage_type_attr)
        if not hasattr(self, "persistency_attr"):
            persistency = "UNKNOWN"
        else:
            persistency = str(self.persistency_attr)
        return f"Directory(path={path},site={site},storage_type={storage_type},persistency={persistency})"

    def __getstate__(self) -> dict[str, Any]:
        """Get the state of the Directory object for the pickle protocol."""
        return self.__dict__

    def __setstate__(self, state: dict[str, Any]) -> None:
        """Set the state of the Directory object for the pickle protocol."""
        self.__dict__.update(state)

        if not hasattr(self, "files"):
            self.files = set()

        if hasattr(self, "site") and hasattr(self, "identifier"):
            if self.identifier not in self.site:
                self.site[self.identifier] = self
