import typing
from collections.abc import Sequence

import mypy.nodes

from puya import log
from puya.awst.nodes import IntrinsicCall, MethodConstant
from puya.errors import CodeError
from puya.parse import SourceLocation
from puyapy.awst_build import pytypes
from puyapy.awst_build.eb import _expect as expect
from puyapy.awst_build.eb._base import FunctionBuilder
from puyapy.awst_build.eb._literals import LiteralBuilderImpl
from puyapy.awst_build.eb._utils import dummy_value
from puyapy.awst_build.eb.bytes import BytesExpressionBuilder
from puyapy.awst_build.eb.factories import builder_for_instance, builder_for_type
from puyapy.awst_build.eb.interface import (
    InstanceBuilder,
    LiteralBuilder,
    NodeBuilder,
    TypeBuilder,
)
from puyapy.awst_build.intrinsic_models import (
    FunctionOpMapping,
    OpMappingWithOverloads,
    PropertyOpMapping,
)

logger = log.get_logger(__name__)

class Arc4SignatureBuilder(FunctionBuilder):
    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        arg = expect.exactly_one_arg(args, location, default=expect.default_none)
        if arg is None:
            return dummy_value(pytypes.BytesType, location)
        str_value = expect.simple_string_literal(arg, default=expect.default_fixed_value(""))
        return BytesExpressionBuilder(
            MethodConstant(
                value=str_value,
                source_location=location,
            )
        )

class IntrinsicEnumTypeBuilder(TypeBuilder[pytypes.IntrinsicEnumType]):
    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        raise CodeError("cannot instantiate enumeration type", location)

    @typing.override
    def member_access(self, name: str, location: SourceLocation) -> NodeBuilder:
        produces = self.produces()
        value = produces.members.get(name)
        if value is None:
            return super().member_access(name, location)
        return LiteralBuilderImpl(value=value, source_location=location)

class IntrinsicNamespaceTypeBuilder(TypeBuilder[pytypes.IntrinsicNamespaceType]):
    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        raise CodeError("cannot instantiate namespace type", location)

    @typing.override
    def member_access(self, name: str, location: SourceLocation) -> NodeBuilder:
        produces = self.produces()
        mapping = produces.members.get(name)
        if mapping is None:
            return super().member_access(name, location)
        if isinstance(mapping, PropertyOpMapping):
            intrinsic_expr = IntrinsicCall(
                op_code=mapping.op_code,
                immediates=[mapping.immediate],
                wtype=mapping.wtype,
                source_location=location,
            )
            return builder_for_instance(mapping.typ, intrinsic_expr)
        else:
            fullname = ".".join((produces.name, name))
            return IntrinsicFunctionExpressionBuilder(fullname, mapping, location)

class IntrinsicFunctionExpressionBuilder(FunctionBuilder):
    def __init__(
        self, fullname: str, mapping: OpMappingWithOverloads, location: SourceLocation
    ) -> None:
        self._fullname = fullname
        self._mapping = mapping
        super().__init__(location)

    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        if not expect.exactly_n_args(args, location, self._mapping.arity):
            return dummy_value(self._mapping.result, location)
        intrinsic_expr = _map_call(self._mapping, args=args, location=location)
        return builder_for_instance(self._mapping.result, intrinsic_expr)

def _best_op_mapping(
    op_mappings: OpMappingWithOverloads, args: Sequence[NodeBuilder]
) -> FunctionOpMapping:
    """Find op mapping that matches as many arguments to immediate args as possible"""
    literal_arg_positions = {
        arg_idx
        for arg_idx, arg in enumerate(args)
        # we can't handle any form of dynamism for immediates, such as `1 if foo else 2`,
        # so we must check for LiteralBuilder only
        if isinstance(arg, LiteralBuilder)
    }
    for op_mapping in sorted(
        op_mappings.overloads, key=lambda om: len(om.literal_arg_positions), reverse=True
    ):
        if literal_arg_positions.issuperset(op_mapping.literal_arg_positions):
            return op_mapping
    # fall back to first, let argument mapping handle logging errors
    return op_mappings.overloads[0]

def _map_call(
    ast_mapper: OpMappingWithOverloads,
    args: Sequence[NodeBuilder],
    location: SourceLocation,
) -> IntrinsicCall:
    if len(ast_mapper.overloads) == 1:
        (op_mapping,) = ast_mapper.overloads
    else:
        op_mapping = _best_op_mapping(ast_mapper, args)

    immediates = list(op_mapping.immediates)
    stack_args = list[InstanceBuilder]()
    for arg, arg_data in zip(args, op_mapping.args, strict=True):
        arg_in = expect.instance_builder(arg, default=expect.default_none)
        if arg_in is None:
            pass
        elif isinstance(arg_data, int):
            immediates_index = arg_data
            literal_type = typing.cast(type[str | int], immediates[immediates_index])
            if not (
                isinstance(arg_in, LiteralBuilder)  # see note in _best_op_mapping
                and isinstance(arg_value := arg_in.value, literal_type)
            ):
                logger.error(
                    f"argument must be a literal {literal_type.__name__} value",
                    location=arg_in.source_location,
                )
                immediates[immediates_index] = literal_type()
            else:
                assert isinstance(arg_value, int | str)
                immediates[immediates_index] = arg_value
        else:
            allowed_pytypes = arg_data
            if isinstance(arg_in.pytype, pytypes.LiteralOnlyType) or (
                isinstance(arg_in.pytype, pytypes.UnionType)
                and any(isinstance(t, pytypes.LiteralOnlyType) for t in arg_in.pytype.types)
            ):
                for allowed_type in allowed_pytypes:
                    type_builder = builder_for_type(allowed_type, arg_in.source_location)
                    if isinstance(type_builder, TypeBuilder):
                        converted = arg_in.try_resolve_literal(type_builder)
                        if converted is not None:
                            arg_in = converted
                            break
            stack_args.append(expect.argument_of_type_else_dummy(arg_in, *allowed_pytypes))
    return IntrinsicCall(
        op_code=op_mapping.op_code,
        wtype=ast_mapper.result_wtype,
        immediates=typing.cast(list[str | int], immediates),
        stack_args=[a.resolve() for a in stack_args],
        source_location=location,
    )


import typing
from collections.abc import Sequence

import mypy.nodes

from puya import log
from puya.awst.nodes import BytesConstant, BytesEncoding, Expression
from puya.parse import SourceLocation
from puyapy.awst_build import intrinsic_factory, pytypes
from puyapy.awst_build.eb import _expect as expect
from puyapy.awst_build.eb._base import FunctionBuilder
from puyapy.awst_build.eb.interface import InstanceBuilder, NodeBuilder
from puyapy.awst_build.eb.none import NoneExpressionBuilder
from puyapy.awst_build.eb.uint64 import UInt64TypeBuilder

logger = log.get_logger(__name__)

class LogBuilder(FunctionBuilder):
    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        empty_utf8: Expression = BytesConstant(
            value=b"", encoding=BytesEncoding.utf8, source_location=location
        )
        args = list(args)
        try:
            sep_index = arg_names.index("sep")
        except ValueError:
            sep = empty_utf8
        else:
            sep_arg = args.pop(sep_index)
            if isinstance(sep_arg, InstanceBuilder) and sep_arg.pytype.is_type_or_subtype(
                pytypes.StringType,
                pytypes.StrLiteralType,
                pytypes.BytesType,
                pytypes.BytesLiteralType,
            ):
                sep = sep_arg.to_bytes(sep_arg.source_location)
            else:
                expect.not_this_type(sep_arg, default=expect.default_none)
                sep = empty_utf8

        bytes_args = []
        for arg in args:
            if not isinstance(arg, InstanceBuilder):
                expect.not_this_type(arg, default=expect.default_none)
            else:
                if arg.pytype == pytypes.IntLiteralType:  # match int exactly, ie exclude bool
                    arg = arg.resolve_literal(UInt64TypeBuilder(arg.source_location))
                # TODO: make to_bytes non-throwing
                bytes_expr = arg.to_bytes(arg.source_location)
                bytes_args.append(bytes_expr)

        if not bytes_args:
            log_value = empty_utf8
        else:
            log_value = bytes_args[0]
            for bytes_expr in bytes_args[1:]:
                arg_plus_sep = intrinsic_factory.concat(log_value, sep, bytes_expr.source_location)
                log_value = intrinsic_factory.concat(
                    arg_plus_sep, bytes_expr, bytes_expr.source_location
                )
        log_expr = intrinsic_factory.log(log_value, location)
        return NoneExpressionBuilder(log_expr)


import typing
from collections.abc import Sequence

import mypy.nodes

from puya.errors import CodeError
from puya.parse import SourceLocation
from puya.program_refs import LogicSigReference
from puyapy.awst_build import pytypes
from puyapy.awst_build.eb.interface import InstanceBuilder, NodeBuilder, TypeBuilder

class LogicSigExpressionBuilder(TypeBuilder):
    def __init__(self, ref: LogicSigReference, location: SourceLocation):
        super().__init__(pytypes.LogicSigType, location)
        self.ref = ref

    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        raise CodeError("cannot instantiate logic signatures", location)


import typing
from collections.abc import Sequence

import mypy.nodes

from puya.awst.nodes import Expression, VoidConstant
from puya.errors import CodeError
from puya.parse import SourceLocation
from puyapy.awst_build import pytypes
from puyapy.awst_build.eb import _expect as expect
from puyapy.awst_build.eb._base import NotIterableInstanceExpressionBuilder
from puyapy.awst_build.eb._utils import constant_bool_and_error
from puyapy.awst_build.eb.interface import InstanceBuilder, NodeBuilder, TypeBuilder

class NoneTypeBuilder(TypeBuilder):
    def __init__(self, location: SourceLocation):
        super().__init__(pytypes.NoneType, location)

    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        expect.no_args(args, location)
        return NoneExpressionBuilder(VoidConstant(location))

class NoneExpressionBuilder(NotIterableInstanceExpressionBuilder):
    def __init__(self, expr: Expression):
        super().__init__(pytypes.NoneType, expr)

    @typing.override
    def to_bytes(self, location: SourceLocation) -> Expression:
        raise CodeError("None is not usable as a value", location)

    @typing.override
    def bool_eval(self, location: SourceLocation, *, negate: bool = False) -> InstanceBuilder:
        return constant_bool_and_error(value=False, location=location, negate=negate)


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

from immutabledict import immutabledict

from puya.awst import wtypes
from puya.awst.nodes import (
    CheckedMaybe,
    Expression,
    IntrinsicCall,
    Not,
    NumericComparison,
    NumericComparisonExpression,
    ReinterpretCast,
)
from puya.parse import SourceLocation
from puyapy.awst_build import intrinsic_factory, pytypes
from puyapy.awst_build.eb._base import NotIterableInstanceExpressionBuilder
from puyapy.awst_build.eb.bool import BoolExpressionBuilder
from puyapy.awst_build.eb.factories import builder_for_instance
from puyapy.awst_build.eb.interface import (
    BuilderComparisonOp,
    InstanceBuilder,
    NodeBuilder,
    TypeBuilder,
)

class ReferenceValueExpressionBuilder(NotIterableInstanceExpressionBuilder, abc.ABC):
    def __init__(
        self,
        expr: Expression,
        *,
        typ: pytypes.PyType,
        native_type: pytypes.RuntimeType,
        native_access_member: str,
        field_mapping: Mapping[str, tuple[str, pytypes.RuntimeType]],
        field_op_code: str,
        field_bool_comment: str,
    ):
        super().__init__(typ, expr)
        self.native_type = native_type
        self.native_access_member = native_access_member
        self.field_mapping = immutabledict[str, tuple[str, pytypes.RuntimeType]](field_mapping)
        self.field_op_code = field_op_code
        self.field_bool_comment = field_bool_comment

    @typing.override
    def member_access(self, name: str, location: SourceLocation) -> NodeBuilder:
        if name == self.native_access_member:
            native_cast = ReinterpretCast(
                expr=self.resolve(), wtype=self.native_type.wtype, source_location=location
            )
            return builder_for_instance(self.native_type, native_cast)
        if name in self.field_mapping:
            immediate, typ = self.field_mapping[name]
            acct_params_get = IntrinsicCall(
                op_code=self.field_op_code,
                immediates=[immediate],
                stack_args=[self.resolve()],
                wtype=wtypes.WTuple((typ.wtype, wtypes.bool_wtype), location),
                source_location=location,
            )
            checked_maybe = CheckedMaybe(acct_params_get, comment=self.field_bool_comment)
            return builder_for_instance(typ, checked_maybe)
        return super().member_access(name, location)

class UInt64BackedReferenceValueExpressionBuilder(ReferenceValueExpressionBuilder):
    def __init__(
        self,
        expr: Expression,
        *,
        typ: pytypes.PyType,
        typ_literal_converter: Callable[[SourceLocation], TypeBuilder],
        native_access_member: str,
        field_mapping: Mapping[str, tuple[str, pytypes.RuntimeType]],
        field_op_code: str,
        field_bool_comment: str,
    ):
        self._typ_literal_converter = typ_literal_converter
        super().__init__(
            expr,
            typ=typ,
            native_type=pytypes.UInt64Type,
            native_access_member=native_access_member,
            field_mapping=field_mapping,
            field_op_code=field_op_code,
            field_bool_comment=field_bool_comment,
        )

    @typing.override
    def to_bytes(self, location: SourceLocation) -> Expression:
        return intrinsic_factory.itob(self.resolve(), location)

    @typing.override
    def bool_eval(self, location: SourceLocation, *, negate: bool = False) -> InstanceBuilder:
        as_bool = ReinterpretCast(
            expr=self.resolve(),
            wtype=wtypes.bool_wtype,
            source_location=self.resolve().source_location,
        )
        if negate:
            expr: Expression = Not(location, as_bool)
        else:
            expr = as_bool
        return BoolExpressionBuilder(expr)

    @typing.override
    def compare(
        self, other: InstanceBuilder, op: BuilderComparisonOp, location: SourceLocation
    ) -> InstanceBuilder:
        other = other.resolve_literal(self._typ_literal_converter(other.source_location))
        if not (
            other.pytype == self.pytype  # can only compare with other of same type?
            and op in (BuilderComparisonOp.eq, BuilderComparisonOp.ne)
        ):
            return NotImplemented
        cmp_expr = NumericComparisonExpression(
            source_location=location,
            lhs=self.resolve(),
            operator=NumericComparison(op.value),
            rhs=other.resolve(),
        )
        return BoolExpressionBuilder(cmp_expr)


import typing
from collections.abc import Sequence

import mypy.nodes

from puya import log, utils
from puya.algo_constants import ENCODED_ADDRESS_LENGTH
from puya.awst import wtypes
from puya.awst.nodes import (
    AddressConstant,
    CheckedMaybe,
    Expression,
    IntrinsicCall,
    NumericComparison,
    NumericComparisonExpression,
    ReinterpretCast,
    TupleItemExpression,
    UInt64Constant,
)
from puya.parse import SourceLocation
from puyapy.awst_build import intrinsic_factory, pytypes
from puyapy.awst_build.eb import _expect as expect
from puyapy.awst_build.eb._base import FunctionBuilder
from puyapy.awst_build.eb._bytes_backed import BytesBackedTypeBuilder
from puyapy.awst_build.eb._utils import (
    cast_to_bytes,
    compare_bytes,
    compare_expr_bytes,
    dummy_value,
)
from puyapy.awst_build.eb.bool import BoolExpressionBuilder
from puyapy.awst_build.eb.interface import (
    BuilderComparisonOp,
    InstanceBuilder,
    LiteralBuilder,
    NodeBuilder,
)
from puyapy.awst_build.eb.reference_types._base import ReferenceValueExpressionBuilder

logger = log.get_logger(__name__)

class AccountTypeBuilder(BytesBackedTypeBuilder):
    def __init__(self, location: SourceLocation):
        super().__init__(pytypes.AccountType, location)

    @typing.override
    def try_convert_literal(
        self, literal: LiteralBuilder, location: SourceLocation
    ) -> InstanceBuilder | None:
        match literal.value:
            case str(str_value):
                if not utils.valid_address(str_value):
                    logger.error(
                        "invalid address value - should have length"
                        f" {ENCODED_ADDRESS_LENGTH} and not include base32 padding",
                        location=literal.source_location,
                    )
                expr = AddressConstant(
                    value=str_value,
                    wtype=wtypes.account_wtype,
                    source_location=location,
                )
                return AccountExpressionBuilder(expr)
        return None

    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        arg = expect.at_most_one_arg(args, location)
        match arg:
            case InstanceBuilder(pytype=pytypes.StrLiteralType):
                return arg.resolve_literal(converter=AccountTypeBuilder(location))
            case None:
                value: Expression = intrinsic_factory.zero_address(location)
            case _:
                arg = expect.argument_of_type_else_dummy(arg, pytypes.BytesType)
                address_bytes_temp = arg.single_eval().resolve()
                is_correct_length = NumericComparisonExpression(
                    operator=NumericComparison.eq,
                    source_location=location,
                    lhs=UInt64Constant(value=32, source_location=location),
                    rhs=intrinsic_factory.bytes_len(address_bytes_temp, location),
                )
                value = CheckedMaybe.from_tuple_items(
                    expr=ReinterpretCast(
                        expr=address_bytes_temp,
                        wtype=wtypes.account_wtype,
                        source_location=location,
                    ),
                    check=is_correct_length,
                    source_location=location,
                    comment="Address length is 32 bytes",
                )
        return AccountExpressionBuilder(value)

class AccountExpressionBuilder(ReferenceValueExpressionBuilder):
    def __init__(self, expr: Expression):
        native_type = pytypes.BytesType
        native_access_member = "bytes"
        field_mapping = {
            "balance": ("AcctBalance", pytypes.UInt64Type),
            "min_balance": ("AcctMinBalance", pytypes.UInt64Type),
            "auth_address": ("AcctAuthAddr", pytypes.AccountType),
            "total_num_uint": ("AcctTotalNumUint", pytypes.UInt64Type),
            "total_num_byte_slice": ("AcctTotalNumByteSlice", pytypes.UInt64Type),
            "total_extra_app_pages": ("AcctTotalExtraAppPages", pytypes.UInt64Type),
            "total_apps_created": ("AcctTotalAppsCreated", pytypes.UInt64Type),
            "total_apps_opted_in": ("AcctTotalAppsOptedIn", pytypes.UInt64Type),
            "total_assets_created": ("AcctTotalAssetsCreated", pytypes.UInt64Type),
            "total_assets": ("AcctTotalAssets", pytypes.UInt64Type),
            "total_boxes": ("AcctTotalBoxes", pytypes.UInt64Type),
            "total_box_bytes": ("AcctTotalBoxBytes", pytypes.UInt64Type),
        }
        field_op_code = "acct_params_get"
        field_bool_comment = "account funded"
        super().__init__(
            expr,
            typ=pytypes.AccountType,
            native_type=native_type,
            native_access_member=native_access_member,
            field_mapping=field_mapping,
            field_op_code=field_op_code,
            field_bool_comment=field_bool_comment,
        )

    @typing.override
    def to_bytes(self, location: SourceLocation) -> Expression:
        return cast_to_bytes(self.resolve(), location)

    @typing.override
    def member_access(self, name: str, location: SourceLocation) -> NodeBuilder:
        if name == "is_opted_in":
            return _IsOptedIn(self.resolve(), location)
        return super().member_access(name, location)

    @typing.override
    def bool_eval(self, location: SourceLocation, *, negate: bool = False) -> InstanceBuilder:
        return compare_expr_bytes(
            lhs=self.resolve(),
            op=BuilderComparisonOp.eq if negate else BuilderComparisonOp.ne,
            rhs=intrinsic_factory.zero_address(location),
            source_location=location,
        )

    @typing.override
    def compare(
        self, other: InstanceBuilder, op: BuilderComparisonOp, location: SourceLocation
    ) -> InstanceBuilder:
        other = other.resolve_literal(converter=AccountTypeBuilder(other.source_location))
        return compare_bytes(self=self, op=op, other=other, source_location=location)

class _IsOptedIn(FunctionBuilder):
    def __init__(self, expr: Expression, source_location: SourceLocation):
        super().__init__(source_location)
        self.expr = expr

    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        arg = expect.exactly_one_arg(args, location, default=expect.default_none)
        if arg is None:
            return dummy_value(pytypes.BoolType, location)
        elif pytypes.AssetType <= arg.pytype:
            return BoolExpressionBuilder(
                TupleItemExpression(
                    base=IntrinsicCall(
                        op_code="asset_holding_get",
                        immediates=["AssetBalance"],
                        stack_args=[self.expr, arg.resolve()],
                        wtype=wtypes.WTuple((wtypes.uint64_wtype, wtypes.bool_wtype), location),
                        source_location=location,
                    ),
                    index=1,
                    source_location=location,
                )
            )
        else:
            arg = expect.argument_of_type_else_dummy(arg, pytypes.ApplicationType)
            return BoolExpressionBuilder(
                IntrinsicCall(
                    op_code="app_opted_in",
                    stack_args=[self.expr, arg.resolve()],
                    source_location=location,
                    wtype=wtypes.bool_wtype,
                )
            )


import typing
from collections.abc import Sequence

import mypy.nodes

from puya import log
from puya.awst import wtypes
from puya.awst.nodes import Expression, ReinterpretCast, UInt64Constant
from puya.parse import SourceLocation
from puyapy.awst_build import pytypes
from puyapy.awst_build.eb import _expect as expect
from puyapy.awst_build.eb.interface import (
    InstanceBuilder,
    LiteralBuilder,
    NodeBuilder,
    TypeBuilder,
)
from puyapy.awst_build.eb.reference_types._base import UInt64BackedReferenceValueExpressionBuilder

logger = log.get_logger(__name__)

class ApplicationTypeBuilder(TypeBuilder):
    def __init__(self, location: SourceLocation):
        super().__init__(pytypes.ApplicationType, location)

    @typing.override
    def try_convert_literal(
        self, literal: LiteralBuilder, location: SourceLocation
    ) -> InstanceBuilder | None:
        match literal.value:
            case int(int_value):
                if int_value < 0 or int_value.bit_length() > 64:  # TODO: should this be 256?
                    logger.error("invalid application ID", location=literal.source_location)
                const = UInt64Constant(value=int_value, source_location=location)
                expr = ReinterpretCast(
                    expr=const, wtype=wtypes.application_wtype, source_location=location
                )
                return ApplicationExpressionBuilder(expr)
        return None

    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        arg = expect.at_most_one_arg(args, location)
        match arg:
            case InstanceBuilder(pytype=pytypes.IntLiteralType):
                return arg.resolve_literal(ApplicationTypeBuilder(location))
            case None:
                uint64_expr: Expression = UInt64Constant(value=0, source_location=location)
            case _:
                arg = expect.argument_of_type_else_dummy(arg, pytypes.UInt64Type)
                uint64_expr = arg.resolve()
        app_expr = ReinterpretCast(
            source_location=location, wtype=wtypes.application_wtype, expr=uint64_expr
        )
        return ApplicationExpressionBuilder(app_expr)

class ApplicationExpressionBuilder(UInt64BackedReferenceValueExpressionBuilder):
    def __init__(self, expr: Expression):
        native_access_member = "id"
        field_mapping = {
            "approval_program": ("AppApprovalProgram", pytypes.BytesType),
            "clear_state_program": ("AppClearStateProgram", pytypes.BytesType),
            "global_num_uint": ("AppGlobalNumUint", pytypes.UInt64Type),
            "global_num_bytes": ("AppGlobalNumByteSlice", pytypes.UInt64Type),
            "local_num_uint": ("AppLocalNumUint", pytypes.UInt64Type),
            "local_num_bytes": ("AppLocalNumByteSlice", pytypes.UInt64Type),
            "extra_program_pages": ("AppExtraProgramPages", pytypes.UInt64Type),
            "creator": ("AppCreator", pytypes.AccountType),
            "address": ("AppAddress", pytypes.AccountType),
        }
        field_op_code = "app_params_get"
        field_bool_comment = "application exists"
        super().__init__(
            expr,
            typ=pytypes.ApplicationType,
            typ_literal_converter=ApplicationTypeBuilder,
            native_access_member=native_access_member,
            field_mapping=field_mapping,
            field_op_code=field_op_code,
            field_bool_comment=field_bool_comment,
        )


import typing
from collections.abc import Sequence

import mypy.nodes

from puya import log
from puya.awst import wtypes
from puya.awst.nodes import (
    CheckedMaybe,
    Expression,
    IntrinsicCall,
    ReinterpretCast,
    UInt64Constant,
)
from puya.parse import SourceLocation
from puyapy.awst_build import pytypes
from puyapy.awst_build.eb import _expect as expect
from puyapy.awst_build.eb._base import FunctionBuilder
from puyapy.awst_build.eb.factories import builder_for_instance
from puyapy.awst_build.eb.interface import (
    InstanceBuilder,
    LiteralBuilder,
    NodeBuilder,
    TypeBuilder,
)
from puyapy.awst_build.eb.reference_types._base import UInt64BackedReferenceValueExpressionBuilder

logger = log.get_logger(__name__)

class AssetTypeBuilder(TypeBuilder[pytypes.RuntimeType]):
    def __init__(self, location: SourceLocation):
        super().__init__(pytypes.AssetType, location)

    @typing.override
    def try_convert_literal(
        self, literal: LiteralBuilder, location: SourceLocation
    ) -> InstanceBuilder | None:
        match literal.value:
            case int(int_value):
                if int_value < 0 or int_value.bit_length() > 64:  # TODO: should this be 256?
                    logger.error("invalid asset ID", location=literal.source_location)
                const = UInt64Constant(value=int_value, source_location=location)
                expr = ReinterpretCast(
                    expr=const, wtype=self.produces().wtype, source_location=location
                )
                return AssetExpressionBuilder(expr)
        return None

    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        arg = expect.at_most_one_arg(args, location)
        match arg:
            case InstanceBuilder(pytype=pytypes.IntLiteralType):
                return arg.resolve_literal(AssetTypeBuilder(location))
            case None:
                uint64_expr: Expression = UInt64Constant(value=0, source_location=location)
            case _:
                arg = expect.argument_of_type_else_dummy(arg, pytypes.UInt64Type)
                uint64_expr = arg.resolve()
        asset_expr = ReinterpretCast(
            expr=uint64_expr, wtype=wtypes.asset_wtype, source_location=location
        )
        return AssetExpressionBuilder(asset_expr)

ASSET_HOLDING_FIELD_MAPPING: typing.Final = {
    "balance": ("AssetBalance", pytypes.UInt64Type),
    "frozen": ("AssetFrozen", pytypes.BoolType),
}

class AssetHoldingExpressionBuilder(FunctionBuilder):
    def __init__(self, asset: Expression, holding_field: str, location: SourceLocation):
        self.asset = asset
        self.holding_field = holding_field
        super().__init__(location)

    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        arg = expect.exactly_one_arg_of_type_else_dummy(args, pytypes.AccountType, location)
        account_expr = arg.resolve()
        immediate, typ = ASSET_HOLDING_FIELD_MAPPING[self.holding_field]
        asset_params_get = IntrinsicCall(
            source_location=location,
            wtype=wtypes.WTuple((typ.wtype, wtypes.bool_wtype), location),
            op_code="asset_holding_get",
            immediates=[immediate],
            stack_args=[account_expr, self.asset],
        )
        return builder_for_instance(
            typ, CheckedMaybe(asset_params_get, comment="account opted into asset")
        )

class AssetExpressionBuilder(UInt64BackedReferenceValueExpressionBuilder):
    def __init__(self, expr: Expression):
        native_access_member = "id"
        field_mapping = {
            "total": ("AssetTotal", pytypes.UInt64Type),
            "decimals": ("AssetDecimals", pytypes.UInt64Type),
            "default_frozen": ("AssetDefaultFrozen", pytypes.BoolType),
            "unit_name": ("AssetUnitName", pytypes.BytesType),
            "name": ("AssetName", pytypes.BytesType),
            "url": ("AssetURL", pytypes.BytesType),
            "metadata_hash": ("AssetMetadataHash", pytypes.BytesType),
            "manager": ("AssetManager", pytypes.AccountType),
            "reserve": ("AssetReserve", pytypes.AccountType),
            "freeze": ("AssetFreeze", pytypes.AccountType),
            "clawback": ("AssetClawback", pytypes.AccountType),
            "creator": ("AssetCreator", pytypes.AccountType),
        }
        field_op_code = "asset_params_get"
        field_bool_comment = "asset exists"
        super().__init__(
            expr,
            typ=pytypes.AssetType,
            typ_literal_converter=AssetTypeBuilder,
            native_access_member=native_access_member,
            field_mapping=field_mapping,
            field_op_code=field_op_code,
            field_bool_comment=field_bool_comment,
        )

    @typing.override
    def member_access(self, name: str, location: SourceLocation) -> NodeBuilder:
        if name in ASSET_HOLDING_FIELD_MAPPING:
            return AssetHoldingExpressionBuilder(self.resolve(), name, location)
        return super().member_access(name, location)


from puyapy.awst_build.eb.storage.box import (
    BoxGenericTypeExpressionBuilder,
    BoxProxyExpressionBuilder,
    BoxTypeBuilder,
)
from puyapy.awst_build.eb.storage.box_map import (
    BoxMapGenericTypeExpressionBuilder,
    BoxMapProxyExpressionBuilder,
    BoxMapTypeBuilder,
)
from puyapy.awst_build.eb.storage.box_ref import (
    BoxRefProxyExpressionBuilder,
    BoxRefTypeBuilder,
)
from puyapy.awst_build.eb.storage.global_state import (
    GlobalStateExpressionBuilder,
    GlobalStateGenericTypeBuilder,
    GlobalStateTypeBuilder,
)
from puyapy.awst_build.eb.storage.local_state import (
    LocalStateExpressionBuilder,
    LocalStateGenericTypeBuilder,
    LocalStateTypeBuilder,
)

__all__ = [
    "BoxTypeBuilder",
    "BoxGenericTypeExpressionBuilder",
    "BoxProxyExpressionBuilder",
    "BoxRefTypeBuilder",
    "BoxRefProxyExpressionBuilder",
    "BoxMapTypeBuilder",
    "BoxMapGenericTypeExpressionBuilder",
    "BoxMapProxyExpressionBuilder",
    "GlobalStateTypeBuilder",
    "GlobalStateGenericTypeBuilder",
    "GlobalStateExpressionBuilder",
    "LocalStateTypeBuilder",
    "LocalStateGenericTypeBuilder",
    "LocalStateExpressionBuilder",
]


import abc
import typing
from collections.abc import Sequence

import mypy.nodes

from puya.awst.nodes import (
    BoxValueExpression,
    ExpressionStatement,
    StateDelete,
    StateGet,
    StateGetEx,
    Statement,
)
from puya.parse import SourceLocation
from puyapy.awst_build import pytypes
from puyapy.awst_build.eb import _expect as expect
from puyapy.awst_build.eb._base import FunctionBuilder
from puyapy.awst_build.eb._utils import cast_to_bytes
from puyapy.awst_build.eb.factories import builder_for_instance
from puyapy.awst_build.eb.interface import InstanceBuilder, NodeBuilder
from puyapy.awst_build.eb.storage._util import box_length_checked, index_box_bytes, slice_box_bytes
from puyapy.awst_build.eb.storage._value_proxy import ValueProxyExpressionBuilder
from puyapy.awst_build.eb.tuple import TupleExpressionBuilder
from puyapy.awst_build.eb.uint64 import UInt64ExpressionBuilder

class _BoxKeyExpressionIntermediateExpressionBuilder(FunctionBuilder, abc.ABC):
    def __init__(self, box: BoxValueExpression, content_type: pytypes.PyType) -> None:
        super().__init__(box.source_location)
        self.box = box
        self.content_type = content_type

class BoxGetExpressionBuilder(_BoxKeyExpressionIntermediateExpressionBuilder):
    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        default_arg_inst = expect.exactly_one_arg_of_type_else_dummy(
            args, self.content_type, location
        )
        default_expr = default_arg_inst.resolve()
        return builder_for_instance(
            self.content_type,
            StateGet(field=self.box, default=default_expr, source_location=location),
        )

class BoxMaybeExpressionBuilder(_BoxKeyExpressionIntermediateExpressionBuilder):
    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        expect.no_args(args, location)
        result_type = pytypes.GenericTupleType.parameterise(
            [self.content_type, pytypes.BoolType], location
        )
        return TupleExpressionBuilder(
            StateGetEx(field=self.box, source_location=location),
            result_type,
        )

class BoxValueExpressionBuilder(ValueProxyExpressionBuilder[pytypes.PyType, BoxValueExpression]):
    """This is used to intercept operations on Box and BoxMap values to use more efficient
    ops (i.e. box_extract, box_length) where possible and provide support for operations like
    delete"""

    @typing.override
    def delete(self, location: SourceLocation) -> Statement:
        return ExpressionStatement(
            expr=StateDelete(field=self.resolve(), source_location=location)
        )

    @typing.override
    def member_access(self, name: str, location: SourceLocation) -> NodeBuilder:
        if name == "bytes":
            return _ValueBytes(self.resolve(), location)
        else:
            return super().member_access(name, location)

    @typing.override
    def index(self, index: InstanceBuilder, location: SourceLocation) -> InstanceBuilder:
        if pytypes.BytesType <= self.pytype:
            return index_box_bytes(self.resolve(), index, location)
        return super().index(index, location)

    @typing.override
    def slice_index(
        self,
        begin_index: InstanceBuilder | None,
        end_index: InstanceBuilder | None,
        stride: InstanceBuilder | None,
        location: SourceLocation,
    ) -> InstanceBuilder:
        if pytypes.BytesType <= self.pytype:
            return slice_box_bytes(self.resolve(), begin_index, end_index, stride, location)
        return super().slice_index(begin_index, end_index, stride, location)

class _ValueBytes(ValueProxyExpressionBuilder):
    def __init__(self, expr: BoxValueExpression, location: SourceLocation) -> None:
        self._typed = expr
        super().__init__(pytypes.BytesType, cast_to_bytes(expr, location))

    @typing.override
    def member_access(self, name: str, location: SourceLocation) -> NodeBuilder:
        match name:
            case "length":
                return UInt64ExpressionBuilder(box_length_checked(self._typed, location))
            case _:
                return super().member_access(name, location)

    @typing.override
    def index(self, index: InstanceBuilder, location: SourceLocation) -> InstanceBuilder:
        return index_box_bytes(self._typed, index, location)

    @typing.override
    def slice_index(
        self,
        begin_index: InstanceBuilder | None,
        end_index: InstanceBuilder | None,
        stride: InstanceBuilder | None,
        location: SourceLocation,
    ) -> InstanceBuilder:
        return slice_box_bytes(self._typed, begin_index, end_index, stride, location)


import typing

import attrs

from puya import log
from puya.awst import wtypes
from puya.awst.nodes import (
    BinaryBooleanOperator,
    BytesConstant,
    BytesEncoding,
    Expression,
    Lvalue,
    ReinterpretCast,
    Statement,
)
from puya.errors import CodeError
from puya.parse import SourceLocation
from puyapy.awst_build import pytypes
from puyapy.awst_build.eb import _expect as expect
from puyapy.awst_build.eb.interface import (
    BuilderBinaryOp,
    BuilderComparisonOp,
    BuilderUnaryOp,
    InstanceBuilder,
    NodeBuilder,
    StorageProxyConstructorArgs,
    StorageProxyConstructorResult,
    TypeBuilder,
)

logger = log.get_logger(__name__)

@typing.final
class StorageProxyDefinitionBuilder(StorageProxyConstructorResult):
    def __init__(
        self,
        args: StorageProxyConstructorArgs,
        typ: pytypes.StorageProxyType | pytypes.StorageMapProxyType,
        location: SourceLocation,
    ):
        super().__init__(location)
        self._typ = typ
        self._args = args

    @typing.override
    @property
    def args(self) -> StorageProxyConstructorArgs:
        return self._args

    @typing.override
    @property
    def pytype(self) -> pytypes.StorageProxyType | pytypes.StorageMapProxyType:
        return self._typ

    @typing.override
    def to_bytes(self, location: SourceLocation) -> Expression:
        return self._assign_first(location)

    @typing.override
    def resolve(self) -> Expression:
        return self._assign_first(self.source_location)

    @typing.override
    def resolve_literal(self, converter: TypeBuilder) -> InstanceBuilder:
        return self.try_resolve_literal(converter)

    @typing.override
    def try_resolve_literal(self, converter: TypeBuilder) -> InstanceBuilder:
        return self

    @typing.override
    def resolve_lvalue(self) -> Lvalue:
        raise CodeError(f"{self._typ} is not valid as an assignment target", self.source_location)

    @typing.override
    def delete(self, location: SourceLocation) -> Statement:
        raise self._assign_first(location)

    @typing.override
    def bool_eval(self, location: SourceLocation, *, negate: bool = False) -> InstanceBuilder:
        return self._assign_first(location)

    @typing.override
    def unary_op(self, op: BuilderUnaryOp, location: SourceLocation) -> InstanceBuilder:
        return self._assign_first(location)

    @typing.override
    def contains(self, item: InstanceBuilder, location: SourceLocation) -> InstanceBuilder:
        return self._assign_first(location)

    @typing.override
    def index(self, index: InstanceBuilder, location: SourceLocation) -> InstanceBuilder:
        return self._assign_first(location)

    @typing.override
    def slice_index(
        self,
        begin_index: InstanceBuilder | None,
        end_index: InstanceBuilder | None,
        stride: InstanceBuilder | None,
        location: SourceLocation,
    ) -> InstanceBuilder:
        return self._assign_first(location)

    @typing.override
    def iterate(self) -> Expression:
        return self._assign_first(self.source_location)

    @typing.override
    def iterable_item_type(self) -> pytypes.PyType:
        return self._assign_first(self.source_location)

    @typing.override
    def member_access(self, name: str, location: SourceLocation) -> NodeBuilder:
        return self._assign_first(location)

    @typing.override
    def compare(
        self, other: InstanceBuilder, op: BuilderComparisonOp, location: SourceLocation
    ) -> InstanceBuilder:
        return self._assign_first(location)

    @typing.override
    def binary_op(
        self,
        other: InstanceBuilder,
        op: BuilderBinaryOp,
        location: SourceLocation,
        *,
        reverse: bool,
    ) -> InstanceBuilder:
        return self._assign_first(location)

    @typing.override
    def bool_binary_op(
        self, other: InstanceBuilder, op: BinaryBooleanOperator, location: SourceLocation
    ) -> InstanceBuilder:
        return self._assign_first(location)

    @typing.override
    def augmented_assignment(
        self, op: BuilderBinaryOp, rhs: InstanceBuilder, location: SourceLocation
    ) -> Statement:
        return self._assign_first(location)

    def _assign_first(self, location: SourceLocation) -> typing.Never:
        raise CodeError(
            f"{self._typ} with inferred"
            f" key{'_prefix' if isinstance(self._typ, pytypes.StorageMapProxyType) else ''}"
            " must be assigned to an instance variable before being used",
            location,
        )

    @typing.override
    def single_eval(self) -> InstanceBuilder:
        if self._args.initial_value is None:
            return self
        return StorageProxyDefinitionBuilder(
            attrs.evolve(self._args, initial_value=self._args.initial_value.single_eval()),
            typ=self._typ,
            location=self.source_location,
        )

def extract_key_override(
    key_arg: NodeBuilder | None, location: SourceLocation, *, typ: wtypes.WType
) -> Expression | None:
    if key_arg is None:
        return None
    if isinstance(key_arg, InstanceBuilder) and key_arg.pytype.is_type_or_subtype(
        pytypes.StringType,
        pytypes.StrLiteralType,
        pytypes.BytesType,
        pytypes.BytesLiteralType,
    ):
        key_override = key_arg.to_bytes(key_arg.source_location)
        if isinstance(key_override, BytesConstant):
            return attrs.evolve(key_override, wtype=typ)
        else:
            return ReinterpretCast(expr=key_override, wtype=typ, source_location=location)
    else:
        expect.not_this_type(key_arg, default=expect.default_none)
        return BytesConstant(
            value=b"0",
            wtype=typ,
            encoding=BytesEncoding.unknown,
            source_location=key_arg.source_location,
        )

def extract_description(descr_arg: NodeBuilder | None) -> str | None:
    if descr_arg is None:
        return None
    return expect.simple_string_literal(descr_arg, default=expect.default_none)

def parse_storage_proxy_constructor_args(
    arg_mapping: dict[str, NodeBuilder],
    key_wtype: wtypes.WType,
    key_arg_name: str,
    descr_arg_name: str | None,
    location: SourceLocation,
    initial_value: InstanceBuilder | None = None,
) -> StorageProxyConstructorArgs:
    key_arg = arg_mapping.get(key_arg_name)
    key_override = extract_key_override(key_arg, location, typ=key_wtype)

    description = None
    if descr_arg_name is not None:
        descr_arg = arg_mapping.get(descr_arg_name)
        description = extract_description(descr_arg)

    return StorageProxyConstructorArgs(
        key=key_override,
        description=description,
        initial_value=initial_value,
        key_arg_name=key_arg_name,
    )


from puya import log
from puya.awst import wtypes
from puya.awst.nodes import (
    BoxValueExpression,
    CheckedMaybe,
    Expression,
    IntrinsicCall,
    NumericComparison,
    NumericComparisonExpression,
    TupleItemExpression,
    UInt64BinaryOperation,
    UInt64BinaryOperator,
    UInt64Constant,
)
from puya.parse import SourceLocation
from puyapy.awst_build import intrinsic_factory, pytypes
from puyapy.awst_build.eb import _expect as expect
from puyapy.awst_build.eb._utils import resolve_negative_literal_index
from puyapy.awst_build.eb.bytes import BytesExpressionBuilder
from puyapy.awst_build.eb.interface import (
    BuilderBinaryOp,
    InstanceBuilder,
    LiteralBuilder,
    NodeBuilder,
)
from puyapy.awst_build.eb.uint64 import UInt64ExpressionBuilder

logger = log.get_logger(__name__)

def index_box_bytes(
    box: BoxValueExpression,
    index: InstanceBuilder,
    location: SourceLocation,
) -> InstanceBuilder:
    box_length = box_length_unchecked(box, location)
    begin_index = resolve_negative_literal_index(index, box_length, location)
    return BytesExpressionBuilder(
        IntrinsicCall(
            op_code="box_extract",
            stack_args=[
                box.key,
                begin_index.resolve(),
                UInt64Constant(value=1, source_location=location),
            ],
            source_location=location,
            wtype=wtypes.bytes_wtype,
        )
    )

def slice_box_bytes(
    box: BoxValueExpression,
    begin_index: InstanceBuilder | None,
    end_index: InstanceBuilder | None,
    stride: InstanceBuilder | None,
    location: SourceLocation,
) -> InstanceBuilder:
    if stride:
        logger.error("stride is not supported when slicing boxes", location=location)
    len_expr = box_length_unchecked(box, location).single_eval().resolve()

    begin_index_expr = _eval_slice_component(len_expr, begin_index, location) or UInt64Constant(
        value=0, source_location=location
    )
    end_index_expr = _eval_slice_component(len_expr, end_index, location) or len_expr
    length_expr = (
        UInt64ExpressionBuilder(end_index_expr)
        .binary_op(
            UInt64ExpressionBuilder(begin_index_expr), BuilderBinaryOp.sub, location, reverse=False
        )
        .resolve()
    )

    return BytesExpressionBuilder(
        IntrinsicCall(
            op_code="box_extract",
            stack_args=[box.key, begin_index_expr, length_expr],
            source_location=location,
            wtype=wtypes.bytes_wtype,
        )
    )

def box_length_unchecked(box: BoxValueExpression, location: SourceLocation) -> InstanceBuilder:
    box_len_expr = _box_len(box.key, location)
    box_length = TupleItemExpression(
        base=box_len_expr,
        index=0,
        source_location=location,
    )
    return UInt64ExpressionBuilder(box_length)

def box_length_checked(box: BoxValueExpression, location: SourceLocation) -> Expression:
    box_len_expr = _box_len(box.key, location)
    return CheckedMaybe(box_len_expr, comment=box.exists_assertion_message or "box exists")

def _box_len(box_key: Expression, location: SourceLocation) -> IntrinsicCall:
    assert box_key.wtype == wtypes.box_key
    return IntrinsicCall(
        op_code="box_len",
        wtype=wtypes.WTuple([wtypes.uint64_wtype, wtypes.bool_wtype], source_location=location),
        stack_args=[box_key],
        source_location=location,
    )

def _eval_slice_component(
    len_expr: Expression, val: NodeBuilder | None, location: SourceLocation
) -> Expression | None:
    if val is None:
        return None

    if not isinstance(val, LiteralBuilder):
        # no negatives to deal with here, easy
        temp_index = (
            expect.argument_of_type_else_dummy(val, pytypes.UInt64Type).single_eval().resolve()
        )
        return intrinsic_factory.select(
            false=len_expr,
            true=temp_index,
            condition=NumericComparisonExpression(
                lhs=temp_index,
                operator=NumericComparison.lt,
                rhs=len_expr,
                source_location=location,
            ),
            loc=location,
        )

    int_lit = val.value
    if not isinstance(int_lit, int):
        logger.error(f"Invalid literal for slicing: {int_lit!r}", location=val.source_location)
        int_lit = 0

    # take the min of abs(int_lit) and len(self.expr)
    abs_lit_expr = UInt64Constant(value=abs(int_lit), source_location=val.source_location)
    trunc_value_expr = intrinsic_factory.select(
        false=len_expr,
        true=abs_lit_expr,
        condition=NumericComparisonExpression(
            lhs=abs_lit_expr,
            operator=NumericComparison.lt,
            rhs=len_expr,
            source_location=location,
        ),
        loc=location,
    )
    return (
        trunc_value_expr
        if int_lit >= 0
        else UInt64BinaryOperation(
            left=len_expr,
            op=UInt64BinaryOperator.sub,
            right=trunc_value_expr,
            source_location=location,
        )
    )


import typing

import typing_extensions

from puya.awst.nodes import Expression, Statement
from puya.parse import SourceLocation
from puyapy.awst_build import pytypes
from puyapy.awst_build.eb._base import (
    InstanceExpressionBuilder,
)
from puyapy.awst_build.eb.factories import builder_for_instance
from puyapy.awst_build.eb.interface import (
    BuilderBinaryOp,
    BuilderComparisonOp,
    BuilderUnaryOp,
    InstanceBuilder,
    NodeBuilder,
)

_TPyType_co = typing_extensions.TypeVar(
    "_TPyType_co", bound=pytypes.PyType, default=pytypes.PyType, covariant=True
)
_TExpression_co = typing_extensions.TypeVar(
    "_TExpression_co", bound=Expression, default=Expression, covariant=True
)

class ValueProxyExpressionBuilder(InstanceExpressionBuilder[_TPyType_co, _TExpression_co]):
    def __init__(self, typ: _TPyType_co, expr: _TExpression_co):
        super().__init__(typ, expr)

    @property
    def _proxied(self) -> InstanceBuilder:
        return builder_for_instance(self.pytype, self.resolve())

    @typing.override
    def to_bytes(self, location: SourceLocation) -> Expression:
        return self._proxied.to_bytes(location)

    @typing.override
    def delete(self, location: SourceLocation) -> Statement:
        return self._proxied.delete(location)

    @typing.override
    def bool_eval(self, location: SourceLocation, *, negate: bool = False) -> InstanceBuilder:
        return self._proxied.bool_eval(location, negate=negate)

    @typing.override
    def unary_op(self, op: BuilderUnaryOp, location: SourceLocation) -> InstanceBuilder:
        return self._proxied.unary_op(op, location)

    @typing.override
    def contains(self, item: InstanceBuilder, location: SourceLocation) -> InstanceBuilder:
        return self._proxied.contains(item, location)

    @typing.override
    def compare(
        self, other: InstanceBuilder, op: BuilderComparisonOp, location: SourceLocation
    ) -> InstanceBuilder:
        return self._proxied.compare(other, op, location)

    @typing.override
    def binary_op(
        self,
        other: InstanceBuilder,
        op: BuilderBinaryOp,
        location: SourceLocation,
        *,
        reverse: bool,
    ) -> InstanceBuilder:
        return self._proxied.binary_op(other, op, location, reverse=reverse)

    @typing.override
    def augmented_assignment(
        self, op: BuilderBinaryOp, rhs: InstanceBuilder, location: SourceLocation
    ) -> Statement:
        return self._proxied.augmented_assignment(op, rhs, location)

    @typing.override
    def index(self, index: InstanceBuilder, location: SourceLocation) -> InstanceBuilder:
        return self._proxied.index(index, location)

    @typing.override
    def slice_index(
        self,
        begin_index: InstanceBuilder | None,
        end_index: InstanceBuilder | None,
        stride: InstanceBuilder | None,
        location: SourceLocation,
    ) -> InstanceBuilder:
        return self._proxied.slice_index(begin_index, end_index, stride, location)

    @typing.override
    def member_access(self, name: str, location: SourceLocation) -> NodeBuilder:
        return self._proxied.member_access(name, location)

    @typing.override
    def iterate(self) -> Expression:
        return self._proxied.iterate()

    @typing.override
    def iterable_item_type(self) -> pytypes.PyType:
        return self._proxied.iterable_item_type()


import typing
from collections.abc import Sequence

import mypy.nodes

from puya import log
from puya.awst import wtypes
from puya.awst.nodes import BoxValueExpression, Expression, Not, StateExists
from puya.errors import CodeError
from puya.parse import SourceLocation
from puyapy.awst_build import pytypes
from puyapy.awst_build.eb._base import GenericTypeBuilder, NotIterableInstanceExpressionBuilder
from puyapy.awst_build.eb._bytes_backed import BytesBackedInstanceExpressionBuilder
from puyapy.awst_build.eb.bool import BoolExpressionBuilder
from puyapy.awst_build.eb.interface import (
    InstanceBuilder,
    NodeBuilder,
    StorageProxyConstructorArgs,
    StorageProxyConstructorResult,
    TypeBuilder,
)
from puyapy.awst_build.eb.storage._common import (
    BoxGetExpressionBuilder,
    BoxMaybeExpressionBuilder,
    BoxValueExpressionBuilder,
)
from puyapy.awst_build.eb.storage._storage import (
    StorageProxyDefinitionBuilder,
    parse_storage_proxy_constructor_args,
)
from puyapy.awst_build.eb.storage._util import box_length_checked
from puyapy.awst_build.eb.uint64 import UInt64ExpressionBuilder
from puyapy.awst_build.utils import get_arg_mapping

logger = log.get_logger(__name__)

class BoxTypeBuilder(TypeBuilder[pytypes.StorageProxyType]):
    def __init__(self, typ: pytypes.PyType, location: SourceLocation) -> None:
        assert isinstance(typ, pytypes.StorageProxyType)
        assert typ.generic == pytypes.GenericBoxType
        super().__init__(typ, location)

    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        return _init(args, arg_names, location, result_type=self.produces())

class BoxGenericTypeExpressionBuilder(GenericTypeBuilder):
    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        return _init(args, arg_names, location, result_type=None)

def _init(
    args: Sequence[NodeBuilder],
    arg_names: list[str | None],
    location: SourceLocation,
    *,
    result_type: pytypes.StorageProxyType | None,
) -> InstanceBuilder:
    type_arg_name = "type_"
    key_arg_name = "key"
    arg_mapping, _ = get_arg_mapping(
        required_positional_names=[type_arg_name],
        optional_kw_only=[key_arg_name],
        args=args,
        arg_names=arg_names,
        call_location=location,
        raise_on_missing=True,
    )
    type_arg = arg_mapping[type_arg_name]
    match type_arg.pytype:
        case pytypes.TypeType(typ=content):
            pass
        case _:
            raise CodeError("first argument must be a type reference", location)

    if result_type is None:
        result_type = pytypes.GenericBoxType.parameterise([content], location)
    elif result_type.content != content:
        logger.error(
            "explicit type annotation does not match first argument"
            " - suggest to remove the explicit type annotation, it shouldn't be required",
            location=location,
        )

    typed_args = parse_storage_proxy_constructor_args(
        arg_mapping,
        key_wtype=wtypes.box_key,
        key_arg_name=key_arg_name,
        descr_arg_name=None,
        location=location,
    )

    if typed_args.key is None:
        return StorageProxyDefinitionBuilder(typed_args, result_type, location)
    return _BoxProxyExpressionBuilderFromConstructor(typed_args, result_type)

class BoxProxyExpressionBuilder(
    NotIterableInstanceExpressionBuilder[pytypes.StorageProxyType],
    BytesBackedInstanceExpressionBuilder[pytypes.StorageProxyType],
    bytes_member="key",
):
    def __init__(self, expr: Expression, typ: pytypes.PyType, member_name: str | None = None):
        assert isinstance(typ, pytypes.StorageProxyType)
        assert typ.generic == pytypes.GenericBoxType
        self._member_name = member_name
        super().__init__(typ, expr)

    def _box_key_expr(self, location: SourceLocation) -> BoxValueExpression:
        if self._member_name:
            exists_assertion_message = f"check self.{self._member_name} exists"
        else:
            exists_assertion_message = "check Box exists"
        return BoxValueExpression(
            key=self.resolve(),
            wtype=self.pytype.content_wtype,
            exists_assertion_message=exists_assertion_message,
            source_location=location,
        )

    def _get_value(self, location: SourceLocation) -> BoxValueExpressionBuilder:
        return BoxValueExpressionBuilder(self.pytype.content, self._box_key_expr(location))

    @typing.override
    def member_access(self, name: str, location: SourceLocation) -> NodeBuilder:
        match name:
            case "value":
                return self._get_value(location)
            case "get":
                return BoxGetExpressionBuilder(
                    self._box_key_expr(location), content_type=self.pytype.content
                )
            case "maybe":
                return BoxMaybeExpressionBuilder(
                    self._box_key_expr(location), content_type=self.pytype.content
                )
            case "length":
                return UInt64ExpressionBuilder(
                    box_length_checked(self._get_value(location).resolve(), location)
                )
        return super().member_access(name, location)

    @typing.override
    def bool_eval(self, location: SourceLocation, *, negate: bool = False) -> InstanceBuilder:
        box_exists = StateExists(
            field=self._box_key_expr(location),
            source_location=location,
        )
        return BoolExpressionBuilder(
            Not(expr=box_exists, source_location=location) if negate else box_exists
        )

class _BoxProxyExpressionBuilderFromConstructor(
    BoxProxyExpressionBuilder, StorageProxyConstructorResult
):
    def __init__(self, args: StorageProxyConstructorArgs, typ: pytypes.StorageProxyType):
        assert args.key is not None
        super().__init__(args.key, typ, member_name=None)
        self._args = args

    @typing.override
    @property
    def args(self) -> StorageProxyConstructorArgs:
        return self._args


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

import mypy.nodes

from puya import log
from puya.awst import wtypes
from puya.awst.nodes import (
    BoxValueExpression,
    Expression,
    StateExists,
    StateGet,
    StateGetEx,
)
from puya.errors import CodeError
from puya.parse import SourceLocation
from puyapy.awst_build import intrinsic_factory, pytypes
from puyapy.awst_build.eb import _expect as expect
from puyapy.awst_build.eb._base import FunctionBuilder, GenericTypeBuilder
from puyapy.awst_build.eb._bytes_backed import BytesBackedInstanceExpressionBuilder
from puyapy.awst_build.eb._utils import dummy_value
from puyapy.awst_build.eb.bool import BoolExpressionBuilder
from puyapy.awst_build.eb.factories import builder_for_instance
from puyapy.awst_build.eb.interface import (
    InstanceBuilder,
    NodeBuilder,
    StorageProxyConstructorArgs,
    StorageProxyConstructorResult,
    TypeBuilder,
)
from puyapy.awst_build.eb.storage._common import BoxValueExpressionBuilder
from puyapy.awst_build.eb.storage._storage import (
    StorageProxyDefinitionBuilder,
    parse_storage_proxy_constructor_args,
)
from puyapy.awst_build.eb.storage._util import box_length_checked
from puyapy.awst_build.eb.tuple import TupleExpressionBuilder
from puyapy.awst_build.eb.uint64 import UInt64ExpressionBuilder
from puyapy.awst_build.utils import get_arg_mapping

logger = log.get_logger(__name__)

class BoxMapTypeBuilder(TypeBuilder[pytypes.StorageMapProxyType]):
    def __init__(self, typ: pytypes.PyType, location: SourceLocation) -> None:
        assert isinstance(typ, pytypes.StorageMapProxyType)
        assert typ.generic == pytypes.GenericBoxMapType
        super().__init__(typ, location)

    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        return _init(args, arg_names, location, result_type=self.produces())

class BoxMapGenericTypeExpressionBuilder(GenericTypeBuilder):
    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        return _init(args, arg_names, location, result_type=None)

def _init(
    args: Sequence[NodeBuilder],
    arg_names: list[str | None],
    location: SourceLocation,
    *,
    result_type: pytypes.StorageMapProxyType | None,
) -> InstanceBuilder:
    key_type_arg_name = "key_type"
    value_type_arg_name = "value_type"
    key_prefix_arg_name = "key_prefix"

    arg_mapping, _ = get_arg_mapping(
        required_positional_names=[key_type_arg_name, value_type_arg_name],
        optional_kw_only=[key_prefix_arg_name],
        args=args,
        arg_names=arg_names,
        call_location=location,
        raise_on_missing=True,
    )
    key_type_arg = arg_mapping[key_type_arg_name]
    value_type_arg = arg_mapping[value_type_arg_name]
    match key_type_arg.pytype, value_type_arg.pytype:
        case pytypes.TypeType(typ=key), pytypes.TypeType(typ=content):
            pass
        case _:
            raise CodeError("first and second arguments must be type references", location)

    if result_type is None:
        result_type = pytypes.GenericBoxMapType.parameterise([key, content], location)
    elif not (result_type.key == key and result_type.content == content):
        logger.error(
            "explicit type annotation does not match first argument"
            " - suggest to remove the explicit type annotation, it shouldn't be required",
            location=location,
        )
    # the type of the key is not retained in the AWST, so to
    wtypes.validate_persistable(result_type.key_wtype, location)

    typed_args = parse_storage_proxy_constructor_args(
        arg_mapping,
        key_wtype=wtypes.box_key,
        key_arg_name=key_prefix_arg_name,
        descr_arg_name=None,
        location=location,
    )

    if typed_args.key is None:
        return StorageProxyDefinitionBuilder(typed_args, result_type, location)
    return _BoxMapProxyExpressionBuilderFromConstructor(typed_args, result_type)

class BoxMapProxyExpressionBuilder(
    BytesBackedInstanceExpressionBuilder[pytypes.StorageMapProxyType], bytes_member="key_prefix"
):
    def __init__(self, expr: Expression, typ: pytypes.PyType, member_name: str | None = None):
        assert isinstance(typ, pytypes.StorageMapProxyType)
        assert typ.generic == pytypes.GenericBoxMapType
        self._member_name = member_name
        super().__init__(typ, expr)

    def _build_box_value(
        self, key: InstanceBuilder, location: SourceLocation
    ) -> BoxValueExpression:
        key_data = key.to_bytes(location)
        key_prefix = self.resolve()
        full_key = intrinsic_factory.concat(
            key_prefix, key_data, location, result_type=wtypes.box_key
        )
        if self._member_name:
            exists_assertion_message = f"check self.{self._member_name} entry exists"
        else:
            exists_assertion_message = "check BoxMap entry exists"
        return BoxValueExpression(
            key=full_key,
            wtype=self.pytype.content_wtype,
            exists_assertion_message=exists_assertion_message,
            source_location=location,
        )

    @typing.override
    def index(self, index: InstanceBuilder, location: SourceLocation) -> InstanceBuilder:
        return BoxValueExpressionBuilder(
            self.pytype.content, self._build_box_value(index, location)
        )

    @typing.override
    def member_access(self, name: str, location: SourceLocation) -> NodeBuilder:
        match name:
            case "length":
                return _Length(location, self._build_box_value, self.pytype)
            case "maybe":
                return _Maybe(location, self._build_box_value, self.pytype)
            case "get":
                return _Get(location, self._build_box_value, self.pytype)
            case _:
                return super().member_access(name, location)

    @typing.override
    def contains(self, item: InstanceBuilder, location: SourceLocation) -> InstanceBuilder:
        box_exists = StateExists(
            field=self._build_box_value(item, location), source_location=location
        )
        return BoolExpressionBuilder(box_exists)

    @typing.override
    def slice_index(
        self,
        begin_index: InstanceBuilder | None,
        end_index: InstanceBuilder | None,
        stride: InstanceBuilder | None,
        location: SourceLocation,
    ) -> InstanceBuilder:
        raise CodeError("slicing of BoxMap is not supported", location)

    @typing.override
    def iterate(self) -> typing.Never:  # pragma: no cover
        raise CodeError("iteration of BoxMap is not supported", self.source_location)

    @typing.override
    def iterable_item_type(self) -> typing.Never:
        self.iterate()

    @typing.override
    def bool_eval(self, location: SourceLocation, *, negate: bool = False) -> InstanceBuilder:
        raise CodeError("cannot determine if a BoxMap is empty or not", location)

class _BoxMapProxyExpressionBuilderFromConstructor(
    BoxMapProxyExpressionBuilder, StorageProxyConstructorResult
):
    def __init__(self, args: StorageProxyConstructorArgs, typ: pytypes.StorageMapProxyType):
        assert args.key is not None
        super().__init__(args.key, typ, member_name=None)
        self._args = args

    @typing.override
    @property
    def args(self) -> StorageProxyConstructorArgs:
        return self._args

BoxValueBuilder = Callable[[InstanceBuilder, SourceLocation], BoxValueExpression]

class _MethodBase(FunctionBuilder, abc.ABC):
    def __init__(
        self,
        location: SourceLocation,
        box_value_builder: BoxValueBuilder,
        box_type: pytypes.StorageMapProxyType,
    ) -> None:
        super().__init__(location)
        self.build_box_value = box_value_builder
        self.box_type = box_type

class _Length(_MethodBase):
    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        key_arg_name = "key"
        args_map, any_missing = get_arg_mapping(
            required_positional_names=[key_arg_name],
            args=args,
            arg_names=arg_names,
            call_location=location,
            raise_on_missing=False,
        )
        if any_missing:
            return dummy_value(pytypes.UInt64Type, location)

        key_arg = expect.argument_of_type_else_dummy(
            args_map[key_arg_name], self.box_type.key, resolve_literal=True
        )
        key = self.build_box_value(key_arg, location)
        return UInt64ExpressionBuilder(box_length_checked(key, location))

class _Get(_MethodBase):
    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        key_arg_name = "key"
        default_arg_name = "default"
        args_map, any_missing = get_arg_mapping(
            required_positional_names=[key_arg_name],
            required_kw_only=[default_arg_name],
            args=args,
            arg_names=arg_names,
            call_location=location,
            raise_on_missing=False,
        )
        if any_missing:
            return dummy_value(self.box_type.content, location)
        key_arg = expect.argument_of_type_else_dummy(args_map[key_arg_name], self.box_type.key)
        default_arg = expect.argument_of_type_else_dummy(
            args_map[default_arg_name], self.box_type.content
        )
        key = self.build_box_value(key_arg, location)
        result_expr = StateGet(default=default_arg.resolve(), field=key, source_location=location)
        return builder_for_instance(self.box_type.content, result_expr)

class _Maybe(_MethodBase):
    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        result_typ = pytypes.GenericTupleType.parameterise(
            [self.box_type.content, pytypes.BoolType], location
        )
        key_arg_name = "key"
        args_map, any_missing = get_arg_mapping(
            required_positional_names=[key_arg_name],
            args=args,
            arg_names=arg_names,
            call_location=location,
            raise_on_missing=False,
        )
        if any_missing:
            return dummy_value(result_typ, location)
        item_key_inst = expect.argument_of_type_else_dummy(
            args_map[key_arg_name], self.box_type.key
        )
        key = self.build_box_value(item_key_inst, location)
        return TupleExpressionBuilder(StateGetEx(field=key, source_location=location), result_typ)


import typing
from collections.abc import Sequence

import mypy.nodes

from puya.awst import wtypes
from puya.awst.nodes import BoxValueExpression, Expression, IntrinsicCall, Not, StateExists
from puya.parse import SourceLocation
from puyapy.awst_build import pytypes
from puyapy.awst_build.eb import _expect as expect
from puyapy.awst_build.eb._base import FunctionBuilder, NotIterableInstanceExpressionBuilder
from puyapy.awst_build.eb._bytes_backed import BytesBackedInstanceExpressionBuilder
from puyapy.awst_build.eb._utils import dummy_value
from puyapy.awst_build.eb.bool import BoolExpressionBuilder
from puyapy.awst_build.eb.factories import builder_for_instance
from puyapy.awst_build.eb.interface import (
    InstanceBuilder,
    NodeBuilder,
    StorageProxyConstructorArgs,
    StorageProxyConstructorResult,
    TypeBuilder,
)
from puyapy.awst_build.eb.none import NoneExpressionBuilder
from puyapy.awst_build.eb.storage._common import BoxGetExpressionBuilder, BoxMaybeExpressionBuilder
from puyapy.awst_build.eb.storage._storage import (
    StorageProxyDefinitionBuilder,
    parse_storage_proxy_constructor_args,
)
from puyapy.awst_build.eb.storage._util import box_length_checked
from puyapy.awst_build.eb.uint64 import UInt64ExpressionBuilder
from puyapy.awst_build.utils import get_arg_mapping

class BoxRefTypeBuilder(TypeBuilder[pytypes.StorageProxyType]):
    def __init__(self, location: SourceLocation) -> None:
        super().__init__(pytypes.BoxRefType, location)

    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        key_arg_name = "key"
        arg_mapping, _ = get_arg_mapping(
            optional_kw_only=[key_arg_name],
            args=args,
            arg_names=arg_names,
            call_location=location,
            raise_on_missing=False,
        )
        typed_args = parse_storage_proxy_constructor_args(
            arg_mapping,
            key_wtype=wtypes.box_key,
            key_arg_name=key_arg_name,
            descr_arg_name=None,
            location=location,
        )

        if typed_args.key is None:
            return StorageProxyDefinitionBuilder(typed_args, self.produces(), location)
        return _BoxRefProxyExpressionBuilderFromConstructor(typed_args)

class BoxRefProxyExpressionBuilder(
    NotIterableInstanceExpressionBuilder[pytypes.StorageProxyType],
    BytesBackedInstanceExpressionBuilder[pytypes.StorageProxyType],
    bytes_member="key",
):
    def __init__(self, expr: Expression, member_name: str | None = None):
        super().__init__(pytypes.BoxRefType, expr)
        self._member_name = member_name

    def _box_key_expr(self, location: SourceLocation) -> BoxValueExpression:
        if self._member_name:
            exists_assertion_message = f"check self.{self._member_name} exists"
        else:
            exists_assertion_message = "check BoxRef exists"
        return BoxValueExpression(
            key=self.resolve(),
            wtype=wtypes.bytes_wtype,
            exists_assertion_message=exists_assertion_message,
            source_location=location,
        )

    @typing.override
    def bool_eval(self, location: SourceLocation, *, negate: bool = False) -> InstanceBuilder:
        box_exists = StateExists(
            field=self._box_key_expr(location),
            source_location=location,
        )
        return BoolExpressionBuilder(
            Not(expr=box_exists, source_location=location) if negate else box_exists
        )

    @typing.override
    def member_access(self, name: str, location: SourceLocation) -> NodeBuilder:
        match name:
            case "create":
                return _Create(location, box_proxy=self.resolve())
            case "delete":
                return _IntrinsicMethod(
                    location,
                    box_proxy=self.resolve(),
                    op_code="box_del",
                    args={},
                    return_type=pytypes.BoolType,
                )
            case "extract":
                return _IntrinsicMethod(
                    location,
                    box_proxy=self.resolve(),
                    op_code="box_extract",
                    args={"start_index": pytypes.UInt64Type, "length": pytypes.UInt64Type},
                    return_type=pytypes.BytesType,
                )
            case "resize":
                return _IntrinsicMethod(
                    location,
                    box_proxy=self.resolve(),
                    op_code="box_resize",
                    args={"new_size": pytypes.UInt64Type},
                    return_type=pytypes.NoneType,
                )
            case "replace":
                return _IntrinsicMethod(
                    location,
                    box_proxy=self.resolve(),
                    op_code="box_replace",
                    args={"start_index": pytypes.UInt64Type, "value": pytypes.BytesType},
                    return_type=pytypes.NoneType,
                )
            case "splice":
                return _IntrinsicMethod(
                    location,
                    box_proxy=self.resolve(),
                    op_code="box_splice",
                    args={
                        "start_index": pytypes.UInt64Type,
                        "length": pytypes.UInt64Type,
                        "value": pytypes.BytesType,
                    },
                    return_type=pytypes.NoneType,
                )

            case "get":
                return BoxGetExpressionBuilder(
                    self._box_key_expr(location), content_type=pytypes.BytesType
                )
            case "put":
                return _Put(location, box_proxy=self.resolve())
            case "maybe":
                return BoxMaybeExpressionBuilder(
                    self._box_key_expr(location), content_type=pytypes.BytesType
                )
            case "length":
                return UInt64ExpressionBuilder(
                    box_length_checked(self._box_key_expr(location), location)
                )
            case _:
                return super().member_access(name, location)

class _BoxRefProxyExpressionBuilderFromConstructor(
    BoxRefProxyExpressionBuilder, StorageProxyConstructorResult
):
    def __init__(self, args: StorageProxyConstructorArgs):
        assert args.key is not None
        super().__init__(args.key, member_name=None)
        self._args = args

    @typing.override
    @property
    def args(self) -> StorageProxyConstructorArgs:
        return self._args

class _IntrinsicMethod(FunctionBuilder):
    def __init__(
        self,
        location: SourceLocation,
        *,
        box_proxy: Expression,
        op_code: str,
        args: dict[str, pytypes.PyType],
        return_type: pytypes.RuntimeType,
    ) -> None:
        super().__init__(location)
        self.box_proxy = box_proxy
        self.op_code = op_code
        self.args = args
        self.return_type = return_type

    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        args_map, any_missing = get_arg_mapping(
            required_positional_names=list(self.args.keys()),
            args=args,
            arg_names=arg_names,
            call_location=location,
            raise_on_missing=False,
        )
        if any_missing:
            return dummy_value(self.return_type, location)
        stack_args = [
            expect.argument_of_type_else_dummy(
                args_map[arg_name], arg_type, resolve_literal=True
            ).resolve()
            for arg_name, arg_type in self.args.items()
        ]
        result_expr = IntrinsicCall(
            op_code=self.op_code,
            stack_args=[self.box_proxy, *stack_args],
            wtype=self.return_type.wtype,
            source_location=location,
        )
        return builder_for_instance(self.return_type, result_expr)

class _Create(FunctionBuilder):
    def __init__(self, location: SourceLocation, *, box_proxy: Expression) -> None:
        super().__init__(location)
        self.box_proxy = box_proxy

    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        arg = expect.exactly_one_arg_of_type_else_dummy(
            args, pytypes.UInt64Type, location, resolve_literal=True
        )
        size = arg.resolve()
        return BoolExpressionBuilder(
            IntrinsicCall(
                op_code="box_create",
                stack_args=[self.box_proxy, size],
                wtype=wtypes.bool_wtype,
                source_location=location,
            )
        )

class _Put(FunctionBuilder):
    def __init__(self, location: SourceLocation, *, box_proxy: Expression) -> None:
        super().__init__(location)
        self.box_proxy = box_proxy

    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        arg = expect.exactly_one_arg_of_type_else_dummy(
            args, pytypes.BytesType, location, resolve_literal=True
        )
        data = arg.resolve()
        return NoneExpressionBuilder(
            IntrinsicCall(
                op_code="box_put",
                stack_args=[self.box_proxy, data],
                wtype=wtypes.void_wtype,
                source_location=location,
            )
        )


import abc
import typing
from collections.abc import Sequence

import mypy.nodes

from puya import log
from puya.awst import wtypes
from puya.awst.nodes import (
    AppStateExpression,
    Expression,
    ExpressionStatement,
    Not,
    StateDelete,
    StateExists,
    StateGet,
    StateGetEx,
    Statement,
)
from puya.errors import CodeError
from puya.parse import SourceLocation
from puyapy.awst_build import pytypes
from puyapy.awst_build.eb import _expect as expect
from puyapy.awst_build.eb._base import (
    FunctionBuilder,
    GenericTypeBuilder,
    NotIterableInstanceExpressionBuilder,
)
from puyapy.awst_build.eb._bytes_backed import BytesBackedInstanceExpressionBuilder
from puyapy.awst_build.eb.bool import BoolExpressionBuilder
from puyapy.awst_build.eb.factories import builder_for_instance
from puyapy.awst_build.eb.interface import (
    InstanceBuilder,
    NodeBuilder,
    StorageProxyConstructorArgs,
    StorageProxyConstructorResult,
    TypeBuilder,
)
from puyapy.awst_build.eb.storage._storage import (
    StorageProxyDefinitionBuilder,
    parse_storage_proxy_constructor_args,
)
from puyapy.awst_build.eb.storage._value_proxy import ValueProxyExpressionBuilder
from puyapy.awst_build.eb.tuple import TupleExpressionBuilder
from puyapy.awst_build.utils import get_arg_mapping

logger = log.get_logger(__name__)

class GlobalStateTypeBuilder(TypeBuilder[pytypes.StorageProxyType]):
    def __init__(self, typ: pytypes.PyType, location: SourceLocation) -> None:
        assert isinstance(typ, pytypes.StorageProxyType)
        assert typ.generic == pytypes.GenericGlobalStateType
        self._typ = typ
        super().__init__(typ, location)

    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        return _init(args, arg_names, location, result_type=self._typ)

class GlobalStateGenericTypeBuilder(GenericTypeBuilder):
    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        return _init(args, arg_names, location, result_type=None)

def _init(
    args: Sequence[NodeBuilder],
    arg_names: list[str | None],
    location: SourceLocation,
    *,
    result_type: pytypes.StorageProxyType | None,
) -> InstanceBuilder:
    type_or_value_arg_name = "type_or_initial_value"
    key_arg_name = "key"
    descr_arg_name = "description"
    arg_mapping, _ = get_arg_mapping(
        required_positional_names=[type_or_value_arg_name],
        optional_kw_only=[key_arg_name, descr_arg_name],
        args=args,
        arg_names=arg_names,
        call_location=location,
        raise_on_missing=True,
    )

    match arg_mapping[type_or_value_arg_name]:
        case NodeBuilder(pytype=pytypes.TypeType(typ=content)):
            iv_builder = None
            if result_type is None:
                result_type = pytypes.GenericGlobalStateType.parameterise([content], location)
            elif result_type.content != content:
                logger.error(
                    "explicit type annotation does not match first argument"
                    " - suggest to remove the explicit type annotation, it shouldn't be required",
                    location=location,
                )
        case value_arg if result_type is not None:
            iv_builder = expect.argument_of_type_else_dummy(value_arg, result_type.content)
        case InstanceBuilder(pytype=content) as iv_builder:
            result_type = pytypes.GenericGlobalStateType.parameterise([content], location)
        case _:
            raise CodeError(
                "first argument must be a type reference or an initial value", location
            )

    typed_args = parse_storage_proxy_constructor_args(
        arg_mapping,
        key_wtype=wtypes.state_key,
        key_arg_name=key_arg_name,
        descr_arg_name=descr_arg_name,
        initial_value=iv_builder,
        location=location,
    )

    if typed_args.key is None:
        return StorageProxyDefinitionBuilder(typed_args, result_type, location)
    return _GlobalStateExpressionBuilderFromConstructor(typed_args, result_type)

class GlobalStateExpressionBuilder(
    NotIterableInstanceExpressionBuilder[pytypes.StorageProxyType],
    BytesBackedInstanceExpressionBuilder[pytypes.StorageProxyType],
    bytes_member="key",
):
    def __init__(self, expr: Expression, typ: pytypes.PyType, member_name: str | None = None):
        assert isinstance(typ, pytypes.StorageProxyType)
        assert typ.generic == pytypes.GenericGlobalStateType
        self.member_name: typing.Final = member_name
        super().__init__(typ, expr)

    @typing.override
    def bool_eval(self, location: SourceLocation, *, negate: bool = False) -> InstanceBuilder:
        field = _build_field(self, location)
        exists_expr = StateExists(field=field, source_location=location)
        if negate:
            expr: Expression = Not(location, exists_expr)
        else:
            expr = exists_expr
        return BoolExpressionBuilder(expr)

    @typing.override
    def member_access(self, name: str, location: SourceLocation) -> NodeBuilder:
        match name:
            case "value":
                field = _build_field(self, location)
                return _Value(self.pytype.content, field)
            case "get":
                return _Get(self, location)
            case "maybe":
                return _Maybe(self, location)
            case _:
                return super().member_access(name, location)

def _build_field(
    self: GlobalStateExpressionBuilder, location: SourceLocation
) -> AppStateExpression:
    if self.member_name:
        exists_assertion_message = f"check self.{self.member_name} exists"
    else:
        exists_assertion_message = "check GlobalState exists"
    return AppStateExpression(
        key=self.resolve(),
        wtype=self.pytype.content_wtype,
        exists_assertion_message=exists_assertion_message,
        source_location=location,
    )

class _GlobalStateExpressionBuilderFromConstructor(
    GlobalStateExpressionBuilder, StorageProxyConstructorResult
):
    def __init__(self, args: StorageProxyConstructorArgs, typ: pytypes.StorageProxyType):
        assert args.key is not None
        super().__init__(args.key, typ, member_name=None)
        self._args = args

    @typing.override
    @property
    def args(self) -> StorageProxyConstructorArgs:
        return self._args

    @typing.override
    def resolve(self) -> Expression:
        if self._args.initial_value is not None:
            logger.error(
                "providing an initial value is only allowed when assigning to a member variable",
                location=self.source_location,
            )
        return super().resolve()

class _MemberFunction(FunctionBuilder, abc.ABC):
    def __init__(self, base: GlobalStateExpressionBuilder, location: SourceLocation):
        super().__init__(location)
        self._base = base

class _Maybe(_MemberFunction):
    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        content_typ = self._base.pytype.content
        expect.no_args(args, location)
        field = _build_field(self._base, self.source_location)
        expr = StateGetEx(field=field, source_location=location)
        result_typ = pytypes.GenericTupleType.parameterise(
            [content_typ, pytypes.BoolType], location
        )
        return TupleExpressionBuilder(expr, result_typ)

class _Get(_MemberFunction):
    @typing.override
    def call(
        self,
        args: Sequence[NodeBuilder],
        arg_kinds: list[mypy.nodes.ArgKind],
        arg_names: list[str | None],
        location: SourceLocation,
    ) -> InstanceBuilder:
        content_typ = self._base.pytype.content
        default_arg = expect.exactly_one_arg_of_type_else_dummy(args, content_typ, location)
        default_expr = default_arg.resolve()
        field = _build_field(self._base, self.source_location)
        expr = StateGet(field=field, default=default_expr, source_location=location)
        return builder_for_instance(content_typ, expr)

class _Value(ValueProxyExpressionBuilder[pytypes.PyType, AppStateExpression]):
    @typing.override
    def delete(self, location: SourceLocation) -> Statement:
        return ExpressionStatement(StateDelete(field=self.resolve(), source_location=location))