# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import dataclasses
from pathlib import Path
from typing import Collection, List, Mapping, Optional, Union

import libcst as cst
from libcst._metadata_dependent import LazyValue, MetadataDependent
from libcst.helpers.module import calculate_module_and_package, ModuleNameAndPackage
from libcst.metadata.base_provider import BatchableMetadataProvider
from libcst.metadata.scope_provider import (
    QualifiedName,
    QualifiedNameSource,
    ScopeProvider,
)


class QualifiedNameProvider(BatchableMetadataProvider[Collection[QualifiedName]]):
    """
    Compute possible qualified names of a variable CSTNode
    (extends `PEP-3155 <https://www.python.org/dev/peps/pep-3155/>`_).
    It uses the
    :func:`~libcst.metadata.Scope.get_qualified_names_for` underlying to get qualified names.
    Multiple qualified names may be returned, such as when we have conditional imports or an
    import shadows another. E.g., the provider finds ``a.b``, ``d.e`` and
    ``f.g`` as possible qualified names of ``c``::

        >>> wrapper = MetadataWrapper(
        >>>     cst.parse_module(dedent(
        >>>     '''
        >>>         if something:
        >>>             from a import b as c
        >>>         elif otherthing:
        >>>             from d import e as c
        >>>         else:
        >>>             from f import g as c
        >>>         c()
        >>>     '''
        >>>     ))
        >>> )
        >>> call = wrapper.module.body[1].body[0].value
        >>> wrapper.resolve(QualifiedNameProvider)[call],
        {
            QualifiedName(name="a.b", source=QualifiedNameSource.IMPORT),
            QualifiedName(name="d.e", source=QualifiedNameSource.IMPORT),
            QualifiedName(name="f.g", source=QualifiedNameSource.IMPORT),
        }

    For qualified name of a variable in a function or a comprehension, please refer
    :func:`~libcst.metadata.Scope.get_qualified_names_for` for more detail.
    """

    METADATA_DEPENDENCIES = (ScopeProvider,)

    def visit_Module(self, node: cst.Module) -> Optional[bool]:
        visitor = QualifiedNameVisitor(self)
        node.visit(visitor)

    @staticmethod
    def has_name(
        visitor: MetadataDependent, node: cst.CSTNode, name: Union[str, QualifiedName]
    ) -> bool:
        """Check if any of qualified name has the str name or :class:`~libcst.metadata.QualifiedName` name."""
        qualified_names = visitor.get_metadata(QualifiedNameProvider, node, set())
        if isinstance(name, str):
            return any(qn.name == name for qn in qualified_names)
        else:
            return any(qn == name for qn in qualified_names)


class QualifiedNameVisitor(cst.CSTVisitor):
    def __init__(self, provider: "QualifiedNameProvider") -> None:
        self.provider: QualifiedNameProvider = provider

    def on_visit(self, node: cst.CSTNode) -> bool:
        scope = self.provider.get_metadata(ScopeProvider, node, None)
        if scope:
            self.provider.set_metadata(
                node, LazyValue(lambda: scope.get_qualified_names_for(node))
            )
        else:
            self.provider.set_metadata(node, set())
        super().on_visit(node)
        return True


class FullyQualifiedNameProvider(BatchableMetadataProvider[Collection[QualifiedName]]):
    """
    Provide fully qualified names for CST nodes. Like :class:`QualifiedNameProvider`,
    but the provided :class:`QualifiedName` instances have absolute identifier names
    instead of local to the current module.

    This provider is initialized with the current module's fully qualified name, and can
    be used with :class:`~libcst.metadata.FullRepoManager`. The module's fully qualified
    name itself is stored as a metadata of the :class:`~libcst.Module` node. Compared to
    :class:`QualifiedNameProvider`, it also resolves relative imports.

    Example usage::

        >>> mgr = FullRepoManager(".", {"dir/a.py"}, {FullyQualifiedNameProvider})
        >>> wrapper = mgr.get_metadata_wrapper_for_path("dir/a.py")
        >>> fqnames = wrapper.resolve(FullyQualifiedNameProvider)
        >>> {type(k): v for (k, v) in fqnames.items()}
        {<class 'libcst._nodes.module.Module'>: {QualifiedName(name='dir.a', source=<QualifiedNameSource.LOCAL: 3>)}}

    """

    METADATA_DEPENDENCIES = (QualifiedNameProvider,)

    @classmethod
    def gen_cache(
        cls, root_path: Path, paths: List[str], timeout: Optional[int] = None
    ) -> Mapping[str, ModuleNameAndPackage]:
        cache = {path: calculate_module_and_package(root_path, path) for path in paths}
        return cache

    def __init__(self, cache: ModuleNameAndPackage) -> None:
        super().__init__(cache)
        self.module_name: str = cache.name
        self.package_name: str = cache.package

    def visit_Module(self, node: cst.Module) -> bool:
        visitor = FullyQualifiedNameVisitor(self, self.module_name, self.package_name)
        node.visit(visitor)
        self.set_metadata(
            node,
            {QualifiedName(name=self.module_name, source=QualifiedNameSource.LOCAL)},
        )
        return True


class FullyQualifiedNameVisitor(cst.CSTVisitor):
    @staticmethod
    def _fully_qualify_local(module_name: str, package_name: str, name: str) -> str:
        abs_name = name.lstrip(".")
        num_dots = len(name) - len(abs_name)
        # handle relative import
        if num_dots > 0:
            name = abs_name
            # see importlib._bootstrap._resolve_name
            # https://github.com/python/cpython/blob/3.10/Lib/importlib/_bootstrap.py#L902
            bits = package_name.rsplit(".", num_dots - 1)
            if len(bits) < num_dots:
                raise ImportError("attempted relative import beyond top-level package")
            module_name = bits[0]

        return f"{module_name}.{name}"

    @staticmethod
    def _fully_qualify(
        module_name: str, package_name: str, qname: QualifiedName
    ) -> QualifiedName:
        if qname.source == QualifiedNameSource.BUILTIN:
            # builtins are already fully qualified
            return qname
        name = qname.name
        if qname.source == QualifiedNameSource.IMPORT and not name.startswith("."):
            # non-relative imports are already fully qualified
            return qname
        new_name = FullyQualifiedNameVisitor._fully_qualify_local(
            module_name, package_name, qname.name
        )
        return dataclasses.replace(qname, name=new_name)

    def __init__(
        self, provider: FullyQualifiedNameProvider, module_name: str, package_name: str
    ) -> None:
        self.module_name = module_name
        self.package_name = package_name
        self.provider = provider

    def on_visit(self, node: cst.CSTNode) -> bool:
        qnames = self.provider.get_metadata(QualifiedNameProvider, node)
        if qnames is not None:
            self.provider.set_metadata(
                node,
                {
                    FullyQualifiedNameVisitor._fully_qualify(
                        self.module_name, self.package_name, qname
                    )
                    for qname in qnames
                },
            )
        return True