import contextlib
import typing
from collections import Counter
from collections.abc import Iterable, Iterator, Mapping, Sequence
from operator import itemgetter

import attrs
from immutabledict import immutabledict

import puya.artifact_metadata as models
from puya import algo_constants, log
from puya.avm import AVMType
from puya.awst import (
    nodes as awst_nodes,
    wtypes,
)
from puya.awst.arc4_types import maybe_avm_to_arc4_equivalent_type, wtype_to_arc4
from puya.awst.function_traverser import FunctionTraverser
from puya.errors import InternalError
from puya.ir._arc4_default_args import convert_default_args
from puya.ir.context import IRBuildContext
from puya.parse import SourceLocation
from puya.utils import StableSet, set_add, unique

__all__ = [
    "build_contract_metadata",
]

logger = log.get_logger(__name__)

def build_contract_metadata(
    ctx: IRBuildContext, contract: awst_nodes.Contract
) -> tuple[
    models.ContractMetaData,
    dict[awst_nodes.ContractMethod, models.ARC4Method],
]:
    global_state = dict[str, models.ContractState]()
    local_state = dict[str, models.ContractState]()
    boxes = dict[str, models.ContractState]()
    for state in contract.app_state:
        translated = _translate_state(state)
        match state.kind:
            case awst_nodes.AppStorageKind.app_global:
                global_state[state.member_name] = translated
            case awst_nodes.AppStorageKind.account_local:
                local_state[state.member_name] = translated
            case awst_nodes.AppStorageKind.box:
                boxes[state.member_name] = translated
            case unexpected:
                typing.assert_never(unexpected)
    state_totals = _build_state_totals(
        contract.state_totals,
        global_state=global_state,
        local_state=local_state,
        location=contract.source_location,
    )
    arc4_method_data, type_refs = _extract_arc4_methods_and_type_refs(ctx, contract)
    structs = _extract_structs(type_refs)
    template_var_types = _TemplateVariableTypeCollector.collect(ctx, contract, arc4_method_data)
    metadata = models.ContractMetaData(
        description=contract.description,
        name=contract.name,
        ref=contract.id,
        arc4_methods=list(arc4_method_data.values()),
        global_state=immutabledict(global_state),
        local_state=immutabledict(local_state),
        boxes=immutabledict(boxes),
        state_totals=state_totals,
        structs=immutabledict(structs),
        template_variable_types=immutabledict(template_var_types),
    )
    return metadata, arc4_method_data

def _translate_state(
    state: awst_nodes.AppStorageDefinition,
) -> models.ContractState:
    storage_type = wtypes.persistable_stack_type(state.storage_wtype, state.source_location)
    if state.key_wtype is not None:
        if state.kind is not awst_nodes.AppStorageKind.box:
            raise InternalError(
                f"maps of {state.kind} are not supported by IR backend yet", state.source_location
            )
        arc56_key_type = _get_arc56_type(state.key_wtype, state.source_location)
        is_map = True
    else:
        arc56_key_type = (
            "AVMString" if state.key.encoding == awst_nodes.BytesEncoding.utf8 else "AVMBytes"
        )
        is_map = False
    arc56_value_type = _get_arc56_type(state.storage_wtype, state.source_location)
    return models.ContractState(
        name=state.member_name,
        source_location=state.source_location,
        key_or_prefix=state.key.value,
        arc56_key_type=arc56_key_type,
        arc56_value_type=arc56_value_type,
        storage_type=storage_type,
        description=state.description,
        is_map=is_map,
    )

def _build_state_totals(
    declared_totals: awst_nodes.StateTotals | None,
    *,
    global_state: Mapping[str, models.ContractState],
    local_state: Mapping[str, models.ContractState],
    location: SourceLocation,
) -> models.StateTotals:
    global_by_type = Counter(s.storage_type for s in global_state.values())
    local_by_type = Counter(s.storage_type for s in local_state.values())
    merged = models.StateTotals(
        global_uints=global_by_type[AVMType.uint64],
        global_bytes=global_by_type[AVMType.bytes],
        local_uints=local_by_type[AVMType.uint64],
        local_bytes=local_by_type[AVMType.bytes],
    )
    if declared_totals is not None:
        insufficient_fields = []
        declared_dict = attrs.asdict(declared_totals, filter=attrs.filters.include(int))
        for field, declared in declared_dict.items():
            calculated = getattr(merged, field)
            if declared < calculated:
                insufficient_fields.append(f"{field}: {declared=}, {calculated=}")
            merged = attrs.evolve(merged, **{field: declared})
        if insufficient_fields:
            logger.warning(
                f"State totals declared on the class are less than totals calculated from"
                f" explicitly declared properties: {', '.join(sorted(insufficient_fields))}.",
                location=location,
            )
    global_total = merged.global_uints + merged.global_bytes
    local_total = merged.local_uints + merged.local_bytes
    if global_total > algo_constants.MAX_GLOBAL_STATE_KEYS:
        logger.warning(
            f"Total global state key count of {global_total}"
            f" exceeds consensus parameter value {algo_constants.MAX_GLOBAL_STATE_KEYS}",
            location=location,
        )
    if local_total > algo_constants.MAX_LOCAL_STATE_KEYS:
        logger.warning(
            f"Total local state key count of {local_total}"
            f" exceeds consensus parameter value {algo_constants.MAX_LOCAL_STATE_KEYS}",
            location=location,
        )
    return merged

class _TemplateVariableTypeCollector(FunctionTraverser):
    def __init__(self, context: IRBuildContext) -> None:
        self.context = context
        self.vars = dict[str, awst_nodes.TemplateVar]()
        self._seen_functions = set[awst_nodes.Function]()
        self._func_stack = list[awst_nodes.Function]()

    def process_func(self, func: awst_nodes.Function) -> None:
        if set_add(self._seen_functions, func):
            with self._enter_func(func):
                func.body.accept(self)

    @classmethod
    def collect(
        cls,
        context: IRBuildContext,
        contract: awst_nodes.Contract,
        routable_methods: Iterable[awst_nodes.ContractMethod],
    ) -> dict[str, str]:
        collector = cls(context)
        for function in (contract.approval_program, contract.clear_program, *routable_methods):
            collector.process_func(function)
        return {
            name: _get_arc56_type(var.wtype, var.source_location)
            for name, var in collector.vars.items()
        }

    def visit_template_var(self, var: awst_nodes.TemplateVar) -> None:
        try:
            existing = self.vars[var.name]
        except KeyError:
            self.vars[var.name] = var
        else:
            if existing.wtype != var.wtype:
                logger.error(
                    "inconsistent types specified for template var",
                    location=var.source_location,
                )
                logger.info("other template var", location=existing.source_location)

    @contextlib.contextmanager
    def _enter_func(self, func: awst_nodes.Function) -> Iterator[None]:
        self._func_stack.append(func)
        try:
            yield
        finally:
            self._func_stack.pop()

    @property
    def current_func(self) -> awst_nodes.Function:
        return self._func_stack[-1]

    def visit_subroutine_call_expression(self, expr: awst_nodes.SubroutineCallExpression) -> None:
        target = self.context.resolve_function_reference(
            expr.target,
            expr.source_location,
            caller=self.current_func,
        )
        self.process_func(target)

def _extract_arc4_methods_and_type_refs(
    ctx: IRBuildContext, contract: awst_nodes.Contract
) -> tuple[dict[awst_nodes.ContractMethod, models.ARC4Method], list[wtypes.WType]]:
    event_collector = _EventCollector(ctx)
    type_refs = [
        typ
        for state in contract.app_state
        for typ in (state.key_wtype, state.storage_wtype)
        if typ is not None
    ]
    methods = dict[awst_nodes.ContractMethod, models.ARC4Method]()
    for method_name in unique(m.member_name for m in contract.methods):
        m = contract.resolve_contract_method(method_name)
        assert m is not None  # shouldn't logically be possible
        match m.arc4_method_config:
            case None:
                pass
            case awst_nodes.ARC4BareMethodConfig() as bare_method_config:
                methods[m] = models.ARC4BareMethod(
                    id=m.full_name, desc=m.documentation.description, config=bare_method_config
                )
            case abi_method_config:
                event_wtypes = event_collector.process_func(m)
                events = list(map(_wtype_to_struct, event_wtypes))
                methods[m] = _abi_method_metadata(ctx, contract, m, abi_method_config, events)
                if m.return_type is not None:
                    type_refs.append(m.return_type)
                type_refs.extend(arg.wtype for arg in m.args)
                for event_struct in event_wtypes:
                    type_refs.extend(event_struct.types)
    return methods, type_refs

def _extract_structs(
    typ_refs: Sequence[wtypes.WType],
) -> dict[str, models.ARC4Struct]:
    # Produce a unique mapping of struct names to ARC4Struct definitions.
    # Will recursively include any structs referenced in fields
    struct_wtypes = list(filter(_is_arc4_struct, typ_refs))
    struct_results = dict[str, models.ARC4Struct]()
    while struct_wtypes:
        struct = struct_wtypes.pop()
        if struct.name in struct_results:
            continue
        struct_wtypes.extend(
            wtype
            for wtype in struct.fields.values()
            if _is_arc4_struct(wtype) and wtype.name not in struct_results
        )
        struct_results[struct.name] = _wtype_to_struct(struct)
    return dict(sorted(struct_results.items(), key=itemgetter(0)))

def _abi_method_metadata(
    ctx: IRBuildContext,
    contract: awst_nodes.Contract,
    m: awst_nodes.ContractMethod,
    config: awst_nodes.ARC4ABIMethodConfig,
    events: Sequence[models.ARC4Struct],
) -> models.ARC4ABIMethod:
    assert config is m.arc4_method_config
    default_args = convert_default_args(ctx, contract, m, config)
    args = [
        models.ARC4MethodArg(
            name=a.name,
            type_=wtype_to_arc4(a.wtype),
            struct=_get_arc4_struct_name(a.wtype),
            desc=m.documentation.args.get(a.name),
            client_default=default_args[a],
        )
        for a in m.args
    ]
    returns = models.ARC4Returns(
        desc=m.documentation.returns,
        type_=wtype_to_arc4(m.return_type),
        struct=_get_arc4_struct_name(m.return_type),
    )
    return models.ARC4ABIMethod(
        id=m.full_name,
        desc=m.documentation.description,
        args=args,
        returns=returns,
        events=events,
        config=config,
    )

def _is_arc4_struct(
    wtype: wtypes.WType | None,
) -> typing.TypeGuard[wtypes.ARC4Struct | wtypes.WTuple]:
    return isinstance(wtype, wtypes.ARC4Struct | wtypes.WTuple) and bool(wtype.fields)

def _get_arc4_struct_name(wtype: wtypes.WType) -> str | None:
    return wtype.name if _is_arc4_struct(wtype) else None

def _wtype_to_struct(s: wtypes.ARC4Struct | wtypes.WTuple) -> models.ARC4Struct:
    fields = []
    assert s.fields
    for field_name, field_wtype in s.fields.items():
        if not isinstance(field_wtype, wtypes.ARC4Type):
            maybe_arc4_field_wtype = maybe_avm_to_arc4_equivalent_type(field_wtype)
            if maybe_arc4_field_wtype is None:
                raise InternalError("expected ARC4 type")
            field_wtype = maybe_arc4_field_wtype
        fields.append(
            models.ARC4StructField(
                name=field_name,
                type=field_wtype.arc4_name,
                struct=_get_arc4_struct_name(field_wtype),
            )
        )
    return models.ARC4Struct(fullname=s.name, desc=s.desc, fields=fields)

@attrs.frozen
class _EventCollector(FunctionTraverser):
    context: IRBuildContext
    emits: dict[awst_nodes.Function, StableSet[wtypes.ARC4Struct]] = attrs.field(
        factory=dict, init=False
    )
    _func_stack: list[awst_nodes.Function] = attrs.field(factory=list)

    def process_func(self, func: awst_nodes.Function) -> StableSet[wtypes.ARC4Struct]:
        if func not in self.emits:
            self.emits[func] = StableSet[wtypes.ARC4Struct]()
            with self._enter_func(func):
                func.body.accept(self)
        return self.emits[func]

    @contextlib.contextmanager
    def _enter_func(self, func: awst_nodes.Function) -> Iterator[None]:
        self._func_stack.append(func)
        try:
            yield
        finally:
            self._func_stack.pop()

    @property
    def current_func(self) -> awst_nodes.Function:
        return self._func_stack[-1]

    def visit_emit(self, emit: awst_nodes.Emit) -> None:
        assert isinstance(emit.value.wtype, wtypes.ARC4Struct)
        self.emits[self.current_func].add(emit.value.wtype)

    def visit_subroutine_call_expression(self, expr: awst_nodes.SubroutineCallExpression) -> None:
        target = self.context.resolve_function_reference(
            expr.target,
            expr.source_location,
            caller=self.current_func,
        )
        self.emits[self.current_func] |= self.process_func(target)

def _get_arc56_type(wtype: wtypes.WType, loc: SourceLocation) -> str:
    if isinstance(wtype, wtypes.ARC4Struct):
        return wtype.name
    if isinstance(wtype, wtypes.ARC4Type):
        return wtype.arc4_name
    if wtype == wtypes.string_wtype:
        return "AVMString"
    storage_type = wtypes.persistable_stack_type(wtype, loc)
    match storage_type:
        case AVMType.uint64:
            return "AVMUint64"
        case AVMType.bytes:
            return "AVMBytes"


from collections import deque
from collections.abc import Iterator

from puya.awst import (
    nodes as awst_nodes,
    wtypes,
)
from puya.ir import models
from puya.ir.models import Parameter, Subroutine
from puya.ir.types_ import wtype_to_ir_type, wtype_to_ir_types
from puya.ir.utils import format_tuple_index
from puya.parse import SourceLocation
from puya.utils import set_add

def bfs_block_order(start: models.BasicBlock) -> Iterator[models.BasicBlock]:
    q = deque((start,))
    visited = {start}
    while q:
        block = q.popleft()
        yield block
        q.extend(succ for succ in block.successors if set_add(visited, succ))

def make_subroutine(func: awst_nodes.Function, *, allow_implicits: bool) -> Subroutine:
    """Pre-construct subroutine with an empty body"""
    parameters = [
        param
        for arg in func.args
        for param in _expand_tuple_parameters(
            arg.name,
            arg.wtype,
            allow_implicits=allow_implicits,
            source_location=arg.source_location,
        )
    ]
    returns = wtype_to_ir_types(func.return_type)
    return Subroutine(
        id=func.full_name,
        short_name=func.short_name,
        parameters=parameters,
        returns=returns,
        body=[],
        inline=func.inline,
        source_location=func.source_location,
    )

def _expand_tuple_parameters(
    name: str, typ: wtypes.WType, *, allow_implicits: bool, source_location: SourceLocation | None
) -> Iterator[Parameter]:
    if isinstance(typ, wtypes.WTuple):
        for item_idx, item_type in enumerate(typ.types):
            item_name = format_tuple_index(typ, name, item_idx)
            yield from _expand_tuple_parameters(
                item_name,
                item_type,
                allow_implicits=allow_implicits,
                source_location=source_location,
            )
    else:
        yield Parameter(
            name=name,
            ir_type=wtype_to_ir_type(typ),
            version=0,
            implicit_return=allow_implicits and not typ.immutable,
            source_location=source_location,
        )


import typing
from collections.abc import Iterable, Mapping, Sequence

import attrs

from puya import (
    artifact_metadata as md,
    log,
)
from puya.avm import OnCompletionAction
from puya.awst import (
    nodes as awst_nodes,
    wtypes,
)
from puya.awst.arc4_types import maybe_avm_to_arc4_equivalent_type
from puya.errors import CodeError
from puya.parse import SourceLocation
from puya.utils import set_add

__all__ = [
    "create_abi_router",
    "AWSTContractMethodSignature",
]

logger = log.get_logger(__name__)

ALL_VALID_APPROVAL_ON_COMPLETION_ACTIONS = {
    OnCompletionAction.NoOp,
    OnCompletionAction.OptIn,
    OnCompletionAction.CloseOut,
    OnCompletionAction.UpdateApplication,
    OnCompletionAction.DeleteApplication,
}

@attrs.frozen(kw_only=True)
class AWSTContractMethodSignature:
    target: awst_nodes.ContractMethodTarget
    parameter_types: Sequence[wtypes.WType]
    return_type: wtypes.WType

def _btoi(
    bytes_arg: awst_nodes.Expression, location: SourceLocation | None = None
) -> awst_nodes.IntrinsicCall:
    return awst_nodes.IntrinsicCall(
        op_code="btoi",
        stack_args=[bytes_arg],
        wtype=wtypes.uint64_wtype,
        source_location=location or bytes_arg.source_location,
    )

def _txn(
    immediate: str, wtype: wtypes.WType, location: SourceLocation
) -> awst_nodes.IntrinsicCall:
    return awst_nodes.IntrinsicCall(
        op_code="txn",
        immediates=[immediate],
        wtype=wtype,
        source_location=location,
    )

def _txn_app_args(index: int, loc: SourceLocation) -> awst_nodes.IntrinsicCall:
    return awst_nodes.IntrinsicCall(
        op_code="txna",
        immediates=["ApplicationArgs", index],
        source_location=loc,
        wtype=wtypes.bytes_wtype,
    )

def create_block(
    location: SourceLocation, comment: str | None, *stmts: awst_nodes.Statement
) -> awst_nodes.Block:
    return awst_nodes.Block(source_location=location, body=stmts, comment=comment)

def call(
    location: SourceLocation, sig: AWSTContractMethodSignature, *args: awst_nodes.Expression
) -> awst_nodes.SubroutineCallExpression:
    return awst_nodes.SubroutineCallExpression(
        target=sig.target,
        args=[awst_nodes.CallArg(name=None, value=arg) for arg in args],
        wtype=sig.return_type,
        source_location=location,
    )

def app_arg(
    index: int,
    wtype: wtypes.WType,
    location: SourceLocation,
) -> awst_nodes.Expression:
    value = _txn_app_args(index, location)
    if wtype == wtypes.bytes_wtype:
        return value
    return awst_nodes.ReinterpretCast(
        source_location=location,
        expr=value,
        wtype=wtype,
    )

def _non_zero(value: awst_nodes.Expression) -> awst_nodes.Expression:
    location = value.source_location
    return awst_nodes.NumericComparisonExpression(
        lhs=value,
        rhs=constant(0, location),
        operator=awst_nodes.NumericComparison.ne,
        source_location=location,
    )

def _is_zero(value: awst_nodes.Expression) -> awst_nodes.Expression:
    location = value.source_location
    return awst_nodes.NumericComparisonExpression(
        lhs=value,
        rhs=constant(0, location),
        operator=awst_nodes.NumericComparison.eq,
        source_location=location,
    )

def return_(value: bool, location: SourceLocation) -> awst_nodes.ReturnStatement:  # noqa: FBT001
    return awst_nodes.ReturnStatement(
        value=awst_nodes.BoolConstant(value=value, source_location=location),
        source_location=location,
    )

def reject(location: SourceLocation) -> awst_nodes.ReturnStatement:
    return return_(False, location)  # noqa: FBT003

def approve(location: SourceLocation) -> awst_nodes.ReturnStatement:
    return return_(True, location)  # noqa: FBT003

def on_completion(location: SourceLocation) -> awst_nodes.Expression:
    return _txn("OnCompletion", wtypes.uint64_wtype, location)

def route_bare_methods(
    location: SourceLocation,
    bare_methods: Mapping[md.ARC4BareMethod, AWSTContractMethodSignature],
) -> awst_nodes.Block | None:
    bare_blocks = dict[OnCompletionAction, awst_nodes.Block]()
    for method, sig in bare_methods.items():
        bare_location = method.config_location
        bare_block = create_block(
            bare_location,
            sig.target.member_name,
            *assert_create_state(method),
            awst_nodes.ExpressionStatement(expr=call(bare_location, sig)),
            approve(bare_location),
        )
        for oca in method.allowed_completion_types:
            if bare_blocks.setdefault(oca, bare_block) is not bare_block:
                logger.error(
                    f"cannot have multiple bare methods handling the same "
                    f"OnCompletionAction: {oca.name}",
                    location=bare_location,
                )
    return create_block(
        location,
        "bare_routing",
        *_maybe_switch(
            on_completion(location),
            {constant(oca.value, location): block for oca, block in bare_blocks.items()},
        ),
    )

def log_arc4_result(
    location: SourceLocation, result_expression: awst_nodes.Expression
) -> awst_nodes.ExpressionStatement:
    abi_log_prefix = awst_nodes.BytesConstant(
        source_location=location,
        value=0x151F7C75.to_bytes(4),
        encoding=awst_nodes.BytesEncoding.base16,
    )
    abi_log = awst_nodes.BytesBinaryOperation(
        source_location=location,
        left=abi_log_prefix,
        op=awst_nodes.BytesBinaryOperator.add,
        right=awst_nodes.ReinterpretCast(
            expr=result_expression,
            wtype=wtypes.bytes_wtype,
            source_location=result_expression.source_location,
        ),
    )
    log_op = awst_nodes.IntrinsicCall(
        op_code="log",
        stack_args=[abi_log],
        wtype=wtypes.void_wtype,
        source_location=location,
    )
    return awst_nodes.ExpressionStatement(log_op)

def assert_create_state(method: md.ARC4Method) -> Sequence[awst_nodes.Statement]:
    app_id = _txn("ApplicationID", wtypes.uint64_wtype, method.config_location)
    match method.create:
        case awst_nodes.ARC4CreateOption.allow:
            # if create is allowed but not required, we don't need to check anything
            return ()
        case awst_nodes.ARC4CreateOption.disallow:
            condition = _non_zero(app_id)
            error_message = "can only call when not creating"
        case awst_nodes.ARC4CreateOption.require:
            condition = _is_zero(app_id)
            error_message = "can only call when creating"
        case invalid:
            typing.assert_never(invalid)
    return [
        awst_nodes.ExpressionStatement(
            awst_nodes.AssertExpression(
                condition=condition,
                error_message=error_message,
                source_location=method.config_location,
            )
        )
    ]

def constant(value: int, location: SourceLocation) -> awst_nodes.Expression:
    return awst_nodes.UInt64Constant(value=value, source_location=location)

def left_shift(value: awst_nodes.Expression, location: SourceLocation) -> awst_nodes.Expression:
    return awst_nodes.UInt64BinaryOperation(
        source_location=location,
        left=constant(1, location),
        op=awst_nodes.UInt64BinaryOperator.lshift,
        right=value,
    )

def bit_and(
    lhs: awst_nodes.Expression, rhs: awst_nodes.Expression, location: SourceLocation
) -> awst_nodes.Expression:
    return awst_nodes.UInt64BinaryOperation(
        source_location=location,
        left=lhs,
        op=awst_nodes.UInt64BinaryOperator.bit_and,
        right=rhs,
    )

def uint64_sub(
    lhs: awst_nodes.Expression, rhs: awst_nodes.Expression, location: SourceLocation
) -> awst_nodes.Expression:
    return awst_nodes.UInt64BinaryOperation(
        source_location=location,
        left=lhs,
        op=awst_nodes.UInt64BinaryOperator.sub,
        right=rhs,
    )

def bit_packed_oca(
    allowed_oca: Iterable[OnCompletionAction], location: SourceLocation
) -> awst_nodes.Expression:
    """Returns an integer constant, where each bit corresponding to an OnCompletionAction is
    set to 1 if that action is allowed. This allows comparing a transaction's on completion value
    against a set of allowed actions using a bitwise and op"""
    bit_packed_value = 0
    for value in allowed_oca:
        bit_packed_value |= 1 << value.value
    return constant(bit_packed_value, location)

def check_allowed_oca(
    allowed_ocas: Sequence[OnCompletionAction], location: SourceLocation
) -> Sequence[awst_nodes.Statement]:
    not_allowed_ocas = sorted(
        a for a in ALL_VALID_APPROVAL_ON_COMPLETION_ACTIONS if a not in allowed_ocas
    )
    if not not_allowed_ocas:
        # all actions are allowed, don't need to check
        return ()
    match allowed_ocas, not_allowed_ocas:
        case [[single_allowed], _]:
            condition: awst_nodes.Expression = awst_nodes.NumericComparisonExpression(
                lhs=on_completion(location),
                rhs=awst_nodes.UInt64Constant(
                    source_location=location,
                    value=single_allowed.value,
                    teal_alias=single_allowed.name,
                ),
                operator=awst_nodes.NumericComparison.eq,
                source_location=location,
            )
        case _, [single_disallowed]:
            condition = awst_nodes.NumericComparisonExpression(
                lhs=on_completion(location),
                rhs=awst_nodes.UInt64Constant(
                    source_location=location,
                    value=single_disallowed.value,
                    teal_alias=single_disallowed.name,
                ),
                operator=awst_nodes.NumericComparison.ne,
                source_location=location,
            )
        case _:
            condition = bit_and(
                left_shift(on_completion(location), location),
                bit_packed_oca(allowed_ocas, location),
                location,
            )
    oca_desc = ", ".join(a.name for a in allowed_ocas)
    if len(allowed_ocas) > 1:
        oca_desc = f"one of {oca_desc}"
    return (
        awst_nodes.ExpressionStatement(
            awst_nodes.AssertExpression(
                condition=condition,
                error_message=f"OnCompletion is not {oca_desc}",
                source_location=location,
            )
        ),
    )

def _map_abi_args(
    arg_types: Sequence[wtypes.WType], location: SourceLocation
) -> Iterable[awst_nodes.Expression]:
    transaction_arg_offset = 0
    incoming_types = []
    for a in arg_types:
        if isinstance(a, wtypes.WGroupTransaction):
            transaction_arg_offset += 1
        else:
            if isinstance(a, wtypes.ARC4Type):
                arc4_type = a
            else:
                converted = maybe_avm_to_arc4_equivalent_type(a)
                if converted is not None:
                    arc4_type = converted
                elif _reference_type_array(a) is not None:
                    arc4_type = wtypes.arc4_byte_alias
                else:
                    raise CodeError(f"not an ARC4 type or native equivalent: {a}", location)
            incoming_types.append(arc4_type)

    if len(incoming_types) > 15:
        unpacked_types, packed_types = incoming_types[:14], incoming_types[14:]
    else:
        unpacked_types, packed_types = incoming_types, []
    abi_args = [
        app_arg(array_index, arg_wtype, location)
        for array_index, arg_wtype in enumerate(unpacked_types, start=1)
    ]
    if packed_types:
        abi_args.extend(
            awst_nodes.TupleItemExpression(
                base=app_arg(
                    15, wtypes.ARC4Tuple(types=packed_types, source_location=location), location
                ),
                index=tuple_index,
                source_location=location,
            )
            for tuple_index, _ in enumerate(packed_types)
        )
    abi_args.reverse()  # reverse so we can pop off end

    for arg in arg_types:
        if isinstance(arg, wtypes.WGroupTransaction):
            transaction_index = uint64_sub(
                _txn("GroupIndex", wtypes.uint64_wtype, location),
                constant(transaction_arg_offset, location),
                location,
            )
            yield awst_nodes.GroupTransactionReference(
                index=transaction_index, wtype=arg, source_location=location
            )
            transaction_arg_offset -= 1
        else:
            abi_arg = abi_args.pop()
            if (ref_array := _reference_type_array(arg)) is not None:
                uint64_index = _btoi(abi_arg, location)
                yield awst_nodes.IntrinsicCall(
                    op_code="txnas",
                    immediates=[ref_array],
                    stack_args=[uint64_index],
                    wtype=arg,
                    source_location=location,
                )
            else:
                if abi_arg.wtype != arg:
                    abi_arg = awst_nodes.ARC4Decode(
                        value=abi_arg, wtype=arg, source_location=location
                    )
                yield abi_arg

def route_abi_methods(
    location: SourceLocation,
    methods: Mapping[md.ARC4ABIMethod, AWSTContractMethodSignature],
) -> awst_nodes.Block:
    method_routing_cases = dict[awst_nodes.Expression, awst_nodes.Block]()
    seen_signatures = set[str]()
    for method, sig in methods.items():
        abi_loc = method.config_location
        abi_args = list(_map_abi_args(sig.parameter_types, location))
        method_result = call(abi_loc, sig, *abi_args)
        match sig.return_type:
            case wtypes.void_wtype:
                call_and_maybe_log = awst_nodes.ExpressionStatement(method_result)
            case wtypes.ARC4Type():
                call_and_maybe_log = log_arc4_result(abi_loc, method_result)
            case _:
                converted_return_type = maybe_avm_to_arc4_equivalent_type(sig.return_type)
                if converted_return_type is None:
                    raise CodeError(f"{sig.return_type} is not a valid ABI return type", abi_loc)
                arc4_encoded = awst_nodes.ARC4Encode(
                    value=method_result,
                    wtype=converted_return_type,
                    source_location=method_result.source_location,
                )
                call_and_maybe_log = log_arc4_result(abi_loc, arc4_encoded)

        arc4_signature = method.signature
        if not set_add(seen_signatures, arc4_signature):
            raise CodeError(
                f"Cannot have duplicate ARC4 method signatures: {arc4_signature}", abi_loc
            )
        method_routing_cases[
            awst_nodes.MethodConstant(source_location=location, value=arc4_signature)
        ] = create_block(
            abi_loc,
            f"{method.name}_route",
            *check_allowed_oca(method.allowed_completion_types, abi_loc),
            *assert_create_state(method),
            call_and_maybe_log,
            approve(abi_loc),
        )
    return create_block(
        location,
        "abi_routing",
        *_maybe_switch(_txn_app_args(0, location), method_routing_cases),
    )

def _maybe_switch(
    value: awst_nodes.Expression, cases: Mapping[awst_nodes.Expression, awst_nodes.Block]
) -> Sequence[awst_nodes.Statement]:
    if not cases:
        return ()
    return [
        awst_nodes.Switch(
            value=value,
            cases=cases,
            default_case=None,
            source_location=value.source_location,
        )
    ]

def create_abi_router(
    contract: awst_nodes.Contract,
    arc4_methods_with_signatures: Mapping[md.ARC4Method, AWSTContractMethodSignature],
) -> awst_nodes.ContractMethod:
    router_location = contract.source_location
    abi_methods = {}
    bare_methods = {}
    for method, sig in arc4_methods_with_signatures.items():
        if isinstance(method, md.ARC4BareMethod):
            bare_methods[method] = sig
        else:
            abi_methods[method] = sig

    abi_routing = route_abi_methods(router_location, abi_methods)
    bare_routing = route_bare_methods(router_location, bare_methods)
    num_app_args = _txn("NumAppArgs", wtypes.uint64_wtype, router_location)
    router = [
        awst_nodes.IfElse(
            condition=_non_zero(num_app_args),
            if_branch=abi_routing,
            else_branch=bare_routing,
            source_location=router_location,
        ),
        reject(router_location),
    ]
    approval_program = awst_nodes.ContractMethod(
        cref=contract.id,
        member_name="__puya_arc4_router__",
        source_location=router_location,
        args=[],
        return_type=wtypes.bool_wtype,
        body=create_block(router_location, None, *router),
        documentation=awst_nodes.MethodDocumentation(),
        arc4_method_config=None,
        inline=True,
    )
    return approval_program

def _reference_type_array(wtype: wtypes.WType) -> str | None:
    match wtype:
        case wtypes.asset_wtype:
            return "Assets"
        case wtypes.account_wtype:
            return "Accounts"
        case wtypes.application_wtype:
            return "Applications"
    return None