from collections.abc import Sequence

import attrs

from puya import log
from puya.awst import (
    nodes as awst_nodes,
    wtypes,
)
from puya.errors import CodeError
from puya.ir.builder._tuple_util import build_tuple_item_names
from puya.ir.builder._utils import assign_targets, new_register_version
from puya.ir.context import IRFunctionBuildContext
from puya.ir.models import InvokeSubroutine, Register, Subroutine, Value, ValueProvider, ValueTuple
from puya.parse import SourceLocation

logger = log.get_logger(__name__)

def visit_subroutine_call_expression(
    context: IRFunctionBuildContext, expr: awst_nodes.SubroutineCallExpression
) -> ValueProvider | None:
    target = context.resolve_subroutine(expr.target, expr.source_location)
    return _call_subroutine(context, target, expr.args, expr.source_location)

def visit_puya_lib_call_expression(
    context: IRFunctionBuildContext, call: awst_nodes.PuyaLibCall
) -> ValueProvider | None:
    try:
        target = context.embedded_funcs_lookup[call.func.value.id]
    except KeyError:
        raise CodeError(f"invalid puya_lib {call.func.name}", call.source_location) from None
    return _call_subroutine(context, target, call.args, call.source_location)

def _call_subroutine(
    context: IRFunctionBuildContext,
    target: Subroutine,
    args: Sequence[awst_nodes.CallArg],
    call_location: SourceLocation,
) -> ValueProvider | None:
    arg_lookup = _build_arg_lookup(context, args, call_location)

    resolved_args = []
    implicit_args = []
    for idx, param in enumerate(target.parameters):
        arg_val = arg_lookup.get(index=idx, param_name=param.name)
        resolved_args.append(arg_val)
        if param.implicit_return:
            if arg_val in implicit_args:
                logger.error(
                    "mutable values cannot be passed more than once to a subroutine",
                    location=arg_val.source_location,
                )
            implicit_args.append(arg_val)
    if not arg_lookup.is_empty:
        raise CodeError("function call arguments do not match signature", call_location) from None
    invoke_expr = InvokeSubroutine(
        source_location=call_location, args=resolved_args, target=target
    )
    if not implicit_args:
        return invoke_expr

    return_values = context.visitor.materialise_value_provider(invoke_expr, target.short_name)
    while implicit_args:
        in_arg = implicit_args.pop()
        out_value = return_values.pop()
        if isinstance(in_arg, Register):
            out_arg = new_register_version(context, in_arg)
            assign_targets(
                context,
                source=out_value,
                targets=[out_arg],
                assignment_location=call_location,
            )

    return (
        ValueTuple(values=return_values, source_location=call_location) if return_values else None
    )

@attrs.define
class _ArgLookup:
    _source_location: SourceLocation
    _positional_args: dict[int, Value] = attrs.field(factory=dict, init=False)
    _named_args: dict[str, Value] = attrs.field(factory=dict, init=False)
    _arg_idx: int = attrs.field(default=0, init=False)

    @property
    def is_empty(self) -> bool:
        return not self._named_args and not self._positional_args

    def add(self, name: str | None, value: Value) -> None:
        if name is None:
            self._positional_args[self._arg_idx] = value
        else:
            self._named_args[name] = value
        self._arg_idx += 1

    def get(self, index: int, param_name: str | None) -> Value:
        if param_name is not None:
            by_name = self._named_args.pop(param_name, None)
            if by_name is not None:
                return by_name
        try:
            return self._positional_args.pop(index)
        except KeyError:
            raise CodeError(
                "function call arguments do not match signature", self._source_location
            ) from None

def _build_arg_lookup(
    context: IRFunctionBuildContext,
    args: Sequence[awst_nodes.CallArg],
    call_location: SourceLocation,
) -> _ArgLookup:
    lookup = _ArgLookup(call_location)
    for expr_arg in args:
        if not isinstance(expr_arg.value.wtype, wtypes.WTuple):
            value = context.visitor.visit_and_materialise_single(expr_arg.value)
            lookup.add(name=expr_arg.name, value=value)
        else:
            values = context.visitor.visit_and_materialise(expr_arg.value)
            if expr_arg.name is None:
                for tup_value in values:
                    lookup.add(name=None, value=tup_value)
            else:
                item_names = build_tuple_item_names(
                    base_name=expr_arg.name,
                    wtype=expr_arg.value.wtype,
                    source_location=call_location,
                )
                for tup_value, (tup_item_name, _) in zip(values, item_names, strict=True):
                    lookup.add(name=tup_item_name, value=tup_value)

    return lookup


from puya import log
from puya.awst import (
    nodes as awst_nodes,
    wtypes,
)
from puya.errors import InternalError
from puya.ir import intrinsic_factory
from puya.ir.builder._tuple_util import build_tuple_registers
from puya.ir.builder._utils import assign_targets, new_register_version
from puya.ir.context import IRFunctionBuildContext
from puya.ir.models import (
    BasicBlock,
    ConditionalBranch,
    Switch,
    Value,
    ValueProvider,
    ValueTuple,
)
from puya.ir.types_ import get_wtype_arity, wtype_to_ir_type
from puya.parse import SourceLocation
from puya.utils import lazy_setdefault

logger = log.get_logger(__name__)

def handle_if_else(context: IRFunctionBuildContext, stmt: awst_nodes.IfElse) -> None:
    if_body = context.block_builder.mkblock(stmt.if_branch, "if_body")
    else_body = None
    if stmt.else_branch:
        else_body = context.block_builder.mkblock(stmt.else_branch, "else_body")
    next_block = context.block_builder.mkblock(stmt.source_location, "after_if_else")

    process_conditional(
        context,
        stmt.condition,
        true=if_body,
        false=else_body or next_block,
        loc=stmt.source_location,
    )
    _branch(context, if_body, stmt.if_branch, next_block)
    if else_body:
        assert stmt.else_branch
        _branch(context, else_body, stmt.else_branch, next_block)
    # activate the "next" block if it is reachable, which  might not be the case
    # if all paths within the "if" and "else" branches return early
    context.block_builder.try_activate_block(next_block)

def handle_switch(context: IRFunctionBuildContext, statement: awst_nodes.Switch) -> None:
    case_blocks = dict[Value, BasicBlock]()
    ir_blocks = dict[awst_nodes.Block | None, BasicBlock]()
    for value, block in statement.cases.items():
        ir_value = context.visitor.visit_and_materialise_single(value)
        if ir_value in case_blocks:
            logger.error("code block is unreachable", location=block.source_location)
        else:
            case_blocks[ir_value] = lazy_setdefault(
                ir_blocks,
                block,
                lambda _: context.block_builder.mkblock(
                    block,  # noqa: B023
                    f"switch_case_{len(ir_blocks)}",
                ),
            )
    default_block = lazy_setdefault(
        ir_blocks,
        statement.default_case,
        lambda b: context.block_builder.mkblock(
            b, "switch_case_default", fallback_location=statement.source_location
        ),
    )
    next_block = context.block_builder.mkblock(statement.source_location, "switch_case_next")

    switch_value = context.visitor.visit_and_materialise_single(statement.value)
    context.block_builder.terminate(
        Switch(
            value=switch_value,
            cases=case_blocks,
            default=default_block,
            source_location=statement.source_location,
        )
    )
    for block_, ir_block in ir_blocks.items():
        _branch(context, ir_block, block_, next_block)

    # activate the "next" block if it is reachable, which  might not be the case
    # if all code paths within the cases return early
    context.block_builder.try_activate_block(next_block)

def _branch(
    context: IRFunctionBuildContext,
    ir_block: BasicBlock,
    ast_block: awst_nodes.Block | None,
    next_ir_block: BasicBlock,
) -> None:
    context.block_builder.activate_block(ir_block)
    if ast_block is not None:
        ast_block.accept(context.visitor)
    context.block_builder.goto(next_ir_block)

def process_conditional(
    context: IRFunctionBuildContext,
    expr: awst_nodes.Expression,
    *,
    true: BasicBlock,
    false: BasicBlock,
    loc: SourceLocation,
) -> None:
    if expr.wtype != wtypes.bool_wtype:
        raise InternalError(
            "_process_conditional should only be used for boolean conditionals", loc
        )
    match expr:
        case awst_nodes.BooleanBinaryOperation(
            op=bool_op, left=lhs, right=rhs, source_location=loc
        ):
            # Short circuit boolean binary operators in a conditional context.
            contd = context.block_builder.mkblock(loc, f"{bool_op}_contd")
            if bool_op == "and":
                process_conditional(context, lhs, true=contd, false=false, loc=loc)
            elif bool_op == "or":
                process_conditional(context, lhs, true=true, false=contd, loc=loc)
            else:
                raise InternalError(
                    f"Unhandled boolean operator for short circuiting: {bool_op}", loc
                )
            context.block_builder.activate_block(contd)
            process_conditional(context, rhs, true=true, false=false, loc=loc)
        case awst_nodes.Not(expr=expr, source_location=loc):
            process_conditional(context, expr, true=false, false=true, loc=loc)
        case _:
            condition_value = context.visitor.visit_and_materialise_single(expr)
            context.block_builder.terminate(
                ConditionalBranch(
                    condition=condition_value,
                    non_zero=true,
                    zero=false,
                    source_location=loc,
                )
            )

def handle_while_loop(context: IRFunctionBuildContext, statement: awst_nodes.WhileLoop) -> None:
    top, next_block = context.block_builder.mkblocks(
        "while_top", "after_while", source_location=statement.source_location
    )
    body = context.block_builder.mkblock(statement.loop_body, "while_body")
    context.block_builder.goto(top)
    with context.block_builder.activate_open_block(top):
        process_conditional(
            context,
            statement.condition,
            true=body,
            false=next_block,
            loc=statement.source_location,
        )
        context.block_builder.activate_block(body)
        with context.block_builder.enter_loop(on_continue=top, on_break=next_block):
            statement.loop_body.accept(context.visitor)
        context.block_builder.goto(top)

    context.block_builder.activate_block(next_block)

def handle_conditional_expression(
    context: IRFunctionBuildContext, expr: awst_nodes.ConditionalExpression
) -> ValueProvider:
    # if lhs and rhs are both guaranteed to not produce side effects, we can use a simple select op
    # TODO: expand detection of side-effect free to include "pure" ops
    if (
        get_wtype_arity(expr.wtype) == 1
        and isinstance(
            expr.true_expr, awst_nodes.VarExpression | awst_nodes.CompileTimeConstantExpression
        )
        and isinstance(
            expr.false_expr, awst_nodes.VarExpression | awst_nodes.CompileTimeConstantExpression
        )
    ):
        false_reg = context.visitor.visit_and_materialise_single(expr.false_expr)
        true_reg = context.visitor.visit_and_materialise_single(expr.true_expr)
        condition_value = context.visitor.visit_and_materialise_single(expr.condition)
        return intrinsic_factory.select(
            condition=condition_value,
            true=true_reg,
            false=false_reg,
            type_=wtype_to_ir_type(expr),
            source_location=expr.source_location,
        )
    true_block, false_block, merge_block = context.block_builder.mkblocks(
        "ternary_true", "ternary_false", "ternary_merge", source_location=expr.source_location
    )
    tmp_var_name = context.next_tmp_name("ternary_result")
    true_registers = build_tuple_registers(context, tmp_var_name, expr.wtype, expr.source_location)

    process_conditional(
        context,
        expr.condition,
        true=true_block,
        false=false_block,
        loc=expr.source_location,
    )

    context.block_builder.activate_block(true_block)
    true_vp = context.visitor.visit_expr(expr.true_expr)
    assign_targets(
        context,
        source=true_vp,
        targets=true_registers,
        assignment_location=expr.true_expr.source_location,
    )
    context.block_builder.goto(merge_block)

    context.block_builder.activate_block(false_block)
    false_vp = context.visitor.visit_expr(expr.false_expr)
    assign_targets(
        context,
        source=false_vp,
        targets=[new_register_version(context, reg) for reg in true_registers],
        assignment_location=expr.false_expr.source_location,
    )
    context.block_builder.goto(merge_block)

    context.block_builder.activate_block(merge_block)
    result = [
        context.ssa.read_variable(variable=r.name, ir_type=r.ir_type, block=merge_block)
        for r in true_registers
    ]
    if len(result) == 1:
        return result[0]
    return ValueTuple(values=result, source_location=expr.source_location)


import typing

from puya import log
from puya.awst import (
    nodes as awst_nodes,
    wtypes,
)
from puya.awst.nodes import Expression
from puya.errors import CodeError, InternalError
from puya.ir.avm_ops import AVMOp
from puya.ir.builder import arc4
from puya.ir.builder._tuple_util import build_tuple_registers, get_tuple_item_values
from puya.ir.builder._utils import (
    assert_value,
    assign_intrinsic_op,
    assign_targets,
    assign_temp,
)
from puya.ir.context import IRFunctionBuildContext
from puya.ir.models import (
    ConditionalBranch,
    GotoNth,
    Intrinsic,
    Register,
    UInt64Constant,
    Value,
    ValueProvider,
)
from puya.ir.types_ import IRType
from puya.ir.utils import lvalue_items
from puya.parse import SourceLocation

logger = log.get_logger(__name__)

class LoopVariables(typing.NamedTuple):
    item_registers: typing.Sequence[Register]
    index_: Register | None

    def refresh_assignment(self, context: IRFunctionBuildContext) -> "LoopVariables":
        item = [_refresh_mutated_variable(context, i) for i in self.item_registers]
        if self.index_ is None:
            index = None
        else:
            index = _refresh_mutated_variable(context, self.index_)
        return LoopVariables(item, index)

class LoopAssigner:
    def __init__(
        self, context: IRFunctionBuildContext, items: awst_nodes.Lvalue, *, has_enumerate: bool
    ):
        self._context: typing.Final = context
        self._items: typing.Final = items
        self.has_enumerate: typing.Final = has_enumerate

    def assign_user_loop_vars(
        self, item_provider: ValueProvider, index_provider: ValueProvider
    ) -> LoopVariables:
        registers = self._build_registers_from_lvalue(self._items)
        if not self.has_enumerate:
            index_register = None
            item_registers = registers
        else:
            (index_register, *item_registers) = registers
        assign_targets(
            self._context,
            source=item_provider,
            targets=item_registers,
            assignment_location=self._items.source_location,
        )
        if index_register:
            assign_targets(
                self._context,
                source=index_provider,
                targets=[index_register],
                assignment_location=self._items.source_location,
            )
        return LoopVariables(item_registers, index_register)

    def _build_registers_from_lvalue(self, target: awst_nodes.Lvalue) -> list[Register]:
        match target:
            case awst_nodes.VarExpression(name=var_name, source_location=var_loc, wtype=var_type):
                return build_tuple_registers(self._context, var_name, var_type, var_loc)
            case awst_nodes.TupleExpression() as tup_expr:
                tuple_items = lvalue_items(tup_expr)
                return [
                    reg for item in tuple_items for reg in self._build_registers_from_lvalue(item)
                ]
            case _:
                raise CodeError(
                    "unsupported assignment target in loop", self._items.source_location
                )

def handle_for_in_loop(context: IRFunctionBuildContext, statement: awst_nodes.ForInLoop) -> None:
    sequence = statement.sequence
    has_enumerate = False
    reverse_items = False
    reverse_index = False

    while True:
        match sequence:
            case awst_nodes.Enumeration():
                if has_enumerate:
                    raise CodeError(
                        "Nested enumeration is not currently supported", sequence.source_location
                    )
                sequence = sequence.expr
                has_enumerate = True
            case awst_nodes.Reversed():
                sequence = sequence.expr
                reverse_items = not reverse_items
                if not has_enumerate:
                    reverse_index = not reverse_index
            case _:
                break

    assign_user_loop_vars = LoopAssigner(
        context,
        items=statement.items,
        has_enumerate=has_enumerate,
    )

    match sequence:
        case awst_nodes.Range(
            start=range_start, stop=range_stop, step=range_step, source_location=range_loc
        ):
            _iterate_urange(
                context,
                loop_body=statement.loop_body,
                assigner=assign_user_loop_vars,
                statement_loc=statement.source_location,
                range_start=range_start,
                range_stop=range_stop,
                range_step=range_step,
                range_loc=range_loc,
                reverse_items=reverse_items,
                reverse_index=reverse_index,
            )
        case awst_nodes.Expression(wtype=wtypes.WTuple(types=item_types)) as tuple_expression:
            if not item_types:
                logger.debug("Skipping ForInStatement which iterates an empty sequence.")
            else:
                _iterate_tuple(
                    context,
                    loop_body=statement.loop_body,
                    assigner=assign_user_loop_vars,
                    tuple_expr=tuple_expression,
                    statement_loc=statement.source_location,
                    reverse_index=reverse_index,
                    reverse_items=reverse_items,
                )
        case awst_nodes.Expression(wtype=wtypes.bytes_wtype):
            bytes_value = context.visitor.visit_and_materialise_single(sequence)
            byte_length = assign_temp(
                context,
                temp_description="bytes_length",
                source=Intrinsic(
                    op=AVMOp.len_,
                    args=[bytes_value],
                    source_location=statement.source_location,
                ),
                source_location=statement.source_location,
            )

            def get_byte_at_index(index_register: Value) -> ValueProvider:
                return Intrinsic(
                    op=AVMOp.extract3,
                    args=[
                        bytes_value,
                        index_register,
                        UInt64Constant(value=1, source_location=None),
                    ],
                    source_location=statement.items.source_location,
                )

            _iterate_indexable(
                context,
                loop_body=statement.loop_body,
                indexable_size=byte_length,
                get_value_at_index=get_byte_at_index,
                assigner=assign_user_loop_vars,
                statement_loc=statement.source_location,
                reverse_index=reverse_index,
                reverse_items=reverse_items,
            )
        case awst_nodes.Expression(wtype=wtypes.ARC4Array() as array_wtype):
            iterator = arc4.build_for_in_array(
                context,
                array_wtype,
                sequence,
                statement.source_location,
            )
            _iterate_indexable(
                context,
                loop_body=statement.loop_body,
                indexable_size=iterator.array_length,
                get_value_at_index=iterator.get_value_at_index,
                assigner=assign_user_loop_vars,
                statement_loc=statement.source_location,
                reverse_index=reverse_index,
                reverse_items=reverse_items,
            )
        case _:
            raise InternalError("Unsupported ForInLoop sequence", statement.source_location)

def _iterate_urange(
    context: IRFunctionBuildContext,
    *,
    loop_body: awst_nodes.Block,
    assigner: LoopAssigner,
    statement_loc: SourceLocation,
    range_start: Expression,
    range_stop: Expression,
    range_step: Expression,
    range_loc: SourceLocation,
    reverse_items: bool,
    reverse_index: bool,
) -> None:
    step = context.visitor.visit_and_materialise_single(range_step)
    stop = context.visitor.visit_and_materialise_single(range_stop)
    start = context.visitor.visit_and_materialise_single(range_start)
    assert_value(context, step, source_location=statement_loc, comment="Step cannot be zero")

    if reverse_items or reverse_index:
        return _iterate_urange_with_reversal(
            context,
            loop_body=loop_body,
            assigner=assigner,
            statement_loc=statement_loc,
            start=start,
            stop=stop,
            step=step,
            range_loc=range_loc,
            reverse_items=reverse_items,
            reverse_index=reverse_index,
        )
    else:
        return _iterate_urange_simple(
            context,
            loop_body=loop_body,
            assigner=assigner,
            statement_loc=statement_loc,
            start=start,
            stop=stop,
            step=step,
            range_loc=range_loc,
        )

def _iterate_urange_simple(
    context: IRFunctionBuildContext,
    *,
    loop_body: awst_nodes.Block,
    assigner: LoopAssigner,
    statement_loc: SourceLocation,
    start: Value,
    stop: Value,
    step: Value,
    range_loc: SourceLocation,
) -> None:
    body = context.block_builder.mkblock(loop_body, "for_body")
    header, footer, next_block = context.block_builder.mkblocks(
        "for_header", "for_footer", "after_for", source_location=statement_loc
    )

    loop_vars = assigner.assign_user_loop_vars(
        start, UInt64Constant(value=0, source_location=None)
    )
    context.block_builder.goto(header)
    with context.block_builder.activate_open_block(header):
        (current_range_item,), current_range_index = loop_vars.refresh_assignment(context)
        continue_looping = assign_intrinsic_op(
            context,
            target="continue_looping",
            op=AVMOp.lt,
            args=[current_range_item, stop],
            source_location=range_loc,
        )
        context.block_builder.terminate(
            ConditionalBranch(
                condition=continue_looping,
                non_zero=body,
                zero=next_block,
                source_location=statement_loc,
            )
        )

        context.block_builder.activate_block(body)
        with context.block_builder.enter_loop(on_continue=footer, on_break=next_block):
            loop_body.accept(context.visitor)

        context.block_builder.goto(footer)
        if context.block_builder.try_activate_block(footer):
            assign_intrinsic_op(
                context,
                target=current_range_item,
                op=AVMOp.add,
                args=[current_range_item, step],
                source_location=range_loc,
            )
            if current_range_index:
                assign_intrinsic_op(
                    context,
                    target=current_range_index,
                    op=AVMOp.add,
                    args=[current_range_index, 1],
                    source_location=range_loc,
                )
            context.block_builder.goto(header)

    context.block_builder.activate_block(next_block)

def _iterate_urange_with_reversal(
    context: IRFunctionBuildContext,
    *,
    loop_body: awst_nodes.Block,
    assigner: LoopAssigner,
    statement_loc: SourceLocation,
    start: Value,
    stop: Value,
    step: Value,
    range_loc: SourceLocation,
    reverse_items: bool,
    reverse_index: bool,
) -> None:
    assert reverse_items or reverse_index
    body = context.block_builder.mkblock(loop_body, "for_body")
    header, footer, increment_block, next_block = context.block_builder.mkblocks(
        "for_header", "for_footer", "for_increment", "after_for", source_location=statement_loc
    )

    # The following code will result in underflow if we don't pre-check the urange
    # params
    should_loop = assign_intrinsic_op(
        context,
        target="should_loop",
        op=AVMOp.lt,
        args=[start, stop],
        source_location=statement_loc,
    )
    context.block_builder.terminate(
        ConditionalBranch(
            condition=should_loop,
            non_zero=header,
            zero=next_block,
            source_location=statement_loc,
        )
    )

    context.block_builder.activate_block(header)
    # iteration_count = ((stop - 1) - start) // step + 1
    # => iteration_count - 1 = (stop - start - 1) // step
    range_length = assign_intrinsic_op(
        context,
        target="range_length",
        op=AVMOp.sub,
        args=[stop, start],
        source_location=range_loc,
    )
    range_length_minus_one = assign_intrinsic_op(
        context,
        target="range_length_minus_one",
        op=AVMOp.sub,
        args=[range_length, 1],
        source_location=range_loc,
    )
    iteration_count_minus_one = assign_intrinsic_op(
        context,
        target="iteration_count_minus_one",
        op=AVMOp.div_floor,
        args=[range_length_minus_one, step],
        source_location=range_loc,
    )
    range_delta = assign_intrinsic_op(
        context,
        target="range_delta",
        op=AVMOp.mul,
        args=[step, iteration_count_minus_one],
        source_location=range_loc,
    )
    max_range_item = assign_intrinsic_op(
        context,
        target="max_range_item",
        op=AVMOp.add,
        args=[start, range_delta],
        source_location=range_loc,
    )
    loop_vars = assigner.assign_user_loop_vars(
        start if not reverse_items else max_range_item,
        (
            UInt64Constant(value=0, source_location=None)
            if not reverse_index
            else iteration_count_minus_one
        ),
    )
    context.block_builder.goto(body)

    with context.block_builder.activate_open_block(body):
        (current_range_item,), current_range_index = loop_vars.refresh_assignment(context)

        with context.block_builder.enter_loop(on_continue=footer, on_break=next_block):
            loop_body.accept(context.visitor)

        context.block_builder.goto(footer)
        if context.block_builder.try_activate_block(footer):
            continue_looping_op = Intrinsic(
                op=AVMOp.lt,
                args=(
                    [current_range_item, max_range_item]
                    if not reverse_items
                    else [start, current_range_item]
                ),
                source_location=range_loc,
            )
            continue_looping = assign_temp(
                context,
                source=continue_looping_op,
                temp_description="continue_looping",
                source_location=range_loc,
            )

            context.block_builder.terminate(
                ConditionalBranch(
                    condition=continue_looping,
                    non_zero=increment_block,
                    zero=next_block,
                    source_location=statement_loc,
                )
            )

            context.block_builder.activate_block(increment_block)
            assign_intrinsic_op(
                context,
                target=current_range_item,
                op=AVMOp.add if not reverse_items else AVMOp.sub,
                args=[current_range_item, step],
                source_location=range_loc,
            )
            if current_range_index:
                assign_intrinsic_op(
                    context,
                    target=current_range_index,
                    op=AVMOp.add if not reverse_index else AVMOp.sub,
                    args=[current_range_index, 1],
                    source_location=range_loc,
                )
            context.block_builder.goto(body)

    context.block_builder.activate_block(next_block)

def _iterate_indexable(
    context: IRFunctionBuildContext,
    *,
    loop_body: awst_nodes.Block,
    assigner: LoopAssigner,
    statement_loc: SourceLocation,
    indexable_size: Value,
    get_value_at_index: typing.Callable[[Value], ValueProvider],
    reverse_items: bool,
    reverse_index: bool,
) -> None:
    body = context.block_builder.mkblock(loop_body, "for_body")
    header, footer, next_block = context.block_builder.mkblocks(
        "for_header", "for_footer", "after_for", source_location=statement_loc
    )

    index_internal = assign_temp(
        context,
        source=UInt64Constant(value=0, source_location=None),
        temp_description="item_index_internal",
        source_location=None,
    )
    reverse_index_internal = assign_temp(
        context,
        source=indexable_size,
        temp_description="reverse_index_internal",
        source_location=None,
    )
    context.block_builder.goto(header)

    with context.block_builder.activate_open_block(header):
        current_index_internal = _refresh_mutated_variable(context, index_internal)
        if not (reverse_items or reverse_index):
            continue_looping = assign_intrinsic_op(
                context,
                target="continue_looping",
                op=AVMOp.lt,
                args=[current_index_internal, indexable_size],
                source_location=statement_loc,
            )
        else:
            continue_looping = assign_intrinsic_op(
                context,
                target="continue_looping",
                op=AVMOp.gt,
                args=[_refresh_mutated_variable(context, reverse_index_internal), 0],
                source_location=statement_loc,
            )
        context.block_builder.terminate(
            ConditionalBranch(
                condition=continue_looping,
                non_zero=body,
                zero=next_block,
                source_location=statement_loc,
            )
        )

        context.block_builder.activate_block(body)
        if reverse_items or reverse_index:
            reverse_index_internal = assign_intrinsic_op(
                context,
                target=reverse_index_internal,
                op=AVMOp.sub,
                args=[_refresh_mutated_variable(context, reverse_index_internal), 1],
                source_location=None,
            )
        assigner.assign_user_loop_vars(
            get_value_at_index(
                reverse_index_internal if reverse_items else current_index_internal
            ),
            reverse_index_internal if reverse_index else current_index_internal,
        )
        with context.block_builder.enter_loop(on_continue=footer, on_break=next_block):
            loop_body.accept(context.visitor)
        context.block_builder.goto(footer)

        if context.block_builder.try_activate_block(footer):
            if not (reverse_items and reverse_index):
                assign_intrinsic_op(
                    context,
                    target=index_internal,
                    op=AVMOp.add,
                    args=[current_index_internal, 1],
                    source_location=None,
                )
            context.block_builder.goto(header)

    context.block_builder.activate_block(next_block)

def _iterate_tuple(
    context: IRFunctionBuildContext,
    *,
    loop_body: awst_nodes.Block,
    assigner: LoopAssigner,
    tuple_expr: awst_nodes.Expression,
    statement_loc: SourceLocation,
    reverse_index: bool,
    reverse_items: bool,
) -> None:
    tuple_values = context.visitor.visit_and_materialise(tuple_expr)
    assert isinstance(tuple_expr.wtype, wtypes.WTuple), "tuple_expr wtype must be WTuple"
    tuple_wtype = tuple_expr.wtype
    max_index = len(tuple_wtype.types) - 1
    loop_counter_name = context.next_tmp_name("loop_counter")

    def assign_counter_and_user_vars(loop_count: int) -> Register:
        counter_reg = context.ssa.new_register(loop_counter_name, IRType.uint64, None)
        assign_targets(
            context,
            source=UInt64Constant(value=loop_count, source_location=None),
            targets=[counter_reg],
            assignment_location=None,
        )
        item_index = loop_count if not reverse_items else (max_index - loop_count)
        item_reg, index_reg = assigner.assign_user_loop_vars(
            get_tuple_item_values(
                tuple_values=tuple_values,
                tuple_wtype=tuple_wtype,
                index=item_index,
                target_wtype=tuple_wtype.types[item_index],
                source_location=statement_loc,
            ),
            UInt64Constant(
                value=loop_count if not reverse_index else (max_index - loop_count),
                source_location=None,
            ),
        )
        if index_reg and not reverse_index:
            return index_reg
        else:
            return counter_reg

    # construct basic blocks
    body = context.block_builder.mkblock(loop_body, "for_body")
    footer, next_block = context.block_builder.mkblocks(
        "for_footer", "after_for", source_location=statement_loc
    )
    headers = {
        idx: context.block_builder.mkblock(statement_loc, f"for_header_{idx}")
        for idx in range(1, len(tuple_wtype.types))
    }

    # first item - assigned in current block
    loop_counter = assign_counter_and_user_vars(0)

    # body
    context.block_builder.goto(body)
    with context.block_builder.activate_open_block(body):
        current_loop_counter = _refresh_mutated_variable(context, loop_counter)
        with context.block_builder.enter_loop(on_continue=footer, on_break=next_block):
            loop_body.accept(context.visitor)

        # footer + follow-up headers, iff the loop body doesn't exit unconditionally on first item
        context.block_builder.goto(footer)
        if context.block_builder.try_activate_block(footer):
            # footer
            context.block_builder.terminate(
                GotoNth(
                    value=current_loop_counter,
                    blocks=list(headers.values()),
                    default=next_block,
                    source_location=statement_loc,
                )
            )

            # headers for remaining items
            for idx, header in headers.items():
                context.block_builder.activate_block(header)
                assign_counter_and_user_vars(idx)
                context.block_builder.goto(body)

    context.block_builder.activate_block(next_block)

def _refresh_mutated_variable(context: IRFunctionBuildContext, reg: Register) -> Register:
    """
    Given a register pointing to an underlying root operand (ie name) that is mutated,
    do an SSA read in the current block.

    This is *only* required when there is control flow involved in the generated IR,
    if it's only the builder that needs to loop then it should usually have an updated
    reference to the most recent assigned register which will still be valid because it's
    within the same block.
    """
    return context.ssa.read_variable(reg.name, reg.ir_type, context.block_builder.active_block)


import typing
from collections.abc import Mapping, Sequence

import attrs

import puya.awst.txn_fields
from puya import log
from puya.awst import (
    nodes as awst_nodes,
    wtypes,
)
from puya.awst.function_traverser import FunctionTraverser
from puya.awst.to_code_visitor import ToCodeVisitor
from puya.awst.wtypes import WInnerTransactionFields
from puya.errors import CodeError, InternalError
from puya.ir.avm_ops import AVMOp
from puya.ir.builder._tuple_util import build_tuple_item_names
from puya.ir.builder._utils import assign, assign_intrinsic_op
from puya.ir.builder.blocks import BlocksBuilder
from puya.ir.context import IRFunctionBuildContext
from puya.ir.models import (
    ConditionalBranch,
    InnerTransactionField,
    Intrinsic,
    ITxnConstant,
    Register,
    UInt64Constant,
    Value,
    ValueProvider,
    ValueTuple,
)
from puya.ir.ssa import BraunSSA
from puya.ir.types_ import IRType, wtype_to_ir_type
from puya.ir.utils import format_tuple_index
from puya.parse import SourceLocation
from puya.utils import StableSet, positive_index

logger = log.get_logger(__name__)

_INNER_TRANSACTION_NON_ARRAY_FIELDS = [f for f in puya.awst.txn_fields.TxnField if not f.is_array]

@attrs.frozen(kw_only=True)
class CreateInnerTransactionFieldData:
    var_name: str
    field: puya.awst.txn_fields.TxnField
    field_counts: set[int] = attrs.field(factory=set)
    """The observed number of values for this field
    For non-array fields this will be either 0 or 1
    For array fields this will be 0 -> N
    Capturing these ranges allows generating much simpler IR
    """
    field_count_register_name: str

    @property
    def value_registers(self) -> dict[int, str]:
        return {idx: self.get_value_register_name(idx) for idx in range(max(self.field_counts))}

    def get_value_register_name(self, index: int) -> str:
        return f"{self.var_name}%%param_{self.field.immediate}_idx_{index}"

@attrs.frozen(kw_only=True)
class CreateInnerTransactionData:
    var_name: str
    fields: dict[puya.awst.txn_fields.TxnField, CreateInnerTransactionFieldData] = attrs.field(
        factory=dict, init=False
    )

    def get_or_add_field_data(
        self, field: puya.awst.txn_fields.TxnField
    ) -> CreateInnerTransactionFieldData:
        try:
            field_data = self.fields[field]
        except KeyError:
            field_data = self.fields[field] = CreateInnerTransactionFieldData(
                var_name=self.var_name,
                field=field,
                field_count_register_name=f"{self.var_name}%%{field.immediate}_length",
            )
        return field_data

class InnerTransactionBuilder:
    def __init__(self, context: IRFunctionBuildContext):
        self.context = context
        self._inner_txn_fields_data = dict[str, CreateInnerTransactionData]()
        self._create_itxn_counter = iter(range(2**64 - 1))

    @property
    def ssa(self) -> BraunSSA:
        return self.context.block_builder.ssa

    @property
    def block_builder(self) -> BlocksBuilder:
        return self.context.block_builder

    def handle_inner_transaction_field_assignments(
        self, stmt: awst_nodes.AssignmentStatement
    ) -> bool:
        # awst_nodes.CreateInnerTransaction is used to create a set of fields used in an inner txn
        # When lowering this to IR field values and sizes are stored to local registers using
        # the top level variable name as a prefix.
        # For unused fields the size is 0
        # For array fields the size goes up to the length of the array
        # Other fields will have a size of 1
        #
        # Then itxn field sets are referenced the values are read using normal SSA functions
        # this allows for complex use cases such as branches and loops to still function

        value = stmt.value
        source_location = stmt.source_location
        target = stmt.target
        match value:
            case awst_nodes.CreateInnerTransaction(fields=fields):
                ((var_name, var_loc),) = _get_assignment_target_local_names(target, 1)
                self._set_inner_transaction_fields(var_name, fields, var_loc)
                return True
            case awst_nodes.Copy(
                value=awst_nodes.Expression(wtype=wtypes.WInnerTransactionFields()) as copy_source
            ):
                ((var_name, var_loc),) = _get_assignment_target_local_names(target, 1)
                src_var_name = self._resolve_inner_txn_params_var_name(copy_source)
                self._copy_inner_transaction_fields(var_name, src_var_name, var_loc)
                return True
            case awst_nodes.TupleExpression(items=tuple_items) as tuple_source if any(
                isinstance(t, WInnerTransactionFields) for t in tuple_source.wtype.types
            ):
                names = _get_assignment_target_local_names(target, len(tuple_items))
                for (item_name, item_loc), item_value in zip(names, tuple_items, strict=True):
                    match item_value:
                        case awst_nodes.CreateInnerTransaction(fields=fields):
                            self._set_inner_transaction_fields(item_name, fields, item_loc)
                        case awst_nodes.Copy(
                            value=awst_nodes.Expression(
                                wtype=wtypes.WInnerTransactionFields()
                            ) as copy_source
                        ):
                            src_var_name = self._resolve_inner_txn_params_var_name(copy_source)
                            self._copy_inner_transaction_fields(item_name, src_var_name, item_loc)
                        case awst_nodes.Expression(wtype=wtypes.WInnerTransactionFields()):
                            raise CodeError(
                                "Unexpected Inner Transaction encountered in tuple", item_loc
                            )
                        case _:
                            value_provider = self.context.visitor.visit_expr(item_value)
                            assign(
                                self.context,
                                value_provider,
                                name=item_name,
                                register_location=item_loc,
                                assignment_location=source_location,
                            )
                return True
            case awst_nodes.Expression(wtype=wtypes.WInnerTransactionFields()):
                raise CodeError(
                    "Inner Transaction params can only be reassigned using copy()",
                    source_location,
                )
            case _:
                return False

    def _visit_submit_expr(self, expr: awst_nodes.Expression) -> Sequence[Value]:
        value_provider = self.context.visitor.visit_expr(expr)
        match value_provider:
            case ValueTuple(values=values):
                return values
            case Value() as value:
                return (value,)
        raise InternalError(
            "Unexpected result for SubmitInnerTransaction expr", expr.source_location
        )

    def add_inner_transaction_submit_result_assignments(
        self,
        targets: Sequence[Value],
        source_expr: awst_nodes.Expression,
        ass_loc: SourceLocation | None,
    ) -> None:
        """Performs special handling for inner transaction related assignments"""
        # awst_nodes.SubmitInnerTransaction is used to submit itxn field sets as an inner
        # transaction group
        # Then all non aray fields are assigned to local registers and cached
        # This allows these fields to be read later even if additional itxn submissions have been
        # performed.
        # Array fields can not be cached in the same way currently, due to the lack of an
        # efficient array type.

        # targets are the assignment results, which will all be registers here,
        # the only case they aren't registers is when assigning to storage, which will
        # never be supported for itxn's because they're ephemeral
        itx_targets = [
            t for t in targets if isinstance(t, Register) and t.ir_type == IRType.itxn_group_idx
        ]
        source_actions = SourceActionExtractor.visit(source_expr)
        if len(itx_targets) != len(source_actions):
            raise CodeError("unsupported inner transaction assignment", ass_loc)
        for itxn_target, source_action in zip(itx_targets, source_actions, strict=True):
            match source_action:
                case _CopySource(var_name=src_var_name):
                    self._copy_submit_inner_transaction_fields(itxn_target.name, src_var_name)
                case _AssignSubmit(is_last=is_last):
                    self._assign_submit_inner_transaction_fields(itxn_target, is_last=is_last)
                case unexpected:
                    typing.assert_never(unexpected)

    def handle_update_inner_transaction(self, call: awst_nodes.UpdateInnerTransaction) -> None:
        var_name = self._resolve_inner_txn_params_var_name(call.itxn)
        self._set_inner_transaction_fields(
            var_name, call.fields, call.source_location, update=True
        )

    def handle_inner_transaction_field(
        self, itxn_field: awst_nodes.InnerTransactionField
    ) -> ValueProvider:
        src_loc = itxn_field.source_location
        field = itxn_field.field
        if field.is_array != bool(itxn_field.array_index):
            raise InternalError(
                "inconsistent array_index for inner transaction field",
                src_loc,
            )

        itxn = self.context.visitor.visit_expr(itxn_field.itxn)
        if not isinstance(itxn, Register | ITxnConstant):
            itxn_field_desc = {itxn_field.itxn.accept(ToCodeVisitor())}
            raise CodeError(
                f"Could not resolve inner transaction group index for {itxn_field_desc}",
                src_loc,
            )

        # use cached field if available
        if isinstance(itxn, Register):
            field_var_name = _get_txn_field_var_name(itxn.name, field.immediate)
            if self.ssa.has_version(field_var_name):
                return self.ssa.read_variable(
                    field_var_name, wtype_to_ir_type(field.wtype), self.block_builder.active_block
                )

        match itxn:
            # use is_last register if it is defined
            case Register(name=itxn_name) if self.ssa.has_version(_get_txn_is_last(itxn_name)):
                is_last_in_group: Value = self.ssa.read_variable(
                    _get_txn_is_last(itxn_name),
                    IRType.bool,
                    self.block_builder.active_block,
                )
            # otherwise infer based on itxn expr
            case _:
                is_last_in_group = UInt64Constant(
                    value=int(_is_last_itxn(itxn_field.itxn)),
                    ir_type=IRType.bool,
                    source_location=src_loc,
                )

        return InnerTransactionField(
            group_index=itxn,
            is_last_in_group=is_last_in_group,
            array_index=(
                self.context.visitor.visit_and_materialise_single(itxn_field.array_index)
                if itxn_field.array_index
                else None
            ),
            field=field.immediate,
            type=wtype_to_ir_type(field.wtype),
            source_location=src_loc,
        )

    def handle_submit_inner_transaction(
        self, submit: awst_nodes.SubmitInnerTransaction
    ) -> Sequence[ITxnConstant]:
        src_loc = submit.source_location

        self.block_builder.add(
            Intrinsic(
                op=AVMOp.itxn_begin,
                source_location=src_loc,
            )
        )

        group_indexes = []
        for group_index, param in enumerate(submit.itxns):
            submit_var_loc = param.source_location
            if group_index > 0:
                self.block_builder.add(
                    Intrinsic(
                        op=AVMOp.itxn_next,
                        source_location=submit_var_loc,
                    )
                )
            param_var_name = self._resolve_inner_txn_params_var_name(param)
            next_txn = self.block_builder.mkblock(submit_var_loc, "next_txn")
            param_data = self._inner_txn_fields_data[param_var_name]

            # with the current implementation, reversing the order itxn_field is called
            # results in less stack manipulations as most values are naturally in the
            # required order when stack allocation occurs
            for field, field_data in reversed(param_data.fields.items()):
                field_value_counts = sorted(field_data.field_counts)
                if not field_value_counts or field_value_counts == [0]:
                    # nothing to do
                    continue

                min_num_values, *remaining_values = field_value_counts
                # values 0 -> min_num_values do not need to test
                # values min_num_values -> max_num_values need to check if they are set
                next_field = self.block_builder.mkblock(submit_var_loc, "next_field")
                self._set_field_values(field_data, 0, min_num_values)

                if remaining_values:
                    last_num_values = min_num_values
                    for next_num_values in remaining_values:
                        set_fields_blk = self.block_builder.mkblock(
                            submit_var_loc,
                            f"set_{field.immediate}_{last_num_values}_to_{next_num_values - 1}",
                        )
                        self.block_builder.terminate(
                            ConditionalBranch(
                                condition=self._get_is_field_count_gte(
                                    field_data, next_num_values
                                ),
                                non_zero=set_fields_blk,
                                zero=next_field,
                                source_location=submit_var_loc,
                            )
                        )
                        self.block_builder.activate_block(set_fields_blk)
                        self._set_field_values(field_data, last_num_values, next_num_values)
                        last_num_values = next_num_values

                    self.block_builder.goto(next_field)
                    self.block_builder.activate_block(next_field)

            group_indexes.append(
                ITxnConstant(
                    value=group_index,
                    source_location=submit_var_loc,
                    ir_type=IRType.itxn_group_idx,
                )
            )

            self.block_builder.goto(next_txn)
            self.block_builder.activate_block(next_txn)

        self.block_builder.add(
            Intrinsic(
                op=AVMOp.itxn_submit,
                source_location=src_loc,
            )
        )

        return group_indexes

    def _assign_submit_inner_transaction_fields(
        self,
        target: Register,
        *,
        is_last: bool,
    ) -> None:
        var_name = target.name
        is_last_in_group = assign(
            self.context,
            source=UInt64Constant(
                value=int(is_last),
                ir_type=IRType.bool,
                source_location=None,
            ),
            name=_get_txn_is_last(var_name),
            register_location=None,
            assignment_location=None,
        )
        for field in _INNER_TRANSACTION_NON_ARRAY_FIELDS:
            field_reg = _get_txn_field_var_name(var_name, field.immediate)
            assign(
                context=self.context,
                source=InnerTransactionField(
                    field=field.immediate,
                    group_index=target,
                    is_last_in_group=is_last_in_group,
                    type=wtype_to_ir_type(field.wtype),
                    array_index=None,
                    source_location=None,
                ),
                name=field_reg,
                register_location=None,
                assignment_location=None,
            )

    def _copy_submit_inner_transaction_fields(self, dest_var_name: str, src_var_name: str) -> None:
        active_block = self.context.block_builder.active_block
        for field in _INNER_TRANSACTION_NON_ARRAY_FIELDS:
            src_field = _get_txn_field_var_name(src_var_name, field.immediate)
            dest_field = _get_txn_field_var_name(dest_var_name, field.immediate)
            assign(
                context=self.context,
                source=self.context.ssa.read_variable(
                    src_field, wtype_to_ir_type(field.wtype), active_block
                ),
                name=dest_field,
                register_location=None,
                assignment_location=None,
            )

    def _set_field_values(
        self,
        field_data: CreateInnerTransactionFieldData,
        idx_from: int,
        idx_to: int,
    ) -> None:
        field = field_data.field
        field_ir_type = wtype_to_ir_type(field.wtype)
        for idx in range(idx_from, idx_to):
            field_value = self.ssa.read_variable(
                field_data.get_value_register_name(idx),
                field_ir_type,
                self.block_builder.active_block,
            )
            self.block_builder.add(
                Intrinsic(
                    op=AVMOp.itxn_field,
                    source_location=None,
                    immediates=[field.immediate],
                    args=[field_value],
                )
            )

    def _get_is_field_count_gte(
        self, field_data: CreateInnerTransactionFieldData, count: int
    ) -> Register:
        field = field_data.field
        len_register = self.ssa.read_variable(
            field_data.field_count_register_name,
            IRType.uint64,
            self.block_builder.active_block,
        )

        is_field_count_gte = assign_intrinsic_op(
            self.context,
            target=f"is_{field.immediate}_count_gte_{count}",
            op=AVMOp.gte,
            args=[len_register, count],
            source_location=None,
        )
        return is_field_count_gte

    def _set_inner_transaction_fields(
        self,
        var_name: str,
        inner_txn_fields: Mapping[puya.awst.txn_fields.TxnField, awst_nodes.Expression],
        var_loc: SourceLocation,
        *,
        update: bool = False,
    ) -> None:
        param_data = self._inner_txn_fields_data.setdefault(
            var_name, CreateInnerTransactionData(var_name=var_name)
        )
        # assign a unique constant to var_name, not used for anything directly, but prevents
        # an undefined variable warning
        assign(
            context=self.context,
            source=ITxnConstant(
                value=next(self._create_itxn_counter),
                source_location=var_loc,
                ir_type=IRType.itxn_field_set,
            ),
            name=var_name,
            assignment_location=var_loc,
        )
        fields = StableSet.from_iter(inner_txn_fields)
        if not update:
            # add missing fields to end
            for field in puya.awst.txn_fields.TxnField:
                if field.is_inner_param and field not in fields:
                    fields.add(field)
        for field in fields:
            field_data = param_data.get_or_add_field_data(field)
            arg_expr = inner_txn_fields.get(field)
            values: Sequence[ValueProvider] = []
            count_loc = arg_expr.source_location if arg_expr else var_loc
            if arg_expr:
                match self.context.visitor.visit_expr(arg_expr):
                    case ValueTuple(values=values):
                        pass
                    case ValueProvider() as vp:
                        values = [vp]

            field_data.field_counts.add(len(values))
            for idx, value in enumerate(values):
                assign(
                    context=self.context,
                    source=value,
                    name=field_data.get_value_register_name(idx),
                    register_location=var_loc,
                    assignment_location=value.source_location,
                )
            assign(
                context=self.context,
                source=UInt64Constant(
                    value=len(values),
                    source_location=count_loc,
                ),
                name=field_data.field_count_register_name,
                register_location=var_loc,
                assignment_location=count_loc,
            )

    def _copy_inner_transaction_fields(
        self, dest_var_name: str, src_var_name: str, var_loc: SourceLocation
    ) -> None:
        src_params_data = self._inner_txn_fields_data[src_var_name]
        dest_params_data = self._inner_txn_fields_data.setdefault(
            dest_var_name, CreateInnerTransactionData(var_name=dest_var_name)
        )
        for field in puya.awst.txn_fields.TxnField:
            if not field.is_inner_param:
                continue
            src_field_data = src_params_data.get_or_add_field_data(field)

            dest_field_data = dest_params_data.get_or_add_field_data(field)
            dest_field_data.field_counts.update(src_field_data.field_counts)
            for idx, src_field_register in src_field_data.value_registers.items():
                dest_field_register = dest_field_data.get_value_register_name(idx)
                assign(
                    context=self.context,
                    source=self.ssa.read_variable(
                        src_field_register,
                        wtype_to_ir_type(field.wtype),
                        self.block_builder.active_block,
                    ),
                    name=dest_field_register,
                    assignment_location=var_loc,
                )
            assign(
                context=self.context,
                source=self.ssa.read_variable(
                    src_field_data.field_count_register_name,
                    IRType.uint64,
                    self.block_builder.active_block,
                ),
                name=dest_field_data.field_count_register_name,
                assignment_location=var_loc,
            )

    def _resolve_inner_txn_params_var_name(self, params: awst_nodes.Expression) -> str:
        match params:
            case awst_nodes.CreateInnerTransaction() as itxn:
                var_name = self.context.next_tmp_name(description="inner_txn_params")
                self._set_inner_transaction_fields(
                    var_name=var_name, inner_txn_fields=itxn.fields, var_loc=itxn.source_location
                )
            case awst_nodes.VarExpression(name=var_name):
                pass
            case awst_nodes.TupleItemExpression(
                base=awst_nodes.VarExpression(name=name, wtype=wtypes.WTuple() as base_wtype),
                index=index,
            ):
                return format_tuple_index(base_wtype, name, index)
            case awst_nodes.Copy(value=value):
                return self._resolve_inner_txn_params_var_name(value)
            case _:
                raise InternalError(
                    "Could not resolve var_name for inner transaction params",
                    params.source_location,
                )
        return var_name

def _get_assignment_target_local_names(
    target: awst_nodes.Expression, expected_number: int
) -> Sequence[tuple[str, SourceLocation]]:
    match target:
        case awst_nodes.VarExpression(name=var_name) if expected_number == 1:
            return [(var_name, target.source_location)]
        case awst_nodes.VarExpression(name=var_name, wtype=wtypes.WTuple() as var_wtype):
            return [
                (format_tuple_index(var_wtype, var_name, idx), target.source_location)
                for idx in range(expected_number)
            ]
        case awst_nodes.TupleExpression(items=items) if expected_number == len(items) and all(
            isinstance(i, awst_nodes.VarExpression) for i in items
        ):
            items = typing.cast(Sequence[awst_nodes.VarExpression], items)
            return [(expr.name, expr.source_location) for expr in items]
        case awst_nodes.TupleItemExpression(
            base=awst_nodes.TupleExpression(wtype=tuple_wtype) as base, index=index
        ):
            tuple_names = _get_assignment_target_local_names(base, len(tuple_wtype.types))
            return [tuple_names[index]]
        case awst_nodes.FieldExpression(
            base=awst_nodes.TupleExpression(wtype=tuple_wtype) as base,
            name=name,
            source_location=name_loc,
        ):
            tuple_names = _get_assignment_target_local_names(base, len(tuple_wtype.types))
            index = tuple_wtype.name_to_index(name, name_loc)
            return [tuple_names[index]]
    raise CodeError(
        "Inner Transactions can only be assigned to local variables",
        target.source_location,
    )

@attrs.frozen
class _CopySource:
    var_name: str

@attrs.frozen
class _AssignSubmit:
    index: int
    is_last: bool

_SourceAction = _CopySource | _AssignSubmit

class SourceActionExtractor(FunctionTraverser):
    def __init__(self) -> None:
        self._actions = list[_SourceAction]()

    @classmethod
    def visit(cls, node: awst_nodes.Expression) -> list[_SourceAction]:
        visitor = cls()
        node.accept(visitor)
        return visitor._actions  # noqa: SLF001

    def visit_submit_inner_transaction(self, call: awst_nodes.SubmitInnerTransaction) -> None:
        itxns = len(call.itxns)
        self._actions.extend(
            _AssignSubmit(index=idx, is_last=idx == itxns - 1) for idx in range(itxns)
        )

    def visit_var_expression(self, expr: awst_nodes.VarExpression) -> None:
        self._actions.extend(
            [
                _CopySource(var_name=name)
                for name, ir_type in build_tuple_item_names(
                    expr.name, expr.wtype, expr.source_location
                )
                if ir_type == IRType.itxn_group_idx
            ]
        )

    def visit_inner_transaction_field(self, itxn_field: awst_nodes.InnerTransactionField) -> None:
        # this will consume any referenced inner transaction, so don't need to traverse it
        pass

    def visit_tuple_item_expression(self, expr: awst_nodes.TupleItemExpression) -> None:
        start_len = len(self._actions)
        super().visit_tuple_item_expression(expr)
        added = self._actions[start_len:]
        # only keep the relevant action
        if isinstance(expr.wtype, wtypes.WInnerTransaction):
            self._actions[start_len:] = [added[expr.index]]

    def visit_slice_expression(self, expr: awst_nodes.SliceExpression) -> None:
        start_len = len(self._actions)
        super().visit_slice_expression(expr)
        added = self._actions[start_len:]
        if not added or not isinstance(expr.base.wtype, wtypes.WTuple):
            return
        # determine constant indexes
        tuple_size = len(added)
        begin_index = 0 if expr.begin_index is None else _get_uint64_const(expr.begin_index)
        if begin_index is None:
            return
        begin_index = positive_index(begin_index, added)
        end_index = tuple_size if expr.end_index is None else _get_uint64_const(expr.end_index)
        if end_index is None:
            return
        end_index = positive_index(end_index, added)

        # include relevant items from sliced tuple
        self._actions[start_len:] = [
            added[idx]
            for idx in range(begin_index, end_index)
            if isinstance(expr.base.wtype.types[idx], wtypes.WInnerTransaction)
        ]

def _get_uint64_const(expr: awst_nodes.Expression) -> int | None:
    if isinstance(expr, awst_nodes.IntegerConstant) and expr.wtype == wtypes.uint64_wtype:
        return expr.value
    return None

def _is_last_itxn(expr: awst_nodes.Expression) -> bool:
    # is last itxn if expr is a submit expr of size 1 OR
    if not isinstance(expr, awst_nodes.TupleItemExpression | awst_nodes.FieldExpression):
        return _is_submit_expr_of_size(expr, 1)

    # if expr is a tuple item expression with an index into the last item of a submit expr
    base = expr.base
    if not isinstance(base.wtype, wtypes.WTuple):
        return False

    index = (
        expr.index
        if isinstance(expr, awst_nodes.TupleItemExpression)
        else base.wtype.name_to_index(expr.name, expr.source_location)
    )
    tuple_size = len(base.wtype.types)
    if index == -1 or (index + 1) == tuple_size:
        return _is_submit_expr_of_size(base, tuple_size)
    else:
        return False

def _is_submit_expr_of_size(expr: awst_nodes.Expression, expected_group_size: int) -> bool:
    match expr:
        case awst_nodes.SubmitInnerTransaction(itxns=itxns) if len(itxns) == expected_group_size:
            return True
        case awst_nodes.SingleEvaluation(source=source):
            return _is_submit_expr_of_size(source, expected_group_size)
        case _:
            return False

def _get_txn_field_var_name(var_name: str, field: str) -> str:
    return f"{var_name}.{field}"

def _get_txn_is_last(var_name: str) -> str:
    return f"{var_name}._is_last"