"""Low-level infrastructure to find modules.

This build on fscache.py; find_sources.py builds on top of this.
"""

import ast
import collections
import functools
import os
import subprocess
import sys
from enum import Enum

from typing import Dict, List, NamedTuple, Optional, Set, Tuple, Union
from typing_extensions import Final

from mypy.defaults import PYTHON3_VERSION_MIN
from mypy.fscache import FileSystemCache
from mypy.options import Options
from mypy import sitepkgs

# Paths to be searched in find_module().
SearchPaths = NamedTuple(
    'SearchPaths',
    [('python_path', Tuple[str, ...]),  # where user code is found
     ('mypy_path', Tuple[str, ...]),  # from $MYPYPATH or config variable
     ('package_path', Tuple[str, ...]),  # from get_site_packages_dirs()
     ('typeshed_path', Tuple[str, ...]),  # paths in typeshed
     ])

# Package dirs are a two-tuple of path to search and whether to verify the module
OnePackageDir = Tuple[str, bool]
PackageDirs = List[OnePackageDir]

PYTHON_EXTENSIONS = ['.pyi', '.py']  # type: Final


# TODO: Consider adding more reasons here?
# E.g. if we deduce a module would likely be found if the user were
# to set the --namespace-packages flag.
class ModuleNotFoundReason(Enum):
    # The module was not found: we found neither stubs nor a plausible code
    # implementation (with or without a py.typed file).
    NOT_FOUND = 0

    # The implementation for this module plausibly exists (e.g. we
    # found a matching folder or *.py file), but either the parent package
    # did not contain a py.typed file or we were unable to find a
    # corresponding *-stubs package.
    FOUND_WITHOUT_TYPE_HINTS = 1

    def error_message_templates(self) -> Tuple[str, str]:
        if self is ModuleNotFoundReason.NOT_FOUND:
            msg = "Cannot find implementation or library stub for module named '{}'"
            note = "See https://mypy.readthedocs.io/en/latest/running_mypy.html#missing-imports"
        elif self is ModuleNotFoundReason.FOUND_WITHOUT_TYPE_HINTS:
            msg = "Skipping analyzing '{}': found module but no type hints or library stubs"
            note = "See https://mypy.readthedocs.io/en/latest/running_mypy.html#missing-imports"
        else:
            assert False
        return msg, note


# If we found the module, returns the path to the module as a str.
# Otherwise, returns the reason why the module wasn't found.
ModuleSearchResult = Union[str, ModuleNotFoundReason]


class BuildSource:
    """A single source file."""

    def __init__(self, path: Optional[str], module: Optional[str],
                 text: Optional[str] = None, base_dir: Optional[str] = None) -> None:
        self.path = path  # File where it's found (e.g. 'xxx/yyy/foo/bar.py')
        self.module = module or '__main__'  # Module name (e.g. 'foo.bar')
        self.text = text  # Source code, if initially supplied, else None
        self.base_dir = base_dir  # Directory where the package is rooted (e.g. 'xxx/yyy')

    def __repr__(self) -> str:
        return '<BuildSource path=%r module=%r has_text=%s base_dir=%s>' % (
            self.path,
            self.module,
            self.text is not None,
            self.base_dir)


class FindModuleCache:
    """Module finder with integrated cache.

    Module locations and some intermediate results are cached internally
    and can be cleared with the clear() method.

    All file system accesses are performed through a FileSystemCache,
    which is not ever cleared by this class. If necessary it must be
    cleared by client code.
    """

    def __init__(self,
                 search_paths: SearchPaths,
                 fscache: Optional[FileSystemCache] = None,
                 options: Optional[Options] = None,
                 ns_packages: Optional[List[str]] = None) -> None:
        self.search_paths = search_paths
        self.fscache = fscache or FileSystemCache()
        # Cache for get_toplevel_possibilities:
        # search_paths -> (toplevel_id -> list(package_dirs))
        self.initial_components = {}  # type: Dict[Tuple[str, ...], Dict[str, List[str]]]
        # Cache find_module: id -> result
        self.results = {}  # type: Dict[str, ModuleSearchResult]
        self.ns_ancestors = {}  # type: Dict[str, str]
        self.options = options
        self.ns_packages = ns_packages or []  # type: List[str]

    def clear(self) -> None:
        self.results.clear()
        self.initial_components.clear()
        self.ns_ancestors.clear()

    def find_lib_path_dirs(self, id: str, lib_path: Tuple[str, ...]) -> PackageDirs:
        """Find which elements of a lib_path have the directory a module needs to exist.

        This is run for the python_path, mypy_path, and typeshed_path search paths."""
        components = id.split('.')
        dir_chain = os.sep.join(components[:-1])  # e.g., 'foo/bar'

        dirs = []
        for pathitem in self.get_toplevel_possibilities(lib_path, components[0]):
            # e.g., '/usr/lib/python3.4/foo/bar'
            dir = os.path.normpath(os.path.join(pathitem, dir_chain))
            if self.fscache.isdir(dir):
                dirs.append((dir, True))
        return dirs

    def get_toplevel_possibilities(self, lib_path: Tuple[str, ...], id: str) -> List[str]:
        """Find which elements of lib_path could contain a particular top-level module.

        In practice, almost all modules can be routed to the correct entry in
        lib_path by looking at just the first component of the module name.

        We take advantage of this by enumerating the contents of all of the
        directories on the lib_path and building a map of which entries in
        the lib_path could contain each potential top-level module that appears.
        """

        if lib_path in self.initial_components:
            return self.initial_components[lib_path].get(id, [])

        # Enumerate all the files in the directories on lib_path and produce the map
        components = {}  # type: Dict[str, List[str]]
        for dir in lib_path:
            try:
                contents = self.fscache.listdir(dir)
            except OSError:
                contents = []
            # False positives are fine for correctness here, since we will check
            # precisely later, so we only look at the root of every filename without
            # any concern for the exact details.
            for name in contents:
                name = os.path.splitext(name)[0]
                components.setdefault(name, []).append(dir)

        self.initial_components[lib_path] = components
        return components.get(id, [])

    def find_module(self, id: str) -> ModuleSearchResult:
        """Return the path of the module source file or why it wasn't found."""
        if id not in self.results:
            self.results[id] = self._find_module(id)
        return self.results[id]

    def _find_module_non_stub_helper(self, components: List[str],
                                     pkg_dir: str) -> Union[OnePackageDir, ModuleNotFoundReason]:
        plausible_match = False
        dir_path = pkg_dir
        for index, component in enumerate(components):
            dir_path = os.path.join(dir_path, component)
            if self.fscache.isfile(os.path.join(dir_path, 'py.typed')):
                return os.path.join(pkg_dir, *components[:-1]), index == 0
            elif not plausible_match and (self.fscache.isdir(dir_path)
                                          or self.fscache.isfile(dir_path + ".py")):
                plausible_match = True
        if plausible_match:
            return ModuleNotFoundReason.FOUND_WITHOUT_TYPE_HINTS
        else:
            return ModuleNotFoundReason.NOT_FOUND

    def _update_ns_ancestors(self, components: List[str], match: Tuple[str, bool]) -> None:
        path, verify = match
        for i in range(1, len(components)):
            pkg_id = '.'.join(components[:-i])
            if pkg_id not in self.ns_ancestors and self.fscache.isdir(path):
                self.ns_ancestors[pkg_id] = path
            path = os.path.dirname(path)

    def _find_module(self, id: str) -> ModuleSearchResult:
        fscache = self.fscache

        # If we're looking for a module like 'foo.bar.baz', it's likely that most of the
        # many elements of lib_path don't even have a subdirectory 'foo/bar'.  Discover
        # that only once and cache it for when we look for modules like 'foo.bar.blah'
        # that will require the same subdirectory.
        components = id.split('.')
        dir_chain = os.sep.join(components[:-1])  # e.g., 'foo/bar'
        # TODO (ethanhs): refactor each path search to its own method with lru_cache

        # We have two sets of folders so that we collect *all* stubs folders and
        # put them in the front of the search path
        third_party_inline_dirs = []  # type: PackageDirs
        third_party_stubs_dirs = []  # type: PackageDirs
        found_possible_third_party_missing_type_hints = False
        # Third-party stub/typed packages
        for pkg_dir in self.search_paths.package_path:
            stub_name = components[0] + '-stubs'
            stub_dir = os.path.join(pkg_dir, stub_name)
            if fscache.isdir(stub_dir):
                stub_typed_file = os.path.join(stub_dir, 'py.typed')
                stub_components = [stub_name] + components[1:]
                path = os.path.join(pkg_dir, *stub_components[:-1])
                if fscache.isdir(path):
                    if fscache.isfile(stub_typed_file):
                        # Stub packages can have a py.typed file, which must include
                        # 'partial\n' to make the package partial
                        # Partial here means that mypy should look at the runtime
                        # package if installed.
                        if fscache.read(stub_typed_file).decode().strip() == 'partial':
                            runtime_path = os.path.join(pkg_dir, dir_chain)
                            third_party_inline_dirs.append((runtime_path, True))
                            # if the package is partial, we don't verify the module, as
                            # the partial stub package may not have a __init__.pyi
                            third_party_stubs_dirs.append((path, False))
                        else:
                            # handle the edge case where people put a py.typed file
                            # in a stub package, but it isn't partial
                            third_party_stubs_dirs.append((path, True))
                    else:
                        third_party_stubs_dirs.append((path, True))
            non_stub_match = self._find_module_non_stub_helper(components, pkg_dir)
            if isinstance(non_stub_match, ModuleNotFoundReason):
                if non_stub_match is ModuleNotFoundReason.FOUND_WITHOUT_TYPE_HINTS:
                    found_possible_third_party_missing_type_hints = True
            else:
                third_party_inline_dirs.append(non_stub_match)
                self._update_ns_ancestors(components, non_stub_match)
        if self.options and self.options.use_builtins_fixtures:
            # Everything should be in fixtures.
            third_party_inline_dirs.clear()
            third_party_stubs_dirs.clear()
            found_possible_third_party_missing_type_hints = False
        python_mypy_path = self.search_paths.mypy_path + self.search_paths.python_path
        candidate_base_dirs = self.find_lib_path_dirs(id, python_mypy_path) + \
            third_party_stubs_dirs + third_party_inline_dirs + \
            self.find_lib_path_dirs(id, self.search_paths.typeshed_path)

        # If we're looking for a module like 'foo.bar.baz', then candidate_base_dirs now
        # contains just the subdirectories 'foo/bar' that actually exist under the
        # elements of lib_path.  This is probably much shorter than lib_path itself.
        # Now just look for 'baz.pyi', 'baz/__init__.py', etc., inside those directories.
        seplast = os.sep + components[-1]  # so e.g. '/baz'
        sepinit = os.sep + '__init__'
        near_misses = []  # Collect near misses for namespace mode (see below).
        for base_dir, verify in candidate_base_dirs:
            base_path = base_dir + seplast  # so e.g. '/usr/lib/python3.4/foo/bar/baz'
            has_init = False
            dir_prefix = base_dir
            for _ in range(len(components) - 1):
                dir_prefix = os.path.dirname(dir_prefix)
            # Prefer package over module, i.e. baz/__init__.py* over baz.py*.
            for extension in PYTHON_EXTENSIONS:
                path = base_path + sepinit + extension
                path_stubs = base_path + '-stubs' + sepinit + extension
                if fscache.isfile_case(path, dir_prefix):
                    has_init = True
                    if verify and not verify_module(fscache, id, path, dir_prefix):
                        near_misses.append((path, dir_prefix))
                        continue
                    return path
                elif fscache.isfile_case(path_stubs, dir_prefix):
                    if verify and not verify_module(fscache, id, path_stubs, dir_prefix):
                        near_misses.append((path_stubs, dir_prefix))
                        continue
                    return path_stubs

            # In namespace mode, register a potential namespace package
            if self.options and self.options.namespace_packages:
                if fscache.isdir(base_path) and not has_init:
                    near_misses.append((base_path, dir_prefix))

            # No package, look for module.
            for extension in PYTHON_EXTENSIONS:
                path = base_path + extension
                if fscache.isfile_case(path, dir_prefix):
                    if verify and not verify_module(fscache, id, path, dir_prefix):
                        near_misses.append((path, dir_prefix))
                        continue
                    return path

        # In namespace mode, re-check those entries that had 'verify'.
        # Assume search path entries xxx, yyy and zzz, and we're
        # looking for foo.bar.baz.  Suppose near_misses has:
        #
        # - xxx/foo/bar/baz.py
        # - yyy/foo/bar/baz/__init__.py
        # - zzz/foo/bar/baz.pyi
        #
        # If any of the foo directories has __init__.py[i], it wins.
        # Else, we look for foo/bar/__init__.py[i], etc.  If there are
        # none, the first hit wins.  Note that this does not take into
        # account whether the lowest-level module is a file (baz.py),
        # a package (baz/__init__.py), or a stub file (baz.pyi) -- for
        # these the first one encountered along the search path wins.
        #
        # The helper function highest_init_level() returns an int that
        # indicates the highest level at which a __init__.py[i] file
        # is found; if no __init__ was found it returns 0, if we find
        # only foo/bar/__init__.py it returns 1, and if we have
        # foo/__init__.py it returns 2 (regardless of what's in
        # foo/bar).  It doesn't look higher than that.
        if self.options and self.options.namespace_packages and near_misses:
            levels = [highest_init_level(fscache, id, path, dir_prefix)
                      for path, dir_prefix in near_misses]
            index = levels.index(max(levels))
            return near_misses[index][0]

        # Finally, we may be asked to produce an ancestor for an
        # installed package with a py.typed marker that is a
        # subpackage of a namespace package.  We only fess up to these
        # if we would otherwise return "not found".
        ancestor = self.ns_ancestors.get(id)
        if ancestor is not None:
            return ancestor

        if found_possible_third_party_missing_type_hints:
            return ModuleNotFoundReason.FOUND_WITHOUT_TYPE_HINTS
        else:
            return ModuleNotFoundReason.NOT_FOUND

    def find_modules_recursive(self, module: str) -> List[BuildSource]:
        module_path = self.find_module(module)
        if isinstance(module_path, ModuleNotFoundReason):
            return []
        result = [BuildSource(module_path, module, None)]
        if module_path.endswith(('__init__.py', '__init__.pyi')):
            # Subtle: this code prefers the .pyi over the .py if both
            # exists, and also prefers packages over modules if both x/
            # and x.py* exist.  How?  We sort the directory items, so x
            # comes before x.py and x.pyi.  But the preference for .pyi
            # over .py is encoded in find_module(); even though we see
            # x.py before x.pyi, find_module() will find x.pyi first.  We
            # use hits to avoid adding it a second time when we see x.pyi.
            # This also avoids both x.py and x.pyi when x/ was seen first.
            hits = set()  # type: Set[str]
            for item in sorted(self.fscache.listdir(os.path.dirname(module_path))):
                abs_path = os.path.join(os.path.dirname(module_path), item)
                if os.path.isdir(abs_path) and \
                        (os.path.isfile(os.path.join(abs_path, '__init__.py')) or
                        os.path.isfile(os.path.join(abs_path, '__init__.pyi'))):
                    hits.add(item)
                    result += self.find_modules_recursive(module + '.' + item)
                elif item != '__init__.py' and item != '__init__.pyi' and \
                        item.endswith(('.py', '.pyi')):
                    mod = item.split('.')[0]
                    if mod not in hits:
                        hits.add(mod)
                        result += self.find_modules_recursive(module + '.' + mod)
        elif os.path.isdir(module_path) and module in self.ns_packages:
            # Even more subtler: handle recursive decent into PEP 420
            # namespace packages that are explicitly listed on the command
            # line with -p/--packages.
            for item in sorted(self.fscache.listdir(module_path)):
                if os.path.isdir(os.path.join(module_path, item)):
                    result += self.find_modules_recursive(module + '.' + item)
        return result


def verify_module(fscache: FileSystemCache, id: str, path: str, prefix: str) -> bool:
    """Check that all packages containing id have a __init__ file."""
    if path.endswith(('__init__.py', '__init__.pyi')):
        path = os.path.dirname(path)
    for i in range(id.count('.')):
        path = os.path.dirname(path)
        if not any(fscache.isfile_case(os.path.join(path, '__init__{}'.format(extension)),
                                       prefix)
                   for extension in PYTHON_EXTENSIONS):
            return False
    return True


def highest_init_level(fscache: FileSystemCache, id: str, path: str, prefix: str) -> int:
    """Compute the highest level where an __init__ file is found."""
    if path.endswith(('__init__.py', '__init__.pyi')):
        path = os.path.dirname(path)
    level = 0
    for i in range(id.count('.')):
        path = os.path.dirname(path)
        if any(fscache.isfile_case(os.path.join(path, '__init__{}'.format(extension)),
                                   prefix)
               for extension in PYTHON_EXTENSIONS):
            level = i + 1
    return level


def mypy_path() -> List[str]:
    path_env = os.getenv('MYPYPATH')
    if not path_env:
        return []
    return path_env.split(os.pathsep)


def default_lib_path(data_dir: str,
                     pyversion: Tuple[int, int],
                     custom_typeshed_dir: Optional[str]) -> List[str]:
    """Return default standard library search paths."""
    # IDEA: Make this more portable.
    path = []  # type: List[str]

    if custom_typeshed_dir:
        typeshed_dir = custom_typeshed_dir
    else:
        auto = os.path.join(data_dir, 'stubs-auto')
        if os.path.isdir(auto):
            data_dir = auto
        typeshed_dir = os.path.join(data_dir, "typeshed")
    if pyversion[0] == 3:
        # We allow a module for e.g. version 3.5 to be in 3.4/. The assumption
        # is that a module added with 3.4 will still be present in Python 3.5.
        versions = ["%d.%d" % (pyversion[0], minor)
                    for minor in reversed(range(PYTHON3_VERSION_MIN[1], pyversion[1] + 1))]
    else:
        # For Python 2, we only have stubs for 2.7
        versions = ["2.7"]
    # E.g. for Python 3.6, try 3.6/, 3.5/, 3.4/, 3/, 2and3/.
    for v in versions + [str(pyversion[0]), '2and3']:
        for lib_type in ['stdlib', 'third_party']:
            stubdir = os.path.join(typeshed_dir, lib_type, v)
            if os.path.isdir(stubdir):
                path.append(stubdir)

    # Add fallback path that can be used if we have a broken installation.
    if sys.platform != 'win32':
        path.append('/usr/local/lib/mypy')
    if not path:
        print("Could not resolve typeshed subdirectories. If you are using mypy\n"
              "from source, you need to run \"git submodule update --init\".\n"
              "Otherwise your mypy install is broken.\nPython executable is located at "
              "{0}.\nMypy located at {1}".format(sys.executable, data_dir), file=sys.stderr)
        sys.exit(1)
    return path


@functools.lru_cache(maxsize=None)
def get_site_packages_dirs(python_executable: Optional[str]) -> Tuple[List[str], List[str]]:
    """Find package directories for given python.

    This runs a subprocess call, which generates a list of the egg directories, and the site
    package directories. To avoid repeatedly calling a subprocess (which can be slow!) we
    lru_cache the results."""
    def make_abspath(path: str, root: str) -> str:
        """Take a path and make it absolute relative to root if not already absolute."""
        if os.path.isabs(path):
            return os.path.normpath(path)
        else:
            return os.path.join(root, os.path.normpath(path))

    if python_executable is None:
        return [], []
    if python_executable == sys.executable:
        # Use running Python's package dirs
        site_packages = sitepkgs.getsitepackages()
    else:
        # Use subprocess to get the package directory of given Python
        # executable
        site_packages = ast.literal_eval(
            subprocess.check_output([python_executable, sitepkgs.__file__],
            stderr=subprocess.PIPE).decode())
    egg_dirs = []
    for dir in site_packages:
        pth = os.path.join(dir, 'easy-install.pth')
        if os.path.isfile(pth):
            with open(pth) as f:
                egg_dirs.extend([make_abspath(d.rstrip(), dir) for d in f.readlines()])
    return egg_dirs, site_packages


def compute_search_paths(sources: List[BuildSource],
                         options: Options,
                         data_dir: str,
                         alt_lib_path: Optional[str] = None) -> SearchPaths:
    """Compute the search paths as specified in PEP 561.

    There are the following 4 members created:
    - User code (from `sources`)
    - MYPYPATH (set either via config or environment variable)
    - installed package directories (which will later be split into stub-only and inline)
    - typeshed
     """
    # Determine the default module search path.
    lib_path = collections.deque(
        default_lib_path(data_dir,
                         options.python_version,
                         custom_typeshed_dir=options.custom_typeshed_dir))

    if options.use_builtins_fixtures:
        # Use stub builtins (to speed up test cases and to make them easier to
        # debug).  This is a test-only feature, so assume our files are laid out
        # as in the source tree.
        # We also need to allow overriding where to look for it. Argh.
        root_dir = os.getenv('MYPY_TEST_PREFIX', None)
        if not root_dir:
            root_dir = os.path.dirname(os.path.dirname(__file__))
        lib_path.appendleft(os.path.join(root_dir, 'test-data', 'unit', 'lib-stub'))
    # alt_lib_path is used by some tests to bypass the normal lib_path mechanics.
    # If we don't have one, grab directories of source files.
    python_path = []  # type: List[str]
    if not alt_lib_path:
        for source in sources:
            # Include directory of the program file in the module search path.
            if source.base_dir:
                dir = source.base_dir
                if dir not in python_path:
                    python_path.append(dir)

        # Do this even if running as a file, for sanity (mainly because with
        # multiple builds, there could be a mix of files/modules, so its easier
        # to just define the semantics that we always add the current director
        # to the lib_path
        # TODO: Don't do this in some cases; for motivation see see
        # https://github.com/python/mypy/issues/4195#issuecomment-341915031
        if options.bazel:
            dir = '.'
        else:
            dir = os.getcwd()
        if dir not in lib_path:
            python_path.insert(0, dir)

    # Start with a MYPYPATH environment variable at the front of the mypy_path, if defined.
    mypypath = mypy_path()

    # Add a config-defined mypy path.
    mypypath.extend(options.mypy_path)

    # If provided, insert the caller-supplied extra module path to the
    # beginning (highest priority) of the search path.
    if alt_lib_path:
        mypypath.insert(0, alt_lib_path)

    egg_dirs, site_packages = get_site_packages_dirs(options.python_executable)
    for site_dir in site_packages:
        assert site_dir not in lib_path
        if (site_dir in mypypath or
                any(p.startswith(site_dir + os.path.sep) for p in mypypath) or
                os.path.altsep and any(p.startswith(site_dir + os.path.altsep) for p in mypypath)):
            print("{} is in the MYPYPATH. Please remove it.".format(site_dir), file=sys.stderr)
            print("See https://mypy.readthedocs.io/en/latest/running_mypy.html"
                  "#how-mypy-handles-imports for more info", file=sys.stderr)
            sys.exit(1)
        elif site_dir in python_path:
            print("{} is in the PYTHONPATH. Please change directory"
                  " so it is not.".format(site_dir),
                  file=sys.stderr)
            sys.exit(1)

    return SearchPaths(tuple(reversed(python_path)),
                       tuple(mypypath),
                       tuple(egg_dirs + site_packages),
                       tuple(lib_path))
