from __future__ import annotations

import dataclasses
import os
from collections import defaultdict
from collections.abc import Callable
from functools import cached_property
from typing import TYPE_CHECKING

from packaging.specifiers import InvalidSpecifier
from packaging.version import InvalidVersion
from resolvelib import AbstractProvider, RequirementsConflicted
from resolvelib.resolvers import Criterion

from pdm.exceptions import CandidateNotFound, InvalidPyVersion, RequirementError
from pdm.models.candidates import Candidate
from pdm.models.repositories import LockedRepository
from pdm.models.requirements import FileRequirement, Requirement, VcsRequirement, parse_requirement, strip_extras
from pdm.models.specifiers import PySpecSet
from pdm.resolver.python import PythonCandidate, PythonRequirement, find_python_matches, is_python_satisfied_by
from pdm.termui import logger
from pdm.utils import (
    deprecation_warning,
    get_requirement_from_override,
    normalize_name,
    parse_version,
    url_without_fragments,
)

if TYPE_CHECKING:
    from collections.abc import Iterable, Iterator, Mapping, Sequence
    from typing import TypeVar

    from resolvelib.resolvers import RequirementInformation

    from pdm._types import Comparable
    from pdm.models.repositories import BaseRepository
    from pdm.models.working_set import WorkingSet

    ProviderT = TypeVar("ProviderT", bound="type[BaseProvider]")


_PROVIDER_REGISTRY: dict[str, type[BaseProvider]] = {}
_CONFLICT_PRIORITY_THRESHOLD = 5


def get_provider(strategy: str) -> type[BaseProvider]:
    return _PROVIDER_REGISTRY[strategy]


def register_provider(strategy: str) -> Callable[[ProviderT], ProviderT]:
    def wrapper(cls: ProviderT) -> ProviderT:
        _PROVIDER_REGISTRY[strategy] = cls
        return cls

    return wrapper


@register_provider("all")
class BaseProvider(AbstractProvider[Requirement, Candidate, str]):
    def __init__(
        self,
        repository: BaseRepository,
        allow_prereleases: bool | None = None,
        overrides: dict[str, str] | None = None,
        direct_minimal_versions: bool = False,
        *,
        locked_repository: LockedRepository | None = None,
    ) -> None:
        if overrides is not None:  # pragma: no cover
            deprecation_warning(
                "The `overrides` argument is deprecated and will be removed in the future.", stacklevel=2
            )
        if allow_prereleases is not None:  # pragma: no cover
            deprecation_warning(
                "The `allow_prereleases` argument is deprecated and will be removed in the future.", stacklevel=2
            )
        project = repository.environment.project
        self.repository = repository
        self.allow_prereleases = project.pyproject.allow_prereleases  # Root allow_prereleases value
        self.fetched_dependencies: dict[tuple[str, str | None], list[Requirement]] = {}
        self.excludes = {normalize_name(k) for k in project.pyproject.resolution.get("excludes", [])}
        self.direct_minimal_versions = direct_minimal_versions
        self.locked_repository = locked_repository
        self._conflict_counts: defaultdict[str, int] = defaultdict(int)
        self._conflict_promoted: set[str] = set()

    def requirement_preference(self, requirement: Requirement) -> Comparable:
        """Return the preference of a requirement to find candidates.

        - Editable requirements are preferred.
        - File links are preferred.
        - The one with narrower specifierset is preferred.
        """
        editable = requirement.editable
        is_named = requirement.is_named
        is_pinned = requirement.is_pinned
        is_prerelease = bool(requirement.prerelease) or bool(requirement.specifier.prereleases)
        specifier_parts = len(requirement.specifier)
        return (not editable, is_named, not is_pinned, not is_prerelease, -specifier_parts)

    def identify(self, requirement_or_candidate: Requirement | Candidate) -> str:
        return requirement_or_candidate.identify()

    def narrow_requirement_selection(
        self,
        identifiers: Iterable[str],
        resolutions: Mapping[str, Candidate],
        candidates: Mapping[str, Iterator[Candidate]],
        information: Mapping[str, Iterator[RequirementInformation]],
        backtrack_causes: Sequence[RequirementInformation],
    ) -> Iterable[str]:
        backtrack_identifiers: set[str] = set()
        for requirement, parent in backtrack_causes:
            names = [requirement.identify()]
            if parent is not None:
                names.append(parent.identify())
            for name in names:
                backtrack_identifiers.add(name)
                if name not in resolutions:
                    self._conflict_counts[name] += 1
                    if self._conflict_counts[name] >= _CONFLICT_PRIORITY_THRESHOLD:
                        self._conflict_promoted.add(name)

        current_backtrack_causes: list[str] = []
        promoted: list[str] = []
        for identifier in identifiers:
            if identifier == "python":
                return [identifier]
            if identifier in backtrack_identifiers:
                current_backtrack_causes.append(identifier)
                continue
            if identifier in self._conflict_promoted:
                promoted.append(identifier)

        if current_backtrack_causes:
            return current_backtrack_causes
        if promoted:
            return promoted
        return identifiers

    def get_preference(
        self,
        identifier: str,
        resolutions: Mapping[str, Candidate],
        candidates: Mapping[str, Iterator[Candidate]],
        information: Mapping[str, Iterator[RequirementInformation]],
        backtrack_causes: Sequence[RequirementInformation],
    ) -> tuple[Comparable, ...]:
        is_top = any(parent is None for _, parent in information[identifier])
        backtrack_identifiers = {req.identify() for req, _ in backtrack_causes} | {
            parent.identify() for _, parent in backtrack_causes if parent is not None
        }
        # Use the REAL identifier as it may be updated after candidate preparation.
        deps: list[Requirement] = []
        for candidate in candidates[identifier]:
            try:
                deps = self.get_dependencies(candidate)
            except RequirementsConflicted:
                continue
            break
        is_backtrack_cause = any(dep.identify() in backtrack_identifiers for dep in deps)
        is_file_or_url = any(not requirement.is_named for requirement, _ in information[identifier])
        operators = [spec.operator for req, _ in information[identifier] for spec in req.specifier]
        is_python = identifier == "python"
        is_pinned = any(op[:2] == "==" for op in operators)
        constraints = len(operators)
        is_conflict_promoted = identifier in self._conflict_promoted
        return (
            not is_python,
            not is_top,
            not is_conflict_promoted,
            not is_file_or_url,
            not is_pinned,
            not is_backtrack_cause,
            -constraints,
            identifier,
        )

    @cached_property
    def locked_candidates(self) -> dict[str, list[Candidate]]:
        return self.locked_repository.all_candidates if self.locked_repository else {}

    @cached_property
    def overrides(self) -> dict[str, Requirement]:
        """A mapping of package name to the requirement for overriding."""
        from pdm.formats.requirements import RequirementParser

        project_overrides: dict[str, str] = {
            normalize_name(k): v
            for k, v in self.repository.environment.project.pyproject.resolution.get("overrides", {}).items()
        }
        requirements: dict[str, Requirement] = {}
        for name, value in project_overrides.items():
            req = get_requirement_from_override(name, value)
            r = parse_requirement(req)
            requirements[r.identify()] = r

        # Read from --override files
        parser = RequirementParser(self.repository.environment.session)
        for override_file in self.repository.environment.project.core.state.overrides:
            parser.parse_file(override_file)
        for r in parser.requirements:
            # There might be duplicates, we only keep the last one
            requirements[r.identify()] = r

        return requirements

    def _is_direct_requirement(self, requirement: Requirement) -> bool:
        from itertools import chain

        project = self.repository.environment.project
        all_dependencies = chain.from_iterable(project.all_dependencies.values())
        return any(r.is_named and requirement.identify() == r.identify() for r in all_dependencies)

    @cached_property
    def workspace_candidates(self) -> dict[str, Candidate]:
        project = self.repository.environment.project
        return {req.identify(): Candidate(req) for req in project.iter_workspace_dependencies()}

    def _find_workspace_candidates(self, identifier: str) -> Iterable[Candidate]:
        if candidate := self.workspace_candidates.get(identifier):
            yield candidate
            return
        name, extras = strip_extras(identifier)
        if extras and (candidate := self.workspace_candidates.get(name)):
            req = dataclasses.replace(candidate.req, extras=extras)
            yield Candidate(req, name=candidate.name, version=candidate.version)

    def _find_candidates(self, requirement: Requirement) -> Iterable[Candidate]:
        if not requirement.is_named and not isinstance(self.repository, LockedRepository):
            can = Candidate(requirement)
            if not can.name:
                can.prepare(self.repository.environment).metadata
            yield can
        else:
            prerelease = requirement.prerelease
            if prerelease is None and requirement.is_pinned and requirement.specifier.prereleases:
                prerelease = True
            if prerelease is None and (key := requirement.identify()) in self.locked_candidates:
                # keep the prerelease if it is locked
                candidates = self.locked_candidates[key]
                for candidate in candidates:
                    if candidate.version is not None:
                        try:
                            parsed_version = parse_version(candidate.version)
                        except InvalidVersion:  # pragma: no cover
                            pass
                        else:
                            if parsed_version.is_prerelease:
                                prerelease = True
                                break
            found = self.repository.find_candidates(
                requirement,
                self.allow_prereleases if prerelease is None else prerelease,
                minimal_version=self.direct_minimal_versions and self._is_direct_requirement(requirement),
            )

            current_version: str | None = None
            collected_wheels: list[Candidate] = []
            collected_others: list[Candidate] = []
            for candidate in found:
                assert candidate.version is not None
                if current_version is None:
                    current_version = candidate.version
                if candidate.version != current_version:
                    # If there are wheels for the given version, we should only return wheels
                    # to avoid build steps.
                    if collected_wheels:
                        yield collected_wheels[0]
                    elif collected_others:
                        yield collected_others[0]
                    current_version = candidate.version
                    collected_wheels.clear()
                    collected_others.clear()
                if candidate.link and candidate.link.is_wheel:
                    collected_wheels.append(candidate)
                else:
                    collected_others.append(candidate)
            if collected_wheels:
                yield collected_wheels[0]
            elif collected_others:
                yield collected_others[0]

    def find_matches(
        self,
        identifier: str,
        requirements: Mapping[str, Iterator[Requirement]],
        incompatibilities: Mapping[str, Iterator[Candidate]],
    ) -> Callable[[], Iterator[Candidate]]:
        def matches_gen() -> Iterator[Candidate]:
            incompat = list(incompatibilities[identifier])
            if identifier == "python":
                candidates = find_python_matches(identifier, requirements)
                return (c for c in candidates if c not in incompat)
            elif identifier in self.overrides:
                return iter(self._find_candidates(self.overrides[identifier]))
            else:
                name, extras = strip_extras(identifier)
                if name in self.overrides:
                    req = dataclasses.replace(self.overrides[name], extras=extras)
                    return iter(self._find_candidates(req))
            reqs = list(requirements[identifier])
            if not reqs:
                return iter(())
            original_req = min(reqs, key=self.requirement_preference)
            bare_name, extras = strip_extras(identifier)
            if extras and bare_name in requirements:
                # We should consider the requirements for both foo and foo[extra]
                reqs.extend(requirements[bare_name])
            reqs.sort(key=self.requirement_preference)
            if workspace_candidates := list(self._find_workspace_candidates(identifier)):
                return (
                    can.copy_with(dataclasses.replace(can.req, groups=original_req.groups))
                    for can in workspace_candidates
                    if can not in incompat and all(self.is_satisfied_by(r, can) for r in reqs)
                )
            candidates = self._find_candidates(reqs[0])
            return (
                # In some cases we will use candidates from the bare requirement,
                # this will miss the extra dependencies if any. So we associate the original
                # requirement back with the candidate since it is used by `get_dependencies()`.
                can.copy_with(original_req) if extras else can
                for can in candidates
                if can not in incompat and all(self.is_satisfied_by(r, can) for r in reqs)
            )

        return matches_gen

    def _compare_file_reqs(self, req1: FileRequirement, req2: FileRequirement) -> bool:
        backend = self.repository.environment.project.backend
        if req1.path and req2.path:
            return os.path.normpath(req1.path.absolute()) == os.path.normpath(req2.path.absolute())
        left = backend.expand_line(url_without_fragments(req1.get_full_url()))
        right = backend.expand_line(url_without_fragments(req2.get_full_url()))
        return left == right

    def is_satisfied_by(self, requirement: Requirement, candidate: Candidate) -> bool:
        if isinstance(requirement, PythonRequirement):
            return is_python_satisfied_by(requirement, candidate)
        elif (name := candidate.identify()) in self.overrides or strip_extras(name)[0] in self.overrides:
            return True
        if not requirement.is_named:
            if candidate.req.is_named:
                return False
            can_req = candidate.req
            if requirement.is_vcs and can_req.is_vcs:
                return can_req.vcs == requirement.vcs and can_req.repo == requirement.repo  # type: ignore[attr-defined]
            return self._compare_file_reqs(requirement, can_req)  # type: ignore[arg-type]
        version = candidate.version
        this_name = self.repository.environment.project.name
        if version is None or candidate.name == this_name:
            # This should be a URL candidate or self package, consider it to be matching
            return True
        # Allow prereleases if: 1) it is not specified in the tool settings or
        # 2) the candidate doesn't come from PyPI index or 3) the requirement is pinned
        allow_prereleases = (
            self.allow_prereleases in (True, None) or not candidate.req.is_named or requirement.is_pinned
        )
        return requirement.specifier.contains(version, allow_prereleases)

    def _get_dependencies_from_repository(self, candidate: Candidate) -> tuple[list[Requirement], PySpecSet, str]:
        return self.repository.get_dependencies(candidate)

    def get_dependencies(self, candidate: Candidate) -> list[Requirement]:
        if isinstance(candidate, PythonCandidate):
            return []
        try:
            deps, requires_python, _ = self._get_dependencies_from_repository(candidate)
        except (RequirementError, InvalidPyVersion, InvalidSpecifier) as e:
            # When the metadata is invalid, skip this candidate by marking it as conflicting.
            # Here we pass an empty criterion so it doesn't provide any info to the resolution.
            logger.error("Invalid metadata in %s: %s", candidate, e)
            raise RequirementsConflicted(Criterion([], [], [])) from None

        if candidate.req.extras:
            # XXX: If the requirement has extras, add the original candidate
            # (without extras) as its dependency. This ensures the same package with
            # different extras resolve to the same version.
            self_req = dataclasses.replace(
                candidate.req.as_pinned_version(candidate.version),
                extras=None,
                marker=None,
            )
            if self_req not in deps:
                deps.insert(0, self_req)
        self.fetched_dependencies[candidate.dep_key] = deps[:]
        # Filter out incompatible dependencies(e.g. functools32) early so that
        # we don't get errors when building wheels.
        valid_deps: list[Requirement] = []
        for dep in deps:
            if (
                dep.requires_python
                & requires_python
                & candidate.req.requires_python
                & PySpecSet(self.repository.env_spec.requires_python)
            ).is_empty():
                continue
            if dep.marker and not dep.marker.matches(self.repository.env_spec):
                continue
            if dep.identify() in self.excludes:
                continue
            dep.requires_python &= candidate.req.requires_python
            valid_deps.append(dep)
        # A candidate contributes to the Python requirements only when:
        # It isn't an optional dependency, or the requires-python doesn't cover
        # the req's requires-python.
        # For example, A v1 requires python>=3.6, it not eligible on a project with
        # requires-python=">=2.7". But it is eligible if A has environment marker
        # A1; python_version>='3.8'
        new_requires_python = candidate.req.requires_python & self.repository.environment.python_requires
        if not (
            candidate.identify() in self.overrides
            or new_requires_python.is_empty()
            or requires_python.is_superset(new_requires_python)
        ):
            valid_deps.append(PythonRequirement.from_pyspec_set(requires_python))
        return valid_deps


@register_provider("reuse")
class ReusePinProvider(BaseProvider):
    """A provider that reuses preferred pins if possible.

    This is used to implement "add", "remove", and "reuse upgrade",
    where already-pinned candidates in lockfile should be preferred.
    """

    def __init__(
        self,
        repository: BaseRepository,
        allow_prereleases: bool | None = None,
        overrides: dict[str, str] | None = None,
        direct_minimal_versions: bool = False,
        *,
        locked_repository: LockedRepository | None = None,
        tracked_names: Iterable[str],
    ) -> None:
        super().__init__(
            repository=repository,
            allow_prereleases=allow_prereleases,
            overrides=overrides,
            direct_minimal_versions=direct_minimal_versions,
            locked_repository=locked_repository,
        )
        self.tracked_names = set(tracked_names)

    def iter_reuse_candidates(self, identifier: str, requirement: Requirement | None) -> Iterable[Candidate]:
        bare_name = strip_extras(identifier)[0]
        if bare_name in self.tracked_names or identifier not in self.locked_candidates:
            return []
        return sorted(self.locked_candidates[identifier], key=lambda c: c.version or "", reverse=True)

    def get_reuse_candidate(self, identifier: str, requirement: Requirement | None) -> Candidate | None:
        deprecation_warning(
            "The get_reuse_candidate method is deprecated, use iter_reuse_candidates instead.", stacklevel=2
        )
        return next(iter(self.iter_reuse_candidates(identifier, requirement)), None)

    def find_matches(
        self,
        identifier: str,
        requirements: Mapping[str, Iterator[Requirement]],
        incompatibilities: Mapping[str, Iterator[Candidate]],
    ) -> Callable[[], Iterator[Candidate]]:
        super_find = super().find_matches(identifier, requirements, incompatibilities)

        def matches_gen() -> Iterator[Candidate]:
            if list(self._find_workspace_candidates(identifier)):
                yield from super_find()
                return
            requested_req = next(filter(lambda r: r.is_named, requirements[identifier]), None)
            for pin in self.iter_reuse_candidates(identifier, requested_req):
                if identifier not in self.overrides and pin.req.is_named:
                    pin = pin.copy_with(min(requirements[identifier], key=self.requirement_preference))
                incompat = list(incompatibilities[identifier])
                pin._preferred = True  # type: ignore[attr-defined]
                if pin not in incompat and all(self.is_satisfied_by(r, pin) for r in requirements[identifier]):
                    yield pin
            yield from super_find()

        return matches_gen

    def _get_dependencies_from_repository(self, candidate: Candidate) -> tuple[list[Requirement], PySpecSet, str]:
        is_stable_metadata = candidate.req.is_named or (
            isinstance(candidate.req, VcsRequirement) and candidate.req.revision
        )
        if self.locked_repository is not None and is_stable_metadata:
            try:
                return self.locked_repository.get_dependencies(candidate)
            except CandidateNotFound:
                pass
        return super()._get_dependencies_from_repository(candidate)


@register_provider("eager")
class EagerUpdateProvider(ReusePinProvider):
    """A specialized provider to handle an "eager" upgrade strategy.

    An eager upgrade tries to upgrade not only packages specified, but also
    their dependencies (recursively). This contrasts to the "only-if-needed"
    default, which only promises to upgrade the specified package, and
    prevents touching anything else if at all possible.

    The provider is implemented as to keep track of all dependencies of the
    specified packages to upgrade, and free their pins when it has a chance.
    """

    def iter_reuse_candidates(self, identifier: str, requirement: Requirement | None) -> Iterable[Candidate]:
        if identifier in self.tracked_names:
            # If this is a tracked package, don't reuse its pinned version, so it can be upgraded.
            return []
        return super().iter_reuse_candidates(identifier, requirement)

    def get_dependencies(self, candidate: Candidate) -> list[Requirement]:
        # If this package is being tracked for upgrade, remove pins of its
        # dependencies, and start tracking these new packages.
        dependencies = super().get_dependencies(candidate)
        if self.identify(candidate) in self.tracked_names:
            for dependency in dependencies:
                if dependency.key:
                    self.tracked_names.add(dependency.key)
        return dependencies

    def get_preference(
        self,
        identifier: str,
        resolutions: Mapping[str, Candidate],
        candidates: Mapping[str, Iterator[Candidate]],
        information: Mapping[str, Iterator[RequirementInformation]],
        backtrack_causes: Sequence[RequirementInformation],
    ) -> tuple[Comparable, ...]:
        # Resolve tracking packages so we have a chance to unpin them first.
        (python, *others) = super().get_preference(identifier, resolutions, candidates, information, backtrack_causes)
        return (python, identifier not in self.tracked_names, *others)


@register_provider("reuse-installed")
class ReuseInstalledProvider(ReusePinProvider):
    """A provider that reuses installed packages if possible."""

    @cached_property
    def installed(self) -> WorkingSet:
        return self.repository.environment.get_working_set()

    def iter_reuse_candidates(self, identifier: str, requirement: Requirement | None) -> Iterable[Candidate]:
        key = strip_extras(identifier)[0]
        if key not in self.installed or requirement is None:
            return super().iter_reuse_candidates(identifier, requirement)
        else:
            dist = self.installed[key]
            return [Candidate(requirement, installed=dist)]
