import asyncio
import contextlib
import glob
import hashlib
import logging
import os
import re
import shlex
import typing
from . import transformations


LOGGER = logging.getLogger(__name__)


class ArtifactFactory(type):
    """
    Metaclass that creates :class:`Artifact` instances or returns existing instances.

    Attributes:
        REGISTRY: Mapping of artifact names to :class:`Artifact` instances. The registry is used to
            keep track of all artifacts and ensure artifact names are unique.
    """
    # Declare using an old-school type hint because Sphinx struggles with circular references.
    REGISTRY = {}  # type: dict[str, Artifact]

    def __init__(self, name, bases, members):
        super(ArtifactFactory, self).__init__(name, bases, members)

    def __call__(self, name, *args, **kwargs):
        # Evaluate the fully qualified name in light of the current group stack.
        name = Group.evaluate_qualified_name(name)
        # Try to retrieve the instance from the registry.
        if instance := self.REGISTRY.get(name):
            if (cls := instance.__class__) is not self:
                raise ValueError(f"cannot create instance of {self} because artifact `{name}` with "
                                 f"type {cls} already exists")
            return instance
        # Create a new instance, register it with the current group if any, and add it to the
        # registry.
        instance = super(ArtifactFactory, self).__call__(name, *args, **kwargs)
        Group.append(instance)
        self.REGISTRY[name] = instance
        return instance


class Artifact(metaclass=ArtifactFactory):
    """
    Artifact generated by a transformation.

    Args:
        name: Unique name of the artifact.

    Raises:
        ValueError: If the given :code:`name` is already in use by another artifact.

    Attributes:
        children: Transformations that consume this artifact; empty if the artifact is a leaf of the
            directed acyclic graph.
        parent: Transformation that generates this artifact; :code:`None` if the artifact is a root
            of the directed acyclic graph.
        digest: Concise summary of the artifact; :code:`None` if the artifact does not exist, cannot
            be summarized, or should always be generated using its :attr:`parent` transform.
    """
    def __init__(self, name: str) -> None:
        self.name = name
        self._parent = None
        self.children = []

    children: typing.Iterable[transformations.Transformation]

    @property
    def parent(self) -> 'transformations.Transformation':
        return self._parent

    @parent.setter
    def parent(self, value: 'transformations.Transformation') -> None:
        if self._parent is not None:
            raise RuntimeError(f"{self} is already associated with a transformation")
        self._parent = value

    @property
    def digest(self) -> bytes:
        return None

    def __repr__(self):
        return f"{self.__class__.__name__}({self.name})"

    async def __call__(self):
        if self.parent is None:
            return
        return await self.parent

    def __await__(self):
        # See https://stackoverflow.com/a/57078217/1150961 for details.
        return (yield from self().__await__())


class Group(Artifact):
    """
    Artifact representing a group of other artifacts.

    Args:
        name: Name of the group. The group name is added as a prefix for all artifacts created
            within the context manager of the group.

    Attributes:
        members: Members of the group.
    """
    def __init__(self, name: str) -> None:
        super().__init__(name)
        self.members = []

    def __enter__(self) -> "Group":
        if self in self.STACK:  # pragma: no cover
            raise RuntimeError(f"{self} is already in the stack")
        self.STACK.append(self)
        return self

    def __exit__(self, *_):
        if (other := self.STACK.pop()) is not self:  # pragma: no cover
            raise RuntimeError(f"expected the last element to be {self} but got {other}")

    async def __call__(self):
        await asyncio.gather(*self.members)

    STACK: list["Group"] = []

    @classmethod
    def evaluate_qualified_name(cls, name: str) -> str:
        """
        Evaluate the fully qualified name of an artifact given its context.

        Args:
            name: Unqualified name.

        Returns:
            name: Qualified name.
        """
        if not cls.STACK:
            return name
        return os.path.join(cls.STACK[-1].name, name)

    @classmethod
    def append(cls, artifact: Artifact) -> None:
        """
        Append an artifact to the current group if available.

        Args:
            artifact: Artifact to append to the current group.
        """
        if not cls.STACK:
            return
        cls.STACK[-1].members.append(artifact)


@contextlib.contextmanager
def group_artifacts(*names: str, squeeze=True) -> list[Group]:
    """
    Group artifacts in a context.

    Args:
        names: Name of nested groups.
        squeeze: Whether to return a single element if there is only one group.

    Returns:
        groups: Artifacts representing the nested groups.
    """
    groups = []
    try:
        for name in names:
            groups.append(Group(name).__enter__())
        yield groups[0] if len(groups) == 1 and squeeze else groups
    except Exception:
        Group.STACK.clear()
        raise
    finally:
        [group.__exit__() for group in reversed(groups)]


class File(Artifact):
    """
    Artifact representing a file.
    """
    def __init__(self, name: str) -> None:
        super().__init__(name)
        if re.search(r"\s", self.name):
            LOGGER.warning("whitespace in `%s` is a recipe for disaster; expect the unexpected",
                           self.name)
        self._last_modified = None
        self._digest = None

    async def __call__(self):
        # Omit directory creation and existence checks during dry run.
        if transformations.Transformation.DRY_RUN:
            await super().__call__()
            return

        # Create the parent directory if necessary.
        if dirname := os.path.dirname(self.name):
            os.makedirs(dirname, exist_ok=True)
        # Generate the artifact and verify the file exists.
        await super().__call__()
        if not os.path.exists(self.name):
            if self._parent:
                message = f"{self._parent} did not generate {self}"
            else:
                message = f"{self} does not exist and cannot be generated"
            raise FileNotFoundError(message)

    @property
    def digest(self):
        try:
            # Only re-evaluate the digest if the file has not been modified since we last computed
            # the hash.
            last_modified = os.stat(self.name).st_mtime
            if self._last_modified is None or self._last_modified < last_modified:
                algorithm = hashlib.sha256()
                with open(self.name, 'rb') as fp:
                    while (chunk := fp.read(4096)):
                        algorithm.update(chunk)
                self._digest = algorithm.digest()
                self._last_modified = last_modified
            return self._digest
        except FileNotFoundError:
            return None

    @classmethod
    def glob(self, pattern: str) -> typing.Iterable["File"]:
        """
        Create a list of file artifacts based on a glob pattern.

        Args:
            pattern: Pattern passed to :func:`glob.glob`.

        Returns:
            artifacts: File artifacts.
        """
        return normalize_artifacts([filename for filename in glob.glob(pattern)])

    def __str__(self):
        return shlex.quote(self.name)


def normalize_artifacts(
    artifacts: typing.Union[str, Artifact, typing.Iterable[typing.Union[str, Artifact]]]) \
        -> typing.Iterable[Artifact]:
    r"""
    Normalize one or more artifacts. Strings are implicitly converted to :class:`File`\ s.

    Args:
        artifacts: One or more artifacts or transformations.

    Returns:
        normalized: Normalized artifacts.
    """
    if artifacts is None:
        return []
    if not isinstance(artifacts, typing.Iterable) or isinstance(artifacts, str):
        artifacts = [artifacts]
    normalized = []
    for artifact in artifacts:
        if isinstance(artifact, str):
            normalized.append(File(artifact))
        elif isinstance(artifact, Artifact):
            normalized.append(artifact)
        elif isinstance(artifact, transformations.Transformation):
            normalized.extend(artifact.outputs)
        else:
            raise TypeError(artifact)
    return normalized


async def gather_artifacts(*artifacts_or_transformations, num_concurrent: int = None) \
       -> typing.Coroutine:
    """
    Gather one or more artifacts and wrap them in a coroutine.

    Args:
        *artifacts_or_transformations: Artifacts to gather and/or transformations whose artifacts to
            gather.
        num_concurrent: Maximum number of concurrent transformations. Defaults to unrestricted.

    Returns:
        gathered: Coroutine that awaits all gathered artifacts.
    """
    try:
        if num_concurrent is not None:
            transformations.Transformation.SEMAPHORE = asyncio.Semaphore(num_concurrent)
        gathered = []
        for item in artifacts_or_transformations:
            if isinstance(item, Artifact):
                gathered.append(item)
            elif isinstance(item, transformations.Transformation):
                gathered.extend(item)
            else:
                raise TypeError(item)
        await asyncio.gather(*gathered)
    finally:
        transformations.Transformation.SEMAPHORE = None
