# type: ignore
# ruff: noqa
import inspect
import typing
from textwrap import dedent

import mypy.build
import mypy.find_sources
import mypy.fscache
import mypy.nodes
import mypy.options
import mypy.types
import pytest
from mypy.nodes import RevealExpr

from puyapy.compile import get_mypy_options

def get_assignment_var_named(mypy_file: mypy.nodes.MypyFile, name: str) -> mypy.nodes.Var:
    for assignment in [
        stmt for stmt in mypy_file.defs if isinstance(stmt, mypy.nodes.AssignmentStmt)
    ]:
        for lvalue in assignment.lvalues:
            if isinstance(lvalue, mypy.nodes.NameExpr) and lvalue.name == name:
                assert isinstance(lvalue.node, mypy.nodes.Var)
                return lvalue.node
    raise Exception(f"Assignment to '{name}' not found")

def get_class_by_name(mypy_file: mypy.nodes.MypyFile, name: str) -> mypy.nodes.ClassDef:
    for cls_def in (stmt for stmt in mypy_file.defs if isinstance(stmt, mypy.nodes.ClassDef)):
        if cls_def.name == name:
            return cls_def
    raise Exception(f"Assignment to '{name}' not found")

def decompile(function: typing.Callable) -> str:
    source = inspect.getsource(function)
    return dedent("\n".join(source.splitlines()[1:]))

TEST_MODULE = "__test__"

def mypy_parse_and_type_check(source: str | typing.Callable) -> mypy.build.BuildResult:
    code = source if isinstance(source, str) else decompile(source)
    options = get_mypy_options()
    options.export_types = True

    sources = [mypy.build.BuildSource(None, TEST_MODULE, text=dedent(code))]

    return mypy.build.build(sources=sources, options=options)

def strip_error_prefixes(br: mypy.build.BuildResult) -> list[str]:
    return [e.split(":", maxsplit=2)[-1].strip() for e in br.errors]

def get_revealed_types(
    br: mypy.build.BuildResult, tree: mypy.nodes.MypyFile
) -> list[mypy.types.Type | None]:
    import mypy.traverser

    types = []

    class MyVisitor(mypy.traverser.TraverserVisitor):
        def visit_reveal_expr(self, o: RevealExpr) -> None:
            types.append(br.types.get(o.expr, mypy.types.UninhabitedType()))

    visitor = MyVisitor()
    visitor.visit_mypy_file(tree)
    return types

def test_ignore_assignment() -> None:
    def test():
        a_lie: int = "bar"  # type: ignore[assignment]

    result = mypy_parse_and_type_check(test)
    assert not result.errors
    module = result.graph[TEST_MODULE]
    assert module.tree
    assert module.tree.ignored_lines == {1: ["assignment"]}

def test_incorrect_assignment() -> None:
    def test():
        wrong: int = "bar"

    result = mypy_parse_and_type_check(test)
    assert result.errors
    assert result.errors[0].endswith("[assignment]")

def test_cast() -> None:
    def test():
        import typing

        wrong: int = typing.cast(int, "bar")

    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    assert tree
    var = get_assignment_var_named(tree, "wrong")
    assert str(var.type) == "builtins.int"

def test_implicit_cast() -> None:
    def test():
        import typing

        wrong_implicit = typing.cast(int, "bar")

    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    assert tree
    var = get_assignment_var_named(tree, "wrong_implicit")
    assert str(var.type) == "builtins.int"

@pytest.mark.xfail(reason="The observed behaviour here changed after updating mypy to 1.5.0")
def test_assert_isinstance_does_not_cast() -> None:
    def test():
        an_int: int = 123
        assert isinstance(an_int, str)

        another_int = an_int  # despite earlier assert, an_int is still seen as a int

    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    assert tree
    var = get_assignment_var_named(tree, "another_int")
    assert str(var.type) == "builtins.int"

def test_assert_isinstance_does_narrow() -> None:
    def test():
        import typing

        an_int: typing.Any = 123
        assert isinstance(an_int, str)

        another_int = an_int  # earlier assert narrows from Any to str

    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    assert tree
    var = get_assignment_var_named(tree, "another_int")
    assert str(var.type) == "builtins.str"

def test_name_defined():
    def test():
        error = missing + 2  # type: ignore[name-defined]

    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    match tree:
        case mypy.nodes.MypyFile(
            defs=[
                mypy.nodes.AssignmentStmt(
                    rvalue=mypy.nodes.OpExpr(left=mypy.nodes.NameExpr(name="missing") as missing),
                )
            ]
        ):
            assert missing.node is None
        case _:
            raise AssertionError

# TODO: override, what does the type information look like when call these?

def test_ignore_override() -> None:
    def test():
        class ABase:
            def method(self, x: int) -> int:
                return x

        class ADerived(ABase):
            def method(self, x: int) -> str:  # type: ignore[override]
                return str(x)

        class BBase:
            def call(self, o: ABase) -> int:
                result = o.method(42)
                return result

        class BDerived(BBase):
            def call(self, o: ADerived) -> str:  # type: ignore[override]
                result = o.method(42)
                return result

        base_base_result = BBase().call(ABase())
        base_derived_result = BBase().call(ADerived())
        derived_base_result = BDerived().call(ABase())
        derived_derived_result = BDerived().call(ADerived())

    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    assert tree

    test()

def test_super_no_args_outside_method() -> None:
    def test():
        def function() -> None:
            super().__init__()

            class Obj:
                pass

            print(super(Obj, Obj()).__str__())

    result = mypy_parse_and_type_check(test)
    assert result.errors == ['<string>:2: error: "super" used outside class  [misc]']

def test_aliased_import() -> None:
    def test():
        from collections import abc as xyz

        hmm = xyz

    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    assert tree.defs[-1].rvalue.fullname == "collections.abc"
    assert tree.names["xyz"].fullname == "collections.abc"

def test_function_redefinition() -> None:
    def test():
        def foo(x: int) -> int:
            return 2 * x

        def foo(x: str) -> str:
            return 2 * x

        bar = foo("1")

    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    assert result.errors == [
        '<string>:4: error: Name "foo" already defined on line 1  [no-redef]',
        '<string>:7: error: Argument 1 to "foo" has incompatible type "str"; expected "int"  [arg-type]',
    ]
    assert tree.defs[-1].rvalue.callee.node is tree.defs[0]

def test_class_members() -> None:
    def test():
        from typing import ClassVar, Final

        class Class:
            class_var: ClassVar = 1
            class_var_implicit = 2
            class_var_implicit_final: Final = 3

            decl_only: int
            decl_and_set_in_init: int
            decl_and_set_in_method: int
            decl_only_final: Final[int]

            def __init__(self) -> None:
                self.decl_and_set_in_init = 100
                self.new_var_in_init = 101
                self.new_final_var_in_init: Final = 102
                self.decl_only_final = 103

            def method(self) -> None:
                self.decl_and_set_in_method = 200
                self.new_var_in_method = 201
                self.new_final_far_in_method: Final = 202

        obj = Class()
        assert not hasattr(Class, "decl_only")
        assert not hasattr(Class, "decl_and_set_in_init")
        assert not hasattr(obj, "decl_only")
        assert hasattr(obj, "decl_and_set_in_init")

    test()
    result = mypy_parse_and_type_check(test)
    assert len(result.errors) == 1
    assert result.errors[0].endswith(
        "error: Can only declare a final attribute in class body or __init__  [misc]"
    )
    tree = result.graph[TEST_MODULE].tree
    assert tree
    cls_def = tree.defs[1]
    assert isinstance(cls_def, mypy.nodes.ClassDef)
    symtable = cls_def.info.names.copy()
    del symtable["__init__"]
    del symtable["method"]
    assert symtable.keys() == {
        "class_var",
        "class_var_implicit",
        "class_var_implicit_final",
        "decl_and_set_in_init",
        "decl_and_set_in_method",
        "decl_only",
        "decl_only_final",
        "new_final_far_in_method",
        "new_final_var_in_init",
        "new_var_in_init",
        "new_var_in_method",
    }
    vars_ = {}
    for name, sym in symtable.items():
        assert sym.kind == mypy.nodes.MDEF  # THE M STANDS FOR MEMBER.. REMEMBER... REMEMMMMBEERRR
        if name.startswith("new_"):
            assert sym.implicit
        else:
            assert not sym.implicit
        assert isinstance(sym.node, mypy.nodes.Var)
        # unfortunate that implicit class vars and not detected as such :(
        assert sym.node.is_classvar == (name == "class_var")
        # can't use is_initialized_in_class either, it is True even when it's just declared
        assert sym.node.is_initialized_in_class == (
            name.startswith("class_var") or name.startswith("decl_")
        )
        vars_[name] = sym.node
    # note this does nothing v
    typing.assert_type(vars_, dict[str, mypy.nodes.Var])

def test_abc() -> None:
    # TODO: FINISH THIS
    def test():
        from abc import ABC, abstractmethod
        import typing

        class Base(ABC):
            def concrete(self) -> int:
                return id(self)

            @abstractmethod
            def explicit_abstract_ellipsis(self) -> int: ...

            @abstractmethod
            def explicit_abstract_raise_not_implemented(self) -> int:
                raise NotImplementedError

            @abstractmethod
            def explicit_abstract_pass(self) -> int:
                pass

            @abstractmethod
            def explicit_abstract_docstring_only(self) -> int:
                """This is a docstring"""

            @abstractmethod
            def explicit_abstract_with_impl(self, x: int) -> int:
                return x + 1

            def implicit_abstract_returning_int(self) -> int:
                raise NotImplementedError

        class DerivedPass(Base):
            pass

        class DerivedABCPass(Base, ABC):
            pass

        class DerivedPartial(Base):
            def explicit_abstract_ellipsis(self) -> int:
                return 42

        if not typing.TYPE_CHECKING:
            import pytest

            with pytest.raises(TypeError, match="Can't instantiate abstract class"):
                base = Base()
            with pytest.raises(TypeError, match="Can't instantiate abstract class"):
                derived_pass = DerivedPass()
            with pytest.raises(TypeError, match="Can't instantiate abstract class"):
                derived_abc_pass = DerivedABCPass()
            with pytest.raises(TypeError, match="Can't instantiate abstract class"):
                derived_partial = DerivedPartial()

    test()
    result = mypy_parse_and_type_check(test)
    # assert not result.errors
    tree = result.graph[TEST_MODULE].tree
    assert tree

# TODO: various member/class var decls, protocol handling, special decorators,
#       metaclass, generics, etc
#

# TODO: Protocol tests

def test_nested_class() -> None:
    def test() -> None:
        class Outer:
            outer_mem: int = 1

            class Inner:
                inner_mem: int = 2

            inner_inst = Inner()

        outer_inst = Outer()
        assert outer_inst

    test()
    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    assert tree

def test_final_method() -> None:
    def test() -> None:
        import typing

        class Class:
            @typing.final
            def meaning_of_life(self) -> int:
                return 42

    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    assert tree
    class_def = tree.defs[-1]
    assert isinstance(class_def, mypy.nodes.ClassDef)
    deco = class_def.defs.body[0]
    assert isinstance(deco, mypy.nodes.Decorator)
    assert deco.decorators == []
    assert deco.is_final
    assert deco.func.is_final

def test_notypecheck_func() -> None:
    def test() -> None:
        import typing

        @typing.no_type_check
        def my_func(x: dict):
            return x + "lol"

    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    assert tree
    deco = tree.defs[-1]
    assert isinstance(deco, mypy.nodes.Decorator)
    assert deco.decorators, "please don't strip the no_type_check decorator mypy..."

def test_chained_assignment_lvalue_order() -> None:
    def test():
        outer = inner = 0

    result = mypy_parse_and_type_check(test)
    assert not result.errors
    tree = result.graph[TEST_MODULE].tree
    assert tree
    assert len(tree.defs) == 1
    assignment_stmt = tree.defs[0]
    assert isinstance(assignment_stmt, mypy.nodes.AssignmentStmt)
    lvalues = assignment_stmt.lvalues
    assert len(lvalues) == 2
    first, second = lvalues
    assert isinstance(first, mypy.nodes.NameExpr) and first.name == "outer"
    assert isinstance(second, mypy.nodes.NameExpr) and second.name == "inner"

def test_list_expr_assignment_target() -> None:
    """MyPy translates assignment targets with list to tuple"""

    def test():
        [a, b] = 1, 2
        [head, *tail] = [1, 2, 3]

    result = mypy_parse_and_type_check(test)
    assert not result.errors
    tree = result.graph[TEST_MODULE].tree
    assert tree and len(tree.defs) == 2
    for stmt in tree.defs:
        assert isinstance(stmt, mypy.nodes.AssignmentStmt)
        lvalues = stmt.lvalues
        assert len(lvalues) == 1
        assert isinstance(lvalues[0], mypy.nodes.TupleExpr)

def test_assignment_statements() -> None:
    def test():
        class Base:
            var: str

        class Obj(Base):
            def __init__(self) -> None:
                self._tmp = "tmp"
                Obj._tmp = "WHY"
                self._tmp2: int
                self._tmp2 = 123
                self._tmp3: int = len(self._tmp)

            def method(self) -> None:
                self.x: int
                self.x = 2
                self.y = 3
                self.z: int

        lst = [1, 2]
        tup0, tup1 = lst
        star0, *star_rest = lst
        foo, (bar, *baz) = (tup0, (tup1, star0, *star_rest))
        x: int
        x = 0
        lst[2:] = [3, 4]
        o = Obj()
        o.var = "a"
        Obj().var = "b"
        super(Obj, o).var = "c"  # type: ignore[misc]
        d = dict[str, int]()
        d["key"] = x
        o.method()
        ox = o.x

    result = mypy_parse_and_type_check(test)
    assert not result.errors
    tree = result.graph[TEST_MODULE].tree
    assert tree

    import mypy.traverser

    class MyVisitor(mypy.traverser.TraverserVisitor):
        def visit_assignment_stmt(self, stmt: mypy.nodes.AssignmentStmt) -> None:
            (lval,) = stmt.lvalues
            assert isinstance(
                lval,
                mypy.nodes.MemberExpr
                | mypy.nodes.NameExpr
                | mypy.nodes.TupleExpr
                | mypy.nodes.SuperExpr
                | mypy.nodes.IndexExpr,
            )
            if isinstance(lval, mypy.nodes.MemberExpr) and lval.def_var is not None:
                assert lval.def_var is lval.node
            if isinstance(stmt.rvalue, mypy.nodes.TempNode) and stmt.rvalue.no_rhs:
                assert stmt.type is not None
                assert lval.is_new_def
            super().visit_assignment_stmt(stmt)

    vis = MyVisitor()
    tree.accept(vis)

def test_super_exprs() -> None:
    def test():
        class Base:
            def __init__(self, x: int) -> None:
                self.x = x
                self.var: int = 0

            def method(self) -> int:
                return self.x

        class Derived(Base):
            def __init__(self, x: int, y: int) -> None:
                super().__init__(x=x)
                super(Base, self).var = 123  # type: ignore[misc]
                self.y = y

            def method(self) -> int:
                return super().method() + self.y

    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    assert tree
    assert not result.errors

def test_special_decorators() -> None:
    def test() -> None:
        import abc
        import typing

        class Class(abc.ABC):
            @property
            def prop_get_set_del(self) -> int:
                return self._x

            @prop_get_set_del.setter
            def prop_get_set_del(self, value: int) -> None:
                self._x = value

            @prop_get_set_del.deleter
            def prop_get_set_del(self) -> None:
                del self._x

            @property
            def prop_get_set(self) -> str:
                return self._y

            @prop_get_set.setter
            def prop_get_set(self, value: str) -> None:
                self._y = value

            @property
            def prop_get(self) -> bytes:
                return b""

            @abc.abstractmethod
            def abstract_method(self, x: str) -> str: ...

            @staticmethod
            def static() -> int:
                return 42

            @classmethod
            def klass(cls) -> int:
                return cls().static()

            @typing.overload
            def overloaded(self, x: int) -> None: ...

            @typing.overload
            def overloaded(self, x: str) -> None: ...

            def overloaded(self, x: int | str) -> None:
                if isinstance(x, int):
                    self.prop_get_set_del = x
                else:
                    self.prop_get_set = x

    test()
    result = mypy_parse_and_type_check(test)
    assert not result.errors
    tree = result.graph[TEST_MODULE].tree
    assert tree
    class_def = tree.defs[-1]
    assert isinstance(class_def, mypy.nodes.ClassDef)

def test_variadic_tuple_type() -> None:
    def test():
        from collections.abc import Sequence

        def func(items: Sequence[int]) -> tuple[int, ...]:
            return tuple(items)

        def func2(items: Sequence[int]) -> tuple[int, int]:
            return items[0], len(items)

        x = func([1, 2, 3])
        y = func2([1, 2, 3])

    test()
    result = mypy_parse_and_type_check(test)
    assert not result.errors
    tree = result.graph[TEST_MODULE].tree
    assert tree
    func = tree.defs[1]
    assert isinstance(func, mypy.nodes.FuncDef)
    assert isinstance(func.type, mypy.types.CallableType)
    assert type(func.type.ret_type) == mypy.types.Instance
    func2 = tree.defs[2]
    assert isinstance(func2, mypy.nodes.FuncDef)
    assert isinstance(func2.type, mypy.types.CallableType)
    assert type(func2.type.ret_type) == mypy.types.TupleType

def test_tuple_covariance() -> None:
    def test():
        class Base:
            pass

        class Derived(Base):
            pass

        def func_out() -> tuple[Base, object]:
            return Derived(), 1

        def func_in(a: tuple[Base, object]) -> None:
            print(a)

        func_in((Derived(), 2))

    test()
    result = mypy_parse_and_type_check(test)
    assert not result.errors
    tree = result.graph[TEST_MODULE].tree
    assert tree

def test_if_else_expr_combined_type() -> None:
    def test():
        import typing

        def random() -> float:
            return 0.1

        class Base:
            pass

        class Derived(Base):
            pass

        typing.reveal_type(Base() if random() > 0.5 else Derived())

        typing.reveal_type("str" if random() < 0.5 else 1)
        okay: str | int = "str" if random() < 0.5 else 1
        typing.reveal_type(okay)
        bad: str = "str" if random() == 1 else 1

    result = mypy_parse_and_type_check(test)
    assert strip_error_prefixes(result) == [
        'note: Revealed type is "__test__.Base"',
        "note: Revealed type is \"Union[Literal['str']?, Literal[1]?]\"",
        'note: Revealed type is "Union[builtins.str, builtins.int]"',
        'error: Incompatible types in assignment (expression has type "str | int",'
        ' variable has type "str")  [assignment]',
    ]
    tree = result.graph[TEST_MODULE].tree
    assert tree
    assert list(map(repr, get_revealed_types(result, tree))) == [
        "__test__.Base",
        "Union[Literal['str']?, Literal[1]?]",
        "Union[builtins.str, builtins.int]",
    ]

def test_or_expr_combined_type() -> None:
    def test():
        import typing

        def random() -> float:
            return 0.1

        class Base:
            def __bool__(self) -> bool:
                return random() < 0.5

        class Derived(Base):
            pass

        typing.reveal_type(Base() or Derived())
        typing.reveal_type(Derived() or Base())
        true = True
        typing.reveal_type(true or Base())
        one = 1
        typing.reveal_type(Base() or one)
        bad: Derived = Base() or Derived()
        okay: Base = Base() or Derived()

    result = mypy_parse_and_type_check(test)
    assert strip_error_prefixes(result) == [
        'note: Revealed type is "__test__.Base"',
        'note: Revealed type is "__test__.Base"',
        'note: Revealed type is "Union[Literal[True], __test__.Base]"',
        'note: Revealed type is "Union[__test__.Base, builtins.int]"',
        'error: Incompatible types in assignment (expression has type "Base",'
        ' variable has type "Derived")  [assignment]',
    ]
    tree = result.graph[TEST_MODULE].tree
    assert tree
    assert list(map(repr, get_revealed_types(result, tree))) == [
        "__test__.Base",
        "__test__.Base",
        "Union[Literal[True], __test__.Base]",
        "Union[__test__.Base, builtins.int]",
    ]

def test_global_types() -> None:
    def test():
        import typing

        a = "a"
        b: typing.Final = "b"
        c = a + b
        d: typing.Final = a + b
        one = 1
        two: typing.Final = 2
        three = one + two
        four: typing.Final = three + one
        true = True
        false: typing.Final = False
        true_and_false = true and false
        true_or_false: typing.Final = true or false
        b_a = b"a"
        b_b: typing.Final = b"b"
        b_c = b_a + b_b
        b_d: typing.Final = 2 * b_b

        def silly_function() -> None:
            print(
                a,
                b,
                c,
                d,
                one,
                two,
                three,
                four,
                true,
                false,
                true_and_false,
                true_or_false,
                b_a,
                b_b,
                b_c,
                b_d,
            )

    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    assert tree
    silly_func_def = tree.defs[-1]
    assert isinstance(silly_func_def, mypy.nodes.FuncDef)
    assert len(silly_func_def.body.body) == 1
    print_stmt = silly_func_def.body.body[0]
    assert isinstance(print_stmt, mypy.nodes.ExpressionStmt)
    print_call = print_stmt.expr
    assert isinstance(print_call, mypy.nodes.CallExpr)
    for arg_name_expr in print_call.args:
        assert isinstance(arg_name_expr, mypy.nodes.NameExpr)
        assert arg_name_expr.kind == mypy.nodes.GDEF
        assert not (
            arg_name_expr.is_alias_rvalue
            or arg_name_expr.is_inferred_def
            or arg_name_expr.is_new_def
            or arg_name_expr.is_special_form
        )
        assert isinstance(arg_name_expr.node, mypy.nodes.Var)
        if isinstance(arg_name_expr.node.type, mypy.types.LiteralType):
            inst_type = arg_name_expr.node.type.fallback
        else:
            assert isinstance(arg_name_expr.node.type, mypy.types.Instance)
            inst_type = arg_name_expr.node.type
        assert inst_type.type.fullname in (
            "builtins.str",
            "builtins.int",
            "builtins.bool",
            "builtins.bytes",
        )

def test_annotated_metadata() -> None:
    """Negative test - this would be a good way to add key overrides to global state storage,
    but unfortunately MyPy strips it entirely from the AST, even the unanalyzed_type"""

    def test() -> None:
        import typing
        import dataclasses

        @dataclasses.dataclass
        class StateKey:
            key: bytes

        counter: typing.Annotated[int, StateKey(b"c")] = 1

    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    assert tree
    ass = tree.defs[-1]
    assert isinstance(ass, mypy.nodes.AssignmentStmt)
    assert isinstance(ass.unanalyzed_type, mypy.types.UnboundType)
    assert len(ass.unanalyzed_type.args) == 2
    _, ann_arg = ass.unanalyzed_type.args
    assert isinstance(ann_arg, mypy.types.RawExpressionType) and ann_arg.literal_value is None

def test_dataclass_transform_frozen() -> None:
    def test() -> None:
        import typing

        @typing.dataclass_transform(
            eq_default=False, order_default=False, kw_only_default=False, field_specifiers=()
        )
        class _StructMeta(type):
            def __new__(
                cls,
                name: str,
                bases: tuple[type, ...],
                namespace: dict[str, object],
                *,
                kw_only: bool = False,
            ) -> "_StructMeta":
                raise NotImplementedError

        class StructBase(metaclass=_StructMeta):
            pass

        class Struct(StructBase):
            field: int

        class FrozenStruct(StructBase, frozen=True):
            field: int

        class UnfrozenStruct(StructBase, frozen=False):
            field: int

        f1 = Struct(1)
        f2 = FrozenStruct(1)
        f3 = UnfrozenStruct(1)
        f1.field = 1
        f2.field = 2
        f3.field = 3

    result = mypy_parse_and_type_check(test)
    assert strip_error_prefixes(result) == [
        'error: Property "field" defined in "FrozenStruct" is read-only  [misc]',
    ]
    tree = result.graph[TEST_MODULE].tree
    assert isinstance(tree, mypy.nodes.MypyFile)
    struct_cls_def = get_class_by_name(tree, "Struct")
    assert struct_cls_def.info.metadata["dataclass"]["frozen"] is False
    frozen_struct_cls_def = get_class_by_name(tree, "FrozenStruct")
    assert frozen_struct_cls_def.info.metadata["dataclass"]["frozen"] is True
    unfrozen_struct_cls_def = get_class_by_name(tree, "UnfrozenStruct")
    assert unfrozen_struct_cls_def.info.metadata["dataclass"]["frozen"] is False

def test_bound_method_types() -> None:
    def test():
        import typing

        def func(
            a: bool,
            /,
            b: int,
            b2: None = None,
            *args: str,
            c: bytes,
            d: float | None = None,
            **kwargs: complex,
        ) -> bool:
            return a

        class Class:
            def __init__(self) -> None:
                pass

            @staticmethod
            def static_method(b: int) -> int:
                return b

            @classmethod
            def class_method(cls, c: str) -> str:
                return c

            def instance_method(self, d: bytes) -> bytes:
                return d

        typing.reveal_type(func)
        typing.reveal_type(Class)
        typing.reveal_type(Class.__init__)
        typing.reveal_type(Class.static_method)
        typing.reveal_type(Class.class_method)
        typing.reveal_type(Class.instance_method)
        inst = Class()
        typing.reveal_type(inst.static_method)
        typing.reveal_type(inst.class_method)
        typing.reveal_type(inst.instance_method)

    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    assert tree
    reveled_types = get_revealed_types(result, tree)
    assert len(reveled_types) == 9
    assert all(isinstance(t, mypy.types.CallableType) for t in reveled_types)

def test_special_indexing() -> None:
    def test():
        import typing

        def func(abc: typing.Literal["a", "b", "c"]) -> str:
            X = typing.Literal[0]
            Y = tuple[typing.Literal[0, 1, 2], str]
            y = tuple[typing.Literal[0, 1, 2], str]((0, abc))
            return abc

    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    assert tree

@pytest.mark.xfail(reason="mypy doesn't consider equality in match-case with value patterns")
def test_match_case_reachability() -> None:
    def test():
        import typing

        class MyInt:
            def __init__(self, value: int):
                self.value = value

            def __eq__(self, other: object) -> bool:
                return (other == self.value) or (
                    isinstance(other, MyInt) and other.value == self.value
                )

        match MyInt(1):
            case 1:
                x = None
                typing.reveal_type(x)
            case _:
                x = "ironically unreachable"
                typing.reveal_type(x)

        return typing.reveal_type(x)

    assert test() is None

    result = mypy_parse_and_type_check(test)
    tree = result.graph[TEST_MODULE].tree
    assert list(map(str, get_revealed_types(result, tree))) == [
        "None",
        "builtins.str",
        "Union[builtins.str, None]",
    ]


from collections.abc import Mapping

import pytest

from puyapy.awst_build import pytypes
from tests import VCS_ROOT

_STUB_SUFFIX = ".pyi"

def stub_class_names_and_predefined_aliases() -> list[str]:
    from mypy import build, find_sources, fscache, nodes

    from puyapy.compile import get_mypy_options

    stubs_dir = (VCS_ROOT / "stubs" / "algopy-stubs").resolve()
    mypy_options = get_mypy_options()
    fs_cache = fscache.FileSystemCache()
    mypy_build_sources = find_sources.create_source_list(
        paths=[str(stubs_dir)], options=mypy_options, fscache=fs_cache
    )
    build_result = build.build(sources=mypy_build_sources, options=mypy_options, fscache=fs_cache)
    result = set()

    algopy_module = build_result.files["algopy"]
    modules_to_visit = [algopy_module]
    seen_modules = set()
    while modules_to_visit:
        module = modules_to_visit.pop()
        if module in seen_modules:
            continue
        seen_modules.add(module)
        for name, symbol in module.names.items():
            if name.startswith("_") or symbol.module_hidden or symbol.kind != nodes.GDEF:
                continue
            match symbol.node:
                case nodes.MypyFile() as new_module:
                    modules_to_visit.append(new_module)
                case nodes.TypeAlias(fullname=alias_name):
                    result.add(alias_name)
                case nodes.TypeInfo(fullname=class_name):
                    result.add(class_name)
    return sorted(result)

@pytest.fixture(scope="session")
def builtins_registry() -> Mapping[str, pytypes.PyType]:
    return pytypes.builtins_registry()

@pytest.mark.parametrize(
    "fullname",
    stub_class_names_and_predefined_aliases(),
    ids=str,
)
def test_stub_class_names_lookup(
    builtins_registry: Mapping[str, pytypes.PyType], fullname: str
) -> None:
    assert fullname in builtins_registry, f"{fullname} is missing from pytypes"


import json
from collections.abc import Iterable, Mapping
from pathlib import Path

import attrs
import mypy.nodes
import mypy.types
import pytest

from puya.awst.txn_fields import TxnField
from puya.parse import SourceLocation
from puyapy.awst_build import pytypes
from puyapy.awst_build.context import type_to_pytype
from puyapy.awst_build.eb.transaction.itxn_args import PYTHON_ITXN_ARGUMENTS
from puyapy.awst_build.eb.transaction.txn_fields import PYTHON_TXN_FIELDS
from tests import EXAMPLES_DIR, VCS_ROOT
from tests.utils import get_awst_cache

# the need to use approval / clear_state pages is abstracted away by
# allowing a tuple of pages in the stubs layer
_MAPPED_INNER_TXN_FIELDS = {
    TxnField.ApprovalProgramPages: TxnField.ApprovalProgram,
    TxnField.ClearStateProgramPages: TxnField.ClearStateProgram,
}
_INTENTIONALLY_OMITTED_INNER_TXN_FIELDS = {
    *_MAPPED_INNER_TXN_FIELDS.values(),
    # only allow enum version of type
    TxnField.Type,
}

@attrs.frozen
class FieldType:
    is_array: bool
    field_type: pytypes.PyType

def _get_type_infos(type_names: Iterable[str]) -> Iterable[mypy.nodes.TypeInfo]:
    awst_cache = get_awst_cache(EXAMPLES_DIR)

    for type_name in type_names:
        module_id, symbol_name = type_name.rsplit(".", maxsplit=1)
        module = awst_cache.parse_result.ordered_modules[module_id].node

        symbol = module.names[symbol_name]
        node = symbol.node
        assert isinstance(node, mypy.nodes.TypeInfo), f"{type_name} not found in {module}"
        yield node

@pytest.fixture(scope="session")
def builtins_registry() -> Mapping[str, pytypes.PyType]:
    return pytypes.builtins_registry()

def test_group_transaction_members() -> None:
    gtxn_types = [t.name for t in pytypes.GroupTransactionTypes.values()]
    gtxn_types.append(pytypes.GroupTransactionBaseType.name)
    for type_info in _get_type_infos(gtxn_types):
        unknown = sorted(set(type_info.protocol_members) - PYTHON_TXN_FIELDS.keys())
        assert not unknown, f"{type_info.fullname}: Unknown TxnField members: {unknown}"

def test_field_vs_argument_name_consistency() -> None:
    itxn_args = {
        (_MAPPED_INNER_TXN_FIELDS.get(params.field, params.field), name)
        for name, params in PYTHON_ITXN_ARGUMENTS.items()
    }
    txn_fields = {(f.field, name) for name, f in PYTHON_TXN_FIELDS.items()}
    bad_itxn_args = itxn_args - txn_fields
    assert not bad_itxn_args

def test_inner_transaction_field_setters() -> None:
    unmapped = {
        tf for tf in TxnField if tf.is_inner_param
    } - _INTENTIONALLY_OMITTED_INNER_TXN_FIELDS
    for type_info in _get_type_infos(
        t.name for t in pytypes.InnerTransactionFieldsetTypes.values()
    ):
        init_args: set[str] | None = None
        for member in ("__init__", "set"):
            func_def = type_info.names[member].node
            assert isinstance(func_def, mypy.nodes.FuncDef)
            arg_names = {a for a in func_def.arg_names if a is not None}
            arg_names.remove("self")
            unknown = sorted(arg_names - PYTHON_ITXN_ARGUMENTS.keys())
            assert not unknown, f"{type_info.fullname}: Unknown TxnField param members: {unknown}"
            unmapped -= {
                PYTHON_ITXN_ARGUMENTS[arg_name].field
                for arg_name in arg_names
                if arg_name not in unknown
            }

            if init_args is None:
                init_args = arg_names
            else:
                difference = init_args.symmetric_difference(arg_names)
                assert (
                    not difference
                ), f"{type_info.fullname}.{member} field difference: {difference}"
    assert not unmapped, f"Unmapped inner param fields: {sorted(f.immediate for f in unmapped)}"

def test_inner_transaction_members() -> None:
    for type_info in _get_type_infos(t.name for t in pytypes.InnerTransactionResultTypes.values()):
        unknown = sorted(set(type_info.protocol_members) - PYTHON_TXN_FIELDS.keys())
        assert not unknown, f"{type_info.fullname}: Unknown TxnField members: {unknown}"

_FAKE_SOURCE_LOCATION = SourceLocation(file=Path(__file__).resolve(), line=1)

def test_txn_fields(builtins_registry: Mapping[str, pytypes.PyType]) -> None:
    # collect all fields that are protocol members
    txn_types = [t.name for t in pytypes.GroupTransactionTypes.values()]
    txn_types.append(pytypes.GroupTransactionBaseType.name)
    txn_types.extend(t.name for t in pytypes.InnerTransactionResultTypes.values())
    seen_fields = set[str]()
    invalid_types = ""
    for type_info in _get_type_infos(txn_types):
        for member in type_info.protocol_members:
            seen_fields.add(member)
            txn_field_data = PYTHON_TXN_FIELDS[member]
            field_type = FieldType(
                is_array=txn_field_data.field.is_array, field_type=txn_field_data.type
            )
            member_mypy_type = type_info[member].type
            assert member_mypy_type is not None, f"Expected {type_info}.{member} to have a type"
            member_type = _member_to_field_type(builtins_registry, member_mypy_type)
            assert field_type == member_type

    # add fields that are arguments
    for type_info in _get_type_infos(
        t.name for t in pytypes.InnerTransactionFieldsetTypes.values()
    ):
        for member in ("__init__", "set"):
            func_def = type_info.names[member].node
            assert isinstance(func_def, mypy.nodes.FuncDef)
            assert func_def.type is not None
            func_type = type_to_pytype(
                builtins_registry, func_def.type, source_location=_FAKE_SOURCE_LOCATION
            )
            assert isinstance(func_type, pytypes.FuncType)
            for arg in func_type.args:
                assert arg.name is not None
                if arg.name == "self":
                    continue
                seen_fields.add(arg.name)
                txn_field_param = PYTHON_ITXN_ARGUMENTS[arg.name]
                txn_field = txn_field_param.field
                if isinstance(arg.type, pytypes.UnionType):
                    arg_types = arg.type.types
                else:
                    arg_types = (arg.type,)
                assert set(txn_field_param.literal_overrides.keys()).issubset(arg_types)
                if txn_field.is_array:
                    arg_types = tuple(
                        vt.items for vt in arg_types if isinstance(vt, pytypes.VariadicTupleType)
                    )
                if txn_field_param.auto_serialize_bytes:
                    assert arg_types == (pytypes.ObjectType,)
                else:
                    non_literal_arg_types = {
                        at for at in arg_types if not isinstance(at, pytypes.LiteralOnlyType)
                    }
                    assert non_literal_arg_types == {
                        txn_field_param.type,
                        *txn_field_param.additional_types,
                    }

    # anything missing is an error
    missing_fields = sorted(PYTHON_TXN_FIELDS.keys() - seen_fields)
    assert not missing_fields, f"Txn Fields not mapped: {missing_fields}"

    # any invalid_types is an error
    assert not invalid_types, f"Invalid field types: {invalid_types}"

def test_mismatched_langspec_txn_fields() -> None:
    langspec_path = VCS_ROOT / "langspec.puya.json"
    langspec = json.loads(langspec_path.read_text())
    arg_enums = langspec["arg_enums"]
    all_txn_fields = {field["name"] for field in arg_enums["txn"]}
    txn_array_fields = {field["name"] for field in arg_enums["txna"]}
    txn_single_fields = all_txn_fields - txn_array_fields
    inner_txn_fields = {field["name"] for field in arg_enums["itxn_field"]}

    assert not _set_difference(
        all_txn_fields, [f.immediate for f in TxnField]
    ), "txn field mismatch"

    assert not _set_difference(
        txn_single_fields, [f.immediate for f in TxnField if not f.is_array]
    ), "single txn field mismatch"

    assert not _set_difference(
        txn_array_fields, [f.immediate for f in TxnField if f.is_array]
    ), "array txn field mismatch"

    assert not _set_difference(
        inner_txn_fields, [f.immediate for f in TxnField if f.is_inner_param]
    ), "inner txn field mismatch"

def _set_difference(expected: set[str], actual: list[str]) -> list[str]:
    return list(expected.symmetric_difference(actual))

def _member_to_field_type(
    builtins_registry: Mapping[str, pytypes.PyType], typ: mypy.types.Type
) -> FieldType:
    is_array = False
    if isinstance(typ, mypy.types.CallableType):
        is_array = len(typ.arg_names) > 1
        typ = typ.ret_type

    field_type = type_to_pytype(builtins_registry, typ, source_location=_FAKE_SOURCE_LOCATION)

    return FieldType(is_array=is_array, field_type=field_type)


import functools
import typing
from collections.abc import Collection, Iterable, Mapping, Sequence
from pathlib import Path

import attrs

from puya.awst import nodes as awst_nodes
from puya.compilation_artifacts import CompilationArtifact
from puya.compile import awst_to_teal
from puya.errors import CodeError
from puya.log import Log, LogLevel, logging_context
from puya.parse import SourceLocation
from puya.program_refs import ContractReference, LogicSigReference
from puya.utils import pushd
from puyapy.awst_build.main import transform_ast
from puyapy.compile import output_awst, parse_with_mypy, write_arc4_clients
from puyapy.options import PuyaPyOptions
from puyapy.parse import ParseResult, SourceDiscoveryMechanism
from puyapy.template import parse_template_key_value
from puyapy.utils import determine_out_dir
from tests import EXAMPLES_DIR, TEST_CASES_DIR

APPROVAL_EXTENSIONS = frozenset((".teal", ".awst", ".ir", ".mir", ".arc32.json", ".arc56.json"))
UNSTABLE_LOG_PREFIXES = {
    LogLevel.debug: (
        "Building AWST for ",
        "Skipping algopy stub ",
        "Skipping typeshed stub ",
        "Skipping stdlib stub ",
        "Discovered user module ",
        # ignore platform specific paths
        "Using python executable: ",
        "Using python site-packages: ",
        "Found algopy: ",
    ),
}

def _get_root_dir(path: Path) -> Path:
    if path.is_relative_to(EXAMPLES_DIR):
        return EXAMPLES_DIR
    if path.is_relative_to(TEST_CASES_DIR):
        return TEST_CASES_DIR
    return path if path.is_dir() else path.parent

class _CompileCache(typing.NamedTuple):
    parse_result: ParseResult
    module_awst: awst_nodes.AWST
    compilation_set: Sequence[ContractReference | LogicSigReference]
    logs: list[Log]

@functools.cache
def get_awst_cache(root_dir: Path) -> _CompileCache:
    # note that this caching assumes that AWST is the same across all
    # optimisation and debug levels, which is currently true.
    # if this were to no longer be true, this test speedup strategy would need to be revisited
    with pushd(root_dir), logging_context() as log_ctx:
        parse_result = parse_with_mypy([root_dir])
        awst, compilation_set = transform_ast(parse_result)
    return _CompileCache(parse_result, awst, compilation_set, log_ctx.logs)

@attrs.frozen(kw_only=True)
class CompilationResult:
    module_awst: awst_nodes.AWST
    logs: list[Log]
    teal: list[CompilationArtifact]
    src_path: Path
    """original source path"""
    root_dir: Path
    """examples or test_cases path"""

def narrowed_parse_result(parse_result: ParseResult, src_path: Path) -> ParseResult:
    filtered_ordered_modules = {
        name: sm
        for name, sm in parse_result.ordered_modules.items()
        if sm.path.resolve().is_relative_to(src_path.resolve())
        or sm.discovery_mechanism == SourceDiscoveryMechanism.dependency
    }
    return ParseResult(
        mypy_options=parse_result.mypy_options,
        ordered_modules=filtered_ordered_modules,
    )

def narrowed_compile_context(
    parse_result: ParseResult,
    src_path: Path,
    awst: awst_nodes.AWST,
    compilation_set: Collection[ContractReference | LogicSigReference],
    puyapy_options: PuyaPyOptions,
) -> tuple[
    Mapping[Path, Sequence[str] | None],
    Mapping[ContractReference | LogicSigReference, Path],
]:
    narrowed_sources_by_path = {
        sm.path: sm.lines
        for sm in parse_result.ordered_modules.values()
        if sm.path.resolve().is_relative_to(src_path)
        and sm.discovery_mechanism != SourceDiscoveryMechanism.dependency
    }
    awst_lookup = {n.id: n for n in awst}
    compilation_set = {
        target_id: determine_out_dir(loc.file.parent, puyapy_options)
        for target_id, loc in ((t, awst_lookup[t].source_location) for t in compilation_set)
        if loc.file in narrowed_sources_by_path
    }
    return narrowed_sources_by_path, compilation_set

def _filter_logs(logs: list[Log], root_dir: Path, src_path: Path) -> list[Log]:
    root_dir = root_dir.resolve()
    relative_src_root = src_path.relative_to(root_dir)
    result = []
    for log in logs:
        # ignore logs that come from files outside of src_path as these are
        # logs emitted during the cached AWST parsing step
        relative_path = get_relative_path(log.location, root_dir)
        if relative_path and relative_src_root not in (
            relative_path,
            *relative_path.parents,
        ):
            continue

        # ignore logs that are not output in a consistent order
        log_prefixes_to_ignore = UNSTABLE_LOG_PREFIXES.get(log.level)
        if log_prefixes_to_ignore and log.message.startswith(log_prefixes_to_ignore):
            continue

        result.append(log)
    return result

def get_relative_path(location: SourceLocation | None, root_dir: Path) -> Path | None:
    if not location or not location.file:
        return None
    try:
        return location.file.relative_to(root_dir)
    except ValueError:
        return None

def _get_log_errors(logs: Iterable[Log]) -> str:
    return "\n".join(str(log) for log in logs if log.level == LogLevel.error)

def compile_src_from_options(options: PuyaPyOptions) -> CompilationResult:
    if options.out_dir is None:
        options = attrs.evolve(options, out_dir=Path("out_tests"))
    (src_path,) = options.paths
    root_dir = _get_root_dir(src_path)
    src_path = src_path.resolve()
    parse_result, awst, compilation_set, awst_logs = get_awst_cache(root_dir)
    awst_logs = _filter_logs(awst_logs, root_dir, src_path)

    awst_errors = _get_log_errors(awst_logs)
    if awst_errors:
        raise CodeError(awst_errors)
    # create a new context from cache and specified src
    with logging_context() as log_ctx:
        narrow_sources_by_path, narrow_compilation_set = narrowed_compile_context(
            parse_result, src_path, awst, compilation_set, options
        )
        # this should be correct and exhaustive but relies on independence of examples
        narrowed_awst = [n for n in awst if n.source_location.file in narrow_sources_by_path]

        with pushd(root_dir):
            if options.output_awst and options.out_dir:
                output_awst(narrowed_awst, options.out_dir)

            try:
                teal = awst_to_teal(
                    log_ctx, options, narrow_compilation_set, narrow_sources_by_path, narrowed_awst
                )
            except SystemExit as ex:
                raise CodeError(_get_log_errors(log_ctx.logs)) from ex

            filtered_teal = [t for t in teal if t.id in narrow_compilation_set]

            if options.output_client:
                write_arc4_clients(
                    options.template_vars_prefix, narrow_compilation_set, filtered_teal
                )

        return CompilationResult(
            module_awst=narrowed_awst,
            logs=awst_logs + log_ctx.logs,
            teal=filtered_teal,
            root_dir=root_dir,
            src_path=src_path,
        )

@attrs.frozen
class PuyaTestCase:
    path: Path

    @property
    def root(self) -> Path:
        return self.path.parent

    @property
    def name(self) -> str:
        return self.path.name

    @property
    def template_vars_path(self) -> Path | None:
        template_vars_path = self.path / "template.vars"
        return template_vars_path.resolve() if template_vars_path.exists() else None

    @property
    def id(self) -> str:
        return f"{self.root.stem}/{self.name}"

def load_template_vars(path: Path | None) -> tuple[str, dict[str, int | bytes]]:
    prefix = "TMPL_"
    prefix_prefix = "prefix="

    result = {}
    if path is not None:
        for line in path.read_text("utf8").splitlines():
            if line.startswith(prefix_prefix):
                prefix = line.removeprefix(prefix_prefix)
            else:
                key, value = parse_template_key_value(line)
                result[key] = value
    return prefix, result


import hashlib
import itertools

def sha_256_raw(b: bytes) -> bytes:
    return hashlib.sha256(b).digest()

class MerkleTree:
    def __init__(self, data: list[bytes]):
        self._tree = [[sha_256_raw(d) for d in data]]

        while len(self._tree[-1]) != 1:
            next_level = list[bytes]()
            for pair in itertools.batched(self._tree[-1], 2):
                a, b = pair if len(pair) == 2 else (pair[0], pair[0])
                a_b_hash = self.hash_pair(a, b)
                next_level.append(a_b_hash)
            self._tree.append(next_level)

    @property
    def root(self) -> bytes:
        return self._tree[-1][0]

    @staticmethod
    def hash_pair(a: bytes, b: bytes) -> bytes:
        return sha_256_raw(a + b if a < b else b + a)

    def get_proof(self, data: bytes) -> list[bytes]:
        proof = list[bytes]()
        current = sha_256_raw(data)
        for level in self._tree:
            index = level.index(current)
            sibling_index = index - 1 if index % 2 == 1 else index + 1
            if sibling_index >= len(level):
                proof.append(current)
            else:
                proof.append(level[sibling_index])
            current = self.hash_pair(current, proof[-1])
        return proof[:-1]

    def verify_proof(self, data: bytes, proof: list[bytes]) -> bool:
        result = sha_256_raw(data)
        for item in proof:
            result = self.hash_pair(result, item)
        return result == self.root

    def __repr__(self) -> str:
        return "\n".join([str(level) for level in self._tree])