# Copyright Spack Project Developers. See COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)
import enum
import functools
from typing import Any, Callable, List, Mapping

import spack.binary_distribution
import spack.config
import spack.environment
import spack.llnl.path
import spack.repo
import spack.spec
import spack.store
import spack.traverse
from spack.externals import (
    ExternalSpecsParser,
    complete_architecture,
    complete_variants_and_architecture,
    extract_dicts_from_configuration,
)

from .runtimes import all_libcs


class SpecFilter:
    """Given a method to produce a list of specs, this class can filter them according to
    different criteria.
    """

    def __init__(
        self,
        factory: Callable[[], List[spack.spec.Spec]],
        is_usable: Callable[[spack.spec.Spec], bool],
        include: List[str],
        exclude: List[str],
    ) -> None:
        """
        Args:
            factory: factory to produce a list of specs
            is_usable: predicate that takes a spec in input and returns False if the spec
                should not be considered for this filter, True otherwise.
            include: if present, a "good" spec must match at least one entry in the list
            exclude: if present, a "good" spec must not match any entry in the list
        """
        self.factory = factory
        self.is_usable = is_usable
        self.include = include
        self.exclude = exclude

    def is_selected(self, s: spack.spec.Spec) -> bool:
        if not self.is_usable(s):
            return False

        if self.include and not any(s.satisfies(c) for c in self.include):
            return False

        if self.exclude and any(s.satisfies(c) for c in self.exclude):
            return False

        return True

    def selected_specs(self) -> List[spack.spec.Spec]:
        return [s for s in self.factory() if self.is_selected(s)]

    @staticmethod
    def from_store(configuration, *, packages_with_externals, include, exclude) -> "SpecFilter":
        """Constructs a filter that takes the specs from the current store."""
        is_reusable = functools.partial(
            _is_reusable, packages_with_externals=packages_with_externals, local=True
        )
        factory = functools.partial(_specs_from_store, configuration=configuration)
        return SpecFilter(factory=factory, is_usable=is_reusable, include=include, exclude=exclude)

    @staticmethod
    def from_buildcache(*, packages_with_externals, include, exclude) -> "SpecFilter":
        """Constructs a filter that takes the specs from the configured buildcaches."""
        is_reusable = functools.partial(
            _is_reusable, packages_with_externals=packages_with_externals, local=False
        )
        return SpecFilter(
            factory=_specs_from_mirror, is_usable=is_reusable, include=include, exclude=exclude
        )

    @staticmethod
    def from_environment(*, packages_with_externals, include, exclude, env) -> "SpecFilter":
        is_reusable = functools.partial(
            _is_reusable, packages_with_externals=packages_with_externals, local=True
        )
        factory = functools.partial(_specs_from_environment, env=env)
        return SpecFilter(factory=factory, is_usable=is_reusable, include=include, exclude=exclude)

    @staticmethod
    def from_environment_included_concrete(
        *,
        packages_with_externals,
        include: List[str],
        exclude: List[str],
        env: spack.environment.Environment,
        included_concrete: str,
    ) -> "SpecFilter":
        is_reusable = functools.partial(
            _is_reusable, packages_with_externals=packages_with_externals, local=True
        )
        factory = functools.partial(
            _specs_from_environment_included_concrete, env=env, included_concrete=included_concrete
        )
        return SpecFilter(factory=factory, is_usable=is_reusable, include=include, exclude=exclude)

    @staticmethod
    def from_packages_yaml(
        *, external_parser: ExternalSpecsParser, packages_with_externals, include, exclude
    ) -> "SpecFilter":
        is_reusable = functools.partial(
            _is_reusable, packages_with_externals=packages_with_externals, local=True
        )
        return SpecFilter(
            external_parser.all_specs, is_usable=is_reusable, include=include, exclude=exclude
        )


def _has_runtime_dependencies(spec: spack.spec.Spec) -> bool:
    # TODO (compiler as nodes): this function contains specific names from builtin, and should
    # be made more general
    if "gcc" in spec and "gcc-runtime" not in spec:
        return False

    if "intel-oneapi-compilers" in spec and "intel-oneapi-runtime" not in spec:
        return False

    return True


def _is_reusable(spec: spack.spec.Spec, packages_with_externals, local: bool) -> bool:
    """A spec is reusable if it's not a dev spec, it's imported from the cray manifest, it's not
    external, or it's external with matching packages.yaml entry. The latter prevents two issues:

    1. Externals in build caches: avoid installing an external on the build machine not
       available on the target machine
    2. Local externals: avoid reusing an external if the local config changes. This helps in
       particular when a user removes an external from packages.yaml, and expects that that
       takes effect immediately.

    Arguments:
        spec: the spec to check
        packages_with_externals: the pre-processed packages configuration
    """
    if "dev_path" in spec.variants:
        return False

    if spec.name == "compiler-wrapper":
        return False

    if not spec.external:
        return _has_runtime_dependencies(spec)

    # Cray external manifest externals are always reusable
    if local:
        _, record = spack.store.STORE.db.query_by_spec_hash(spec.dag_hash())
        if record and record.origin == "external-db":
            return True

    try:
        provided = spack.repo.PATH.get(spec).provided_virtual_names()
    except spack.repo.RepoError:
        provided = []

    for name in {spec.name, *provided}:
        for entry in packages_with_externals.get(name, {}).get("externals", []):
            expected_prefix = entry.get("prefix")
            if expected_prefix is not None:
                expected_prefix = spack.llnl.path.path_to_os_path(expected_prefix)[0]
            if (
                spec.satisfies(entry["spec"])
                and spec.external_path == expected_prefix
                and spec.external_modules == entry.get("modules")
            ):
                return True

    return False


def _specs_from_store(configuration):
    store = spack.store.create(configuration)
    with store.db.read_transaction():
        return store.db.query(installed=True)


def _specs_from_mirror():
    try:
        return spack.binary_distribution.update_cache_and_get_specs()
    except (spack.binary_distribution.FetchCacheError, IndexError):
        # this is raised when no mirrors had indices.
        # TODO: update mirror configuration so it can indicate that the
        # TODO: source cache (or any mirror really) doesn't have binaries.
        return []


def _specs_from_environment(env):
    """Return all concrete specs from the environment. This includes all included concrete"""
    if env:
        return list(spack.traverse.traverse_nodes([s for _, s in env.concretized_specs()]))
    else:
        return []


def _specs_from_environment_included_concrete(env, included_concrete):
    """Return only concrete specs from the environment included from the included_concrete"""
    if env:
        assert included_concrete in env.included_concrete_envs
        return [concrete for concrete in env.included_specs_by_hash[included_concrete].values()]
    else:
        return []


class ReuseStrategy(enum.Enum):
    ROOTS = enum.auto()
    DEPENDENCIES = enum.auto()
    NONE = enum.auto()


def create_external_parser(
    packages_with_externals: Any, completion_mode: str
) -> ExternalSpecsParser:
    """Get externals from a pre-processed packages.yaml (with implicit externals)."""
    external_dicts = extract_dicts_from_configuration(packages_with_externals)
    if completion_mode == "default_variants":
        complete_fn = complete_variants_and_architecture
    elif completion_mode == "architecture_only":
        complete_fn = complete_architecture
    else:
        raise ValueError(
            f"Unknown value for concretizer:externals:completion: {completion_mode!r}"
        )
    return ExternalSpecsParser(external_dicts, complete_node=complete_fn)


class ReusableSpecsSelector:
    """Selects specs that can be reused during concretization."""

    def __init__(
        self,
        configuration: spack.config.Configuration,
        external_parser: ExternalSpecsParser,
        packages_with_externals: Any,
    ) -> None:
        self.configuration = configuration
        self.store = spack.store.create(configuration)
        self.reuse_strategy = ReuseStrategy.ROOTS

        reuse_yaml = self.configuration.get("concretizer:reuse", False)
        self.reuse_sources = []
        if not isinstance(reuse_yaml, Mapping):
            self.reuse_sources.append(
                SpecFilter.from_packages_yaml(
                    external_parser=external_parser,
                    packages_with_externals=packages_with_externals,
                    include=[],
                    exclude=[],
                )
            )
            if reuse_yaml is False:
                self.reuse_strategy = ReuseStrategy.NONE
                return

            if reuse_yaml == "dependencies":
                self.reuse_strategy = ReuseStrategy.DEPENDENCIES
            self.reuse_sources.extend(
                [
                    SpecFilter.from_store(
                        configuration=self.configuration,
                        packages_with_externals=packages_with_externals,
                        include=[],
                        exclude=[],
                    ),
                    SpecFilter.from_buildcache(
                        packages_with_externals=packages_with_externals, include=[], exclude=[]
                    ),
                    SpecFilter.from_environment(
                        packages_with_externals=packages_with_externals,
                        include=[],
                        exclude=[],
                        env=spack.environment.active_environment(),  # with all concrete includes
                    ),
                ]
            )
        else:
            has_external_source = False
            roots = reuse_yaml.get("roots", True)
            if roots is True:
                self.reuse_strategy = ReuseStrategy.ROOTS
            else:
                self.reuse_strategy = ReuseStrategy.DEPENDENCIES
            default_include = reuse_yaml.get("include", [])
            default_exclude = reuse_yaml.get("exclude", [])
            default_sources = [{"type": "local"}, {"type": "buildcache"}]
            for source in reuse_yaml.get("from", default_sources):
                include = source.get("include", default_include)
                exclude = source.get("exclude", default_exclude)
                if source["type"] == "environment" and "path" in source:
                    env_dir = spack.environment.as_env_dir(source["path"])
                    active_env = spack.environment.active_environment()
                    if active_env and env_dir in active_env.included_concrete_envs:
                        # If the environment is included as a concrete environment, use the
                        # local copy of specs in the active environment.
                        # note: included concrete environments are only updated at concretization
                        #       time, and reuse needs to match the included specs.
                        self.reuse_sources.append(
                            SpecFilter.from_environment_included_concrete(
                                packages_with_externals=packages_with_externals,
                                include=include,
                                exclude=exclude,
                                env=active_env,
                                included_concrete=env_dir,
                            )
                        )
                    else:
                        # If the environment is not included as a concrete environment, use the
                        # current specs from its lockfile.
                        self.reuse_sources.append(
                            SpecFilter.from_environment(
                                packages_with_externals=packages_with_externals,
                                include=include,
                                exclude=exclude,
                                env=spack.environment.environment_from_name_or_dir(env_dir),
                            )
                        )
                elif source["type"] == "environment":
                    # reusing from the current environment implicitly reuses from all of the
                    # included concrete environments
                    self.reuse_sources.append(
                        SpecFilter.from_environment(
                            packages_with_externals=packages_with_externals,
                            include=include,
                            exclude=exclude,
                            env=spack.environment.active_environment(),
                        )
                    )
                elif source["type"] == "local":
                    self.reuse_sources.append(
                        SpecFilter.from_store(
                            self.configuration,
                            packages_with_externals=packages_with_externals,
                            include=include,
                            exclude=exclude,
                        )
                    )
                elif source["type"] == "buildcache":
                    self.reuse_sources.append(
                        SpecFilter.from_buildcache(
                            packages_with_externals=packages_with_externals,
                            include=include,
                            exclude=exclude,
                        )
                    )
                elif source["type"] == "external":
                    has_external_source = True
                    if include:
                        # Since libcs are implicit externals, we need to implicitly include them
                        include = include + sorted(all_libcs())  # type: ignore[type-var]
                    self.reuse_sources.append(
                        SpecFilter.from_packages_yaml(
                            external_parser=external_parser,
                            packages_with_externals=packages_with_externals,
                            include=include,
                            exclude=exclude,
                        )
                    )

            # If "external" is not specified, we assume that all externals have to be included
            if not has_external_source:
                self.reuse_sources.append(
                    SpecFilter.from_packages_yaml(
                        external_parser=external_parser,
                        packages_with_externals=packages_with_externals,
                        include=[],
                        exclude=[],
                    )
                )

    def reusable_specs(self, specs: List[spack.spec.Spec]) -> List[spack.spec.Spec]:
        result = []
        for reuse_source in self.reuse_sources:
            result.extend(reuse_source.selected_specs())
        # If we only want to reuse dependencies, remove the root specs
        if self.reuse_strategy == ReuseStrategy.DEPENDENCIES:
            result = [spec for spec in result if not any(root in spec for root in specs)]

        return result
