from mypy.nodes import (
    REVEAL_TYPE,
    AssertStmt,
    AssertTypeExpr,
    AssignmentExpr,
    AssignmentStmt,
    AwaitExpr,
    Block,
    CallExpr,
    CastExpr,
    ClassDef,
    ComparisonExpr,
    ConditionalExpr,
    Decorator,
    DelStmt,
    DictExpr,
    DictionaryComprehension,
    ExpressionStmt,
    ForStmt,
    FuncDef,
    FuncItem,
    GeneratorExpr,
    IfStmt,
    Import,
    ImportFrom,
    IndexExpr,
    LambdaExpr,
    ListComprehension,
    ListExpr,
    MatchStmt,
    MemberExpr,
    MypyFile,
    Node,
    OperatorAssignmentStmt,
    OpExpr,
    OverloadedFuncDef,
    RaiseStmt,
    ReturnStmt,
    RevealExpr,
    SetComprehension,
    SetExpr,
    SliceExpr,
    StarExpr,
    SuperExpr,
    TryStmt,
    TupleExpr,
    TypeApplication,
    UnaryExpr,
    WhileStmt,
    WithStmt,
    YieldExpr,
    YieldFromExpr,
)
from mypy.patterns import (
    AsPattern,
    ClassPattern,
    MappingPattern,
    OrPattern,
    SequencePattern,
    StarredPattern,
    ValuePattern,
)


# an inversion of mypy's visitor pattern to allow compatibility in mypy's compiled mode.
# instead of node.accept(visitor), which breaks when the visitor is an extension of any compiled mypy visitor class,
# we call traverser.visit_node(node).
# Implementations for each node type are taken from mypy.traverser.TraverserVisitor, and adapted to this pattern.
class DirectTraverser:
    # Implementation for each case taken from the accept() method for each node class
    def visit_node(self, node: Node) -> None:
        match node:
            case MypyFile():
                self.visit_mypy_file(node)
            case Block():
                self.visit_block(node)
            case FuncItem():
                self.visit_func(node)
            case FuncDef():
                self.visit_func_def(node)
            case OverloadedFuncDef():
                self.visit_overloaded_func_def(node)
            case ClassDef():
                self.visit_class_def(node)
            case Decorator():
                self.visit_decorator(node)
            case ExpressionStmt():
                self.visit_expression_stmt(node)
            case AssignmentStmt():
                self.visit_assignment_stmt(node)
            case OperatorAssignmentStmt():
                self.visit_operator_assignment_stmt(node)
            case WhileStmt():
                self.visit_while_stmt(node)
            case ForStmt():
                self.visit_for_stmt(node)
            case ReturnStmt():
                self.visit_return_stmt(node)
            case AssertStmt():
                self.visit_assert_stmt(node)
            case DelStmt():
                self.visit_del_stmt(node)
            case IfStmt():
                self.visit_if_stmt(node)
            case RaiseStmt():
                self.visit_raise_stmt(node)
            case TryStmt():
                self.visit_try_stmt(node)
            case WithStmt():
                self.visit_with_stmt(node)
            case MatchStmt():
                self.visit_match_stmt(node)
            case MemberExpr():
                self.visit_member_expr(node)
            case YieldFromExpr():
                self.visit_yield_from_expr(node)
            case YieldExpr():
                self.visit_yield_expr(node)
            case CallExpr():
                self.visit_call_expr(node)
            case OpExpr():
                self.visit_op_expr(node)
            case ComparisonExpr():
                self.visit_comparison_expr(node)
            case SliceExpr():
                self.visit_slice_expr(node)
            case CastExpr():
                self.visit_cast_expr(node)
            case AssertTypeExpr():
                self.visit_assert_type_expr(node)
            case RevealExpr():
                self.visit_reveal_expr(node)
            case AssignmentExpr():
                self.visit_assignment_expr(node)
            case UnaryExpr():
                self.visit_unary_expr(node)
            case ListExpr():
                self.visit_list_expr(node)
            case TupleExpr():
                self.visit_tuple_expr(node)
            case DictExpr():
                self.visit_dict_expr(node)
            case SetExpr():
                self.visit_set_expr(node)
            case IndexExpr():
                self.visit_index_expr(node)
            case GeneratorExpr():
                self.visit_generator_expr(node)
            case DictionaryComprehension():
                self.visit_dictionary_comprehension(node)
            case ListComprehension():
                self.visit_list_comprehension(node)
            case SetComprehension():
                self.visit_set_comprehension(node)
            case ConditionalExpr():
                self.visit_conditional_expr(node)
            case TypeApplication():
                self.visit_type_application(node)
            case LambdaExpr():
                self.visit_lambda_expr(node)
            case StarExpr():
                self.visit_star_expr(node)
            case AwaitExpr():
                self.visit_await_expr(node)
            case SuperExpr():
                self.visit_super_expr(node)
            case AsPattern():
                self.visit_as_pattern(node)
            case OrPattern():
                self.visit_or_pattern(node)
            case ValuePattern():
                self.visit_value_pattern(node)
            case SequencePattern():
                self.visit_sequence_pattern(node)
            case StarredPattern():
                self.visit_starred_pattern(node)
            case MappingPattern():
                self.visit_mapping_pattern(node)
            case ClassPattern():
                self.visit_class_pattern(node)
            case Import():
                self.visit_import(node)
            case ImportFrom():
                self.visit_import_from(node)

            # if we hit a node type that TraverserVisitor doesn't implement, just pass
            # as the node is either effectively or explicitly a leaf node in the AST
            case _:
                pass

    def visit_mypy_file(self, o: MypyFile) -> None:
        for d in o.defs:
            self.visit_node(d)

    def visit_block(self, block: Block) -> None:
        for s in block.body:
            self.visit_node(s)

    def visit_func(self, o: FuncItem) -> None:
        if o.arguments is not None:
            for arg in o.arguments:
                init = arg.initializer
                if init is not None:
                    self.visit_node(init)

        self.visit_node(o.body)

    def visit_func_def(self, o: FuncDef) -> None:
        self.visit_func(o)

    def visit_overloaded_func_def(self, o: OverloadedFuncDef) -> None:
        for item in o.items:
            self.visit_node(item)
        if o.impl:
            self.visit_node(o.impl)

    def visit_class_def(self, o: ClassDef) -> None:
        for d in o.decorators:
            self.visit_node(d)
        for base in o.base_type_exprs:
            self.visit_node(base)
        if o.metaclass:
            self.visit_node(o.metaclass)
        for v in o.keywords.values():
            self.visit_node(v)
        self.visit_node(o.defs)
        if o.analyzed:
            self.visit_node(o.analyzed)

    def visit_decorator(self, o: Decorator) -> None:
        self.visit_node(o.func)
        self.visit_node(o.var)
        for decorator in o.decorators:
            self.visit_node(decorator)

    def visit_expression_stmt(self, o: ExpressionStmt) -> None:
        self.visit_node(o.expr)

    def visit_assignment_stmt(self, o: AssignmentStmt) -> None:
        self.visit_node(o.rvalue)
        for l in o.lvalues:
            self.visit_node(l)

    def visit_operator_assignment_stmt(self, o: OperatorAssignmentStmt) -> None:
        self.visit_node(o.rvalue)
        self.visit_node(o.lvalue)

    def visit_while_stmt(self, o: WhileStmt) -> None:
        self.visit_node(o.expr)
        self.visit_node(o.body)
        if o.else_body:
            self.visit_node(o.else_body)

    def visit_for_stmt(self, o: ForStmt) -> None:
        self.visit_node(o.index)
        self.visit_node(o.expr)
        self.visit_node(o.body)
        if o.else_body:
            self.visit_node(o.else_body)

    def visit_return_stmt(self, o: ReturnStmt) -> None:
        if o.expr is not None:
            self.visit_node(o.expr)

    def visit_assert_stmt(self, o: AssertStmt) -> None:
        if o.expr is not None:
            self.visit_node(o.expr)
        if o.msg is not None:
            self.visit_node(o.msg)

    def visit_del_stmt(self, o: DelStmt) -> None:
        if o.expr is not None:
            self.visit_node(o.expr)

    def visit_if_stmt(self, o: IfStmt) -> None:
        for e in o.expr:
            self.visit_node(e)
        for b in o.body:
            self.visit_node(b)
        if o.else_body:
            self.visit_node(o.else_body)

    def visit_raise_stmt(self, o: RaiseStmt) -> None:
        if o.expr is not None:
            self.visit_node(o.expr)
        if o.from_expr is not None:
            self.visit_node(o.from_expr)

    def visit_try_stmt(self, o: TryStmt) -> None:
        self.visit_node(o.body)
        for i in range(len(o.types)):
            tp = o.types[i]
            if tp is not None:
                self.visit_node(tp)
            self.visit_node(o.handlers[i])
        for v in o.vars:
            if v is not None:
                self.visit_node(v)
        if o.else_body is not None:
            self.visit_node(o.else_body)
        if o.finally_body is not None:
            self.visit_node(o.finally_body)

    def visit_with_stmt(self, o: WithStmt) -> None:
        for i in range(len(o.expr)):
            self.visit_node(o.expr[i])
            targ = o.target[i]
            if targ is not None:
                self.visit_node(targ)
        self.visit_node(o.body)

    def visit_match_stmt(self, o: MatchStmt) -> None:
        self.visit_node(o.subject)
        for i in range(len(o.patterns)):
            self.visit_node(o.patterns[i])
            guard = o.guards[i]
            if guard is not None:
                self.visit_node(guard)
            self.visit_node(o.bodies[i])

    def visit_member_expr(self, o: MemberExpr) -> None:
        self.visit_node(o.expr)

    def visit_yield_from_expr(self, o: YieldFromExpr) -> None:
        self.visit_node(o.expr)

    def visit_yield_expr(self, o: YieldExpr) -> None:
        if o.expr:
            self.visit_node(o.expr)

    def visit_call_expr(self, o: CallExpr) -> None:
        self.visit_node(o.callee)
        for a in o.args:
            self.visit_node(a)
        if o.analyzed:
            self.visit_node(o.analyzed)

    def visit_op_expr(self, o: OpExpr) -> None:
        self.visit_node(o.left)
        self.visit_node(o.right)
        if o.analyzed is not None:
            self.visit_node(o.analyzed)

    def visit_comparison_expr(self, o: ComparisonExpr) -> None:
        for operand in o.operands:
            self.visit_node(operand)

    def visit_slice_expr(self, o: SliceExpr) -> None:
        if o.begin_index is not None:
            self.visit_node(o.begin_index)
        if o.end_index is not None:
            self.visit_node(o.end_index)
        if o.stride is not None:
            self.visit_node(o.stride)

    def visit_cast_expr(self, o: CastExpr) -> None:
        self.visit_node(o.expr)

    def visit_assert_type_expr(self, o: AssertTypeExpr) -> None:
        self.visit_node(o.expr)

    def visit_reveal_expr(self, o: RevealExpr) -> None:
        if o.kind == REVEAL_TYPE:
            assert o.expr is not None
            self.visit_node(o.expr)
        else:
            # RevealLocalsExpr doesn't have an inner expression
            pass

    def visit_assignment_expr(self, o: AssignmentExpr) -> None:
        self.visit_node(o.target)
        self.visit_node(o.value)

    def visit_unary_expr(self, o: UnaryExpr) -> None:
        self.visit_node(o.expr)

    def visit_list_expr(self, o: ListExpr) -> None:
        for item in o.items:
            self.visit_node(item)

    def visit_tuple_expr(self, o: TupleExpr) -> None:
        for item in o.items:
            self.visit_node(item)

    def visit_dict_expr(self, o: DictExpr) -> None:
        for k, v in o.items:
            if k is not None:
                self.visit_node(k)
            self.visit_node(v)

    def visit_set_expr(self, o: SetExpr) -> None:
        for item in o.items:
            self.visit_node(item)

    def visit_index_expr(self, o: IndexExpr) -> None:
        self.visit_node(o.base)
        self.visit_node(o.index)
        if o.analyzed:
            self.visit_node(o.analyzed)

    def visit_generator_expr(self, o: GeneratorExpr) -> None:
        for index, sequence, conditions in zip(o.indices, o.sequences, o.condlists):
            self.visit_node(sequence)
            self.visit_node(index)
            for cond in conditions:
                self.visit_node(cond)
        self.visit_node(o.left_expr)

    def visit_dictionary_comprehension(self, o: DictionaryComprehension) -> None:
        for index, sequence, conditions in zip(o.indices, o.sequences, o.condlists):
            self.visit_node(sequence)
            self.visit_node(index)
            for cond in conditions:
                self.visit_node(cond)
        self.visit_node(o.key)
        self.visit_node(o.value)

    def visit_list_comprehension(self, o: ListComprehension) -> None:
        self.visit_node(o.generator)

    def visit_set_comprehension(self, o: SetComprehension) -> None:
        self.visit_node(o.generator)

    def visit_conditional_expr(self, o: ConditionalExpr) -> None:
        self.visit_node(o.cond)
        self.visit_node(o.if_expr)
        self.visit_node(o.else_expr)

    def visit_type_application(self, o: TypeApplication) -> None:
        self.visit_node(o.expr)

    def visit_lambda_expr(self, o: LambdaExpr) -> None:
        self.visit_func(o)

    def visit_star_expr(self, o: StarExpr) -> None:
        self.visit_node(o.expr)

    def visit_await_expr(self, o: AwaitExpr) -> None:
        self.visit_node(o.expr)

    def visit_super_expr(self, o: SuperExpr) -> None:
        self.visit_node(o.call)

    def visit_as_pattern(self, o: AsPattern) -> None:
        if o.pattern is not None:
            self.visit_node(o.pattern)
        if o.name is not None:
            self.visit_node(o.name)

    def visit_or_pattern(self, o: OrPattern) -> None:
        for p in o.patterns:
            self.visit_node(p)

    def visit_value_pattern(self, o: ValuePattern) -> None:
        self.visit_node(o.expr)

    def visit_sequence_pattern(self, o: SequencePattern) -> None:
        for p in o.patterns:
            self.visit_node(p)

    def visit_starred_pattern(self, o: StarredPattern) -> None:
        if o.capture is not None:
            self.visit_node(o.capture)

    def visit_mapping_pattern(self, o: MappingPattern) -> None:
        for key in o.keys:
            self.visit_node(key)
        for value in o.values:
            self.visit_node(value)
        if o.rest is not None:
            self.visit_node(o.rest)

    def visit_class_pattern(self, o: ClassPattern) -> None:
        self.visit_node(o.class_ref)
        for p in o.positionals:
            self.visit_node(p)
        for v in o.keyword_values:
            self.visit_node(v)

    def visit_import(self, o: Import) -> None:
        for a in o.assignments:
            self.visit_node(a)

    def visit_import_from(self, o: ImportFrom) -> None:
        for a in o.assignments:
            self.visit_node(a)