# Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
# For details: https://github.com/PyCQA/astroid/blob/main/LICENSE
"""
Astroid hook for the dataclasses library

Support both built-in dataclasses and pydantic.dataclasses. References:
- https://docs.python.org/3/library/dataclasses.html
- https://pydantic-docs.helpmanual.io/usage/dataclasses/
"""
from typing import Generator, List, Optional, Tuple

from astroid import context, inference_tip
from astroid.builder import parse
from astroid.const import PY37_PLUS, PY39_PLUS
from astroid.exceptions import AstroidSyntaxError, InferenceError, MroError
from astroid.manager import AstroidManager
from astroid.nodes.node_classes import (
    AnnAssign,
    AssignName,
    Attribute,
    Call,
    Name,
    NodeNG,
    Subscript,
    Unknown,
)
from astroid.nodes.scoped_nodes import ClassDef, FunctionDef
from astroid.util import Uninferable

DATACLASSES_DECORATORS = frozenset(("dataclass",))
FIELD_NAME = "field"
DATACLASS_MODULE = "dataclasses"
DEFAULT_FACTORY = "_HAS_DEFAULT_FACTORY"  # based on typing.py


def is_decorated_with_dataclass(node, decorator_names=DATACLASSES_DECORATORS):
    """Return True if a decorated node has a `dataclass` decorator applied."""
    if not isinstance(node, ClassDef) or not node.decorators:
        return False

    return any(
        _looks_like_dataclass_decorator(decorator_attribute, decorator_names)
        for decorator_attribute in node.decorators.nodes
    )


def dataclass_transform(node: ClassDef) -> None:
    """Rewrite a dataclass to be easily understood by pylint"""

    for assign_node in _get_dataclass_attributes(node):
        name = assign_node.target.name

        rhs_node = Unknown(
            lineno=assign_node.lineno,
            col_offset=assign_node.col_offset,
            parent=assign_node,
        )
        rhs_node = AstroidManager().visit_transforms(rhs_node)
        node.instance_attrs[name] = [rhs_node]

    if not _check_generate_dataclass_init(node):
        return

    try:
        reversed_mro = reversed(node.mro())
    except MroError:
        reversed_mro = [node]

    field_assigns = {}
    field_order = []
    for klass in (k for k in reversed_mro if is_decorated_with_dataclass(k)):
        for assign_node in _get_dataclass_attributes(klass, init=True):
            name = assign_node.target.name
            if name not in field_assigns:
                field_order.append(name)
            field_assigns[name] = assign_node

    init_str = _generate_dataclass_init([field_assigns[name] for name in field_order])
    try:
        init_node = parse(init_str)["__init__"]
    except AstroidSyntaxError:
        pass
    else:
        init_node.parent = node
        init_node.lineno, init_node.col_offset = None, None
        node.locals["__init__"] = [init_node]

        root = node.root()
        if DEFAULT_FACTORY not in root.locals:
            new_assign = parse(f"{DEFAULT_FACTORY} = object()").body[0]
            new_assign.parent = root
            root.locals[DEFAULT_FACTORY] = [new_assign.targets[0]]


def _get_dataclass_attributes(node: ClassDef, init: bool = False) -> Generator:
    """Yield the AnnAssign nodes of dataclass attributes for the node.

    If init is True, also include InitVars, but exclude attributes from calls to
    field where init=False.
    """
    for assign_node in node.body:
        if not isinstance(assign_node, AnnAssign) or not isinstance(
            assign_node.target, AssignName
        ):
            continue

        if _is_class_var(assign_node.annotation):
            continue

        if init:
            value = assign_node.value
            if (
                isinstance(value, Call)
                and _looks_like_dataclass_field_call(value, check_scope=False)
                and any(
                    keyword.arg == "init" and not keyword.value.bool_value()
                    for keyword in value.keywords
                )
            ):
                continue
        elif _is_init_var(assign_node.annotation):
            continue

        yield assign_node


def _check_generate_dataclass_init(node: ClassDef) -> bool:
    """Return True if we should generate an __init__ method for node.

    This is True when:
        - node doesn't define its own __init__ method
        - the dataclass decorator was called *without* the keyword argument init=False
    """
    if "__init__" in node.locals:
        return False

    found = None

    for decorator_attribute in node.decorators.nodes:
        if not isinstance(decorator_attribute, Call):
            continue

        if _looks_like_dataclass_decorator(decorator_attribute):
            found = decorator_attribute

    if found is None:
        return True

    # Check for keyword arguments of the form init=False
    return all(
        keyword.arg != "init" or keyword.value.bool_value()
        for keyword in found.keywords
    )


def _generate_dataclass_init(assigns: List[AnnAssign]) -> str:
    """Return an init method for a dataclass given the targets."""
    target_names = []
    params = []
    assignments = []

    for assign in assigns:
        name, annotation, value = assign.target.name, assign.annotation, assign.value
        target_names.append(name)

        if _is_init_var(annotation):
            init_var = True
            if isinstance(annotation, Subscript):
                annotation = annotation.slice
            else:
                # Cannot determine type annotation for parameter from InitVar
                annotation = None
            assignment_str = ""
        else:
            init_var = False
            assignment_str = f"self.{name} = {name}"

        if annotation:
            param_str = f"{name}: {annotation.as_string()}"
        else:
            param_str = name

        if value:
            if isinstance(value, Call) and _looks_like_dataclass_field_call(
                value, check_scope=False
            ):
                result = _get_field_default(value)

                default_type, default_node = result
                if default_type == "default":
                    param_str += f" = {default_node.as_string()}"
                elif default_type == "default_factory":
                    param_str += f" = {DEFAULT_FACTORY}"
                    assignment_str = (
                        f"self.{name} = {default_node.as_string()} "
                        f"if {name} is {DEFAULT_FACTORY} else {name}"
                    )
            else:
                param_str += f" = {value.as_string()}"

        params.append(param_str)
        if not init_var:
            assignments.append(assignment_str)

    params = ", ".join(["self"] + params)
    assignments = "\n    ".join(assignments) if assignments else "pass"
    return f"def __init__({params}) -> None:\n    {assignments}"


def infer_dataclass_attribute(
    node: Unknown, ctx: context.InferenceContext = None
) -> Generator:
    """Inference tip for an Unknown node that was dynamically generated to
    represent a dataclass attribute.

    In the case that a default value is provided, that is inferred first.
    Then, an Instance of the annotated class is yielded.
    """
    assign = node.parent
    if not isinstance(assign, AnnAssign):
        yield Uninferable
        return

    annotation, value = assign.annotation, assign.value
    if value is not None:
        yield from value.infer(context=ctx)
    if annotation is not None:
        yield from _infer_instance_from_annotation(annotation, ctx=ctx)
    else:
        yield Uninferable


def infer_dataclass_field_call(
    node: AssignName, ctx: context.InferenceContext = None
) -> Generator:
    """Inference tip for dataclass field calls."""
    field_call = node.parent.value
    default_type, default = _get_field_default(field_call)
    if not default_type:
        yield Uninferable
    elif default_type == "default":
        yield from default.infer(context=ctx)
    else:
        new_call = parse(default.as_string()).body[0].value
        new_call.parent = field_call.parent
        yield from new_call.infer(context=ctx)


def _looks_like_dataclass_decorator(
    node: NodeNG, decorator_names: List[str] = DATACLASSES_DECORATORS
) -> bool:
    """Return True if node looks like a dataclass decorator.

    Uses inference to lookup the value of the node, and if that fails,
    matches against specific names. (Currently inference fails for dataclass import
    from pydantic.)
    """
    if isinstance(node, Call):  # decorator with arguments
        node = node.func
    try:
        inferred = next(node.infer())
    except (InferenceError, StopIteration):
        inferred = Uninferable

    if inferred is Uninferable:
        if isinstance(node, Name):
            return node.name in decorator_names
        if isinstance(node, Attribute):
            return node.attrname in decorator_names

        return False

    return (
        isinstance(inferred, FunctionDef)
        and inferred.name in decorator_names
        and inferred.root().name == DATACLASS_MODULE
    )


def _looks_like_dataclass_attribute(node: Unknown) -> bool:
    """Return True if node was dynamically generated as the child of an AnnAssign
    statement.
    """
    parent = node.parent
    scope = parent.scope()
    return (
        isinstance(parent, AnnAssign)
        and isinstance(scope, ClassDef)
        and is_decorated_with_dataclass(scope)
    )


def _looks_like_dataclass_field_call(node: Call, check_scope: bool = True) -> bool:
    """Return True if node is calling dataclasses field or Field
    from an AnnAssign statement directly in the body of a ClassDef.

    If check_scope is False, skips checking the statement and body.
    """
    if check_scope:
        stmt = node.statement()
        scope = stmt.scope()
        if not (
            isinstance(stmt, AnnAssign)
            and isinstance(scope, ClassDef)
            and is_decorated_with_dataclass(scope)
        ):
            return False

    try:
        inferred = next(node.func.infer())
    except (InferenceError, StopIteration):
        return False

    if not isinstance(inferred, FunctionDef):
        return False

    return inferred.name == FIELD_NAME and inferred.root().name == DATACLASS_MODULE


def _get_field_default(field_call: Call) -> Tuple[str, Optional[NodeNG]]:
    """Return a the default value of a field call, and the corresponding keyword argument name.

    field(default=...) results in the ... node
    field(default_factory=...) results in a Call node with func ... and no arguments

    If neither or both arguments are present, return ("", None) instead,
    indicating that there is not a valid default value.
    """
    default, default_factory = None, None
    for keyword in field_call.keywords:
        if keyword.arg == "default":
            default = keyword.value
        elif keyword.arg == "default_factory":
            default_factory = keyword.value

    if default is not None and default_factory is None:
        return "default", default

    if default is None and default_factory is not None:
        new_call = Call(
            lineno=field_call.lineno,
            col_offset=field_call.col_offset,
            parent=field_call.parent,
        )
        new_call.postinit(func=default_factory)
        return "default_factory", new_call

    return "", None


def _is_class_var(node: NodeNG) -> bool:
    """Return True if node is a ClassVar, with or without subscripting."""
    if PY39_PLUS:
        try:
            inferred = next(node.infer())
        except (InferenceError, StopIteration):
            return False

        return getattr(inferred, "name", "") == "ClassVar"

    # Before Python 3.9, inference returns typing._SpecialForm instead of ClassVar.
    # Our backup is to inspect the node's structure.
    return isinstance(node, Subscript) and (
        isinstance(node.value, Name)
        and node.value.name == "ClassVar"
        or isinstance(node.value, Attribute)
        and node.value.attrname == "ClassVar"
    )


def _is_init_var(node: NodeNG) -> bool:
    """Return True if node is an InitVar, with or without subscripting."""
    try:
        inferred = next(node.infer())
    except (InferenceError, StopIteration):
        return False

    return getattr(inferred, "name", "") == "InitVar"


# Allowed typing classes for which we support inferring instances
_INFERABLE_TYPING_TYPES = frozenset(
    (
        "Dict",
        "FrozenSet",
        "List",
        "Set",
        "Tuple",
    )
)


def _infer_instance_from_annotation(
    node: NodeNG, ctx: context.InferenceContext = None
) -> Generator:
    """Infer an instance corresponding to the type annotation represented by node.

    Currently has limited support for the typing module.
    """
    klass = None
    try:
        klass = next(node.infer(context=ctx))
    except (InferenceError, StopIteration):
        yield Uninferable
    if not isinstance(klass, ClassDef):
        yield Uninferable
    elif klass.root().name in (
        "typing",
        "_collections_abc",
        "",
    ):  # "" because of synthetic nodes in brain_typing.py
        if klass.name in _INFERABLE_TYPING_TYPES:
            yield klass.instantiate_class()
        else:
            yield Uninferable
    else:
        yield klass.instantiate_class()


if PY37_PLUS:
    AstroidManager().register_transform(
        ClassDef, dataclass_transform, is_decorated_with_dataclass
    )

    AstroidManager().register_transform(
        Call,
        inference_tip(infer_dataclass_field_call, raise_on_overwrite=True),
        _looks_like_dataclass_field_call,
    )

    AstroidManager().register_transform(
        Unknown,
        inference_tip(infer_dataclass_attribute, raise_on_overwrite=True),
        _looks_like_dataclass_attribute,
    )
