"""
API for analyzing Python stubs using mypy.
"""
import os
from os.path import abspath
from typing import Iterable, Optional, Set

from mypy.build import BuildResult, State, build
from mypy.main import process_options
from mypy.nodes import (
    GDEF,
    MDEF,
    Decorator,
    FuncDef,
    MypyFile,
    OverloadedFuncDef,
    SymbolNode,
    TypeAlias,
    TypeInfo,
    TypeVarExpr,
    Var,
)

from .types import RelevantSymbolNode

IGNORED_MODULE_SYMBOLS = [
    # because docstrings are not generated by stubgen
    "__doc__",
    # created by the import system, do not need to be handwritten:
    "__name__",
    "__package__",
    "__file__",
]
"""
Module level definitions that will not be collected by collect_types.
"""


def _mypy_analyze(
    mypy_conf_path: str, root_path: str, stubs_path: Optional[str] = None
) -> BuildResult:
    """
    Parse and analyze the types of the code in root_path.

    :param mypy_conf_path: path to a mypy.ini
    :param root_path: path to the code directory where the type analysis is started
    :param stubs_path: path to the directory of stubs for mypy to use
    :returns: Mypy's analysis result
    """
    # The call to `build.build` is inspired by `mypy/mypy/main.py::main`
    # `build` is not a documented public API
    args = [
        "--config-file",
        mypy_conf_path,
        "--no-incremental",
        "--cache-dir=" + ("nul" if os.name == "nt" else "/dev/null"),
        root_path,
    ]
    sources, options = process_options(args)
    if stubs_path is not None:
        options = options.apply_changes({"mypy_path": [stubs_path]})

    return build(sources, options, None, None)


def is_stubbed_module(module: State) -> bool:
    """
    Check if a module's types were loaded from a stub.

    :param module: The module to check
    """
    return module.path is not None and module.path.endswith(".pyi")


def collect_types(
    symbol_node: SymbolNode, collected_types: Optional[Set[str]] = None
) -> Iterable[RelevantSymbolNode]:
    """
    Collect all relevant type definitions of the symbols in the given node.

    :param symbol_node: any symbol node, e.g. MypyFile (BuildResult.graph.tree)
    :param collected_types: used to avoid collecting duplicates
    """
    if not collected_types:
        collected_types = set()

    # ignore builtins because we don't provide stubs for them
    if "builtins" in symbol_node.fullname():
        return

    # do not collect types twice
    if symbol_node.fullname() in collected_types:
        return
    collected_types.add(symbol_node.fullname())

    if isinstance(symbol_node, MypyFile):
        # the symbol node represents a Python module
        for symbol in symbol_node.names.values():
            # only global and class member definitions are interesting
            if symbol.kind not in [GDEF, MDEF]:
                continue

            if symbol.node and symbol.node.name() in IGNORED_MODULE_SYMBOLS:
                continue

            if symbol.node and symbol.module_public:
                yield from collect_types(symbol.node, collected_types)
    elif isinstance(symbol_node, TypeInfo):
        # the symbol represents a class definition
        yield symbol_node
        for class_member in symbol_node.names.values():
            if class_member.node:
                yield from collect_types(class_member.node, collected_types)
    elif isinstance(
        symbol_node,
        (Decorator, FuncDef, OverloadedFuncDef, Var, TypeAlias, TypeVarExpr),
    ):
        # the symbol represents a function definition,
        # variable, type alias or generic TypeVar
        yield symbol_node
    else:
        assert False, f"Unexpected symbol type {type(symbol_node)}"


def get_stub_types(
    stubs_path: str, mypy_conf_path: str, root_path: Optional[str] = None
) -> Iterable[RelevantSymbolNode]:
    """
    Analyze the stub files in stubs_path and return module
    and class definitions of stubs as symbol nodes.

    Only relevant symbol nodes (e.g. for variables, functions, classes, methods)
    are returned.
    They contain the type annotation information.

    :param stubs_path: where all the stub files are located
    :param mypy_conf_path: path to mypy.ini
    :param root_path: path to the code directory where the type analysis is started
    """
    stubs_path = abspath(stubs_path)

    if root_path:
        build_result = _mypy_analyze(mypy_conf_path, root_path, stubs_path)
    else:
        build_result = _mypy_analyze(mypy_conf_path, stubs_path)

    stubbed_modules = {
        module
        for module in build_result.graph.values()
        if module.path
        and is_stubbed_module(module)
        and module.path.startswith(stubs_path)
    }

    for module in stubbed_modules:
        if module.tree:
            yield from collect_types(module.tree)
