from __future__ import annotations

from mypy.nodes import (
    AssertTypeExpr,
    AssignmentStmt,
    CastExpr,
    ClassDef,
    ForStmt,
    FuncItem,
    NamedTupleExpr,
    NewTypeExpr,
    PromoteExpr,
    TypeAlias,
    TypeAliasExpr,
    TypeAliasStmt,
    TypeApplication,
    TypedDictExpr,
    TypeFormExpr,
    TypeInfo,
    TypeVarExpr,
    Var,
    WithStmt,
)
from mypy.traverser import TraverserVisitor
from mypy.types import Type
from mypy.typetraverser import TypeTraverserVisitor


class MixedTraverserVisitor(TraverserVisitor, TypeTraverserVisitor):
    """Recursive traversal of both Node and Type objects."""

    def __init__(self) -> None:
        self.in_type_alias_expr = False

    # Symbol nodes

    def visit_var(self, var: Var, /) -> None:
        self.visit_optional_type(var.type)

    def visit_func(self, o: FuncItem, /) -> None:
        super().visit_func(o)
        self.visit_optional_type(o.type)

    def visit_class_def(self, o: ClassDef, /) -> None:
        super().visit_class_def(o)
        if o.info:
            self.process_type_info(o.info)

    def process_type_info(self, info: TypeInfo) -> None:
        # TODO: Should we visit generated methods/variables as well?
        # We should for methods generated by us (see below). But it is less clear for
        # 3rd party plugin generated methods (since we don't want to emit errors there).
        for base in info.bases:
            base.accept(self)
        if info.special_alias:
            # We need to accept all types that are conceptually identical like special
            # alias target and corresponding tuple_type or typeddict_type, since those
            # may be copies, and not the same object.
            info.special_alias.accept(self)
        if info.tuple_type:
            info.tuple_type.accept(self)
        if info.typeddict_type:
            info.typeddict_type.accept(self)
        if info.is_named_tuple or info.is_newtype:
            for sym in info.names.values():
                if sym.plugin_generated and sym.node:
                    sym.node.accept(self)

    def visit_type_alias_expr(self, o: TypeAliasExpr, /) -> None:
        super().visit_type_alias_expr(o)
        o.node.accept(self)

    def visit_type_var_expr(self, o: TypeVarExpr, /) -> None:
        super().visit_type_var_expr(o)
        o.upper_bound.accept(self)
        o.default.accept(self)
        for value in o.values:
            value.accept(self)

    def visit_typeddict_expr(self, o: TypedDictExpr, /) -> None:
        super().visit_typeddict_expr(o)
        self.process_type_info(o.info)

    def visit_namedtuple_expr(self, o: NamedTupleExpr, /) -> None:
        super().visit_namedtuple_expr(o)
        self.process_type_info(o.info)

    def visit__promote_expr(self, o: PromoteExpr, /) -> None:
        super().visit__promote_expr(o)
        o.type.accept(self)

    def visit_newtype_expr(self, o: NewTypeExpr, /) -> None:
        super().visit_newtype_expr(o)
        if o.info:
            self.process_type_info(o.info)
        self.visit_optional_type(o.old_type)

    # Statements

    def visit_assignment_stmt(self, o: AssignmentStmt, /) -> None:
        super().visit_assignment_stmt(o)
        self.visit_optional_type(o.type)

    def visit_type_alias_stmt(self, o: TypeAliasStmt, /) -> None:
        super().visit_type_alias_stmt(o)
        if o.alias_node is not None:
            o.alias_node.accept(self)

    def visit_type_alias(self, o: TypeAlias, /) -> None:
        super().visit_type_alias(o)
        self.in_type_alias_expr = True
        o.target.accept(self)
        self.in_type_alias_expr = False

    def visit_for_stmt(self, o: ForStmt, /) -> None:
        super().visit_for_stmt(o)
        self.visit_optional_type(o.index_type)

    def visit_with_stmt(self, o: WithStmt, /) -> None:
        super().visit_with_stmt(o)
        for typ in o.analyzed_types:
            typ.accept(self)

    # Expressions

    def visit_cast_expr(self, o: CastExpr, /) -> None:
        super().visit_cast_expr(o)
        o.type.accept(self)

    def visit_type_form_expr(self, o: TypeFormExpr, /) -> None:
        super().visit_type_form_expr(o)
        o.type.accept(self)

    def visit_assert_type_expr(self, o: AssertTypeExpr, /) -> None:
        super().visit_assert_type_expr(o)
        o.type.accept(self)

    def visit_type_application(self, o: TypeApplication, /) -> None:
        super().visit_type_application(o)
        for t in o.types:
            t.accept(self)

    # Helpers

    def visit_optional_type(self, t: Type | None, /) -> None:
        if t:
            t.accept(self)