import abc
import typing
from collections.abc import Callable, Sequence

import attrs
import mypy.nodes
import mypy.patterns
import mypy.types
import mypy.visitor

from puya import log
from puya.awst.nodes import (
    AppStateExpression,
    AppStorageDefinition,
    ARC4MethodConfig,
    AssertExpression,
    AssignmentExpression,
    AssignmentStatement,
    BinaryBooleanOperator,
    Block,
    BooleanBinaryOperation,
    BytesConstant,
    BytesEncoding,
    ConditionalExpression,
    ContractMethod,
    Expression,
    ExpressionStatement,
    ForInLoop,
    Goto,
    IfElse,
    InstanceSuperMethodTarget,
    Label,
    LoopContinue,
    LoopExit,
    Lvalue,
    ReturnStatement,
    Statement,
    Subroutine,
    SubroutineArgument,
    SubroutineID,
    Switch,
    VarExpression,
    WhileLoop,
)
from puya.errors import CodeError, InternalError
from puya.parse import SourceLocation
from puya.program_refs import ContractReference, LogicSigReference
from puyapy.awst_build import constants, pytypes
from puyapy.awst_build.base_mypy_visitor import BaseMyPyExpressionVisitor, BaseMyPyStatementVisitor
from puyapy.awst_build.context import ASTConversionModuleContext
from puyapy.awst_build.eb import _expect as expect
from puyapy.awst_build.eb._literals import LiteralBuilderImpl
from puyapy.awst_build.eb.arc4 import ARC4ClientTypeBuilder
from puyapy.awst_build.eb.conditional_literal import ConditionalLiteralBuilder
from puyapy.awst_build.eb.contracts import (
    ContractSelfExpressionBuilder,
    ContractTypeExpressionBuilder,
)
from puyapy.awst_build.eb.dict_ import DictLiteralBuilder
from puyapy.awst_build.eb.factories import (
    builder_for_instance,
    builder_for_type,
    try_get_builder_for_func,
)
from puyapy.awst_build.eb.interface import (
    BuilderBinaryOp,
    BuilderComparisonOp,
    BuilderUnaryOp,
    CallableBuilder,
    InstanceBuilder,
    LiteralBuilder,
    NodeBuilder,
    StorageProxyConstructorResult,
)
from puyapy.awst_build.eb.logicsig import LogicSigExpressionBuilder
from puyapy.awst_build.eb.none import NoneTypeBuilder
from puyapy.awst_build.eb.subroutine import SubroutineInvokerExpressionBuilder
from puyapy.awst_build.utils import (
    extract_bytes_literal_from_mypy,
    get_decorators_by_fullname,
    get_unaliased_fullname,
    maybe_resolve_literal,
    require_callable_type,
)
from puyapy.models import ContractFragmentBase
from puyapy.parse import parse_docstring

logger = log.get_logger(__name__)

@attrs.frozen
class ContractMethodInfo:
    fragment: ContractFragmentBase
    contract_type: pytypes.ContractType
    is_abstract: bool
    arc4_method_config: ARC4MethodConfig | None

class ExpressionASTConverter(BaseMyPyExpressionVisitor[NodeBuilder], abc.ABC):
    @abc.abstractmethod
    def resolve_local_type(
        self, var_name: str, expr_loc: SourceLocation
    ) -> pytypes.PyType | None: ...

    @abc.abstractmethod
    def builder_for_self(self, expr_loc: SourceLocation) -> NodeBuilder: ...

    # Expressions
    def _visit_ref_expr(self, expr: mypy.nodes.MemberExpr | mypy.nodes.NameExpr) -> NodeBuilder:
        expr_loc = self._location(expr)
        # Do a straight forward lookup at the RefExpr level handle the cases of:
        #  - type aliases
        #  - simple (non-generic) types
        #  - generic type that has not been parameterised
        #    (e.g. when calling a constructor which can infer the type from its arguments)
        # For parameterised generics, these are resolved at the IndexExpr level, without
        # descending into IndexExpr.base.
        # By doing a simple lookup instead of resolving the PyType of expr,
        # we can side step complex constructs in the stubs that we don't support in user code,
        # such as overloads.
        if py_typ := self.context.lookup_pytype(expr.fullname):
            if isinstance(py_typ, pytypes.ContractType):
                if fragment := self.context.contract_fragments.get(py_typ.name):
                    return ContractTypeExpressionBuilder(py_typ, fragment, expr_loc)
            elif pytypes.ARC4ClientBaseType < py_typ:  # provides type info only
                if fragment := self.context.contract_fragments.get(ContractReference(py_typ.name)):
                    return ARC4ClientTypeBuilder(py_typ, expr_loc, fragment)
            elif py_typ == pytypes.LogicSigType:
                ref = LogicSigReference(get_unaliased_fullname(expr))
                return LogicSigExpressionBuilder(ref, expr_loc)
            else:
                return builder_for_type(py_typ, expr_loc)

        if expr.name == "__all__":
            # special case here, we allow __all__ at the module level for it's "public vs private"
            # control implications w.r.t linting etc, but we do so by ignoring it.
            # so this is here just in case someone tries to reference __all__ inside a function,
            # to give a more useful error message.
            raise CodeError("__all__ cannot be referenced inside functions", expr_loc)

        fullname = get_unaliased_fullname(expr)
        if fullname.startswith("builtins."):
            return self._visit_ref_expr_of_builtins(fullname, expr_loc)
        if func_builder := try_get_builder_for_func(fullname, expr_loc):
            return func_builder
        match expr:
            case mypy.nodes.NameExpr(node=mypy.nodes.Var(is_self=True)):
                return self.builder_for_self(expr_loc)
            case mypy.nodes.RefExpr(
                node=mypy.nodes.Decorator() as dec
            ) if constants.SUBROUTINE_HINT in get_decorators_by_fullname(self.context, dec):
                self._precondition(
                    expr.fullname != expr.name,
                    "unqualified name found in call to function",
                    expr_loc,
                )
                mypy_func_type = require_callable_type(dec, expr_loc)
                func_type = self.context.type_to_pytype(mypy_func_type, source_location=expr_loc)
                if not isinstance(func_type, pytypes.FuncType):
                    raise CodeError("decorated function has non-function type", expr_loc)
                return SubroutineInvokerExpressionBuilder(
                    target=SubroutineID(expr.fullname),
                    func_type=func_type,
                    location=expr_loc,
                )
            case mypy.nodes.RefExpr(node=mypy.nodes.FuncDef()):
                raise CodeError(
                    f"Cannot invoke {fullname} as it is not "
                    f"decorated with {constants.SUBROUTINE_HINT_ALIAS}",
                    expr_loc,
                )
            # TODO: is this enough to filter to only global variable references?
            #       it seems like it should be...
            case mypy.nodes.RefExpr(kind=mypy.nodes.GDEF, node=mypy.nodes.Var()):
                self._precondition(
                    not (
                        expr.is_new_def
                        or expr.is_inferred_def
                        or expr.is_alias_rvalue
                        or (isinstance(expr, mypy.nodes.NameExpr) and expr.is_special_form)
                    ),
                    "global variable reference with unexpected flags",
                    expr_loc,
                )
                try:
                    constant_value = self.context.constants[expr.fullname]
                except KeyError as ex:
                    raise CodeError(
                        "Unable to resolve global constant reference", expr_loc
                    ) from ex
                else:
                    return LiteralBuilderImpl(source_location=expr_loc, value=constant_value)
            case (
                mypy.nodes.NameExpr(
                    kind=mypy.nodes.LDEF, node=mypy.nodes.Var(), name=var_name
                ) as name_expr
            ):
                self._precondition(
                    not name_expr.is_special_form,
                    "special form lvalues should only appear"
                    " as a singular lvalue in an assignment statement",
                    expr_loc,
                )
                local_type = self.resolve_local_type(var_name, expr_loc)
                if local_type is None:
                    raise CodeError(
                        "local variable type is unknown - possible use before definition?",
                        expr_loc,
                    )
                var_expr = VarExpression(
                    name=var_name,
                    wtype=local_type.checked_wtype(expr_loc),
                    source_location=expr_loc,
                )
                return builder_for_instance(local_type, var_expr)
        scope = {
            mypy.nodes.LDEF: "local",
            mypy.nodes.MDEF: "member",
            mypy.nodes.GDEF: "global",
            None: "unknown",
        }.get(expr.kind)
        # this can happen in otherwise well-formed code that is just missing a reference
        raise CodeError(
            f"Unable to resolve reference to {expr.fullname or expr.name!r}, {scope=}",
            expr_loc,
        )

    @staticmethod
    def _visit_ref_expr_of_builtins(fullname: str, location: SourceLocation) -> NodeBuilder:
        assert fullname.startswith("builtins.")
        rest_of_name = fullname.removeprefix("builtins.")
        match rest_of_name:
            case "True":
                return LiteralBuilderImpl(source_location=location, value=True)
            case "False":
                return LiteralBuilderImpl(source_location=location, value=False)
            case "None":
                raise CodeError("None is not supported as a value, only a return type", location)
            case "len":
                raise CodeError(
                    "len() is not supported -"
                    " types with a length will have a .length property instead",
                    location,
                )
            case "range":
                raise CodeError("range() is not supported - use algopy.urange() instead", location)
            case "enumerate":
                raise CodeError(
                    "enumerate() is not supported - use algopy.uenumerate() instead", location
                )
            case _:
                raise CodeError(f"Unsupported builtin: {rest_of_name}", location)

    @typing.override
    def visit_name_expr(self, expr: mypy.nodes.NameExpr) -> NodeBuilder:
        return self._visit_ref_expr(expr)

    @typing.override
    def visit_member_expr(self, expr: mypy.nodes.MemberExpr) -> NodeBuilder:
        expr_loc = self._location(expr)
        if isinstance(expr.expr, mypy.nodes.RefExpr) and isinstance(
            expr.expr.node, mypy.nodes.MypyFile
        ):
            # special case for module attribute access
            return self._visit_ref_expr(expr)

        base = expr.expr.accept(self)
        return base.member_access(expr.name, expr_loc)

    @typing.override
    def visit_call_expr(self, call: mypy.nodes.CallExpr) -> NodeBuilder:
        if call.analyzed is not None:
            return self._visit_special_call_expr(call, analyzed=call.analyzed)

        callee = call.callee.accept(self)
        if not isinstance(callee, CallableBuilder):
            raise CodeError("not a callable expression", self._location(call.callee))

        args = [arg.accept(self) for arg in call.args]
        return callee.call(
            args=args,
            arg_kinds=call.arg_kinds,
            arg_names=call.arg_names,
            location=self._location(call),
        )

    def _visit_special_call_expr(
        self, call: mypy.nodes.CallExpr, *, analyzed: mypy.nodes.Expression
    ) -> NodeBuilder:
        match analyzed:
            case mypy.nodes.CastExpr(expr=inner_expr):
                self.context.warning(
                    "use of typing.cast, output may be invalid or insecure TEAL", call
                )
            case mypy.nodes.AssertTypeExpr(expr=inner_expr):
                # just FYI... in case the user thinks this has a runtime effect
                # (it doesn't, not even in Python)
                self.context.warning(
                    "use of typing.assert_type has no effect on compilation", call
                )
            case mypy.nodes.RevealExpr(expr=mypy.nodes.Expression() as inner_expr):
                pass
            case _:
                raise CodeError(
                    f"Unsupported special function call"
                    f" of analyzed type {type(analyzed).__name__}",
                    self._location(call),
                )

        return inner_expr.accept(self)

    @typing.override
    def visit_unary_expr(self, node: mypy.nodes.UnaryExpr) -> NodeBuilder:
        expr_loc = self._location(node)
        builder_or_literal = node.expr.accept(self)
        match node.op:
            case "not":
                return builder_or_literal.bool_eval(expr_loc, negate=True)
            case op_str if op_str in BuilderUnaryOp:
                builder_op = BuilderUnaryOp(op_str)
                return require_instance_builder(builder_or_literal).unary_op(builder_op, expr_loc)
            case _:
                # guard against future python unary operators
                raise InternalError(f"Unable to interpret unary operator '{node.op}'", expr_loc)

    @typing.override
    def visit_op_expr(self, node: mypy.nodes.OpExpr) -> NodeBuilder:
        node_loc = self._location(node)
        lhs = require_instance_builder(node.left.accept(self))
        rhs = require_instance_builder(node.right.accept(self))

        # mypy combines ast.BoolOp and ast.BinOp, but they're kinda different...
        if node.op in BinaryBooleanOperator:
            bool_op = BinaryBooleanOperator(node.op)
            return lhs.bool_binary_op(rhs, bool_op, node_loc)

        try:
            op = BuilderBinaryOp(node.op)
        except ValueError as ex:
            raise InternalError(f"Unknown binary operator: {node.op}") from ex

        result: NodeBuilder = NotImplemented
        if isinstance(lhs, NodeBuilder):
            result = lhs.binary_op(other=rhs, op=op, location=node_loc, reverse=False)
        if result is NotImplemented and isinstance(rhs, NodeBuilder):
            result = rhs.binary_op(other=lhs, op=op, location=node_loc, reverse=True)
        if result is NotImplemented:
            raise CodeError(f"unsupported operation {op.value} between types", node_loc)
        return result

    @typing.override
    def visit_index_expr(self, expr: mypy.nodes.IndexExpr) -> NodeBuilder:
        expr_location = self._location(expr)
        if isinstance(expr.method_type, mypy.types.CallableType):
            result_pytyp = self.context.type_to_pytype(
                expr.method_type.ret_type, source_location=expr_location
            )
            if isinstance(result_pytyp, pytypes.PseudoGenericFunctionType):
                return builder_for_type(result_pytyp, expr_location)
        match expr.analyzed:
            case None:
                pass
            case mypy.nodes.TypeAliasExpr():
                raise CodeError("type aliases are not supported inside subroutines", expr_location)
            case mypy.nodes.TypeApplication(types=type_args):
                # TODO: resolve expr to generic type builder instead
                assert isinstance(
                    expr.base, mypy.nodes.RefExpr
                ), "expect base of type application to be RefExpr"
                pytyp = self.context.lookup_pytype(expr.base.fullname)
                if pytyp is None:
                    raise CodeError("unknown base type", expr_location)
                param_typ = pytyp.parameterise(
                    [
                        self.context.type_to_pytype(
                            t, in_type_args=True, source_location=expr_location
                        )
                        for t in type_args
                    ],
                    expr_location,
                )
                return builder_for_type(param_typ, expr_location)
            case _:
                typing.assert_never(expr.analyzed)

        base_builder = require_instance_builder(expr.base.accept(self))
        match expr.index:
            # special case handling of SliceExpr, so we don't need to handle slice Literal's
            # or some such everywhere
            # TODO: SliceBuilder?
            case mypy.nodes.SliceExpr(begin_index=begin, end_index=end, stride=stride):
                return base_builder.slice_index(
                    begin_index=(require_instance_builder(begin.accept(self)) if begin else None),
                    end_index=(require_instance_builder(end.accept(self)) if end else None),
                    stride=(require_instance_builder(stride.accept(self)) if stride else None),
                    location=expr_location,
                )

        index_builder = require_instance_builder(expr.index.accept(self))
        return base_builder.index(index=index_builder, location=expr_location)

    @typing.override
    def visit_conditional_expr(self, expr: mypy.nodes.ConditionalExpr) -> NodeBuilder:
        expr_loc = self._location(expr)
        condition = expr.cond.accept(self).bool_eval(self._location(expr.cond))
        true_b = require_instance_builder(expr.if_expr.accept(self))
        false_b = require_instance_builder(expr.else_expr.accept(self))
        if true_b.pytype <= false_b.pytype:
            expr_pytype = true_b.pytype
        elif false_b.pytype < true_b.pytype:
            expr_pytype = false_b.pytype
        else:
            raise CodeError(
                "incompatible result types for 'true' and 'false' expressions", expr_loc
            )

        if (
            isinstance(expr_pytype, pytypes.LiteralOnlyType)
            and isinstance(true_b, LiteralBuilder)
            and isinstance(false_b, LiteralBuilder)
        ):
            return ConditionalLiteralBuilder(
                true_literal=true_b,
                false_literal=false_b,
                condition=condition,
                location=expr_loc,
            )
        else:
            cond_expr = ConditionalExpression(
                condition=condition.resolve(),
                true_expr=true_b.resolve(),
                false_expr=false_b.resolve(),
                source_location=expr_loc,
            )
            return builder_for_instance(expr_pytype, cond_expr)

    @typing.override
    def visit_comparison_expr(self, expr: mypy.nodes.ComparisonExpr) -> NodeBuilder:
        from puyapy.awst_build.eb.bool import BoolExpressionBuilder

        expr_loc = self._location(expr)
        self._precondition(
            len(expr.operands) == (len(expr.operators) + 1),
            "operands don't match with operators",
            expr_loc,
        )
        # we can decompose this into a series of AND operations, but potentially need to generate
        # assignment expressions for everything except first and last operands, e.g.:
        #   a < b < c
        # becomes:
        #   (a < (tmp := b)) AND (tmp < c)
        # TODO: what about comparing different types that can never be equal? according to Python
        #       type signatures that should be possible, and we can always know it's value at
        #       compile time, but it would always result in a constant ...

        operands = [require_instance_builder(o.accept(self)) for o in expr.operands]
        operands[1:-1] = [operand.single_eval() for operand in list(operands)[1:-1]]

        comparisons = [
            self._build_compare(operator=operator, lhs=lhs, rhs=rhs).resolve()
            for operator, lhs, rhs in zip(expr.operators, operands, operands[1:], strict=False)
        ]

        result = prev = comparisons[0]
        for curr in comparisons[1:]:
            result = BooleanBinaryOperation(
                left=result,
                op=BinaryBooleanOperator.and_,
                right=curr,
                source_location=prev.source_location.try_merge(curr.source_location),
            )
            prev = curr
        return BoolExpressionBuilder(result)

    @staticmethod
    def _build_compare(
        operator: str, lhs: InstanceBuilder, rhs: InstanceBuilder
    ) -> InstanceBuilder:
        cmp_loc = lhs.source_location.try_merge(rhs.source_location)
        match operator:
            case "not in":
                return rhs.contains(lhs, cmp_loc).bool_eval(cmp_loc, negate=True)
            case "in":
                return rhs.contains(lhs, cmp_loc)

        result: InstanceBuilder = NotImplemented
        op = BuilderComparisonOp(operator)
        if isinstance(lhs, NodeBuilder):
            result = lhs.compare(other=rhs, op=op, location=cmp_loc)
        if result is NotImplemented and isinstance(rhs, NodeBuilder):
            result = rhs.compare(other=lhs, op=op.reversed(), location=cmp_loc)
        if result is NotImplemented:
            raise CodeError(f"unsupported comparison {operator!r} between types", cmp_loc)
        return result

    @typing.override
    def visit_int_expr(self, expr: mypy.nodes.IntExpr) -> LiteralBuilder:
        return LiteralBuilderImpl(value=expr.value, source_location=self._location(expr))

    @typing.override
    def visit_str_expr(self, expr: mypy.nodes.StrExpr) -> LiteralBuilder:
        return LiteralBuilderImpl(value=expr.value, source_location=self._location(expr))

    @typing.override
    def visit_bytes_expr(self, expr: mypy.nodes.BytesExpr) -> LiteralBuilder:
        bytes_const = extract_bytes_literal_from_mypy(expr)
        return LiteralBuilderImpl(value=bytes_const, source_location=self._location(expr))

    @typing.override
    def visit_tuple_expr(self, mypy_expr: mypy.nodes.TupleExpr) -> NodeBuilder:
        from puyapy.awst_build.eb.tuple import TupleLiteralBuilder

        location = self._location(mypy_expr)
        item_builders = [
            require_instance_builder(mypy_item.accept(self)) for mypy_item in mypy_expr.items
        ]
        return TupleLiteralBuilder(item_builders, location)

    @typing.override
    def visit_dict_expr(self, expr: mypy.nodes.DictExpr) -> NodeBuilder:
        location = self._location(expr)
        mappings = dict[str, InstanceBuilder]()
        for mypy_name, mypy_value in expr.items:
            if mypy_name is None:
                logger.error("None is not usable as a value", location=location)
                continue
            key_node = mypy_name.accept(self)
            value_node = mypy_value.accept(self)
            key = expect.simple_string_literal(key_node, default=expect.default_none)
            value = require_instance_builder(value_node)
            if key is not None:
                mappings[key] = value
        return DictLiteralBuilder(mappings, location)

    # unsupported expressions

    def visit_list_comprehension(self, expr: mypy.nodes.ListComprehension) -> NodeBuilder:
        raise CodeError("List comprehensions are not supported", self._location(expr))

    def visit_slice_expr(self, expr: mypy.nodes.SliceExpr) -> NodeBuilder:
        raise CodeError("Slices are not supported outside of indexing", self._location(expr))

    def visit_lambda_expr(self, expr: mypy.nodes.LambdaExpr) -> NodeBuilder:
        raise CodeError("lambda functions are not supported", self._location(expr))

    def visit_ellipsis(self, expr: mypy.nodes.EllipsisExpr) -> NodeBuilder:
        raise CodeError("ellipsis expressions are not supported", self._location(expr))

    def visit_list_expr(self, expr: mypy.nodes.ListExpr) -> NodeBuilder:
        raise CodeError("Python lists are not supported", self._location(expr))

class FunctionASTConverter(
    BaseMyPyStatementVisitor[Statement | Sequence[Statement] | None],
    ExpressionASTConverter,
):
    def __init__(
        self,
        context: ASTConversionModuleContext,
        func_def: mypy.nodes.FuncDef,
        contract_method_info: ContractMethodInfo | None,
        source_location: SourceLocation,
        *,
        inline: bool | None,
    ):
        super().__init__(context=context)
        func_loc = context.node_location(func_def)  # TODO: why not source_location ??
        self.contract_method_info = contract_method_info
        self.func_def = func_def
        self._precondition(
            not func_def.is_trivial_body,
            "trivial functions should be skipped at a higher level",
            func_loc,
        )
        type_info = func_def.type
        self._precondition(
            isinstance(type_info, mypy.types.CallableType),
            "should only receive functions/methods with a type of CallableType",
            func_loc,
        )
        assert isinstance(type_info, mypy.types.CallableType)
        self._precondition(
            len(type_info.arg_types) == len(func_def.arguments),
            "FuncDef argument type list length should match arguments list length",
            func_loc,
        )
        self._break_label_stack = list[Label | None]()
        # convert the return type
        self._return_type = context.type_to_pytype(type_info.ret_type, source_location=func_loc)
        return_wtype = self._return_type.checked_wtype(func_loc)
        # check & convert the arguments
        mypy_args = func_def.arguments
        mypy_arg_types = type_info.arg_types
        if type_info.def_extras.get("first_arg"):
            # function is a method
            if not mypy_args:
                logger.error("method declaration is missing 'self' argument", location=func_loc)
            else:
                self._precondition(
                    mypy_args[0].variable.is_self,
                    "if function is a method, first variable should be self-like",
                    func_loc,
                )
                mypy_args = mypy_args[1:]
                mypy_arg_types = mypy_arg_types[1:]
        elif mypy_args:
            self._precondition(
                not mypy_args[0].variable.is_self,
                "if function is not a method, first variable should be self-like",
                func_loc,
            )
        self._symtable = dict[str, pytypes.PyType]()
        args = list[SubroutineArgument]()
        for arg, arg_type in zip(mypy_args, mypy_arg_types, strict=True):
            arg_loc = context.node_location(arg)
            if arg.kind.is_star():
                raise CodeError("variadic functions are not supported", arg_loc)
            if arg.initializer is not None:
                context.error(
                    "default function argument values are not supported yet", arg.initializer
                )
            pytyp = context.type_to_pytype(arg_type, source_location=arg_loc)
            wtype = pytyp.checked_wtype(arg_loc)
            arg_name = arg.variable.name
            args.append(SubroutineArgument(name=arg_name, wtype=wtype, source_location=arg_loc))
            self._symtable[arg_name] = pytyp
        # translate body
        translated_body = self.visit_block(func_def.body)
        # build result
        self.result: Subroutine | ContractMethod
        documentation = parse_docstring(func_def.docstring)
        if contract_method_info is None:
            self.result = Subroutine(
                id=func_def.fullname,
                name=func_def.name,
                source_location=source_location,
                args=args,
                return_type=return_wtype,
                body=translated_body,
                documentation=documentation,
                inline=inline,
            )
        else:
            self.result = ContractMethod(
                cref=contract_method_info.fragment.id,
                member_name=func_def.name,
                source_location=source_location,
                args=args,
                return_type=return_wtype,
                body=translated_body,
                documentation=documentation,
                arc4_method_config=contract_method_info.arc4_method_config,
                inline=inline,
            )

    @classmethod
    @typing.overload
    def convert(
        cls,
        context: ASTConversionModuleContext,
        func_def: mypy.nodes.FuncDef,
        source_location: SourceLocation,
        *,
        inline: bool | None,
    ) -> Subroutine: ...

    @classmethod
    @typing.overload
    def convert(
        cls,
        context: ASTConversionModuleContext,
        func_def: mypy.nodes.FuncDef,
        source_location: SourceLocation,
        contract_method_info: ContractMethodInfo,
        *,
        inline: bool | None,
    ) -> ContractMethod: ...

    @classmethod
    def convert(
        cls,
        context: ASTConversionModuleContext,
        func_def: mypy.nodes.FuncDef,
        source_location: SourceLocation,
        contract_method_info: ContractMethodInfo | None = None,
        *,
        inline: bool | None,
    ) -> Subroutine | ContractMethod:
        return cls(
            context=context,
            func_def=func_def,
            contract_method_info=contract_method_info,
            inline=inline,
            source_location=source_location,
        ).result

    @typing.override
    def builder_for_self(self, expr_loc: SourceLocation) -> NodeBuilder:
        if self.contract_method_info is None:
            raise InternalError("variable is inferred as self outside of contract scope", expr_loc)
        return ContractSelfExpressionBuilder(
            fragment=self.contract_method_info.fragment,
            pytype=self.contract_method_info.contract_type,
            location=expr_loc,
        )

    @typing.override
    def resolve_local_type(self, var_name: str, expr_loc: SourceLocation) -> pytypes.PyType | None:
        return self._symtable.get(var_name)

    # expressions

    @typing.override
    def visit_assignment_expr(self, expr: mypy.nodes.AssignmentExpr) -> NodeBuilder:
        expr_loc = self._location(expr)
        self._precondition(
            not isinstance(expr, mypy.nodes.TupleExpr | mypy.nodes.ListExpr),
            "Python doesn't support tuple unpacking in assignment expressions",
            expr_loc,
        )
        # TODO: test if self. assignment?
        source = require_instance_builder(expr.value.accept(self))
        if not self._assign_type(expr.target, source.pytype):
            # typing error, just return the source to continue with
            return source
        value = source.resolve()
        target = self.resolve_lvalue(expr.target)
        result = AssignmentExpression(source_location=expr_loc, value=value, target=target)
        result_typ = source.pytype
        return builder_for_instance(result_typ, result)

    @typing.override
    def visit_super_expr(self, super_expr: mypy.nodes.SuperExpr) -> NodeBuilder:
        super_loc = self._location(super_expr)
        if self.contract_method_info is None:
            raise CodeError("super() expression should not occur outside of class", super_loc)
        if super_expr.call.args:
            raise CodeError(
                "only the zero-arguments version of super() is supported",
                self._location(super_expr.call),
            )
        if (
            super_expr.name.startswith("__")
            and super_expr.name.endswith("__")
            and super_expr.name != "__init__"
        ):
            raise CodeError("the only dunder method supported for calling is init", super_loc)
        for base in self.contract_method_info.fragment.mro:
            try:
                base_sym_typ = base.symbols[super_expr.name]
            except KeyError:
                pass
            else:
                if not isinstance(base_sym_typ, pytypes.FuncType):
                    raise CodeError("super() is only supported for calling functions", super_loc)
                func_type = base_sym_typ
                base_method = base.resolve_method(super_expr.name, include_inherited=False)
                # the check for class abstractness is because in theory the next method in the MRO
                # when compiling a derived class could still be non-trivial
                if (
                    base_method is not None
                    and base_method.is_trivial
                    and not self.contract_method_info.is_abstract
                ):
                    raise CodeError("call to trivial method via super()", super_loc)
                break
        else:
            if super_expr.name == "__init__":
                # support fall through to object __init__, it's the only method from
                # non-member bases that is callable, and it's a no-op, but it can be useful
                # in the case of multiple inheritance
                return NoneTypeBuilder(super_loc)
            raise CodeError(
                f"unable to locate method {super_expr.name}"
                f" in bases of {self.contract_method_info.fragment.id}",
                super_loc,
            )

        super_target = InstanceSuperMethodTarget(member_name=super_expr.name)
        return SubroutineInvokerExpressionBuilder(
            target=super_target, func_type=func_type, location=super_loc
        )

    # statements

    def visit_block(self, block: mypy.nodes.Block) -> Block:
        translated_body = []
        for stmt in block.body:
            translated_stmt = stmt.accept(self)
            if translated_stmt is None:
                pass
            elif isinstance(translated_stmt, Statement):
                translated_body.append(translated_stmt)
            else:
                translated_body.extend(translated_stmt)
        return Block(
            source_location=self._location(block),
            body=translated_body,
        )

    def visit_expression_stmt(self, stmt: mypy.nodes.ExpressionStmt) -> ExpressionStatement | None:
        stmt_loc = self._location(stmt)
        if isinstance(stmt.expr, mypy.nodes.StrExpr):
            if stmt.expr.value != self.func_def.docstring:
                logger.warning(
                    "String literal is not assigned and is not part of docstring",
                    location=stmt_loc,
                )
            return None
        if isinstance(stmt.expr, mypy.nodes.TupleExpr) and len(stmt.expr.items) == 1:
            raise CodeError(
                "Tuple being constructed without assignment,"
                " check for a stray comma at the end of the statement",
                stmt_loc,
            )
        expr_builder = require_instance_builder(stmt.expr.accept(self))
        if expr_builder.pytype != pytypes.NoneType:
            if isinstance(stmt.expr, mypy.nodes.CallExpr) and isinstance(
                stmt.expr.analyzed, mypy.nodes.RevealExpr | mypy.nodes.AssertTypeExpr
            ):
                # special case to ignore ignoring the result of typing.reveal_type/assert_type
                pass
            elif isinstance(expr_builder.pytype, pytypes.InnerTransactionResultType) or (
                isinstance(expr_builder.pytype, pytypes.TupleLikeType)
                and any(
                    isinstance(i, pytypes.InnerTransactionResultType)
                    for i in expr_builder.pytype.items
                )
            ):
                # special case to ignore inner transaction result types
                # could maybe expand this check to consider whether an expression has known
                # side-effects
                pass
            else:
                self.context.warning("expression result is ignored", stmt_loc)
        return ExpressionStatement(expr=expr_builder.resolve())

    def visit_assignment_stmt(self, stmt: mypy.nodes.AssignmentStmt) -> Sequence[Statement]:
        stmt_loc = self._location(stmt)
        match stmt.lvalues:
            case [mypy.nodes.NameExpr(is_special_form=True)]:
                self._error(
                    "type aliases, type vars, and type constructors"
                    " are not supported inside functions",
                    stmt_loc,
                )
                return []
        match stmt.rvalue:
            case mypy.nodes.TempNode(no_rhs=True):
                assert (
                    len(stmt.lvalues) == 1
                ), "mypy should guarantee that an annotation assigment has only a single lvalue"
                (lvalue,) = stmt.lvalues
                assert (
                    stmt.type is not None
                ), "mypy should guarantee that an annotation assignment has a type set"
                annotated_type = self.context.type_to_pytype(stmt.type, source_location=stmt_loc)
                self._assign_type(lvalue, annotated_type)
                # forward type-declaration only, no-op
                return []
        rvalue = require_instance_builder(stmt.rvalue.accept(self))
        if not all(self._assign_type(lvalue, rvalue.pytype) for lvalue in stmt.lvalues):
            # already a typing error, don't need a second one
            return []

        if isinstance(rvalue, StorageProxyConstructorResult):
            try:
                (lvalue,) = stmt.lvalues
            except ValueError:
                # this is true regardless of whether it's a self assignment or not,
                # these are objects so aliasing is an issue in terms of semantic compatibility
                raise CodeError(
                    f"{rvalue.pytype} can only be assigned to a single variable", stmt_loc
                ) from None
            if is_self_member(lvalue):
                return self._handle_proxy_assignment(lvalue, rvalue, stmt_loc)
            elif rvalue.args.initial_value is not None:
                raise CodeError(
                    "providing an initial value is only allowed"
                    " when assigning to a member variable",
                    stmt_loc,
                )
        elif isinstance(
            rvalue.pytype, pytypes.StorageProxyType | pytypes.StorageMapProxyType
        ) and any(is_self_member(lvalue) for lvalue in stmt.lvalues):
            raise CodeError("unsupported usage of storage type", stmt_loc)
        elif len(stmt.lvalues) > 1:
            rvalue = rvalue.single_eval()

        return [
            AssignmentStatement(
                value=rvalue.resolve(),
                target=self.resolve_lvalue(lvalue),
                source_location=stmt_loc,
            )
            for lvalue in reversed(stmt.lvalues)
        ]

    def _assign_type(
        self,
        lvalue: mypy.nodes.Expression,
        typ: pytypes.PyType,
    ) -> bool:
        match lvalue:
            case mypy.nodes.NameExpr(name=var_name):
                if var_name == "_":
                    raise CodeError(
                        "_ is not currently supported as a variable name", self._location(lvalue)
                    )
                lvalue_loc = self._location(lvalue)
                symbol_type = self._symtable.setdefault(var_name, typ)
                if not (symbol_type <= typ):
                    logger.error(
                        f"{var_name!r} already has type {symbol_type}"
                        f" which is not compatible with {typ}",
                        location=lvalue_loc,
                    )
                    return False
                return True
            case mypy.nodes.TupleExpr(items=lval_items) | mypy.nodes.ListExpr(items=lval_items):
                if star_expr := next(
                    (
                        lval_item
                        for lval_item in lval_items
                        if isinstance(lval_item, mypy.nodes.StarExpr)
                    ),
                    None,
                ):
                    raise CodeError(
                        "star expressions are not supported", self._location(star_expr)
                    )
                match typ:
                    case pytypes.SequenceType(items=homogenous_type):
                        tuple_item_types = (homogenous_type,) * len(lval_items)
                    case pytypes.TupleLikeType(items=tuple_item_types):
                        if len(tuple_item_types) != len(lval_items):
                            raise CodeError(
                                f"length mismatch source size of {len(tuple_item_types)}"
                                f" and target size of {len(lval_items)}",
                                self._location(lvalue),
                            )
                    case _:
                        raise CodeError(
                            "source type not supported for unpacking", self._location(lvalue)
                        )
                return all(
                    self._assign_type(lval_item, inner_typ)
                    for lval_item, inner_typ in zip(lval_items, tuple_item_types, strict=True)
                )
            case _:
                return True

    def _handle_proxy_assignment(
        self,
        lvalue: mypy.nodes.MemberExpr,
        rvalue: StorageProxyConstructorResult,
        stmt_loc: SourceLocation,
    ) -> Sequence[Statement]:
        if self.contract_method_info is None:
            raise InternalError("Assignment to self outside of a contract class", stmt_loc)
        pytype = rvalue.pytype
        if self.func_def.name != "__init__":
            raise CodeError(
                f"{pytype.generic or pytype}"
                " can only be assigned to a member variable in the __init__ method",
                stmt_loc,
            )
        member_loc = self._location(lvalue)
        member_name = lvalue.name
        storage = self.contract_method_info.fragment.resolve_storage(member_name)
        if storage is None:
            raise CodeError("unable to resolve class storage member", member_loc)
        if storage.definition is not None:
            logger.info(
                f"previous definition of {member_name} was here",
                location=storage.definition.source_location,
            )
            logger.error(
                f"redefinition of {member_name}",
                location=stmt_loc,
            )

        key = rvalue.args.key
        if key is None:
            key = BytesConstant(
                value=member_name.encode("utf8"),
                wtype=pytype.wtype,
                encoding=BytesEncoding.utf8,
                source_location=member_loc,
            )
        elif not isinstance(key, BytesConstant):
            logger.error(
                f"assigning {pytype} to a member variable"
                f" requires a constant value for {rvalue.args.key_arg_name}",
                location=stmt_loc,
            )
            key = BytesConstant(
                value=b"0",
                wtype=key.wtype,
                encoding=BytesEncoding.unknown,
                source_location=key.source_location,
            )
        storage.definition = AppStorageDefinition(
            source_location=member_loc,
            member_name=member_name,
            kind=storage.kind,
            storage_wtype=pytype.content_wtype,
            key_wtype=(
                pytype.key_wtype if isinstance(pytype, pytypes.StorageMapProxyType) else None
            ),
            key=key,
            description=rvalue.args.description,
        )

        if rvalue.args.initial_value is None:
            return []
        match pytype:
            case pytypes.StorageProxyType(
                generic=pytypes.GenericGlobalStateType, content_wtype=content_wtype
            ):
                global_state_target = AppStateExpression(
                    key=key,
                    exists_assertion_message=None,  # this is a write, not a read
                    wtype=content_wtype,
                    source_location=member_loc,
                )
                return [
                    AssignmentStatement(
                        target=global_state_target,
                        value=rvalue.args.initial_value.resolve(),
                        source_location=stmt_loc,
                    )
                ]
            case unhandled:
                raise InternalError(
                    f"Don't know how to do initialise-on-declaration for {unhandled}", stmt_loc
                )

    def resolve_lvalue(self, lvalue: mypy.nodes.Expression) -> Lvalue:
        builder_or_literal = lvalue.accept(self)
        builder = require_instance_builder(builder_or_literal)
        return builder.resolve_lvalue()

    @typing.override
    def empty_statement(self, stmt: mypy.nodes.Statement) -> None:
        return None

    def visit_operator_assignment_stmt(self, stmt: mypy.nodes.OperatorAssignmentStmt) -> Statement:
        stmt_loc = self._location(stmt)
        builder = require_instance_builder(stmt.lvalue.accept(self))
        rhs = require_instance_builder(stmt.rvalue.accept(self))
        try:
            op = BuilderBinaryOp(stmt.op)
        except ValueError as ex:
            raise InternalError(f"Unknown binary operator: {stmt.op}") from ex
        return builder.augmented_assignment(op=op, rhs=rhs, location=stmt_loc)

    def visit_if_stmt(self, stmt: mypy.nodes.IfStmt) -> IfElse:
        self._precondition(
            len(stmt.body) == len(stmt.expr),
            "mismatch between if statement condition list length and body list length",
            stmt,
        )
        self._precondition(
            len(stmt.expr) == 1,
            "python ast module should produce normalised if statements",
            stmt,
        )
        (expr,) = stmt.expr
        (if_body,) = stmt.body
        condition = expr.accept(self).bool_eval(self._location(expr))
        if_branch = self.visit_block(if_body)
        else_branch = self.visit_block(stmt.else_body) if stmt.else_body else None
        return IfElse(
            source_location=self._location(stmt),
            condition=condition.resolve(),
            if_branch=if_branch,
            else_branch=else_branch,
        )

    @typing.overload
    def _build_loop(
        self,
        stmt: mypy.nodes.WhileStmt,
        ctor: Callable[[Block, SourceLocation], WhileLoop],
    ) -> WhileLoop | Block: ...

    @typing.overload
    def _build_loop(
        self,
        stmt: mypy.nodes.ForStmt,
        ctor: Callable[[Block, SourceLocation], ForInLoop],
    ) -> ForInLoop | Block: ...

    def _build_loop(
        self,
        stmt: mypy.nodes.WhileStmt | mypy.nodes.ForStmt,
        ctor: Callable[[Block, SourceLocation], WhileLoop | ForInLoop],
    ) -> WhileLoop | ForInLoop | Block:
        loop_loc = self._location(stmt)
        if stmt.else_body is None:
            break_label = None
        else:
            break_label = Label(f"after_loop_L{loop_loc.line}")

        self._break_label_stack.append(break_label)
        try:
            block = self.visit_block(stmt.body)
        finally:
            self._break_label_stack.pop()
        loop = ctor(block, loop_loc)
        if stmt.else_body is None:
            return loop
        else_body = self.visit_block(stmt.else_body)
        # empty block, exists so we have a goto target that skips the else body
        after_block = Block(label=break_label, body=[], source_location=loop_loc)
        return Block(
            body=[loop, else_body, after_block],
            comment=f"loop_with_else_L{loop_loc.line}",
            source_location=loop_loc,
        )

    def visit_while_stmt(self, stmt: mypy.nodes.WhileStmt) -> WhileLoop | Block:
        condition = stmt.expr.accept(self).bool_eval(self._location(stmt.expr))
        return self._build_loop(
            stmt,
            lambda loop_body, loop_loc: WhileLoop(
                condition=condition.resolve(),
                loop_body=loop_body,
                source_location=loop_loc,
            ),
        )

    def visit_for_stmt(self, stmt: mypy.nodes.ForStmt) -> ForInLoop | Block:
        sequence_builder = require_instance_builder(stmt.expr.accept(self))
        sequence = sequence_builder.iterate()
        iter_item_type = sequence_builder.iterable_item_type()
        self._assign_type(stmt.index, iter_item_type)
        items = self.resolve_lvalue(stmt.index)
        return self._build_loop(
            stmt,
            lambda loop_body, loop_loc: ForInLoop(
                items=items,
                sequence=sequence,
                loop_body=loop_body,
                source_location=loop_loc,
            ),
        )

    def visit_break_stmt(self, stmt: mypy.nodes.BreakStmt) -> LoopExit | Goto:
        stmt_loc = self._location(stmt)
        break_label = self._break_label_stack[-1]
        if break_label is None:
            return LoopExit(stmt_loc)
        else:
            return Goto(target=break_label, source_location=stmt_loc)

    def visit_continue_stmt(self, stmt: mypy.nodes.ContinueStmt) -> LoopContinue:
        return LoopContinue(self._location(stmt))

    def visit_assert_stmt(self, stmt: mypy.nodes.AssertStmt) -> ExpressionStatement:
        error_message: str | None = None
        if stmt.msg is not None:
            msg = stmt.msg.accept(self)
            match msg:
                case LiteralBuilder(value=str(error_message)):
                    pass
                case _:
                    self._error("only literal strings are supported as assertion messages", stmt)
        condition = stmt.expr.accept(self).bool_eval(self._location(stmt.expr))
        return ExpressionStatement(
            AssertExpression(
                condition=condition.resolve(),
                error_message=error_message,
                source_location=self._location(stmt),
            )
        )

    def visit_del_stmt(self, stmt: mypy.nodes.DelStmt) -> Statement:
        stmt_expr = stmt.expr.accept(self)
        del_item = require_instance_builder(stmt_expr)
        return del_item.delete(self._location(stmt))

    def visit_return_stmt(self, stmt: mypy.nodes.ReturnStmt) -> ReturnStatement | None:
        loc = self._location(stmt)
        return_expr = stmt.expr
        if return_expr is None:
            if self._return_type is not pytypes.NoneType:
                self._error(
                    "function is typed as returning a value, so a value must be returned", loc
                )
            return ReturnStatement(source_location=loc, value=None)

        returning_builder = require_instance_builder(return_expr.accept(self))
        if not (self._return_type <= returning_builder.pytype):
            self._error(
                f"invalid return type of {returning_builder.pytype}, expected {self._return_type}",
                loc,
            )
        return ReturnStatement(source_location=loc, value=returning_builder.resolve())

    def visit_match_stmt(self, stmt: mypy.nodes.MatchStmt) -> Switch | None:
        loc = self._location(stmt)
        subject = require_instance_builder(stmt.subject.accept(self))
        case_block_map = dict[Expression, Block]()
        default_block: Block | None = None
        for pattern, guard, block in zip(stmt.patterns, stmt.guards, stmt.bodies, strict=True):
            if guard is not None:
                self._error("guard clauses are not supported", guard)
            if default_block is not None:
                self._error("default case already encountered", pattern)
                continue
            case_block = self.visit_block(block)
            pattern_loc = self._location(pattern)
            case_value_builder: InstanceBuilder | None = None
            match pattern:
                case mypy.patterns.AsPattern(name=None, pattern=None):
                    default_block = case_block
                case mypy.patterns.ValuePattern(expr=case_expr):
                    case_value_builder = require_instance_builder(case_expr.accept(self))
                    case_value_builder = maybe_resolve_literal(case_value_builder, subject.pytype)
                case mypy.patterns.SingletonPattern(value=bool(bool_literal)):
                    case_value_builder = LiteralBuilderImpl(
                        value=bool_literal, source_location=pattern_loc
                    )
                    case_value_builder = maybe_resolve_literal(case_value_builder, subject.pytype)
                case mypy.patterns.ClassPattern() as cls_pattern:
                    class_builder = cls_pattern.class_ref.accept(self)
                    if not isinstance(class_builder, CallableBuilder):
                        logger.error("expected a type", location=class_builder.source_location)
                    elif cls_pattern.keyword_values:
                        self._error(
                            "attribute matching is not supported", cls_pattern.keyword_values[0]
                        )
                    else:
                        cls_args = []
                        has_error = False
                        for pos_case in cls_pattern.positionals:
                            if isinstance(pos_case, mypy.patterns.ValuePattern):
                                cls_args.append(pos_case.expr.accept(self))
                            else:
                                has_error = True
                                self._error("expected a value pattern", pos_case)
                        if not has_error:
                            case_value_builder = class_builder.call(
                                args=cls_args,
                                arg_kinds=[mypy.nodes.ARG_POS] * len(cls_args),
                                arg_names=[None] * len(cls_args),
                                location=pattern_loc,
                            )

                case _:
                    logger.error("unsupported case pattern", location=pattern_loc)
            if case_value_builder is not None:
                if case_value_builder.pytype != subject.pytype:
                    # TODO: what about other comparable types, or subtypes?
                    logger.error(
                        "type mismatch,"
                        " case values must be the exact same type as the subject type",
                        location=pattern_loc,
                    )
                else:
                    case_block_map[case_value_builder.resolve()] = case_block

        return Switch(
            source_location=loc,
            value=subject.resolve(),
            cases=case_block_map,
            default_case=default_block,
        )

    # Unsupported statements
    @typing.override
    def visit_function(self, fdef: mypy.nodes.FuncDef, _: mypy.nodes.Decorator | None) -> None:
        self._error("nested functions are not supported", fdef)

    @typing.override
    def visit_class_def(self, cdef: mypy.nodes.ClassDef) -> None:
        self._error("classes nested inside functions are not supported", cdef)

    @typing.override
    def visit_type_alias_stmt(self, o: mypy.nodes.TypeAliasStmt) -> None:
        self._error("type alias statements are not supported", o)

def is_self_member(
    expr: mypy.nodes.Expression,
) -> typing.TypeGuard[mypy.nodes.MemberExpr]:
    match expr:
        case mypy.nodes.MemberExpr(expr=mypy.nodes.NameExpr(node=mypy.nodes.Var(is_self=True))):
            return True
    return False

def require_instance_builder(builder_or_literal: NodeBuilder) -> InstanceBuilder:
    return expect.instance_builder(builder_or_literal, default=expect.default_raise)


import contextlib
import operator
import re
import typing
import warnings
from collections.abc import Callable, Iterator, Mapping, Sequence
from itertools import zip_longest

import mypy.build
import mypy.nodes
import mypy.types
from mypy.types import get_proper_type, is_named_instance

from puya import log
from puya.errors import CodeError, InternalError
from puya.parse import SourceLocation
from puya.utils import unique
from puyapy.awst_build import pytypes
from puyapy.awst_build.context import ASTConversionModuleContext
from puyapy.awst_build.eb.factories import builder_for_type
from puyapy.awst_build.eb.interface import InstanceBuilder, NodeBuilder, TypeBuilder
from puyapy.models import ConstantValue

logger = log.get_logger(__name__)

def refers_to_fullname(ref_expr: mypy.nodes.RefExpr, *fullnames: str) -> bool:
    """Is node a name or member expression with the given full name?"""
    if isinstance(ref_expr.node, mypy.nodes.TypeAlias):
        return is_named_instance(ref_expr.node.target, fullnames)
    else:
        return ref_expr.fullname in fullnames

def get_unaliased_fullname(ref_expr: mypy.nodes.RefExpr) -> str:
    alias = get_aliased_instance(ref_expr)
    if alias:
        return alias.type.fullname
    return ref_expr.fullname

def get_aliased_instance(ref_expr: mypy.nodes.RefExpr) -> mypy.types.Instance | None:
    if isinstance(ref_expr.node, mypy.nodes.TypeAlias):
        t = get_proper_type(ref_expr.node.target)
        if isinstance(t, mypy.types.Instance):
            return t
        if (
            isinstance(t, mypy.types.TupleType)
            and t.partial_fallback.type.fullname != "builtins.tuple"
        ):
            return t.partial_fallback
    return None

def get_decorators_by_fullname(
    ctx: ASTConversionModuleContext,
    decorator: mypy.nodes.Decorator,
    *,
    original: bool = False,
) -> dict[str, mypy.nodes.Expression]:
    result = dict[str, mypy.nodes.Expression]()
    for d in decorator.original_decorators if original else decorator.decorators:
        if isinstance(d, mypy.nodes.RefExpr):
            full_name = get_unaliased_fullname(d)
            result[get_unaliased_fullname(d)] = d
        elif isinstance(d, mypy.nodes.CallExpr) and isinstance(d.callee, mypy.nodes.RefExpr):
            full_name = get_unaliased_fullname(d.callee)
        else:
            raise CodeError("Unsupported decorator usage", ctx.node_location(d))
        result[full_name] = d
    return result

UNARY_OPS: typing.Final[Mapping[str, Callable[[typing.Any], typing.Any]]] = {
    "~": operator.inv,
    "not": operator.not_,
    "+": operator.pos,
    "-": operator.neg,
}

@contextlib.contextmanager
def _log_warnings(location: SourceLocation) -> Iterator[None]:
    folding_warnings = []
    try:
        with warnings.catch_warnings(record=True, action="always") as folding_warnings:
            yield
    finally:
        for warning in folding_warnings:
            logger.warning(warning.message, location=location)

def fold_unary_expr(location: SourceLocation, op: str, expr: ConstantValue) -> ConstantValue:
    if not (func := UNARY_OPS.get(op)):
        raise InternalError(f"Unhandled unary operator: {op}", location)
    if op == "~" and isinstance(expr, int):
        logger.warning(
            "due to Python ints being signed, bitwise inversion yield a negative number",
            location=location,
        )
    try:
        with _log_warnings(location):
            result = func(expr)
    except Exception as ex:
        raise CodeError(str(ex), location) from ex
    if not isinstance(result, ConstantValue):
        raise CodeError(f"unsupported result type of {type(result).__name__}", location)
    return result

BINARY_OPS: typing.Final[Mapping[str, Callable[[typing.Any, typing.Any], typing.Any]]] = {
    "+": operator.add,
    "-": operator.sub,
    "*": operator.mul,
    "@": operator.matmul,
    "/": operator.truediv,
    "%": operator.mod,
    "**": operator.pow,
    "<<": operator.lshift,
    ">>": operator.rshift,
    "|": operator.or_,
    "^": operator.xor,
    "&": operator.and_,
    "//": operator.floordiv,
    ">": operator.gt,
    "<": operator.lt,
    "==": operator.eq,
    ">=": operator.ge,
    "<=": operator.le,
    "!=": operator.ne,
    "is": operator.is_,
    "is not": operator.is_not,
    "in": lambda a, b: a in b,
    "not in": lambda a, b: a not in b,
    "and": lambda a, b: a and b,
    "or": lambda a, b: a or b,
}

def fold_binary_expr(
    location: SourceLocation,
    op: str,
    lhs: ConstantValue,
    rhs: ConstantValue,
) -> ConstantValue:
    if not (func := BINARY_OPS.get(op)):
        raise InternalError(f"Unhandled binary operator: {op}", location)
    try:
        with _log_warnings(location):
            result = func(lhs, rhs)
    except Exception as ex:
        raise CodeError(str(ex), location) from ex
    if not isinstance(result, ConstantValue):
        raise CodeError(f"unsupported result type of {type(result).__name__}", location)
    return result

def extract_bytes_literal_from_mypy(expr: mypy.nodes.BytesExpr) -> bytes:
    import ast

    bytes_str = expr.value
    # mypy doesn't preserve the value of the quotes, need to reverse it
    if '"' in bytes_str:
        bytes_literal = "b'" + bytes_str + "'"
    else:
        bytes_literal = 'b"' + bytes_str + '"'
    bytes_const: bytes = ast.literal_eval(bytes_literal)
    return bytes_const

TBuilder = typing.TypeVar("TBuilder", bound=NodeBuilder)

def get_arg_mapping(
    *,
    required_positional_names: Sequence[str] | None = None,
    optional_positional_names: Sequence[str] | None = None,
    required_kw_only: Sequence[str] | None = None,
    optional_kw_only: Sequence[str] | None = None,
    args: Sequence[TBuilder],
    arg_names: Sequence[str | None],
    call_location: SourceLocation,
    raise_on_missing: bool,
) -> tuple[dict[str, TBuilder], bool]:
    required_positional_names = required_positional_names or []
    optional_positional_names = optional_positional_names or []
    required_kw_only = required_kw_only or []
    optional_kw_only = optional_kw_only or []
    all_names = [
        *required_positional_names,
        *optional_positional_names,
        *required_kw_only,
        *optional_kw_only,
    ]
    if len(all_names) != len(set(all_names)):
        raise InternalError("duplicate names", call_location)

    arg_mapping = dict[str, TBuilder]()
    num_positionals = arg_names.count(None)
    positional_names = [*required_positional_names, *optional_positional_names]
    for positional_name, positional_arg in zip_longest(positional_names, args[:num_positionals]):
        if positional_name is None:
            logger.error("unexpected positional argument", location=positional_arg.source_location)
        elif positional_arg is None:
            break  # positional arguments could be named
        else:
            arg_mapping[positional_name] = positional_arg
    for arg_name, kw_arg in zip(arg_names[num_positionals:], args[num_positionals:], strict=True):
        if arg_name is None:
            # shouldn't be possible, might happen with type ignore?
            logger.error(
                "non-keyword argument appears after keywords", location=kw_arg.source_location
            )
        elif arg_name in arg_mapping:
            # again, shouldn't be possible, but might happen with type ignore?
            logger.error("duplicate argument name", location=kw_arg.source_location)
        elif arg_name not in all_names:
            logger.error("unrecognised keyword argument", location=kw_arg.source_location)
        else:
            arg_mapping[arg_name] = kw_arg
    any_missing = False
    if missing_args := set(required_positional_names).difference(arg_mapping.keys()):
        msg = f"missing required positional argument(s): {', '.join(sorted(missing_args))}"
        if raise_on_missing:
            raise CodeError(msg, call_location)
        any_missing = True
        logger.error(msg, location=call_location)
    elif missing_kwargs := set(required_kw_only).difference(arg_mapping.keys()):
        msg = f"missing required keyword argument(s): {', '.join(sorted(missing_kwargs))}"
        if raise_on_missing:
            raise CodeError(msg, call_location)
        any_missing = True
        logger.error(msg, location=call_location)
    return arg_mapping, any_missing

def snake_case(s: str) -> str:
    s = s.replace("-", " ")
    s = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", s)
    s = re.sub(r"([a-z\d])([A-Z])", r"\1_\2", s)
    return re.sub(r"[-\s]", "_", s).lower()

def require_callable_type(
    node: mypy.nodes.FuncBase | mypy.nodes.Decorator, location: SourceLocation
) -> mypy.types.CallableType:
    if isinstance(node, mypy.nodes.Decorator):
        typ = node.var.type
    else:
        typ = node.type
    if typ is None:
        raise InternalError(f"unable to resolve type of {node.fullname}", location)
    if not isinstance(typ, mypy.types.CallableType):
        raise CodeError("expected a callable type", location)
    return typ

def maybe_resolve_literal(
    operand: TBuilder, target_type: pytypes.PyType
) -> TBuilder | InstanceBuilder:
    if not isinstance(operand, InstanceBuilder):
        return operand
    if not (target_type <= operand.pytype):
        target_type_builder = builder_for_type(target_type, operand.source_location)
        if isinstance(target_type_builder, TypeBuilder):
            return operand.resolve_literal(target_type_builder)
    return operand

def resolve_literal(operand: InstanceBuilder, target_type: pytypes.PyType) -> InstanceBuilder:
    target_type_builder = builder_for_type(target_type, operand.source_location)
    assert isinstance(target_type_builder, TypeBuilder)
    return operand.resolve_literal(target_type_builder)

def determine_base_type(
    first: pytypes.PyType, *rest: pytypes.PyType, location: SourceLocation
) -> pytypes.PyType:
    operands = unique((first, *rest))
    if len(operands) == 1:
        return first
    for candidate in operands:
        if all(candidate <= operand for operand in operands):
            return candidate
    return pytypes.UnionType(operands, location)

def extract_decorator_args(
    decorator: mypy.nodes.Expression, location: SourceLocation
) -> list[tuple[str | None, mypy.nodes.Expression]]:
    match decorator:
        case mypy.nodes.RefExpr():
            return []
        case mypy.nodes.CallExpr(args=args, arg_names=arg_names):
            return list(zip(arg_names, args, strict=True))
        case unexpected_node:
            raise InternalError(f"unexpected decorator node: {unexpected_node}", location)

def get_subroutine_decorator_inline_arg(
    context: ASTConversionModuleContext, decorator: mypy.nodes.Expression
) -> bool | None:
    args = extract_decorator_args(decorator, context.node_location(decorator))
    for name, expr in args:
        match name, expr:
            case "inline", mypy.nodes.NameExpr(fullname="builtins.True"):
                return True
            case "inline", mypy.nodes.NameExpr(fullname="builtins.False"):
                return False
            case "inline", mypy.nodes.NameExpr(fullname="builtins.None"):
                return None
            case _:
                logger.error("unexpected argument", location=context.node_location(expr))
    return None


import argparse
import json
import typing
from collections.abc import Iterable, Sequence
from pathlib import Path

import attrs
from cattrs.preconf.json import make_converter

from puya import (
    arc56_models as arc56,
    log,
)
from puya.arc32 import OCA_ARC32_MAPPING
from puya.arc56 import allowed_call_oca, allowed_create_oca
from puya.artifact_metadata import ARC4MethodArg, ARC4Returns
from puya.errors import PuyaError
from puyapy.awst_build.arc4_client_gen import write_arc4_client

logger = log.get_logger(__name__)
ARC32_OCA_MAPPING = {v: k for k, v in OCA_ARC32_MAPPING.items()}

@attrs.define(kw_only=True)
class PuyaGenOptions:
    paths: Sequence[Path] = attrs.field(default=(), repr=lambda p: str(list(map(str, p))))
    log_level: log.LogLevel = log.LogLevel.info
    out_dir: Path | None = None

def main() -> None:
    parser = argparse.ArgumentParser(
        prog="puyapy-clientgen",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Output algopy contract client for typed ARC4 ABI calls from an "
        "ARC-32 or ARC-56 application spec",
    )
    parser.add_argument(
        "--out-dir",
        type=Path,
        help="Path for outputting client, defaults to app spec folder",
        default=False,
    )
    parser.add_argument("paths", type=Path, nargs="+", metavar="PATH")
    options = PuyaGenOptions()
    parser.parse_args(namespace=options)
    log.configure_logging(min_log_level=options.log_level)
    output_stubs(options.paths, options.out_dir)

def output_stubs(paths: Sequence[Path], out_dir: Path | None) -> None:
    try:
        app_spec_paths = resolve_app_specs(paths)
        for app_spec_path in app_spec_paths:
            app_spec_json = app_spec_path.read_text("utf8")
            if app_spec_path.name.endswith(".arc56.json"):
                app_spec = parse_arc56(app_spec_json)
            else:
                # TODO: use algokit_utils to do this conversion when it is available?
                app_spec = _convert_arc32_to_arc56(app_spec_json)
            write_arc4_client(app_spec, out_dir or app_spec_path.parent)
    except PuyaError as ex:
        logger.error(str(ex))  # noqa: TRY400

def _convert_arc32_to_arc56(app_spec_json: str) -> arc56.Contract:
    name, structs, methods = _parse_arc32_app_spec_methods(app_spec_json)
    return arc56.Contract(
        name=name,
        structs=structs,
        methods=methods,
    )

def resolve_app_specs(paths: Sequence[Path]) -> Sequence[Path]:
    app_specs = list[Path]()
    for path in paths:
        if not path.is_dir():
            app_specs.append(path)
        else:
            patterns = ["application.json", "*.arc32.json"]

            app_specs = []
            for pattern in patterns:
                app_specs.extend(path.rglob(pattern))

    app_specs = sorted(set(app_specs))
    if not app_specs:
        raise PuyaError("No app specs found")
    return app_specs

def parse_arc56(app_spec_json: str) -> arc56.Contract:
    converter = make_converter(omit_if_default=True)
    return converter.loads(app_spec_json, arc56.Contract)

def _parse_arc32_app_spec_methods(
    app_spec_json: str,
) -> tuple[str, dict[str, Sequence[arc56.StructField]], Sequence[arc56.Method]]:
    # only need to parse a limited subset of ARC32 for client generation
    # i.e. ABI methods, their OCA parameters, and any struct info
    # default args are ignored as they aren't supported for on chain calls currently

    app_spec = json.loads(app_spec_json)
    contract = app_spec["contract"]
    hints = app_spec["hints"]
    contract_name = contract["name"]
    methods = list[arc56.Method]()
    arc4_methods = {m.signature: m for m in _parse_methods(contract["methods"])}
    known_structs = dict[str, Sequence[arc56.StructField]]()
    for arc4_method in arc4_methods.values():
        method_hints = hints[str(arc4_method.signature)]
        arg_to_struct = dict[str, str]()
        for param, struct_name, struct in _parse_structs(method_hints.get("structs", {})):
            for known_struct_name, known_struct in known_structs.items():
                if known_struct == struct:
                    arg_to_struct[param] = known_struct_name
                    break
            else:
                known_structs[struct_name] = struct
                arg_to_struct[param] = struct_name
        methods.append(
            arc56.Method(
                name=arc4_method.signature.name,
                desc=arc4_method.desc,
                args=[
                    arc56.MethodArg(
                        name=a.name,
                        type=a.type_,
                        struct=arg_to_struct.get(a.name),
                    )
                    for a in arc4_method.signature.args
                ],
                returns=arc56.MethodReturns(
                    type=arc4_method.signature.returns.type_,
                    desc=arc4_method.desc,
                    struct=arg_to_struct.get("output"),
                ),
                readonly=bool(method_hints.get("read_only") or arc4_method.readonly),
                actions=_parse_call_config(method_hints["call_config"]),
            )
        )
    return contract_name, known_structs, methods

@attrs.frozen(kw_only=True)
class _MethodSignature:
    name: str
    args: tuple[ARC4MethodArg, ...]
    returns: ARC4Returns

    def __str__(self) -> str:
        return f"{self.name}({','.join(a.type_ for a in self.args)}){self.returns.type_}"

@attrs.frozen(kw_only=True)
class _Method:
    signature: _MethodSignature
    desc: str | None
    readonly: bool | None

def _parse_methods(methods: list[dict[str, typing.Any]]) -> Iterable[_Method]:
    for method in methods:
        signature = _parse_signature(method)
        yield _Method(
            signature=signature,
            desc=method.get("desc"),
            readonly=method.get("readonly"),
        )

def _parse_signature(method: dict[str, typing.Any]) -> _MethodSignature:
    returns = method["returns"]
    return _MethodSignature(
        name=method["name"],
        args=tuple(
            ARC4MethodArg(
                name=arg["name"],
                type_=arg["type"],
                desc=arg.get("desc"),
                struct=arg.get("struct"),
                client_default=None,  # not supported on chain
            )
            for arg in method["args"]
        ),
        returns=ARC4Returns(
            type_=returns["type"],
            desc=returns.get("desc"),
            struct=returns.get("struct"),
        ),
    )

def _parse_call_config(method_call_config: dict[str, str]) -> arc56.MethodActions:
    create = []
    call = []
    for oca, call_config in method_call_config.items():
        action = ARC32_OCA_MAPPING[oca].name
        match call_config:
            case "CREATE" if allowed_create_oca(action):
                create.append(action)
            case "CALL" if allowed_call_oca(action):
                call.append(action)
            # allowed creates is narrower than calls so only need to check that
            case "ALL" if allowed_create_oca(action):
                create.append(action)
                call.append(action)
            case invalid:
                raise PuyaError(f"invalid call config option: {invalid}")

    return arc56.MethodActions(
        create=create,
        call=call,
    )

def _parse_structs(
    structs: dict[str, dict[str, typing.Any]],
) -> Iterable[tuple[str, str, Sequence[arc56.StructField]]]:
    for param, struct_config in structs.items():
        yield (
            param,
            struct_config["name"],
            [arc56.StructField(name=f[0], type=f[1]) for f in struct_config["elements"]],
        )

if __name__ == "__main__":
    main()


import os
import shutil
import subprocess
import sysconfig
from collections.abc import Callable, Mapping, Sequence
from importlib import metadata
from pathlib import Path

import mypy.build
import mypy.errors
import mypy.find_sources
import mypy.fscache
import mypy.modulefinder
import mypy.nodes
import mypy.options
import mypy.util
from packaging import version

from puya import log
from puya.arc56 import create_arc56_json
from puya.awst.nodes import AWST
from puya.awst.serialize import awst_to_json
from puya.awst.to_code_visitor import ToCodeVisitor
from puya.compilation_artifacts import CompilationArtifact, CompiledContract
from puya.compile import awst_to_teal
from puya.errors import log_exceptions
from puya.program_refs import ContractReference, LogicSigReference
from puya.utils import make_path_relative_to_cwd
from puyapy.awst_build.arc4_client_gen import write_arc4_client
from puyapy.awst_build.main import transform_ast
from puyapy.client_gen import parse_arc56
from puyapy.options import PuyaPyOptions
from puyapy.parse import TYPESHED_PATH, ParseResult, parse_and_typecheck
from puyapy.utils import determine_out_dir

# this should contain the lowest version number that this compiler does NOT support
# i.e. the next minor version after what is defined in stubs/pyproject.toml:tool.poetry.version
MAX_SUPPORTED_ALGOPY_VERSION_EX = version.parse("2.7.0")
MIN_SUPPORTED_ALGOPY_VERSION = version.parse(f"{MAX_SUPPORTED_ALGOPY_VERSION_EX.major}.0.0")

logger = log.get_logger(__name__)

def compile_to_teal(puyapy_options: PuyaPyOptions) -> None:
    """Drive the actual core compilation step."""
    with log.logging_context() as log_ctx, log_exceptions():
        logger.debug(puyapy_options)
        try:
            parse_result = parse_with_mypy(puyapy_options.paths)
            log_ctx.sources_by_path = parse_result.sources_by_path
            log_ctx.exit_if_errors()
            awst, compilation_targets = transform_ast(parse_result)
        except mypy.errors.CompileError:
            # the placement of this catch is probably overly conservative,
            # but in parse_with_mypy there is a piece copied from mypyc, around setting
            # the location during mypy callbacks in case errors are produced.
            # also this error should have already been logged
            assert log_ctx.num_errors > 0, "expected mypy errors to be logged"
        log_ctx.exit_if_errors()
        if puyapy_options.output_awst or puyapy_options.output_awst_json:
            awst_out_dir = (
                puyapy_options.out_dir or Path.cwd()  # TODO: maybe make this defaulted on init?
            )
            nodes = [
                n for n in awst if n.source_location.file in parse_result.explicit_source_paths
            ]
            if puyapy_options.output_awst:
                output_awst(nodes, awst_out_dir)
            if puyapy_options.output_awst_json:
                output_awst_json(nodes, awst_out_dir)
        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_targets
            )
            if loc.file
        }
        teal = awst_to_teal(
            log_ctx, puyapy_options, compilation_set, parse_result.sources_by_path, awst
        )
        log_ctx.exit_if_errors()
        if puyapy_options.output_client:
            write_arc4_clients(puyapy_options.template_vars_prefix, compilation_set, teal)
    # needs to be outside the with block
    log_ctx.exit_if_errors()

def write_arc4_clients(
    template_prefix: str,
    compilation_set: Mapping[ContractReference | LogicSigReference, Path],
    artifacts: Sequence[CompilationArtifact],
) -> None:
    for artifact in artifacts:
        if isinstance(artifact, CompiledContract) and artifact.metadata.is_arc4:
            contract_out_dir = compilation_set.get(artifact.id)
            if contract_out_dir:
                app_spec_json = create_arc56_json(
                    approval_program=artifact.approval_program,
                    clear_program=artifact.clear_program,
                    metadata=artifact.metadata,
                    template_prefix=template_prefix,
                )
                # use round trip of ARC-56 -> reparse to ensure consistency
                # of client output regardless if generating from ARC-56 or
                # Puya ARC4Contract
                contract = parse_arc56(app_spec_json)
                write_arc4_client(contract, contract_out_dir)

def parse_with_mypy(paths: Sequence[Path]) -> ParseResult:
    mypy_options = get_mypy_options()

    _, ordered_modules = parse_and_typecheck(paths, mypy_options)
    return ParseResult(
        mypy_options=mypy_options,
        ordered_modules=ordered_modules,
    )

def get_mypy_options() -> mypy.options.Options:
    mypy_opts = mypy.options.Options()

    # improve mypy parsing performance by using a cut-down typeshed
    mypy_opts.custom_typeshed_dir = str(TYPESHED_PATH)
    mypy_opts.abs_custom_typeshed_dir = str(TYPESHED_PATH.resolve())

    # set python_executable so third-party packages can be found
    mypy_opts.python_executable = _get_python_executable()

    mypy_opts.preserve_asts = True
    mypy_opts.include_docstrings = True
    # next two options disable caching entirely.
    # slows things down but prevents intermittent failures.
    mypy_opts.incremental = False
    mypy_opts.cache_dir = os.devnull

    # strict mode flags, need to review these and all others too
    mypy_opts.disallow_any_generics = True
    mypy_opts.disallow_subclassing_any = True
    mypy_opts.disallow_untyped_calls = True
    mypy_opts.disallow_untyped_defs = True
    mypy_opts.disallow_incomplete_defs = True
    mypy_opts.check_untyped_defs = True
    mypy_opts.disallow_untyped_decorators = True
    mypy_opts.warn_redundant_casts = True
    mypy_opts.warn_unused_ignores = True
    mypy_opts.warn_return_any = True
    mypy_opts.strict_equality = True
    mypy_opts.strict_concatenate = True

    # disallow use of any
    mypy_opts.disallow_any_unimported = True
    mypy_opts.disallow_any_expr = True
    mypy_opts.disallow_any_decorated = True
    mypy_opts.disallow_any_explicit = True

    mypy_opts.pretty = True  # show source in output

    return mypy_opts

def _get_python_executable() -> str | None:
    prefix = _get_prefix()
    if not prefix:
        logger.warning("Could not determine python prefix or algopy version")
        return None
    logger.info(f"Found python prefix: {prefix}")
    venv_paths = sysconfig.get_paths(vars={"base": prefix})

    python_exe = None
    for python in ("python3", "python"):
        python_exe = shutil.which(python, path=venv_paths["scripts"])
        if python_exe:
            logger.debug(f"Using python executable: {python_exe}")
            break
    else:
        logger.warning("Found a python prefix, but could not find the expected python interpreter")
    # use glob here, as we don't want to assume the python version
    discovered_site_packages = list(
        Path(prefix).glob(str(Path("[Ll]ib") / "**" / "site-packages"))
    )
    try:
        (site_packages,) = discovered_site_packages
    except ValueError:
        logger.warning(
            "Found a prefix, but could not find the expected"
            f" site-packages: {prefix=}, {discovered_site_packages=}"
        )
    else:
        logger.debug(f"Using python site-packages: {site_packages}")
        _check_algopy_version(site_packages)

    return python_exe

def _get_prefix() -> str | None:
    # look for VIRTUAL_ENV as we want the venv puyapy is being run against (i.e. the project),
    # if no venv is active, then fallback to the ambient python prefix
    venv = os.getenv("VIRTUAL_ENV")
    if venv:
        return venv
    for python in ("python3", "python"):
        prefix_result = subprocess.run(  # noqa: S602
            f"{python} -c 'import sys; print(sys.prefix or sys.base_prefix)'",
            shell=True,
            text=True,
            capture_output=True,
            check=False,
        )
        if prefix_result.returncode == 0 and (maybe_prefix := prefix_result.stdout.strip()):
            return maybe_prefix
    return None

_STUBS_PACKAGE_NAME = "algorand-python"

def _check_algopy_version(site_packages: Path) -> None:
    pkgs = metadata.Distribution.discover(name=_STUBS_PACKAGE_NAME, path=[str(site_packages)])
    try:
        (algopy,) = pkgs
    except ValueError:
        logger.warning("Could not determine algopy version")
        return
    algopy_version = version.parse(algopy.version)
    logger.debug(f"Found algopy: {algopy_version}")

    if not (MIN_SUPPORTED_ALGOPY_VERSION <= algopy_version < MAX_SUPPORTED_ALGOPY_VERSION_EX):
        logger.warning(
            f"{_STUBS_PACKAGE_NAME} version {algopy_version} is outside the supported range:"
            f" >={MIN_SUPPORTED_ALGOPY_VERSION}, <{MAX_SUPPORTED_ALGOPY_VERSION_EX}",
            important=True,
            related_lines=[
                "This will cause typing errors if there are incompatibilities in the API used.",
                "Please update your algorand-python package to be in the supported range.",
            ],
        )

def output_awst(awst: AWST, awst_out_dir: Path) -> None:
    _output_awst_any(awst, ToCodeVisitor().visit_module, awst_out_dir, ".awst")

def output_awst_json(awst: AWST, awst_out_dir: Path) -> None:
    _output_awst_any(awst, awst_to_json, awst_out_dir, ".awst.json")

def _output_awst_any(
    awst: AWST, formatter: Callable[[AWST], str], awst_out_dir: Path, suffix: str
) -> None:
    out_text = formatter(awst)
    awst_out_dir.mkdir(exist_ok=True)
    output_path = awst_out_dir / f"module{suffix}"
    logger.info(f"writing {make_path_relative_to_cwd(output_path)}")
    output_path.write_text(out_text, "utf-8")


import abc
import contextlib
import typing
from collections.abc import Iterator, Mapping, Sequence, Set
from functools import cached_property

import attrs

from puya import log
from puya.avm import OnCompletionAction
from puya.awst.nodes import (
    AppStorageDefinition,
    AppStorageKind,
    ARC4ABIMethodConfig,
    ARC4BareMethodConfig,
    ARC4CreateOption,
    ContractMethod,
    StateTotals,
)
from puya.errors import CodeError
from puya.parse import SourceLocation
from puya.program_refs import ContractReference
from puyapy.awst_build import pytypes

logger = log.get_logger(__name__)

@attrs.frozen
class ContractClassOptions:
    name_override: str | None
    scratch_slot_reservations: Set[int] | None
    state_totals: StateTotals | None
    avm_version: int | None

@attrs.frozen
class ARC4BareMethodData:
    pytype: pytypes.FuncType
    source_location: SourceLocation
    member_name: str
    config: ARC4BareMethodConfig
    is_bare: bool = attrs.field(default=True, init=False)

@attrs.frozen
class ARC4ABIMethodData:
    pytype: pytypes.FuncType
    source_location: SourceLocation
    member_name: str
    config: ARC4ABIMethodConfig
    is_bare: bool = attrs.field(default=False, init=False)
    _signature: dict[str, pytypes.PyType]
    _arc4_signature: Mapping[str, pytypes.PyType] = attrs.field(init=False)

    @_arc4_signature.default
    def _arc4_signature_default(self) -> Mapping[str, pytypes.PyType]:
        from puyapy.awst_build.arc4_utils import pytype_to_arc4_pytype  # TODO: resolve circularity

        def on_error(bad_type: pytypes.PyType) -> typing.Never:
            raise CodeError(
                f"invalid type for an ARC4 method: {bad_type}", self.config.source_location
            )

        return {k: pytype_to_arc4_pytype(v, on_error=on_error) for k, v in self._signature.items()}

    @property
    def signature(self) -> Mapping[str, pytypes.PyType]:
        return self._signature

    @cached_property
    def return_type(self) -> pytypes.PyType:
        return self._signature["output"]

    @cached_property
    def arc4_return_type(self) -> pytypes.PyType:
        return self._arc4_signature["output"]

    @cached_property
    def argument_types(self) -> Sequence[pytypes.PyType]:
        names, types = zip(*self._signature.items(), strict=True)
        assert names[-1] == "output"
        return tuple(types[:-1])

    @cached_property
    def arc4_argument_types(self) -> Sequence[pytypes.PyType]:
        names, types = zip(*self._arc4_signature.items(), strict=True)
        assert names[-1] == "output"
        return tuple(types[:-1])

ARC4MethodData: typing.TypeAlias = ARC4BareMethodData | ARC4ABIMethodData

@attrs.define
class ContractFragmentStorage:
    member_name: str = attrs.field(on_setattr=attrs.setters.frozen)
    source_location: SourceLocation = attrs.field(on_setattr=attrs.setters.frozen)
    kind: AppStorageKind = attrs.field(on_setattr=attrs.setters.frozen)
    definition: AppStorageDefinition | None

@attrs.define(kw_only=True)
class ContractFragmentMethod:
    member_name: str = attrs.field(on_setattr=attrs.setters.frozen)
    source_location: SourceLocation = attrs.field(on_setattr=attrs.setters.frozen)
    metadata: ARC4MethodData | None = attrs.field(on_setattr=attrs.setters.frozen)
    is_trivial: bool = attrs.field(on_setattr=attrs.setters.frozen)
    synthetic: bool = attrs.field(default=False, on_setattr=attrs.setters.frozen)
    inheritable: bool = attrs.field(default=True, on_setattr=attrs.setters.frozen)
    implementation: ContractMethod | None = attrs.field(default=None)

    @implementation.validator
    def _implementation_validate(self, _attr: object, value: ContractMethod | None) -> None:
        if value is None:
            return
        assert not self.is_trivial, "trivial methods should not have an implementation"
        assert value.member_name == self.member_name

class ContractFragmentBase(abc.ABC):
    @property
    @abc.abstractmethod
    def id(self) -> ContractReference: ...

    @property
    @abc.abstractmethod
    def mro(self) -> Sequence["ContractFragmentBase"]: ...

    @property
    @abc.abstractmethod
    def symbols(self) -> Mapping[str, pytypes.PyType | None]: ...

    @typing.final
    def resolve_symbol(self, name: str) -> pytypes.PyType | None:
        for frag in (self, *self.mro):
            with contextlib.suppress(KeyError):
                return frag.symbols[name]
        return None

    @abc.abstractmethod
    def resolve_method(
        self, name: str, *, include_inherited: bool = True
    ) -> ContractFragmentMethod | None: ...

    @abc.abstractmethod
    def methods(
        self, *, include_inherited: bool = True, include_overridden: bool = False
    ) -> Iterator[ContractFragmentMethod]: ...

    @abc.abstractmethod
    def resolve_storage(
        self, name: str, *, include_inherited: bool = True
    ) -> ContractFragmentStorage | None: ...

    @abc.abstractmethod
    def state(self, *, include_inherited: bool = True) -> Iterator[ContractFragmentStorage]: ...

    @typing.final
    def find_arc4_method_metadata(
        self,
        *,
        bare: bool | None = None,
        oca: OnCompletionAction | None = None,
        can_create: bool | None = None,
    ) -> Iterator[ARC4MethodData]:
        for method in self.methods():
            md = method.metadata
            if md is not None:
                bare_matches = bare is None or bare == md.is_bare
                oca_matches = oca is None or oca in md.config.allowed_completion_types
                can_create_matches = can_create is None or (
                    can_create != (md.config.create == ARC4CreateOption.disallow)
                )
                if bare_matches and oca_matches and can_create_matches:
                    yield md

ConstantValue: typing.TypeAlias = int | str | bytes | bool


from collections.abc import Sequence
from pathlib import Path

import attrs

from puya.log import LogLevel
from puya.options import PuyaOptions

@attrs.frozen(kw_only=True)
class PuyaPyOptions(PuyaOptions):
    paths: Sequence[Path] = attrs.field(default=(), repr=lambda p: str(list(map(str, p))))
    output_awst: bool = False
    output_awst_json: bool = False
    output_client: bool = False
    out_dir: Path | None = attrs.field(default=None, repr=False)
    log_level: LogLevel = LogLevel.info