from itertools import product
from typing import List, Literal, Optional, cast

import pytest
from dataclasses import dataclass

import pyteal as pt
from pyteal.ast.subroutine import ABIReturnSubroutine, evaluate_subroutine

options = pt.CompileOptions(version=5)


def test_subroutine_definition():
    def fn0Args():
        return pt.Return()

    def fn1Args(a1):
        return pt.Return()

    def fn2Args(a1, a2):
        return pt.Return()

    def fn10Args(a1, a2, a3, a4, a5, a6, a7, a8, a9, a10):
        return pt.Return()

    lam0Args = lambda: pt.Return()  # noqa: E731
    lam1Args = lambda a1: pt.Return()  # noqa: E731
    lam2Args = lambda a1, a2: pt.Return()  # noqa: E731
    lam10Args = (
        lambda a1, a2, a3, a4, a5, a6, a7, a8, a9, a10: pt.Return()
    )  # noqa: E731

    def fnWithExprAnnotations(a: pt.Expr, b: pt.Expr) -> pt.Expr:
        return pt.Return()

    def fnWithOnlyReturnExprAnnotations(a, b) -> pt.Expr:
        return pt.Return()

    def fnWithOnlyArgExprAnnotations(a: pt.Expr, b: pt.Expr):
        return pt.Return()

    def fnWithPartialExprAnnotations(a, b: pt.Expr) -> pt.Expr:
        return pt.Return()

    cases = (
        (fn0Args, 0, "fn0Args"),
        (fn1Args, 1, "fn1Args"),
        (fn2Args, 2, "fn2Args"),
        (fn10Args, 10, "fn10Args"),
        (lam0Args, 0, "<lambda>"),
        (lam1Args, 1, "<lambda>"),
        (lam2Args, 2, "<lambda>"),
        (lam10Args, 10, "<lambda>"),
        (fnWithExprAnnotations, 2, "fnWithExprAnnotations"),
        (fnWithOnlyReturnExprAnnotations, 2, "fnWithOnlyReturnExprAnnotations"),
        (fnWithOnlyArgExprAnnotations, 2, "fnWithOnlyArgExprAnnotations"),
        (fnWithPartialExprAnnotations, 2, "fnWithPartialExprAnnotations"),
    )

    for (fn, numArgs, name) in cases:
        definition = pt.SubroutineDefinition(fn, pt.TealType.none)
        assert definition.argument_count() == numArgs
        assert definition.name() == name

        if numArgs > 0:
            with pytest.raises(pt.TealInputError):
                definition.invoke([pt.Int(1)] * (numArgs - 1))

        with pytest.raises(pt.TealInputError):
            definition.invoke([pt.Int(1)] * (numArgs + 1))

        if numArgs > 0:
            with pytest.raises(pt.TealInputError):
                definition.invoke([1] * numArgs)

        args = [pt.Int(1)] * numArgs
        invocation = definition.invoke(args)
        assert isinstance(invocation, pt.SubroutineCall)
        assert invocation.subroutine is definition
        assert invocation.args == args


@dataclass
class ABISubroutineTC:
    definition: pt.ABIReturnSubroutine
    arg_instances: list[pt.Expr | pt.abi.BaseType]
    name: str
    ret_type: str | pt.abi.TypeSpec
    signature: Optional[str]


def test_abi_subroutine_definition():
    @pt.ABIReturnSubroutine
    def fn_0arg_0ret() -> pt.Expr:
        return pt.Return()

    @pt.ABIReturnSubroutine
    def fn_0arg_uint64_ret(*, output: pt.abi.Uint64) -> pt.Expr:
        return output.set(1)

    @pt.ABIReturnSubroutine
    def fn_1arg_0ret(a: pt.abi.Uint64) -> pt.Expr:
        return pt.Return()

    @pt.ABIReturnSubroutine
    def fn_1arg_1ret(a: pt.abi.Uint64, *, output: pt.abi.Uint64) -> pt.Expr:
        return output.set(a)

    @pt.ABIReturnSubroutine
    def fn_2arg_0ret(
        a: pt.abi.Uint64, b: pt.abi.StaticArray[pt.abi.Byte, Literal[10]]
    ) -> pt.Expr:
        return pt.Return()

    @pt.ABIReturnSubroutine
    def fn_2arg_1ret(
        a: pt.abi.Uint64,
        b: pt.abi.StaticArray[pt.abi.Byte, Literal[10]],
        *,
        output: pt.abi.Byte,
    ) -> pt.Expr:
        return output.set(b[a.get() % pt.Int(10)])

    @pt.ABIReturnSubroutine
    def fn_2arg_1ret_with_expr(
        a: pt.Expr,
        b: pt.abi.StaticArray[pt.abi.Byte, Literal[10]],
        *,
        output: pt.abi.Byte,
    ) -> pt.Expr:
        return output.set(b[a % pt.Int(10)])

    @pt.ABIReturnSubroutine
    def fn_w_tuple1arg(
        a: pt.Expr,
        b: pt.abi.Tuple1[pt.abi.Byte],
        *,
        output: pt.abi.Byte,
    ) -> pt.Expr:
        return output.set(pt.Int(1))

    cases = (
        ABISubroutineTC(fn_0arg_0ret, [], "fn_0arg_0ret", "void", "fn_0arg_0ret()void"),
        ABISubroutineTC(
            fn_0arg_uint64_ret,
            [],
            "fn_0arg_uint64_ret",
            pt.abi.Uint64TypeSpec(),
            "fn_0arg_uint64_ret()uint64",
        ),
        ABISubroutineTC(
            fn_1arg_0ret,
            [pt.abi.Uint64()],
            "fn_1arg_0ret",
            "void",
            "fn_1arg_0ret(uint64)void",
        ),
        ABISubroutineTC(
            fn_1arg_1ret,
            [pt.abi.Uint64()],
            "fn_1arg_1ret",
            pt.abi.Uint64TypeSpec(),
            "fn_1arg_1ret(uint64)uint64",
        ),
        ABISubroutineTC(
            fn_2arg_0ret,
            [
                pt.abi.Uint64(),
                pt.abi.StaticArray(
                    pt.abi.StaticArrayTypeSpec(pt.abi.ByteTypeSpec(), 10)
                ),
            ],
            "fn_2arg_0ret",
            "void",
            "fn_2arg_0ret(uint64,byte[10])void",
        ),
        ABISubroutineTC(
            fn_2arg_1ret,
            [
                pt.abi.Uint64(),
                pt.abi.StaticArray(
                    pt.abi.StaticArrayTypeSpec(pt.abi.ByteTypeSpec(), 10)
                ),
            ],
            "fn_2arg_1ret",
            pt.abi.ByteTypeSpec(),
            "fn_2arg_1ret(uint64,byte[10])byte",
        ),
        ABISubroutineTC(
            fn_2arg_1ret_with_expr,
            [
                pt.Int(5),
                pt.abi.StaticArray(
                    pt.abi.StaticArrayTypeSpec(pt.abi.ByteTypeSpec(), 10)
                ),
            ],
            "fn_2arg_1ret_with_expr",
            pt.abi.ByteTypeSpec(),
            None,
        ),
        ABISubroutineTC(
            fn_w_tuple1arg,
            [
                pt.Int(5),
                pt.abi.make(pt.abi.Tuple1[pt.abi.Byte]),
            ],
            "fn_w_tuple1arg",
            pt.abi.ByteTypeSpec(),
            None,
        ),
    )

    for case in cases:
        assert case.definition.subroutine.argument_count() == len(case.arg_instances)
        assert case.definition.name() == case.name

        if len(case.arg_instances) > 0:
            with pytest.raises(pt.TealInputError):
                case.definition(*case.arg_instances[:-1])

        with pytest.raises(pt.TealInputError):
            case.definition(*(case.arg_instances + [pt.abi.Uint64()]))

        assert case.definition.type_of() == case.ret_type
        invoked = case.definition(*case.arg_instances)
        assert isinstance(
            invoked, (pt.Expr if case.ret_type == "void" else pt.abi.ReturnedValue)
        )
        assert case.definition.is_abi_routable() == all(
            map(lambda x: isinstance(x, pt.abi.BaseType), case.arg_instances)
        )

        if case.definition.is_abi_routable():
            assert case.definition.method_signature() == cast(str, case.signature)
        else:
            with pytest.raises(pt.TealInputError):
                case.definition.method_signature()


def test_subroutine_return_reference():
    @ABIReturnSubroutine
    def invalid_ret_type(*, output: pt.abi.Account):
        return output.decode(pt.Bytes(b"\x00"))

    with pytest.raises(pt.TealInputError):
        invalid_ret_type.method_signature()

    @ABIReturnSubroutine
    def invalid_ret_type_collection(
        *, output: pt.abi.Tuple2[pt.abi.Account, pt.abi.Uint64]
    ):
        return output.set(pt.abi.Account(), pt.abi.Uint64())

    with pytest.raises(pt.TealInputError):
        invalid_ret_type_collection.method_signature()

    @ABIReturnSubroutine
    def invalid_ret_type_collection_nested(
        *, output: pt.abi.DynamicArray[pt.abi.Tuple2[pt.abi.Account, pt.abi.Uint64]]
    ):
        return output.set(
            pt.abi.make(
                pt.abi.DynamicArray[pt.abi.Tuple2[pt.abi.Account, pt.abi.Uint64]]
            )
        )

    with pytest.raises(pt.TealInputError):
        invalid_ret_type_collection_nested.method_signature()


def test_subroutine_definition_validate():
    """
    DFS through SubroutineDefinition.validate()'s logic
    """

    def mock_subroutine_definition(implementation, has_abi_output=False):
        mock = pt.SubroutineDefinition(lambda: pt.Return(pt.Int(1)), pt.TealType.uint64)
        mock._validate()  # haven't failed with dummy implementation
        mock.implementation = implementation
        mock.has_abi_output = has_abi_output
        return mock

    not_callable = mock_subroutine_definition("I'm not callable")
    with pytest.raises(pt.TealInputError) as tie:
        not_callable._validate()

    assert tie.value == pt.TealInputError(
        "Input to SubroutineDefinition is not callable"
    )

    # input_types:

    three_params = mock_subroutine_definition(lambda x, y, z: pt.Return(pt.Int(1)))

    params, anns, arg_types, byrefs, abi_args, output_kwarg = three_params._validate()
    assert len(params) == 3
    assert anns == {}
    assert all(at is pt.Expr for at in arg_types)
    assert byrefs == set()
    assert abi_args == {}
    assert output_kwarg == {}

    def bad_return_impl() -> str:
        return pt.Return(pt.Int(1))  # type: ignore

    bad_return = mock_subroutine_definition(bad_return_impl)
    with pytest.raises(pt.TealInputError) as tie:
        bad_return._validate()

    assert tie.value == pt.TealInputError(
        "Function has return of disallowed type <class 'str'>. Only Expr is allowed"
    )

    # now we iterate through the implementation params validating each as we go

    def var_abi_output_impl(*, output: pt.abi.Uint16):
        pt.Return(pt.Int(1))  # this is wrong but ignored

    # raises without abi_output_arg_name:
    var_abi_output_noname = mock_subroutine_definition(var_abi_output_impl)
    with pytest.raises(pt.TealInputError) as tie:
        var_abi_output_noname._validate()

    assert tie.value == pt.TealInputError(
        "Function has a parameter type that is not allowed in a subroutine: parameter output with type KEYWORD_ONLY"
    )

    # copacetic abi output:
    var_abi_output = mock_subroutine_definition(
        var_abi_output_impl, has_abi_output=True
    )
    params, anns, arg_types, byrefs, abi_args, output_kwarg = var_abi_output._validate()
    assert len(params) == 1
    assert anns == {"output": pt.abi.Uint16}
    assert all(at is pt.Expr for at in arg_types)
    assert byrefs == set()
    assert abi_args == {}
    assert output_kwarg == {"output": pt.abi.Uint16TypeSpec()}

    var_positional = mock_subroutine_definition(lambda *args: pt.Return(pt.Int(1)))
    with pytest.raises(pt.TealInputError) as tie:
        var_positional._validate()

    assert tie.value == pt.TealInputError(
        "Function has a parameter type that is not allowed in a subroutine: parameter args with type VAR_POSITIONAL"
    )

    kw_only = mock_subroutine_definition(lambda *, kw: pt.Return(pt.Int(1)))
    with pytest.raises(pt.TealInputError) as tie:
        kw_only._validate()

    assert tie.value == pt.TealInputError(
        "Function has a parameter type that is not allowed in a subroutine: parameter kw with type KEYWORD_ONLY"
    )

    var_keyword = mock_subroutine_definition(lambda **kw: pt.Return(pt.Int(1)))
    with pytest.raises(pt.TealInputError) as tie:
        var_keyword._validate()

    assert tie.value == pt.TealInputError(
        "Function has a parameter type that is not allowed in a subroutine: parameter kw with type VAR_KEYWORD"
    )

    param_default = mock_subroutine_definition(lambda x="niiiice": pt.Return(pt.Int(1)))
    with pytest.raises(pt.TealInputError) as tie:
        param_default._validate()

    assert tie.value == pt.TealInputError(
        "Function has a parameter with a default value, which is not allowed in a subroutine: x"
    )

    # Now we get to _validate_annotation():
    one_vanilla = mock_subroutine_definition(lambda x: pt.Return(pt.Int(1)))

    params, anns, arg_types, byrefs, abi_args, output_kwarg = one_vanilla._validate()
    assert len(params) == 1
    assert anns == {}
    assert all(at is pt.Expr for at in arg_types)
    assert byrefs == set()
    assert abi_args == {}
    assert output_kwarg == {}

    def one_expr_impl(x: pt.Expr):
        return pt.Return(pt.Int(1))

    one_expr = mock_subroutine_definition(one_expr_impl)
    params, anns, arg_types, byrefs, abi_args, output_kwarg = one_expr._validate()
    assert len(params) == 1
    assert anns == {"x": pt.Expr}
    assert all(at is pt.Expr for at in arg_types)
    assert byrefs == set()
    assert abi_args == {}
    assert output_kwarg == {}

    def one_scratchvar_impl(x: pt.ScratchVar):
        return pt.Return(pt.Int(1))

    one_scratchvar = mock_subroutine_definition(one_scratchvar_impl)
    params, anns, arg_types, byrefs, abi_args, output_kwarg = one_scratchvar._validate()
    assert len(params) == 1
    assert anns == {"x": pt.ScratchVar}
    assert all(at is pt.ScratchVar for at in arg_types)
    assert byrefs == {"x"}
    assert abi_args == {}
    assert output_kwarg == {}

    # not is_class()
    def one_nontype_impl(x: "blahBlah"):  # type: ignore # noqa: F821
        return pt.Return(pt.Int(1))

    one_nontype = mock_subroutine_definition(one_nontype_impl)
    with pytest.raises(pt.TealInputError) as tie:
        one_nontype._validate()

    assert tie.value == pt.TealInputError(
        "Function has parameter x of declared type blahBlah which is not a class"
    )

    def one_dynscratchvar_impl(x: pt.DynamicScratchVar):
        return pt.Return(pt.Int(1))

    one_dynscratchvar = mock_subroutine_definition(one_dynscratchvar_impl)
    with pytest.raises(pt.TealInputError) as tie:
        one_dynscratchvar._validate()

    assert tie.value == pt.TealInputError(
        "Function has parameter x of disallowed type <class 'pyteal.DynamicScratchVar'>. Only the types (<class 'pyteal.Expr'>, <class 'pyteal.ScratchVar'>, 'ABI') are allowed"
    )

    # Now we're back to _validate() main body and looking at input_types

    three_params_with_output = mock_subroutine_definition(
        lambda x, y, z, *, output: pt.Return(pt.Int(1)), has_abi_output=True
    )
    four_inputs = [
        pt.TealType.uint64,
        pt.TealType.uint64,
        pt.TealType.bytes,
        pt.TealType.uint64,
    ]

    two_inputs = [pt.TealType.uint64, pt.TealType.bytes]
    with pytest.raises(pt.TealInputError) as tie:
        three_params._validate(input_types=two_inputs)

    assert tie.value == pt.TealInputError(
        "Provided number of input_types (2) does not match detected number of input parameters (3)"
    )

    three_inputs_with_a_wrong_type = [pt.TealType.uint64, pt.Expr, pt.TealType.bytes]

    with pytest.raises(pt.TealInputError) as tie:
        three_params._validate(input_types=three_inputs_with_a_wrong_type)

    assert tie.value == pt.TealInputError(
        "Function has input type <class 'pyteal.Expr'> for parameter y which is not a TealType"
    )

    with pytest.raises(pt.TealInputError) as tie:
        three_params._validate(
            input_types=[pt.TealType.uint64, pt.Expr, pt.TealType.anytype]
        )

    assert tie.value == pt.TealInputError(
        "Function has input type <class 'pyteal.Expr'> for parameter y which is not a TealType"
    )

    with pytest.raises(pt.TealInputError) as tie:
        three_params._validate(
            input_types=[pt.TealType.uint64, None, pt.TealType.anytype]
        )
    assert tie.value == pt.TealInputError(
        "input_type for y is unspecified i.e. None but this is only allowed for ABI arguments"
    )

    # this one gets caught inside of _validate_annotation()
    with pytest.raises(pt.TealInputError) as tie:
        three_params_with_output._validate(input_types=four_inputs)

    # everything should be copacetic
    for x, y, z in product(pt.TealType, pt.TealType, pt.TealType):
        (
            params,
            anns,
            arg_types,
            byrefs,
            abi_args,
            output_kwarg,
        ) = three_params._validate(input_types=[x, y, z])
        assert len(params) == 3
        assert anns == {}
        assert all(at is pt.Expr for at in arg_types)
        assert byrefs == set()
        assert abi_args == {}
        assert output_kwarg == {}

    # annotation / abi type handling:
    abi_annotation_examples = {
        pt.abi.Address: pt.abi.AddressTypeSpec(),
        pt.abi.Bool: pt.abi.BoolTypeSpec(),
        pt.abi.Byte: pt.abi.ByteTypeSpec(),
        pt.abi.DynamicArray[pt.abi.Bool]: pt.abi.DynamicArrayTypeSpec(
            pt.abi.BoolTypeSpec()
        ),
        pt.abi.StaticArray[pt.abi.Uint32, Literal[10]]: pt.abi.StaticArrayTypeSpec(
            pt.abi.Uint32TypeSpec(), 10
        ),
        pt.abi.String: pt.abi.StringTypeSpec(),
        pt.abi.Tuple2[pt.abi.Bool, pt.abi.Uint32]: pt.abi.TupleTypeSpec(
            pt.abi.BoolTypeSpec(), pt.abi.Uint32TypeSpec()
        ),
        pt.abi.Uint8: pt.abi.Uint8TypeSpec(),
        pt.abi.Uint16: pt.abi.Uint16TypeSpec(),
        pt.abi.Uint32: pt.abi.Uint32TypeSpec(),
        pt.abi.Uint64: pt.abi.Uint64TypeSpec(),
    }

    anns = (pt.Expr, pt.ScratchVar) + tuple(abi_annotation_examples.keys())
    for x_ann, z_ann in product(anns, anns):

        def mocker_impl(x: x_ann, y, z: z_ann):
            return pt.Return(pt.Int(1))

        mocker = mock_subroutine_definition(mocker_impl)
        params, anns, arg_types, byrefs, abis, output_kwarg = mocker._validate()
        print(
            f"{x_ann=}, {z_ann=}, {params=}, {anns=}, {arg_types=}, {byrefs=}, {abis=}, {output_kwarg=}"
        )

        assert len(params) == 3

        assert anns == {"x": x_ann, "z": z_ann}

        assert (
            (arg_types[0] is x_ann or arg_types[0] == abi_annotation_examples[x_ann])
            and arg_types[1] is pt.Expr
            and (
                arg_types[2] is z_ann or arg_types[2] == abi_annotation_examples[z_ann]
            )
        ), f"{arg_types[0]} -> {x_ann} and {arg_types[1]} -> {pt.Expr} and {arg_types[2]} -> {z_ann}"

        assert byrefs == set(["x"] if x_ann is pt.ScratchVar else []) | set(
            ["z"] if z_ann is pt.ScratchVar else []
        )
        expected_abis = {}
        if x_ann not in (pt.Expr, pt.ScratchVar):
            expected_abis["x"] = abi_annotation_examples[x_ann]
        if z_ann not in (pt.Expr, pt.ScratchVar):
            expected_abis["z"] = abi_annotation_examples[z_ann]
        assert abis == expected_abis


def test_subroutine_invocation_param_types():
    def fnWithNoAnnotations(a, b):
        return pt.Return()

    def fnWithExprAnnotations(a: pt.Expr, b: pt.Expr) -> pt.Expr:
        return pt.Return()

    def fnWithSVAnnotations(a: pt.ScratchVar, b: pt.ScratchVar):
        return pt.Return()

    def fnWithABIAnnotations(
        a: pt.abi.Byte,
        b: pt.abi.StaticArray[pt.abi.Uint32, Literal[10]],
        c: pt.abi.DynamicArray[pt.abi.Bool],
    ):
        return pt.Return()

    def fnWithMixedAnns1(a: pt.ScratchVar, b: pt.Expr) -> pt.Expr:
        return pt.Return()

    def fnWithMixedAnns2(a: pt.ScratchVar, b) -> pt.Expr:
        return pt.Return()

    def fnWithMixedAnns3(a: pt.Expr, b: pt.ScratchVar):
        return pt.Return()

    def fnWithMixedAnns4(a: pt.ScratchVar, b, c: pt.abi.Uint16) -> pt.Expr:
        return pt.Return()

    sv = pt.ScratchVar()
    x = pt.Int(42)
    s = pt.Bytes("hello")
    av_u16 = pt.abi.Uint16()
    av_bool_dym_arr = pt.abi.DynamicArray(
        pt.abi.DynamicArrayTypeSpec(pt.abi.BoolTypeSpec())
    )
    av_u32_static_arr = pt.abi.StaticArray(
        pt.abi.StaticArrayTypeSpec(pt.abi.Uint32TypeSpec(), 10)
    )
    av_bool = pt.abi.Bool()
    av_byte = pt.abi.Byte()

    cases = [
        ("vanilla 1", fnWithNoAnnotations, [x, s], None),
        ("vanilla 2", fnWithNoAnnotations, [x, x], None),
        ("vanilla no sv's allowed 1", fnWithNoAnnotations, [x, sv], pt.TealInputError),
        ("exprs 1", fnWithExprAnnotations, [x, s], None),
        ("exprs 2", fnWithExprAnnotations, [x, x], None),
        ("exprs no sv's allowed 1", fnWithExprAnnotations, [x, sv], pt.TealInputError),
        ("all sv's 1", fnWithSVAnnotations, [sv, sv], None),
        ("all sv's but strings", fnWithSVAnnotations, [s, s], pt.TealInputError),
        ("all sv's but ints", fnWithSVAnnotations, [x, x], pt.TealInputError),
        (
            "all abi's 1",
            fnWithABIAnnotations,
            [av_byte, av_u32_static_arr, av_bool_dym_arr],
            None,
        ),
        (
            "all abi's but ints 1",
            fnWithABIAnnotations,
            [x, av_u32_static_arr, av_bool_dym_arr],
            pt.TealInputError,
        ),
        (
            "all abi's but ints 2",
            fnWithABIAnnotations,
            [x, av_u32_static_arr, x],
            pt.TealInputError,
        ),
        ("all abi's but ints 3", fnWithABIAnnotations, [x, x, x], pt.TealInputError),
        (
            "all abi's but sv's 1",
            fnWithABIAnnotations,
            [sv, av_u32_static_arr, av_bool_dym_arr],
            pt.TealInputError,
        ),
        (
            "all abi's but sv's 2",
            fnWithABIAnnotations,
            [av_byte, av_u32_static_arr, sv],
            pt.TealInputError,
        ),
        (
            "all abi's but sv's 3",
            fnWithABIAnnotations,
            [av_byte, sv, av_u32_static_arr],
            pt.TealInputError,
        ),
        (
            "all abi's but wrong typed 1",
            fnWithABIAnnotations,
            [av_u32_static_arr, av_u32_static_arr, av_bool_dym_arr],
            pt.TealInputError,
        ),
        (
            "all abi's but wrong typed 2",
            fnWithABIAnnotations,
            [av_bool, av_bool_dym_arr, av_u16],
            pt.TealInputError,
        ),
        (
            "all abi's but wrong typed 3",
            fnWithABIAnnotations,
            [av_u16, av_bool, av_byte],
            pt.TealInputError,
        ),
        ("mixed1 copacetic", fnWithMixedAnns1, [sv, x], None),
        ("mixed1 flipped", fnWithMixedAnns1, [x, sv], pt.TealInputError),
        ("mixed1 missing the sv", fnWithMixedAnns1, [x, s], pt.TealInputError),
        ("mixed1 missing the non-sv", fnWithMixedAnns1, [sv, sv], pt.TealInputError),
        ("mixed2 copacetic", fnWithMixedAnns2, [sv, x], None),
        ("mixed2 flipped", fnWithMixedAnns2, [x, sv], pt.TealInputError),
        ("mixed2 missing the sv", fnWithMixedAnns2, [x, s], pt.TealInputError),
        ("mixed2 missing the non-sv", fnWithMixedAnns2, [sv, sv], pt.TealInputError),
        ("mixed3 copacetic", fnWithMixedAnns3, [s, sv], None),
        ("mixed3 flipped", fnWithMixedAnns3, [sv, x], pt.TealInputError),
        ("mixed3 missing the sv", fnWithMixedAnns3, [x, s], pt.TealInputError),
        ("mixed anno", fnWithMixedAnns4, [sv, x, av_u16], None),
        (
            "mixed anno but wrong typed 1",
            fnWithMixedAnns4,
            [av_byte, x, av_u16],
            pt.TealInputError,
        ),
        (
            "mixed anno but wrong typed 2",
            fnWithMixedAnns4,
            [sv, av_byte, sv],
            pt.TealInputError,
        ),
        (
            "mixed anno but wrong typed 3",
            fnWithMixedAnns4,
            [sv, x, av_byte],
            pt.TealInputError,
        ),
    ]
    for case_name, fn, args, err in cases:
        definition = pt.SubroutineDefinition(fn, pt.TealType.none)
        assert definition.argument_count() == len(args), case_name
        assert definition.name() == fn.__name__, case_name

        if err is None:
            assert len(definition.by_ref_args) == len(
                [x for x in args if isinstance(x, pt.ScratchVar)]
            ), case_name

            invocation = definition.invoke(args)
            assert isinstance(invocation, pt.SubroutineCall), case_name
            assert invocation.subroutine is definition, case_name
            assert invocation.args == args, case_name
            assert invocation.has_return() is False, case_name

        else:
            try:
                with pytest.raises(err):
                    definition.invoke(args)
            except Exception as e:
                assert (
                    not e
                ), f"EXPECTED ERROR of type {err}. encountered unexpected error during invocation case <{case_name}>: {e}"


def test_abi_subroutine_calling_param_types():
    @pt.ABIReturnSubroutine
    def fn_log_add(a: pt.abi.Uint64, b: pt.abi.Uint32) -> pt.Expr:
        return pt.Seq(pt.Log(pt.Itob(a.get() + b.get())), pt.Return())

    @pt.ABIReturnSubroutine
    def fn_ret_add(
        a: pt.abi.Uint64, b: pt.abi.Uint32, *, output: pt.abi.Uint64
    ) -> pt.Expr:
        return output.set(a.get() + b.get() + pt.Int(0xA190))

    @pt.ABIReturnSubroutine
    def fn_abi_annotations_0(
        a: pt.abi.Byte,
        b: pt.abi.StaticArray[pt.abi.Uint32, Literal[10]],
        c: pt.abi.DynamicArray[pt.abi.Bool],
    ) -> pt.Expr:
        return pt.Return()

    @pt.ABIReturnSubroutine
    def fn_abi_annotations_0_with_ret(
        a: pt.abi.Byte,
        b: pt.abi.StaticArray[pt.abi.Uint32, Literal[10]],
        c: pt.abi.DynamicArray[pt.abi.Bool],
        *,
        output: pt.abi.Byte,
    ):
        return output.set(a)

    @pt.ABIReturnSubroutine
    def fn_mixed_annotations_0(a: pt.ScratchVar, b: pt.Expr, c: pt.abi.Byte) -> pt.Expr:
        return pt.Seq(
            a.store(c.get() * pt.Int(0x0FF1CE) * b),
            pt.Return(),
        )

    @pt.ABIReturnSubroutine
    def fn_mixed_annotations_0_with_ret(
        a: pt.ScratchVar, b: pt.Expr, c: pt.abi.Byte, *, output: pt.abi.Uint64
    ) -> pt.Expr:
        return pt.Seq(
            a.store(c.get() * pt.Int(0x0FF1CE) * b),
            output.set(a.load()),
        )

    @pt.ABIReturnSubroutine
    def fn_mixed_annotation_1(
        a: pt.ScratchVar, b: pt.abi.StaticArray[pt.abi.Uint32, Literal[10]]
    ) -> pt.Expr:
        return pt.Seq(
            (intermediate := pt.abi.Uint32()).set(b[a.load() % pt.Int(10)]),
            a.store(intermediate.get()),
            pt.Return(),
        )

    @pt.ABIReturnSubroutine
    def fn_mixed_annotation_1_with_ret(
        a: pt.ScratchVar, b: pt.abi.Uint64, *, output: pt.abi.Bool
    ) -> pt.Expr:
        return output.set((a.load() + b.get()) % pt.Int(2))

    abi_u64 = pt.abi.Uint64()
    abi_u32 = pt.abi.Uint32()
    abi_byte = pt.abi.Byte()
    abi_static_u32_10 = pt.abi.StaticArray(
        pt.abi.StaticArrayTypeSpec(pt.abi.Uint32TypeSpec(), 10)
    )
    abi_dynamic_bool = pt.abi.DynamicArray(
        pt.abi.DynamicArrayTypeSpec(pt.abi.BoolTypeSpec())
    )
    sv = pt.ScratchVar()
    expr_int = pt.Int(1)

    cases = [
        ("vanilla 1", fn_log_add, [abi_u64, abi_u32], "void", None),
        (
            "vanilla 1 with wrong ABI type",
            fn_log_add,
            [abi_u64, abi_u64],
            None,
            pt.TealInputError,
        ),
        (
            "vanilla 1 with ABI return",
            fn_ret_add,
            [abi_u64, abi_u32],
            pt.abi.Uint64TypeSpec(),
            None,
        ),
        (
            "vanilla 1 with ABI return wrong typed",
            fn_ret_add,
            [abi_u32, abi_u64],
            None,
            pt.TealInputError,
        ),
        (
            "full ABI annotations no return",
            fn_abi_annotations_0,
            [abi_byte, abi_static_u32_10, abi_dynamic_bool],
            "void",
            None,
        ),
        (
            "full ABI annotations wrong input 0",
            fn_abi_annotations_0,
            [abi_u64, abi_static_u32_10, abi_dynamic_bool],
            None,
            pt.TealInputError,
        ),
        (
            "full ABI annotations with ABI return",
            fn_abi_annotations_0_with_ret,
            [abi_byte, abi_static_u32_10, abi_dynamic_bool],
            pt.abi.ByteTypeSpec(),
            None,
        ),
        (
            "full ABI annotations with ABI return wrong inputs",
            fn_abi_annotations_0_with_ret,
            [abi_byte, abi_dynamic_bool, abi_static_u32_10],
            None,
            pt.TealInputError,
        ),
        (
            "mixed with ABI annotations 0",
            fn_mixed_annotations_0,
            [sv, expr_int, abi_byte],
            "void",
            None,
        ),
        (
            "mixed with ABI annotations 0 wrong inputs",
            fn_mixed_annotations_0,
            [abi_u64, expr_int, abi_byte],
            None,
            pt.TealInputError,
        ),
        (
            "mixed with ABI annotations 0 with ABI return",
            fn_mixed_annotations_0_with_ret,
            [sv, expr_int, abi_byte],
            pt.abi.Uint64TypeSpec(),
            None,
        ),
        (
            "mixed with ABI annotations 0 with ABI return wrong inputs",
            fn_mixed_annotations_0_with_ret,
            [sv, expr_int, sv],
            None,
            pt.TealInputError,
        ),
        (
            "mixed with ABI annotations 1",
            fn_mixed_annotation_1,
            [sv, abi_static_u32_10],
            "void",
            None,
        ),
        (
            "mixed with ABI annotations 1 with ABI return",
            fn_mixed_annotation_1_with_ret,
            [sv, abi_u64],
            pt.abi.BoolTypeSpec(),
            None,
        ),
        (
            "mixed with ABI annotations 1 with ABI return wrong inputs",
            fn_mixed_annotation_1_with_ret,
            [expr_int, abi_static_u32_10],
            None,
            pt.TealInputError,
        ),
    ]

    for case_name, definition, args, ret_type, err in cases:
        assert definition.subroutine.argument_count() == len(args), case_name
        assert (
            definition.name() == definition.subroutine.implementation.__name__
        ), case_name

        if err is None:
            invocation = definition(*args)
            if ret_type == "void":
                assert isinstance(invocation, pt.SubroutineCall), case_name
                assert not invocation.has_return(), case_name
                assert invocation.args == args, case_name
            else:
                assert isinstance(invocation, pt.abi.ReturnedValue), case_name
                assert invocation.type_spec == ret_type
                assert isinstance(invocation.computation, pt.SubroutineCall), case_name
                assert not invocation.computation.has_return(), case_name
                assert invocation.computation.args == args, case_name
        else:
            try:
                with pytest.raises(err):
                    definition(*args)
            except Exception as e:
                assert (
                    not e
                ), f"EXPECTED ERROR of type {err}. encountered unexpected error during invocation case <{case_name}>: {e}"


def test_subroutine_definition_invalid():
    def fnWithDefaults(a, b=None):
        return pt.Return()

    def fnWithKeywordArgs(a, *, output):
        return pt.Return()

    def fnWithKeywordArgsWrongKWName(a, *, b: pt.abi.Uint64):
        return pt.Return()

    def fnWithMultipleABIKeywordArgs(a, *, b: pt.abi.Byte, c: pt.abi.Bool):
        return pt.Return()

    def fnWithVariableArgs(a, *b):
        return pt.Return()

    def fnWithNonExprReturnAnnotation(a, b) -> pt.TealType.uint64:
        return pt.Return()

    def fnWithNonExprParamAnnotation(a, b: pt.TealType.uint64):
        return pt.Return()

    def fnWithScratchVarSubclass(a, b: pt.DynamicScratchVar):
        return pt.Return()

    def fnReturningExprSubclass(a: pt.ScratchVar, b: pt.Expr) -> pt.Return:
        return pt.Return()

    def fnWithMixedAnns4AndBytesReturn(a: pt.Expr, b: pt.ScratchVar) -> pt.Bytes:
        return pt.Bytes("hello uwu")

    def fnWithMixedAnnsABIRet1(
        a: pt.Expr, b: pt.ScratchVar, c: pt.abi.Uint16
    ) -> pt.abi.StaticArray[pt.abi.Uint32, Literal[10]]:
        return pt.abi.StaticArray(
            pt.abi.StaticArrayTypeSpec(pt.abi.Uint32TypeSpec(), 10)
        )

    def fnWithMixedAnnsABIRet2(
        a: pt.Expr, b: pt.abi.Byte, c: pt.ScratchVar
    ) -> pt.abi.Uint64:
        return pt.abi.Uint64()

    cases = (
        (
            1,
            "TealInputError('Input to SubroutineDefinition is not callable'",
            "TealInputError('Input to ABIReturnSubroutine is not callable'",
        ),
        (
            None,
            "TealInputError('Input to SubroutineDefinition is not callable'",
            "TealInputError('Input to ABIReturnSubroutine is not callable'",
        ),
        (
            fnWithDefaults,
            "TealInputError('Function has a parameter with a default value, which is not allowed in a subroutine: b'",
            "TealInputError('Function has a parameter with a default value, which is not allowed in a subroutine: b'",
        ),
        (
            fnWithKeywordArgs,
            "TealInputError('Function has a parameter type that is not allowed in a subroutine: parameter output with type",
            "TealInputError('ABI return subroutine output-kwarg output must specify ABI type')",
        ),
        (
            fnWithKeywordArgsWrongKWName,
            "TealInputError('Function has a parameter type that is not allowed in a subroutine: parameter b with type",
            "TealInputError('ABI return subroutine output-kwarg name must be `output` at this moment",
        ),
        (
            fnWithMultipleABIKeywordArgs,
            "TealInputError('Function has a parameter type that is not allowed in a subroutine: parameter b with type",
            "multiple output arguments (2) with type annotations",
        ),
        (
            fnWithVariableArgs,
            "TealInputError('Function has a parameter type that is not allowed in a subroutine: parameter b with type",
            "Function has a parameter type that is not allowed in a subroutine: parameter b with type VAR_POSITIONAL",
        ),
        (
            fnWithNonExprReturnAnnotation,
            "Function has return of disallowed type TealType.uint64. Only Expr is allowed",
            "Function has return of disallowed type TealType.uint64. Only Expr is allowed",
        ),
        (
            fnWithNonExprParamAnnotation,
            "Function has parameter b of declared type TealType.uint64 which is not a class",
            "Function has parameter b of declared type TealType.uint64 which is not a class",
        ),
        (
            fnWithScratchVarSubclass,
            "Function has parameter b of disallowed type <class 'pyteal.DynamicScratchVar'>",
            "Function has parameter b of disallowed type <class 'pyteal.DynamicScratchVar'>",
        ),
        (
            fnReturningExprSubclass,
            "Function has return of disallowed type <class 'pyteal.Return'>",
            "Function has return of disallowed type <class 'pyteal.Return'>. Only Expr is allowed",
        ),
        (
            fnWithMixedAnns4AndBytesReturn,
            "Function has return of disallowed type <class 'pyteal.Bytes'>",
            "Function has return of disallowed type <class 'pyteal.Bytes'>. Only Expr is allowed",
        ),
        (
            fnWithMixedAnnsABIRet1,
            "Function has return of disallowed type pyteal.abi.StaticArray[pyteal.abi.Uint32, typing.Literal[10]]. "
            "Only Expr is allowed",
            "Function has return of disallowed type pyteal.abi.StaticArray[pyteal.abi.Uint32, typing.Literal[10]]. "
            "Only Expr is allowed",
        ),
        (
            fnWithMixedAnnsABIRet2,
            "Function has return of disallowed type <class 'pyteal.abi.Uint64'>. Only Expr is allowed",
            "Function has return of disallowed type <class 'pyteal.abi.Uint64'>. Only Expr is allowed",
        ),
    )

    for fn, sub_def_msg, abi_sub_def_msg in cases:
        with pytest.raises(pt.TealInputError) as e:
            print(f"case=[{sub_def_msg}]")
            pt.SubroutineDefinition(fn, pt.TealType.none)

        assert sub_def_msg in str(e), f"failed for case [{fn.__name__}]"

        with pytest.raises(pt.TealInputError) as e:
            print(f"case=[{abi_sub_def_msg}]")
            pt.ABIReturnSubroutine(fn)

        assert abi_sub_def_msg in str(e), f"failed for case[{fn.__name__}]"


def test_subroutine_declaration():
    cases = (
        (pt.TealType.none, pt.Return()),
        (pt.TealType.uint64, pt.Return(pt.Int(1))),
        (pt.TealType.uint64, pt.Int(1)),
        (pt.TealType.bytes, pt.Bytes("value")),
        (pt.TealType.anytype, pt.App.globalGet(pt.Bytes("key"))),
    )

    for (returnType, value) in cases:

        def mySubroutine():
            return value

        definition = pt.SubroutineDefinition(mySubroutine, returnType)

        declaration = pt.SubroutineDeclaration(definition, value)
        assert declaration.type_of() == value.type_of()
        assert declaration.has_return() == value.has_return()

        options.currentSubroutine = definition
        assert declaration.__teal__(options) == value.__teal__(options)
        options.setSubroutine(None)


def test_subroutine_call():
    def mySubroutine():
        return pt.Return()

    returnTypes = (
        pt.TealType.uint64,
        pt.TealType.bytes,
        pt.TealType.anytype,
        pt.TealType.none,
    )

    argCases = (
        [],
        [pt.Int(1)],
        [pt.Int(1), pt.Bytes("value")],
    )

    for returnType in returnTypes:
        definition = pt.SubroutineDefinition(mySubroutine, returnType)

        for args in argCases:
            expr = pt.SubroutineCall(definition, args)

            assert expr.type_of() == returnType
            assert not expr.has_return()

            expected, _ = pt.TealBlock.FromOp(
                options, pt.TealOp(expr, pt.Op.callsub, definition), *args
            )

            actual, _ = expr.__teal__(options)

            assert actual == expected


def test_decorator():
    assert callable(pt.Subroutine)
    assert callable(pt.Subroutine(pt.TealType.anytype))

    @pt.Subroutine(pt.TealType.none)
    def mySubroutine(a):
        return pt.Return()

    assert isinstance(mySubroutine, pt.SubroutineFnWrapper)

    invocation = mySubroutine(pt.Int(1))
    assert isinstance(invocation, pt.SubroutineCall)

    with pytest.raises(pt.TealInputError):
        mySubroutine()

    with pytest.raises(pt.TealInputError):
        mySubroutine(pt.Int(1), pt.Int(2))

    with pytest.raises(pt.TealInputError):
        mySubroutine(pt.Pop(pt.Int(1)))

    with pytest.raises(pt.TealInputError):
        mySubroutine(1)

    with pytest.raises(pt.TealInputError):
        mySubroutine(a=pt.Int(1))


def test_evaluate_subroutine_no_args():
    cases = (
        (pt.TealType.none, pt.Return()),
        (pt.TealType.uint64, pt.Int(1) + pt.Int(2)),
        (pt.TealType.uint64, pt.Return(pt.Int(1) + pt.Int(2))),
        (pt.TealType.bytes, pt.Bytes("value")),
        (pt.TealType.bytes, pt.Return(pt.Bytes("value"))),
    )

    for (returnType, returnValue) in cases:

        def mySubroutine():
            return returnValue

        definition = pt.SubroutineDefinition(mySubroutine, returnType)
        declaration = evaluate_subroutine(definition)

        assert isinstance(declaration, pt.SubroutineDeclaration)
        assert declaration.subroutine is definition

        assert declaration.type_of() == returnValue.type_of()
        assert declaration.has_return() == returnValue.has_return()

        options.setSubroutine(definition)
        expected, _ = pt.Seq([returnValue]).__teal__(options)

        actual, _ = declaration.__teal__(options)
        options.setSubroutine(None)
        assert actual == expected


def test_evaluate_subroutine_1_arg():
    cases = (
        (pt.TealType.none, pt.Return()),
        (pt.TealType.uint64, pt.Int(1) + pt.Int(2)),
        (pt.TealType.uint64, pt.Return(pt.Int(1) + pt.Int(2))),
        (pt.TealType.bytes, pt.Bytes("value")),
        (pt.TealType.bytes, pt.Return(pt.Bytes("value"))),
    )

    for (returnType, returnValue) in cases:
        argSlots: List[pt.ScratchSlot] = []

        def mySubroutine(a1):
            assert isinstance(a1, pt.ScratchLoad)
            argSlots.append(a1.slot)
            return returnValue

        definition = pt.SubroutineDefinition(mySubroutine, returnType)
        declaration = evaluate_subroutine(definition)

        assert isinstance(declaration, pt.SubroutineDeclaration)
        assert declaration.subroutine is definition

        assert declaration.type_of() == returnValue.type_of()
        assert declaration.has_return() == returnValue.has_return()

        assert isinstance(declaration.body, pt.Seq)
        assert len(declaration.body.args) == 2

        assert isinstance(declaration.body.args[0], pt.ScratchStackStore)

        assert declaration.body.args[0].slot is argSlots[-1]

        options.setSubroutine(definition)
        expected, _ = pt.Seq([declaration.body.args[0], returnValue]).__teal__(options)

        actual, _ = declaration.__teal__(options)
        options.setSubroutine(None)
        assert actual == expected


def test_evaluate_subroutine_2_args():
    cases = (
        (pt.TealType.none, pt.Return()),
        (pt.TealType.uint64, pt.Int(1) + pt.Int(2)),
        (pt.TealType.uint64, pt.Return(pt.Int(1) + pt.Int(2))),
        (pt.TealType.bytes, pt.Bytes("value")),
        (pt.TealType.bytes, pt.Return(pt.Bytes("value"))),
    )

    for (returnType, returnValue) in cases:
        argSlots: List[pt.ScratchSlot] = []

        def mySubroutine(a1, a2):
            assert isinstance(a1, pt.ScratchLoad)
            argSlots.append(a1.slot)
            assert isinstance(a2, pt.ScratchLoad)
            argSlots.append(a2.slot)
            return returnValue

        definition = pt.SubroutineDefinition(mySubroutine, returnType)

        declaration = evaluate_subroutine(definition)

        assert isinstance(declaration, pt.SubroutineDeclaration)
        assert declaration.subroutine is definition

        assert declaration.type_of() == returnValue.type_of()
        assert declaration.has_return() == returnValue.has_return()

        assert isinstance(declaration.body, pt.Seq)
        assert len(declaration.body.args) == 3

        assert isinstance(declaration.body.args[0], pt.ScratchStackStore)
        assert isinstance(declaration.body.args[1], pt.ScratchStackStore)

        assert declaration.body.args[0].slot is argSlots[-1]
        assert declaration.body.args[1].slot is argSlots[-2]

        options.setSubroutine(definition)
        expected, _ = pt.Seq(
            [declaration.body.args[0], declaration.body.args[1], returnValue]
        ).__teal__(options)

        actual, _ = declaration.__teal__(options)
        options.setSubroutine(None)
        assert actual == expected


def test_evaluate_subroutine_10_args():
    cases = (
        (pt.TealType.none, pt.Return()),
        (pt.TealType.uint64, pt.Int(1) + pt.Int(2)),
        (pt.TealType.uint64, pt.Return(pt.Int(1) + pt.Int(2))),
        (pt.TealType.bytes, pt.Bytes("value")),
        (pt.TealType.bytes, pt.Return(pt.Bytes("value"))),
    )

    for (returnType, returnValue) in cases:
        argSlots: List[pt.ScratchSlot] = []

        def mySubroutine(a1, a2, a3, a4, a5, a6, a7, a8, a9, a10):
            for a in (a1, a2, a3, a4, a5, a6, a7, a8, a9, a10):
                assert isinstance(a, pt.ScratchLoad)
                argSlots.append(a.slot)
            return returnValue

        definition = pt.SubroutineDefinition(mySubroutine, returnType)
        declaration = evaluate_subroutine(definition)

        assert isinstance(declaration, pt.SubroutineDeclaration)
        assert declaration.subroutine is definition

        assert declaration.type_of() == returnValue.type_of()
        assert declaration.has_return() == returnValue.has_return()

        assert isinstance(declaration.body, pt.Seq)
        assert len(declaration.body.args) == 11

        for i in range(10):
            assert isinstance(declaration.body.args[i], pt.ScratchStackStore)

        for i in range(10):
            assert declaration.body.args[i].slot is argSlots[-i - 1]

        options.setSubroutine(definition)
        expected, _ = pt.Seq(declaration.body.args[:10] + [returnValue]).__teal__(
            options
        )

        actual, _ = declaration.__teal__(options)
        options.setSubroutine(None)
        assert actual == expected
