import typing
from collections.abc import Iterator, Sequence

import attrs

import puya.awst.visitors
import puya.ir.builder.storage
from puya import algo_constants, log, utils
from puya.avm import AVMType
from puya.awst import (
    nodes as awst_nodes,
    wtypes,
)
from puya.awst.nodes import BigUIntBinaryOperator, UInt64BinaryOperator
from puya.awst.to_code_visitor import ToCodeVisitor
from puya.awst.txn_fields import TxnField
from puya.awst.wtypes import WInnerTransaction, WInnerTransactionFields
from puya.errors import CodeError, InternalError
from puya.ir.avm_ops import AVMOp
from puya.ir.builder import arc4, flow_control, storage
from puya.ir.builder._tuple_util import get_tuple_item_values
from puya.ir.builder._utils import (
    OpFactory,
    assert_value,
    assign,
    assign_intrinsic_op,
    assign_targets,
    assign_temp,
    extract_const_int,
    get_implicit_return_is_original,
    get_implicit_return_out,
    mktemp,
)
from puya.ir.builder.arc4 import ARC4_FALSE, ARC4_TRUE
from puya.ir.builder.assignment import (
    handle_assignment,
    handle_assignment_expr,
)
from puya.ir.builder.bytes import (
    visit_bytes_intersection_slice_expression,
    visit_bytes_slice_expression,
)
from puya.ir.builder.callsub import (
    visit_puya_lib_call_expression,
    visit_subroutine_call_expression,
)
from puya.ir.builder.iteration import handle_for_in_loop
from puya.ir.builder.itxn import InnerTransactionBuilder
from puya.ir.context import IRBuildContext
from puya.ir.models import (
    AddressConstant,
    BigUIntConstant,
    BytesConstant,
    CompiledContractReference,
    CompiledLogicSigReference,
    Fail,
    Intrinsic,
    InvokeSubroutine,
    MethodConstant,
    Op,
    ProgramExit,
    Subroutine,
    SubroutineReturn,
    TemplateVar,
    UInt64Constant,
    Value,
    ValueProvider,
    ValueTuple,
)
from puya.ir.types_ import (
    AVMBytesEncoding,
    IRType,
    bytes_enc_to_avm_bytes_enc,
    wtype_to_ir_type,
    wtype_to_ir_types,
)
from puya.ir.utils import format_tuple_index
from puya.parse import SourceLocation

TExpression: typing.TypeAlias = ValueProvider | None
TStatement: typing.TypeAlias = None

logger = log.get_logger(__name__)

class FunctionIRBuilder(
    puya.awst.visitors.ExpressionVisitor[TExpression],
    puya.awst.visitors.StatementVisitor[TStatement],
):
    def __init__(
        self, context: IRBuildContext, function: awst_nodes.Function, subroutine: Subroutine
    ):
        self.context = context.for_function(function, subroutine, self)
        self._itxn = InnerTransactionBuilder(self.context)
        self._single_eval_cache = dict[awst_nodes.SingleEvaluation, TExpression]()
        self._visited_exprs = dict[tuple[int, awst_nodes.Expression], TExpression]()

    @classmethod
    def build_body(
        cls,
        ctx: IRBuildContext,
        function: awst_nodes.Function,
        subroutine: Subroutine,
    ) -> None:
        logger.debug(f"Building IR for function {function.full_name}")
        builder = cls(ctx, function, subroutine)
        func_ctx = builder.context
        with func_ctx.log_exceptions():
            block_builder = func_ctx.block_builder
            for p in subroutine.parameters:
                if p.implicit_return:
                    assign(
                        func_ctx,
                        UInt64Constant(value=1, ir_type=IRType.bool, source_location=None),
                        name=get_implicit_return_is_original(p.name),
                        assignment_location=None,
                    )
                    assign(
                        func_ctx,
                        p,
                        name=get_implicit_return_out(p.name),
                        assignment_location=None,
                    )
            function.body.accept(builder)
            final_block = block_builder.active_block
            if not final_block.terminated:
                if function.return_type != wtypes.void_wtype:
                    raise CodeError("not all paths return a value", function.body.source_location)
                block_builder.terminate(
                    SubroutineReturn(
                        result=[
                            block_builder.ssa.read_variable(
                                get_implicit_return_out(p.name), p.ir_type, final_block
                            )
                            for p in subroutine.parameters
                            if p.implicit_return
                        ],
                        source_location=None,
                    )
                )
            subroutine.body = block_builder.finalise()
            subroutine.validate_with_ssa()

    def visit_copy(self, expr: awst_nodes.Copy) -> TExpression:
        # For reference types, we need to clone the data
        # For value types, we can just visit the expression and the resulting read
        # will effectively be a copy. We assign the copy to a new register in case it is
        # mutated.
        match expr.value.wtype:
            case wtypes.ARC4Type(immutable=False):
                # Arc4 encoded types are value types
                original_value = self.visit_and_materialise_single(expr.value)
                return assign_temp(
                    temp_description="copy",
                    source=original_value,
                    source_location=expr.source_location,
                    context=self.context,
                )
        raise InternalError(
            f"Invalid source wtype for Copy {expr.value.wtype}", expr.source_location
        )

    def visit_arc4_decode(self, expr: awst_nodes.ARC4Decode) -> TExpression:
        return arc4.decode_expr(self.context, expr)

    def visit_arc4_encode(self, expr: awst_nodes.ARC4Encode) -> TExpression:
        return arc4.encode_expr(self.context, expr)

    def visit_compiled_contract(self, expr: awst_nodes.CompiledContract) -> TExpression:
        prefix = self.context.options.template_vars_prefix if expr.prefix is None else expr.prefix
        template_variables = {
            prefix + k: self.visit_and_materialise_single(v)
            for k, v in expr.template_variables.items()
        }
        # TODO: remove implicit coupling
        #       the coupling here is between the order of values in the ValueTuple
        #       and the structure of the high level python type
        #       once we support nested tuples, this coupling can be removed
        #       and instead support names on WTuple, then each value can be accessed and lowered
        #       via a FieldExpression
        program_pages = [
            CompiledContractReference(
                artifact=expr.contract,
                field=field,
                program_page=page,
                ir_type=IRType.bytes,
                source_location=expr.source_location,
                template_variables=template_variables,
            )
            for field in (
                TxnField.ApprovalProgramPages,
                TxnField.ClearStateProgramPages,
            )
            for page in (0, 1)
        ]
        return ValueTuple(
            values=program_pages
            + [
                (
                    self.visit_and_materialise_single(expr.allocation_overrides[field])
                    if field in expr.allocation_overrides
                    else CompiledContractReference(
                        artifact=expr.contract,
                        field=field,
                        ir_type=IRType.uint64,
                        source_location=expr.source_location,
                        template_variables=template_variables,
                    )
                )
                for field in (
                    TxnField.ExtraProgramPages,
                    TxnField.GlobalNumUint,
                    TxnField.GlobalNumByteSlice,
                    TxnField.LocalNumUint,
                    TxnField.LocalNumByteSlice,
                )
            ],
            source_location=expr.source_location,
        )

    def visit_compiled_logicsig(self, expr: awst_nodes.CompiledLogicSig) -> TExpression:
        prefix = self.context.options.template_vars_prefix if expr.prefix is None else expr.prefix
        template_variables = {
            prefix + k: self.visit_and_materialise_single(v)
            for k, v in expr.template_variables.items()
        }
        return ValueTuple(
            values=[
                CompiledLogicSigReference(
                    artifact=expr.logic_sig,
                    ir_type=IRType.bytes,
                    source_location=expr.source_location,
                    template_variables=template_variables,
                )
            ],
            source_location=expr.source_location,
        )

    def visit_assignment_statement(self, stmt: awst_nodes.AssignmentStatement) -> TStatement:
        if not self._itxn.handle_inner_transaction_field_assignments(stmt):
            targets = handle_assignment_expr(
                self.context,
                target=stmt.target,
                value=stmt.value,
                assignment_location=stmt.source_location,
            )
            self._itxn.add_inner_transaction_submit_result_assignments(
                targets, stmt.value, stmt.source_location
            )
        return None

    def visit_assignment_expression(self, expr: awst_nodes.AssignmentExpression) -> TExpression:
        result = handle_assignment_expr(
            self.context,
            target=expr.target,
            value=expr.value,
            assignment_location=expr.source_location,
        )
        if not result:
            # HOW DID YOU GET HERE
            raise CodeError("Assignment expression did not return a result", expr.source_location)
        if len(result) == 1:
            return result[0]
        else:
            return ValueTuple(expr.source_location, list(result))

    def visit_biguint_postfix_unary_operation(
        self, expr: awst_nodes.BigUIntPostfixUnaryOperation
    ) -> TExpression:
        target_value = self.visit_and_materialise_single(expr.target)
        rhs = BigUIntConstant(value=1, source_location=expr.source_location)
        match expr.op:
            case awst_nodes.BigUIntPostfixUnaryOperator.increment:
                binary_op = awst_nodes.BigUIntBinaryOperator.add
            case awst_nodes.BigUIntPostfixUnaryOperator.decrement:
                binary_op = awst_nodes.BigUIntBinaryOperator.sub
            case never:
                typing.assert_never(never)

        new_value = create_biguint_binary_op(binary_op, target_value, rhs, expr.source_location)

        handle_assignment(
            self.context,
            target=expr.target,
            value=new_value,
            is_nested_update=False,
            assignment_location=expr.source_location,
        )
        return target_value

    def visit_uint64_postfix_unary_operation(
        self, expr: awst_nodes.UInt64PostfixUnaryOperation
    ) -> TExpression:
        target_value = self.visit_and_materialise_single(expr.target)
        rhs = UInt64Constant(value=1, source_location=expr.source_location)
        match expr.op:
            case awst_nodes.UInt64PostfixUnaryOperator.increment:
                binary_op = awst_nodes.UInt64BinaryOperator.add
            case awst_nodes.UInt64PostfixUnaryOperator.decrement:
                binary_op = awst_nodes.UInt64BinaryOperator.sub
            case never:
                typing.assert_never(never)

        new_value = create_uint64_binary_op(binary_op, target_value, rhs, expr.source_location)

        handle_assignment(
            self.context,
            target=expr.target,
            value=new_value,
            is_nested_update=False,
            assignment_location=expr.source_location,
        )
        return target_value

    def visit_uint64_binary_operation(self, expr: awst_nodes.UInt64BinaryOperation) -> TExpression:
        left = self.visit_and_materialise_single(expr.left)
        right = self.visit_and_materialise_single(expr.right)
        return create_uint64_binary_op(expr.op, left, right, expr.source_location)

    def visit_biguint_binary_operation(
        self, expr: awst_nodes.BigUIntBinaryOperation
    ) -> TExpression:
        left = self.visit_and_materialise_single(expr.left)
        right = self.visit_and_materialise_single(expr.right)
        return create_biguint_binary_op(expr.op, left, right, expr.source_location)

    def visit_uint64_unary_operation(self, expr: awst_nodes.UInt64UnaryOperation) -> TExpression:
        return Intrinsic(
            op=AVMOp(expr.op),
            args=[self.visit_and_materialise_single(expr.expr)],
            source_location=expr.source_location,
        )

    def visit_bytes_unary_operation(self, expr: awst_nodes.BytesUnaryOperation) -> TExpression:
        return Intrinsic(
            op=AVMOp(f"b{expr.op}"),
            args=[self.visit_and_materialise_single(expr.expr)],
            source_location=expr.source_location,
        )

    def visit_integer_constant(self, expr: awst_nodes.IntegerConstant) -> TExpression:
        match expr.wtype:
            case wtypes.uint64_wtype:
                if expr.value < 0 or expr.value.bit_length() > 64:
                    raise CodeError(f"invalid {expr.wtype} value", expr.source_location)

                return UInt64Constant(
                    value=expr.value,
                    source_location=expr.source_location,
                    teal_alias=expr.teal_alias,
                )
            case wtypes.biguint_wtype:
                if expr.value < 0 or expr.value.bit_length() > algo_constants.MAX_BIGUINT_BITS:
                    raise CodeError(f"invalid {expr.wtype} value", expr.source_location)

                return BigUIntConstant(value=expr.value, source_location=expr.source_location)
            case wtypes.ARC4UIntN(n=bit_size):
                num_bytes = bit_size // 8
                try:
                    arc4_result = expr.value.to_bytes(num_bytes, "big", signed=False)
                except OverflowError:
                    raise CodeError(f"invalid {expr.wtype} value", expr.source_location) from None
                return BytesConstant(
                    value=arc4_result,
                    encoding=AVMBytesEncoding.base16,
                    source_location=expr.source_location,
                )
            case _:
                raise InternalError(
                    f"Unhandled wtype {expr.wtype} for integer constant {expr.value}",
                    expr.source_location,
                )

    def visit_decimal_constant(self, expr: awst_nodes.DecimalConstant) -> TExpression:
        match expr.wtype:
            case wtypes.ARC4UFixedNxM(n=bit_size, m=precision):
                num_bytes = bit_size // 8
                sign, digits, exponent = expr.value.as_tuple()
                adjusted_int = int("".join(map(str, digits)))
                if (
                    sign != 0  # negative
                    or not isinstance(exponent, int)  # infinite
                    or -exponent > precision  # too precise
                    or adjusted_int.bit_length() > bit_size  # too big
                ):
                    raise CodeError(f"invalid {expr.wtype} value", expr.source_location)
                return BytesConstant(
                    source_location=expr.source_location,
                    encoding=AVMBytesEncoding.base16,
                    value=adjusted_int.to_bytes(num_bytes, "big", signed=False),
                )
            case _:
                raise InternalError(
                    f"Unhandled wtype {expr.wtype} for decimal constant {expr.value}",
                    expr.source_location,
                )

    def visit_bool_constant(self, expr: awst_nodes.BoolConstant) -> TExpression:
        match expr.wtype:
            case wtypes.bool_wtype:
                return UInt64Constant(
                    value=int(expr.value),
                    ir_type=IRType.bool,
                    source_location=expr.source_location,
                )
            case wtypes.arc4_bool_wtype:
                return BytesConstant(
                    value=(ARC4_TRUE if expr.value else ARC4_FALSE),
                    encoding=AVMBytesEncoding.base16,
                    ir_type=IRType.bytes,
                    source_location=expr.source_location,
                )
            case _:
                raise InternalError(
                    f"Unexpected wtype {expr.wtype} for BoolConstant", expr.source_location
                )

    def visit_bytes_constant(self, expr: awst_nodes.BytesConstant) -> BytesConstant:
        if len(expr.value) > algo_constants.MAX_BYTES_LENGTH:
            raise CodeError(f"invalid {expr.wtype} value", expr.source_location)
        return BytesConstant(
            value=expr.value,
            encoding=bytes_enc_to_avm_bytes_enc(expr.encoding),
            ir_type=wtype_to_ir_type(expr),
            source_location=expr.source_location,
        )

    def visit_string_constant(self, expr: awst_nodes.StringConstant) -> BytesConstant:
        try:
            value = expr.value.encode("utf8")
        except UnicodeError:
            value = None
        if value is None:
            raise CodeError(f"invalid {expr.wtype} value", expr.source_location)

        match expr.wtype:
            case wtypes.string_wtype:
                encoding = AVMBytesEncoding.utf8
            case wtypes.arc4_string_alias:
                encoding = AVMBytesEncoding.base16
                value = len(value).to_bytes(2) + value
            case _:
                raise InternalError(
                    f"Unexpected wtype {expr.wtype} for StringConstant", expr.source_location
                )

        if len(value) > algo_constants.MAX_BYTES_LENGTH:
            raise CodeError(f"invalid {expr.wtype} value", expr.source_location)

        return BytesConstant(
            value=value,
            encoding=encoding,
            source_location=expr.source_location,
        )

    @typing.override
    def visit_void_constant(self, expr: awst_nodes.VoidConstant) -> TExpression:
        return None

    def visit_address_constant(self, expr: awst_nodes.AddressConstant) -> TExpression:
        if not utils.valid_address(expr.value):
            # TODO: should this be here, or on IR model? there's pros and cons to each
            raise CodeError("invalid Algorand address", expr.source_location)
        return AddressConstant(
            value=expr.value,
            source_location=expr.source_location,
        )

    def visit_numeric_comparison_expression(
        self, expr: awst_nodes.NumericComparisonExpression
    ) -> TExpression:
        left = self.visit_and_materialise_single(expr.lhs)
        right = self.visit_and_materialise_single(expr.rhs)
        if left.atype != right.atype:
            raise InternalError(
                "Numeric comparison between different numeric types", expr.source_location
            )
        op_code = expr.operator.value
        if left.atype == AVMType.bytes:
            op_code = "b" + op_code

        try:
            avm_op = AVMOp(op_code)
        except ValueError as ex:
            raise InternalError(
                f"Unmapped numeric comparison operator {expr.operator}", expr.source_location
            ) from ex

        return Intrinsic(
            op=avm_op,
            args=[left, right],
            source_location=expr.source_location,
        )

    def visit_checked_maybe(self, expr: awst_nodes.CheckedMaybe) -> TExpression:
        value, check = self.visit_and_materialise(expr.expr, ("value", "check"))
        assert_value(
            self.context,
            check,
            comment=expr.comment,
            source_location=expr.source_location,
        )
        return value

    def _expand_tuple_var(self, name: str, wtype: wtypes.WTuple) -> Iterator[Value]:
        for idx, wt in enumerate(wtype.types):
            item_name = format_tuple_index(wtype, name, idx)
            if isinstance(wt, wtypes.WTuple):
                yield from self._expand_tuple_var(item_name, wt)
            else:
                yield self.context.ssa.read_variable(
                    variable=item_name,
                    ir_type=wtype_to_ir_type(wt),
                    block=self.context.block_builder.active_block,
                )

    def visit_var_expression(self, expr: awst_nodes.VarExpression) -> TExpression:
        if isinstance(expr.wtype, wtypes.WTuple):
            values = tuple(self._expand_tuple_var(expr.name, expr.wtype))
            return ValueTuple(values=values, source_location=expr.source_location)
        ir_type = wtype_to_ir_type(expr)
        variable = self.context.ssa.read_variable(
            expr.name, ir_type, self.context.block_builder.active_block
        )
        variable = attrs.evolve(variable, source_location=expr.source_location)
        return variable

    def _add_assert(
        self,
        condition_expr: awst_nodes.Expression | None,
        error_message: str | None,
        loc: SourceLocation,
    ) -> Intrinsic | None:
        condition_value = (
            self.visit_and_materialise_single(condition_expr) if condition_expr else None
        )
        if isinstance(condition_value, UInt64Constant):
            if condition_value.value:
                logger.warning("assertion is always true, ignoring", location=loc)
                return None
            else:
                condition_value = None
        if condition_value is None:
            self.context.block_builder.terminate(
                Fail(source_location=loc, error_message=error_message)
            )
            return None
        else:
            return Intrinsic(
                op=AVMOp("assert"),
                source_location=loc,
                args=[condition_value],
                error_message=error_message,
            )

    def visit_intrinsic_call(self, call: awst_nodes.IntrinsicCall) -> TExpression:
        match call.op_code:
            case "err":
                return self._add_assert(
                    condition_expr=None, error_message=None, loc=call.source_location
                )
            case "return":
                assert not call.immediates, f"return intrinsic had immediates: {call.immediates}"
                (arg_expr,) = call.stack_args
                exit_value = self.visit_and_materialise_single(arg_expr)
                self.context.block_builder.terminate(
                    ProgramExit(source_location=call.source_location, result=exit_value)
                )
                return None
            case "assert":
                (condition_expr,) = call.stack_args
                return self._add_assert(
                    condition_expr=condition_expr, error_message=None, loc=call.source_location
                )
            case _:
                args = [self.visit_and_materialise_single(arg) for arg in call.stack_args]
                return Intrinsic(
                    op=AVMOp(call.op_code),
                    source_location=call.source_location,
                    args=args,
                    immediates=list(call.immediates),
                    types=wtype_to_ir_types(call.wtype),
                )

    def visit_group_transaction_reference(
        self, ref: awst_nodes.GroupTransactionReference
    ) -> TExpression:
        index = self.visit_and_materialise_single(ref.index, "gtxn_idx")
        if (txn_type := ref.wtype.transaction_type) is not None:
            actual_type = assign_intrinsic_op(
                self.context,
                target="gtxn_type",
                op=AVMOp.gtxns,
                immediates=["TypeEnum"],
                args=[index],
                source_location=ref.source_location,
            )
            type_constant = UInt64Constant(
                value=txn_type.value, teal_alias=txn_type.name, source_location=ref.source_location
            )
            type_matches = assign_intrinsic_op(
                self.context,
                target="gtxn_type_matches",
                op=AVMOp.eq,
                args=[actual_type, type_constant],
                source_location=ref.source_location,
            )
            assert_value(
                self.context,
                type_matches,
                comment=f"transaction type is {txn_type.name}",
                source_location=ref.source_location,
            )
        return index

    def visit_create_inner_transaction(self, call: awst_nodes.CreateInnerTransaction) -> None:
        # for semantic compatibility, this is an error, since we don't evaluate the args
        # here (there would be no point, if we hit this node on its own and not as part
        # of a submit or an assigment, it does nothing)
        logger.error(
            "statement has no effect, did you forget to submit?", location=call.source_location
        )

    def visit_submit_inner_transaction(
        self, submit: awst_nodes.SubmitInnerTransaction
    ) -> TExpression:
        result = self._itxn.handle_submit_inner_transaction(submit)
        if len(result) == 1:
            return result[0]
        return ValueTuple(
            values=list(result),
            source_location=submit.source_location,
        )

    def visit_update_inner_transaction(self, call: awst_nodes.UpdateInnerTransaction) -> None:
        self._itxn.handle_update_inner_transaction(call)

    def visit_inner_transaction_field(
        self, itxn_field: awst_nodes.InnerTransactionField
    ) -> TExpression:
        return self._itxn.handle_inner_transaction_field(itxn_field)

    def visit_method_constant(self, expr: awst_nodes.MethodConstant) -> TExpression:
        return MethodConstant(value=expr.value, source_location=expr.source_location)

    def visit_tuple_expression(self, expr: awst_nodes.TupleExpression) -> TExpression:
        items = list[Value]()
        for item in expr.items:
            nested_values = self.visit_and_materialise(item)
            items.extend(nested_values)
        return ValueTuple(
            source_location=expr.source_location,
            values=items,
        )

    def visit_tuple_item_expression(self, expr: awst_nodes.TupleItemExpression) -> TExpression:
        if isinstance(expr.base.wtype, wtypes.WTuple):
            tup = self.visit_and_materialise(expr.base)
            return get_tuple_item_values(
                tuple_values=tup,
                tuple_wtype=expr.base.wtype,
                index=expr.index,
                target_wtype=expr.wtype,
                source_location=expr.source_location,
            )
        elif isinstance(expr.base.wtype, wtypes.ARC4Tuple):
            base = self.visit_and_materialise_single(expr.base)
            return arc4.arc4_tuple_index(
                self.context,
                base=base,
                index=expr.index,
                wtype=expr.base.wtype,
                source_location=expr.source_location,
            )
        else:
            raise InternalError(
                f"Tuple indexing operation IR lowering"
                f" not implemented for base type {expr.base.wtype.name}",
                expr.source_location,
            )

    def visit_field_expression(self, expr: awst_nodes.FieldExpression) -> TExpression:
        if isinstance(expr.base.wtype, wtypes.WStructType):
            raise NotImplementedError
        if isinstance(expr.base.wtype, wtypes.WTuple):
            index = expr.base.wtype.name_to_index(expr.name, expr.source_location)
            tup = self.visit_and_materialise(expr.base)
            return get_tuple_item_values(
                tuple_values=tup,
                tuple_wtype=expr.base.wtype,
                index=index,
                target_wtype=expr.wtype,
                source_location=expr.source_location,
            )
        if isinstance(expr.base.wtype, wtypes.ARC4Struct):
            base = self.visit_and_materialise_single(expr.base)
            index = expr.base.wtype.names.index(expr.name)
            return arc4.arc4_tuple_index(
                self.context,
                base=base,
                index=index,
                wtype=expr.base.wtype,
                source_location=expr.source_location,
            )
        else:
            raise InternalError(
                f"Field access IR lowering"
                f" not implemented for base type {expr.base.wtype.name}",
                expr.source_location,
            )

    def visit_intersection_slice_expression(
        self, expr: awst_nodes.IntersectionSliceExpression
    ) -> TExpression:
        if isinstance(expr.base.wtype, wtypes.WTuple):
            return self._visit_tuple_slice(expr, expr.base.wtype)
        elif expr.base.wtype == wtypes.bytes_wtype:
            return visit_bytes_intersection_slice_expression(self.context, expr)
        else:
            raise InternalError(
                f"IntersectionSlice operation IR lowering not implemented for {expr.wtype.name}",
                expr.source_location,
            )

    def visit_slice_expression(self, expr: awst_nodes.SliceExpression) -> TExpression:
        """Slices an enumerable type."""
        if isinstance(expr.base.wtype, wtypes.WTuple):
            return self._visit_tuple_slice(expr, expr.base.wtype)
        elif expr.base.wtype == wtypes.bytes_wtype:
            return visit_bytes_slice_expression(self.context, expr)
        else:
            raise InternalError(
                f"Slice operation IR lowering not implemented for {expr.wtype.name}",
                expr.source_location,
            )

    def _visit_tuple_slice(
        self,
        expr: awst_nodes.SliceExpression | awst_nodes.IntersectionSliceExpression,
        base_wtype: wtypes.WTuple,
    ) -> TExpression:
        tup = self.visit_and_materialise(expr.base)
        start_i = extract_const_int(expr.begin_index) or 0
        end_i = extract_const_int(expr.end_index)
        return get_tuple_item_values(
            tuple_values=tup,
            tuple_wtype=base_wtype,
            index=(start_i, end_i),
            target_wtype=expr.wtype,
            source_location=expr.source_location,
        )

    def visit_index_expression(self, expr: awst_nodes.IndexExpression) -> TExpression:
        index = self.visit_and_materialise_single(expr.index)
        base = self.visit_and_materialise_single(expr.base)

        if expr.base.wtype == wtypes.bytes_wtype:
            # note: the below works because Bytes is immutable, so this index expression
            # can never appear as an assignment target
            if isinstance(index, UInt64Constant) and index.value <= 255:
                return Intrinsic(
                    op=AVMOp.extract,
                    args=[base],
                    immediates=[index.value, 1],
                    source_location=expr.source_location,
                )
            else:
                return Intrinsic(
                    op=AVMOp.extract3,
                    args=[
                        base,
                        index,
                        UInt64Constant(value=1, source_location=expr.source_location),
                    ],
                    source_location=expr.source_location,
                )
        elif isinstance(expr.base.wtype, wtypes.WArray):
            raise NotImplementedError
        elif isinstance(expr.base.wtype, wtypes.ARC4StaticArray | wtypes.ARC4DynamicArray):
            return arc4.arc4_array_index(
                self.context,
                array_wtype=expr.base.wtype,
                array=base,
                index=index,
                source_location=expr.source_location,
            )
        else:
            raise InternalError(
                f"Indexing operation IR lowering not implemented for {expr.wtype.name}",
                expr.source_location,
            )

    def visit_conditional_expression(self, expr: awst_nodes.ConditionalExpression) -> TExpression:
        return flow_control.handle_conditional_expression(self.context, expr)

    def visit_single_evaluation(self, expr: awst_nodes.SingleEvaluation) -> TExpression:
        try:
            return self._single_eval_cache[expr]
        except KeyError:
            pass
        source = expr.source.accept(self)
        if not (source and source.types):
            result: TExpression = None
        else:
            values = self.materialise_value_provider(source, description="awst_tmp")
            if len(values) == 1:
                (result,) = values
            else:
                result = ValueTuple(values=values, source_location=expr.source_location)
        self._single_eval_cache[expr] = result
        return result

    def visit_app_state_expression(self, expr: awst_nodes.AppStateExpression) -> TExpression:
        return storage.visit_app_state_expression(self.context, expr)

    def visit_app_account_state_expression(
        self, expr: awst_nodes.AppAccountStateExpression
    ) -> TExpression:
        return storage.visit_app_account_state_expression(self.context, expr)

    def visit_box_value_expression(self, expr: awst_nodes.BoxValueExpression) -> TExpression:
        return puya.ir.builder.storage.visit_box_value(self.context, expr)

    def visit_state_get_ex(self, expr: awst_nodes.StateGetEx) -> TExpression:
        return storage.visit_state_get_ex(self.context, expr)

    def visit_state_delete(self, statement: awst_nodes.StateDelete) -> TExpression:
        return storage.visit_state_delete(self.context, statement)

    def visit_state_get(self, expr: awst_nodes.StateGet) -> TExpression:
        return storage.visit_state_get(self.context, expr)

    def visit_state_exists(self, expr: awst_nodes.StateExists) -> TExpression:
        return storage.visit_state_exists(self.context, expr)

    def visit_new_array(self, expr: awst_nodes.NewArray) -> TExpression:
        match expr.wtype:
            case wtypes.ARC4Array():
                return arc4.encode_arc4_array(self.context, expr)
            case wtypes.WArray():
                raise NotImplementedError
            case _:
                typing.assert_never(expr.wtype)

    def visit_bytes_comparison_expression(
        self, expr: awst_nodes.BytesComparisonExpression
    ) -> TExpression:
        left = self.visit_and_materialise_single(expr.lhs)
        right = self.visit_and_materialise_single(expr.rhs)
        op_code = expr.operator.value
        try:
            avm_op = AVMOp(op_code)
        except ValueError as ex:
            raise InternalError(
                f"Unmapped bytes comparison operator {expr.operator}", expr.source_location
            ) from ex

        return Intrinsic(
            op=avm_op,
            args=[left, right],
            source_location=expr.source_location,
        )

    def visit_subroutine_call_expression(
        self, expr: awst_nodes.SubroutineCallExpression
    ) -> TExpression:
        return visit_subroutine_call_expression(self.context, expr)

    def visit_puya_lib_call(self, call: awst_nodes.PuyaLibCall) -> TExpression:
        return visit_puya_lib_call_expression(self.context, call)

    def visit_bytes_binary_operation(self, expr: awst_nodes.BytesBinaryOperation) -> TExpression:
        left = self.visit_and_materialise_single(expr.left)
        right = self.visit_and_materialise_single(expr.right)
        return create_bytes_binary_op(expr.op, left, right, expr.source_location)

    def visit_boolean_binary_operation(
        self, expr: awst_nodes.BooleanBinaryOperation
    ) -> TExpression:
        if not isinstance(expr.right, awst_nodes.VarExpression | awst_nodes.BoolConstant):
            true_block, false_block, merge_block = self.context.block_builder.mkblocks(
                "bool_true", "bool_false", "bool_merge", source_location=expr.source_location
            )
            tmp_name = self.context.next_tmp_name(f"{expr.op}_result")

            flow_control.process_conditional(
                self.context, expr, true=true_block, false=false_block, loc=expr.source_location
            )
            self.context.block_builder.activate_block(true_block)
            assign(
                self.context,
                UInt64Constant(value=1, ir_type=IRType.bool, source_location=None),
                name=tmp_name,
                assignment_location=None,
            )
            self.context.block_builder.goto(merge_block)

            self.context.block_builder.activate_block(false_block)
            assign(
                self.context,
                UInt64Constant(value=0, ir_type=IRType.bool, source_location=None),
                name=tmp_name,
                assignment_location=None,
            )
            self.context.block_builder.goto(merge_block)
            self.context.block_builder.activate_block(merge_block)
            return self.context.ssa.read_variable(
                variable=tmp_name, ir_type=IRType.bool, block=merge_block
            )

        left = self.visit_and_materialise_single(expr.left)
        right = self.visit_and_materialise_single(expr.right)
        match expr.op:
            case "and":
                op = AVMOp.and_
            case "or":
                op = AVMOp.or_
            case _:
                raise InternalError(
                    f"Unexpected/unimplemented boolean operator in IR builder: {expr.op}",
                    expr.source_location,
                )
        return Intrinsic(
            op=op,
            args=[left, right],
            source_location=expr.source_location,
        )

    def visit_not_expression(self, expr: awst_nodes.Not) -> TExpression:
        negated = self.visit_and_materialise_single(expr.expr)
        return Intrinsic(
            op=AVMOp("!"),
            args=[negated],
            source_location=expr.source_location,
        )

    def visit_reinterpret_cast(self, expr: awst_nodes.ReinterpretCast) -> TExpression:
        # should be a no-op for us, but we validate the cast here too
        source = self.visit_expr(expr.expr)
        (inner_ir_type,) = source.types
        outer_ir_type = wtype_to_ir_type(expr)
        # don't need to do anything further if ir types are the same
        if inner_ir_type == outer_ir_type:
            return source
        inner_avm_type = inner_ir_type.avm_type
        outer_avm_type = outer_ir_type.avm_type
        if inner_avm_type != outer_avm_type:
            raise InternalError(
                f"Tried to reinterpret {expr.expr.wtype} as {expr.wtype},"
                " but resulting AVM types are incompatible:"
                f" {inner_avm_type} and {outer_avm_type}, respectively",
                expr.source_location,
            )
        target = mktemp(
            self.context,
            outer_ir_type,
            description=f"reinterpret_{outer_ir_type.name}",
            source_location=expr.source_location,
        )
        assign_targets(
            self.context,
            source=source,
            targets=[target],
            assignment_location=expr.source_location,
        )
        return target

    def visit_block(self, block: awst_nodes.Block) -> TStatement:
        if block.label:
            ir_block = self.context.block_builder.mkblock(block)
            self.context.block_builder.goto(ir_block)
            self.context.block_builder.activate_block(ir_block)

        for stmt in block.body:
            stmt.accept(self)

    def visit_goto(self, statement: awst_nodes.Goto) -> TStatement:
        self.context.block_builder.goto_label(statement.target, statement.source_location)

    def visit_if_else(self, stmt: awst_nodes.IfElse) -> TStatement:
        flow_control.handle_if_else(self.context, stmt)

    def visit_switch(self, statement: awst_nodes.Switch) -> TStatement:
        flow_control.handle_switch(self.context, statement)

    def visit_while_loop(self, statement: awst_nodes.WhileLoop) -> TStatement:
        flow_control.handle_while_loop(self.context, statement)

    def visit_loop_exit(self, statement: awst_nodes.LoopExit) -> TStatement:
        self.context.block_builder.loop_break(statement.source_location)

    def visit_return_statement(self, statement: awst_nodes.ReturnStatement) -> TStatement:
        if statement.value is not None:
            result = list(self.visit_and_materialise(statement.value))
        else:
            result = []

        for param in self.context.subroutine.parameters:
            if param.implicit_return:
                result.append(
                    self.context.ssa.read_variable(
                        param.name,
                        param.ir_type,
                        self.context.block_builder.active_block,
                    )
                )
        return_types = [r.ir_type for r in result]
        if return_types != self.context.subroutine.returns:
            raise CodeError(
                f"invalid return type {return_types}, expected {self.context.subroutine.returns}",
                statement.source_location,
            )
        self.context.block_builder.terminate(
            SubroutineReturn(
                source_location=statement.source_location,
                result=result,
            )
        )

    def visit_assert_expression(self, expr: awst_nodes.AssertExpression) -> TStatement:
        op = self._add_assert(
            condition_expr=expr.condition,
            error_message=expr.error_message,
            loc=expr.source_location,
        )
        if op:
            self.context.block_builder.add(op)

    def visit_template_var(self, expr: awst_nodes.TemplateVar) -> TExpression:
        return TemplateVar(
            name=expr.name,
            ir_type=wtype_to_ir_type(expr.wtype),
            source_location=expr.source_location,
        )

    def visit_loop_continue(self, statement: awst_nodes.LoopContinue) -> TStatement:
        self.context.block_builder.loop_continue(statement.source_location)

    def visit_expression_statement(self, statement: awst_nodes.ExpressionStatement) -> TStatement:
        # NOTE: popping of ignored return values should happen at code gen time
        result = self._visit_and_check_for_double_eval(statement.expr)
        if result is None:
            wtype = statement.expr.wtype
            match wtype:
                case wtypes.void_wtype:
                    pass
                case _ if (isinstance(wtype, WInnerTransaction | WInnerTransactionFields)):
                    # inner transaction wtypes aren't true expressions
                    pass
                case _:
                    raise InternalError(
                        f"Expression statement with type {statement.expr.wtype} "
                        f"generated no result",
                        statement.source_location,
                    )
        elif isinstance(result, Op):
            self.context.block_builder.add(result)
        # If we get a Value (e.g. a Register or some such) it's something that's being
        # discarded effectively.
        # The frontend should have already warned about this

    def visit_uint64_augmented_assignment(
        self, statement: awst_nodes.UInt64AugmentedAssignment
    ) -> TStatement:
        target_value = self.visit_and_materialise_single(statement.target)
        rhs = self.visit_and_materialise_single(statement.value)
        expr = create_uint64_binary_op(statement.op, target_value, rhs, statement.source_location)

        handle_assignment(
            self.context,
            target=statement.target,
            value=expr,
            is_nested_update=False,
            assignment_location=statement.source_location,
        )

    def visit_biguint_augmented_assignment(
        self, statement: awst_nodes.BigUIntAugmentedAssignment
    ) -> TStatement:
        target_value = self.visit_and_materialise_single(statement.target)
        rhs = self.visit_and_materialise_single(statement.value)
        expr = create_biguint_binary_op(statement.op, target_value, rhs, statement.source_location)

        handle_assignment(
            self.context,
            target=statement.target,
            value=expr,
            is_nested_update=False,
            assignment_location=statement.source_location,
        )

    def visit_bytes_augmented_assignment(
        self, statement: awst_nodes.BytesAugmentedAssignment
    ) -> TStatement:
        if statement.target.wtype == wtypes.arc4_string_alias:
            value: ValueProvider = arc4.concat_values(
                self.context, statement.target, statement.value, statement.source_location
            )
        else:
            target_value = self.visit_and_materialise_single(statement.target)
            rhs = self.visit_and_materialise_single(statement.value)
            value = create_bytes_binary_op(
                statement.op, target_value, rhs, statement.source_location
            )

        handle_assignment(
            self.context,
            target=statement.target,
            value=value,
            is_nested_update=False,
            assignment_location=statement.source_location,
        )

    def visit_enumeration(self, expr: awst_nodes.Enumeration) -> TStatement:
        raise CodeError("Nested enumeration is not currently supported", expr.source_location)

    def visit_reversed(self, expr: awst_nodes.Reversed) -> TExpression:
        raise CodeError("Reversed is not valid outside of an enumeration", expr.source_location)

    def visit_for_in_loop(self, statement: awst_nodes.ForInLoop) -> TStatement:
        handle_for_in_loop(self.context, statement)

    def visit_new_struct(self, expr: awst_nodes.NewStruct) -> TExpression:
        match expr.wtype:
            case wtypes.WStructType():
                raise NotImplementedError
            case wtypes.ARC4Struct() as arc4_struct_wtype:
                return arc4.encode_arc4_struct(self.context, expr, arc4_struct_wtype)
            case _:
                typing.assert_never(expr.wtype)

    def visit_array_pop(self, expr: awst_nodes.ArrayPop) -> TExpression:
        source_location = expr.source_location
        match expr.base.wtype:
            case wtypes.ARC4DynamicArray() as array_wtype:
                return arc4.pop_arc4_array(self.context, expr, array_wtype)
            case _:
                raise InternalError(
                    f"Unsupported target for array pop: {expr.base.wtype}", source_location
                )

    def visit_array_concat(self, expr: awst_nodes.ArrayConcat) -> TExpression:
        return arc4.concat_values(
            self.context,
            left_expr=expr.left,
            right_expr=expr.right,
            source_location=expr.source_location,
        )

    def visit_array_extend(self, expr: awst_nodes.ArrayExtend) -> TExpression:
        concat_result = arc4.concat_values(
            self.context,
            left_expr=expr.base,
            right_expr=expr.other,
            source_location=expr.source_location,
        )
        return arc4.handle_arc4_assign(
            self.context,
            target=expr.base,
            value=concat_result,
            is_nested_update=True,
            source_location=expr.source_location,
        )

    def visit_arc4_router(self, expr: awst_nodes.ARC4Router) -> TExpression:
        root = self.context.root
        if not isinstance(root, awst_nodes.Contract):
            raise CodeError(
                "cannot create ARC4 router outside of a contract", expr.source_location
            )

        return InvokeSubroutine(
            target=self.context.routers[root.id],
            args=[],
            source_location=expr.source_location,
        )

    def visit_emit(self, expr: awst_nodes.Emit) -> TExpression:
        factory = OpFactory(self.context, expr.source_location)
        value = self.context.visitor.visit_and_materialise_single(expr.value)
        prefix = MethodConstant(value=expr.signature, source_location=expr.source_location)
        event = factory.concat(prefix, value, "event")

        self.context.block_builder.add(
            Intrinsic(
                op=AVMOp("log"),
                args=[event],
                source_location=expr.source_location,
            )
        )
        return None

    def visit_range(self, node: awst_nodes.Range) -> TExpression:
        raise CodeError("unexpected range location", node.source_location)

    def visit_and_materialise_single(
        self, expr: awst_nodes.Expression, temp_description: str = "tmp"
    ) -> Value:
        """Translate an AWST Expression into a single Value"""
        values = self.visit_and_materialise(expr, temp_description=temp_description)
        try:
            (value,) = values
        except ValueError as ex:
            raise InternalError(
                "visit_and_materialise_single should not be used when"
                f" an expression could be multi-valued, expression was: {expr}",
                expr.source_location,
            ) from ex
        return value

    def visit_and_materialise(
        self, expr: awst_nodes.Expression, temp_description: str | Sequence[str] = "tmp"
    ) -> Sequence[Value]:
        value_seq_or_provider = self._visit_and_check_for_double_eval(
            expr, materialise_as=temp_description
        )
        if value_seq_or_provider is None:
            raise InternalError(
                "No value produced by expression IR conversion", expr.source_location
            )
        return self.materialise_value_provider(value_seq_or_provider, description=temp_description)

    def visit_expr(self, expr: awst_nodes.Expression) -> ValueProvider:
        """Visit the expression and ensure result is not None"""
        value_seq_or_provider = self._visit_and_check_for_double_eval(expr)
        if value_seq_or_provider is None:
            raise InternalError(
                "No value produced by expression IR conversion", expr.source_location
            )
        return value_seq_or_provider

    def _visit_and_check_for_double_eval(
        self, expr: awst_nodes.Expression, *, materialise_as: str | Sequence[str] | None = None
    ) -> ValueProvider | None:
        # explicit SingleEvaluation nodes already handle this
        if isinstance(expr, awst_nodes.SingleEvaluation):
            return expr.accept(self)
        # include the expression in the key to ensure the lifetime of the
        # expression is as long as the cache.
        # Temporary nodes may end up with the same id if nothing is referencing them
        # e.g. such as used in _update_implicit_out_var
        expr_id = id(expr), expr
        try:
            result = self._visited_exprs[expr_id]
        except KeyError:
            pass
        else:
            if isinstance(result, ValueProvider) and not isinstance(result, ValueTuple | Value):
                raise InternalError(
                    "double evaluation of expression without materialization", expr.source_location
                )
            expr_str = expr.accept(ToCodeVisitor())
            logger.debug(
                f"encountered already materialized expression ({expr_str}),"
                f" reusing result: {result!s}",
                location=expr.source_location,
            )
            return result
        source = expr.accept(self)
        if materialise_as is None or not (source and source.types):
            result = source
        else:
            values = self.materialise_value_provider(source, description=materialise_as)
            if len(values) == 1:
                (result,) = values
            else:
                result = ValueTuple(values=values, source_location=expr.source_location)
        self._visited_exprs[expr_id] = result
        return result

    def materialise_value_provider(
        self, provider: ValueProvider, description: str | Sequence[str]
    ) -> list[Value]:
        """
        Given a ValueProvider with arity of N, return a Value sequence of length N.

        Anything which is already a Value is passed through without change.

        Otherwise, the result is assigned to a temporary register, which is returned
        """
        if isinstance(provider, Value):
            return [provider]

        if isinstance(provider, ValueTuple):
            return list(provider.values)

        ir_types = provider.types
        if not ir_types:
            raise InternalError(
                "Attempted to assign from expression that has no result", provider.source_location
            )

        if isinstance(description, str):
            temp_description: Sequence[str] = [description] * len(ir_types)
        else:
            temp_description = description
        targets = [
            mktemp(self.context, ir_type, provider.source_location, description=descr)
            for ir_type, descr in zip(ir_types, temp_description, strict=True)
        ]
        assign_targets(
            context=self.context,
            source=provider,
            targets=targets,
            # TODO: should this be the source location of the site forcing materialisation?
            assignment_location=provider.source_location,
        )
        return list[Value](targets)

def create_uint64_binary_op(
    op: UInt64BinaryOperator, left: Value, right: Value, source_location: SourceLocation
) -> Intrinsic:
    avm_op: AVMOp
    match op:
        case UInt64BinaryOperator.floor_div:
            avm_op = AVMOp.div_floor
        case UInt64BinaryOperator.pow:
            avm_op = AVMOp.exp
        case UInt64BinaryOperator.lshift:
            avm_op = AVMOp.shl
        case UInt64BinaryOperator.rshift:
            avm_op = AVMOp.shr
        case _:
            try:
                avm_op = AVMOp(op.value)
            except ValueError as ex:
                raise InternalError(
                    f"Unhandled uint64 binary operator: {op}", source_location
                ) from ex
    return Intrinsic(op=avm_op, args=[left, right], source_location=source_location)

def create_biguint_binary_op(
    op: BigUIntBinaryOperator, left: Value, right: Value, source_location: SourceLocation
) -> Intrinsic:
    avm_op: AVMOp
    match op:
        case BigUIntBinaryOperator.floor_div:
            avm_op = AVMOp.div_floor_bytes
        case _:
            try:
                avm_op = AVMOp("b" + op.value)
            except ValueError as ex:
                raise InternalError(
                    f"Unhandled uint64 binary operator: {op}", source_location
                ) from ex
    return Intrinsic(
        op=avm_op, args=[left, right], types=(IRType.biguint,), source_location=source_location
    )

def create_bytes_binary_op(
    op: awst_nodes.BytesBinaryOperator, lhs: Value, rhs: Value, source_location: SourceLocation
) -> ValueProvider:
    match op:
        case awst_nodes.BytesBinaryOperator.add:
            return Intrinsic(
                op=AVMOp.concat,
                args=[lhs, rhs],
                source_location=source_location,
            )
        case awst_nodes.BytesBinaryOperator.bit_and:
            return Intrinsic(
                op=AVMOp.bitwise_and_bytes,
                args=[lhs, rhs],
                source_location=source_location,
            )
        case awst_nodes.BytesBinaryOperator.bit_or:
            return Intrinsic(
                op=AVMOp.bitwise_or_bytes,
                args=[lhs, rhs],
                source_location=source_location,
            )
        case awst_nodes.BytesBinaryOperator.bit_xor:
            return Intrinsic(
                op=AVMOp.bitwise_xor_bytes,
                args=[lhs, rhs],
                source_location=source_location,
            )
    raise InternalError("Unsupported BytesBinaryOperator: " + op)

def get_comparison_op_for_wtype(
    numeric_comparison_equivalent: awst_nodes.NumericComparison, wtype: wtypes.WType
) -> AVMOp:
    match wtype:
        case wtypes.biguint_wtype:
            return AVMOp("b" + numeric_comparison_equivalent)
        case wtypes.uint64_wtype:
            return AVMOp(numeric_comparison_equivalent)
        case wtypes.bytes_wtype:
            match numeric_comparison_equivalent:
                case awst_nodes.NumericComparison.eq:
                    return AVMOp.eq
                case awst_nodes.NumericComparison.ne:
                    return AVMOp.neq
    raise InternalError(
        f"unsupported operation of {numeric_comparison_equivalent} on type of {wtype}"
    )


import typing

from puya.avm import AVMType
from puya.awst import (
    nodes as awst_nodes,
    wtypes,
)
from puya.ir import intrinsic_factory
from puya.ir.avm_ops import AVMOp
from puya.ir.builder._utils import assert_value, assign_targets, mktemp
from puya.ir.context import IRFunctionBuildContext
from puya.ir.models import Intrinsic, UInt64Constant, Value, ValueProvider, ValueTuple
from puya.ir.types_ import IRType, wtype_to_ir_type
from puya.parse import SourceLocation

def visit_app_state_expression(
    context: IRFunctionBuildContext, expr: awst_nodes.AppStateExpression
) -> ValueProvider:
    maybe_value, exists = _build_state_get_ex(context, expr, expr.source_location)
    # TODO: add specific (unsafe) optimisation flag to allow skipping this check
    assert_value(
        context,
        value=exists,
        comment=expr.exists_assertion_message or "state exists",
        source_location=expr.source_location,
    )
    return maybe_value

def visit_app_account_state_expression(
    context: IRFunctionBuildContext, expr: awst_nodes.AppAccountStateExpression
) -> ValueProvider:
    maybe_value, exists = _build_state_get_ex(context, expr, expr.source_location)
    # TODO: add specific (unsafe) optimisation flag to allow skipping this check
    assert_value(
        context,
        value=exists,
        comment=expr.exists_assertion_message or "state exists for account",
        source_location=expr.source_location,
    )
    return maybe_value

def visit_box_value(
    context: IRFunctionBuildContext, expr: awst_nodes.BoxValueExpression
) -> ValueProvider:
    maybe_value, exists = _build_state_get_ex(context, expr, expr.source_location)
    # TODO: add specific (unsafe) optimisation flag to allow skipping this check
    assert_value(
        context,
        value=exists,
        comment=expr.exists_assertion_message or "box exists",
        source_location=expr.source_location,
    )
    return maybe_value

def visit_state_exists(
    context: IRFunctionBuildContext, expr: awst_nodes.StateExists
) -> ValueProvider:
    _, exists = _build_state_get_ex(
        context, expr.field, expr.source_location, for_existence_check=True
    )
    return exists

def visit_state_get(context: IRFunctionBuildContext, expr: awst_nodes.StateGet) -> ValueProvider:
    default = context.visitor.visit_and_materialise_single(expr.default)
    maybe_value, exists = _build_state_get_ex(context, expr.field, expr.source_location)
    return intrinsic_factory.select(
        condition=exists,
        true=maybe_value,
        false=default,
        type_=wtype_to_ir_type(expr.wtype),
        source_location=expr.source_location,
    )

def visit_state_get_ex(
    context: IRFunctionBuildContext, expr: awst_nodes.StateGetEx
) -> ValueProvider:
    return ValueTuple(
        values=list(_build_state_get_ex(context, expr.field, expr.source_location)),
        source_location=expr.source_location,
    )

def visit_state_delete(
    context: IRFunctionBuildContext, statement: awst_nodes.StateDelete
) -> ValueProvider | None:
    match statement.field:
        case awst_nodes.BoxValueExpression(key=awst_key):
            op = AVMOp.box_del
            awst_account = None
        case awst_nodes.AppStateExpression(key=awst_key):
            op = AVMOp.app_global_del
            awst_account = None
        case awst_nodes.AppAccountStateExpression(key=awst_key, account=awst_account):
            op = AVMOp.app_local_del
        case _:
            typing.assert_never(statement.field)

    args = []
    if awst_account is not None:
        account_value = context.visitor.visit_and_materialise_single(awst_account)
        args.append(account_value)
    key_value = context.visitor.visit_and_materialise_single(awst_key)
    args.append(key_value)

    context.block_builder.add(
        Intrinsic(op=op, args=args, source_location=statement.source_location)
    )
    return None

def _build_state_get_ex(
    context: IRFunctionBuildContext,
    expr: (
        awst_nodes.AppAccountStateExpression
        | awst_nodes.AppStateExpression
        | awst_nodes.BoxValueExpression
    ),
    source_location: SourceLocation,
    *,
    for_existence_check: bool = False,
) -> tuple[Value, Value]:
    key = context.visitor.visit_and_materialise_single(expr.key)
    args: list[Value]
    true_value_ir_type = get_ex_value_ir_type = wtype_to_ir_type(expr.wtype)
    convert_op: AVMOp | None = None
    if isinstance(expr, awst_nodes.AppStateExpression):
        current_app_offset = UInt64Constant(value=0, source_location=expr.source_location)
        args = [current_app_offset, key]
        op = AVMOp.app_global_get_ex
    elif isinstance(expr, awst_nodes.AppAccountStateExpression):
        current_app_offset = UInt64Constant(value=0, source_location=expr.source_location)
        op = AVMOp.app_local_get_ex
        account = context.visitor.visit_and_materialise_single(expr.account)
        args = [account, current_app_offset, key]
    else:
        args = [key]
        if for_existence_check:
            get_ex_value_ir_type = IRType.uint64
            op = AVMOp.box_len
        else:
            op = AVMOp.box_get
            match wtypes.persistable_stack_type(expr.wtype, source_location):
                case AVMType.uint64:
                    get_ex_value_ir_type = IRType.bytes
                    convert_op = AVMOp.btoi
                case AVMType.bytes:
                    pass
                case invalid:
                    typing.assert_never(invalid)
    get_ex = Intrinsic(
        op=op,
        args=args,
        types=[get_ex_value_ir_type, IRType.bool],
        source_location=source_location,
    )
    value_tmp, did_exist_tmp = context.visitor.materialise_value_provider(
        get_ex, ("maybe_value", "maybe_exists")
    )
    if convert_op is None:
        return value_tmp, did_exist_tmp
    convert = Intrinsic(op=convert_op, args=[value_tmp], source_location=source_location)
    value_tmp_converted = mktemp(
        context,
        ir_type=true_value_ir_type,
        description="maybe_value_converted",
        source_location=expr.source_location,
    )

    assign_targets(
        context, source=convert, targets=[value_tmp_converted], assignment_location=source_location
    )
    return value_tmp_converted, did_exist_tmp


import contextlib
import itertools
import typing
from collections import defaultdict
from collections.abc import Iterator, Mapping
from functools import cached_property

import attrs

import puya.awst.nodes as awst_nodes
from puya.context import CompileContext
from puya.errors import CodeError, log_exceptions
from puya.ir.builder.blocks import BlocksBuilder
from puya.ir.models import Subroutine
from puya.ir.ssa import BraunSSA
from puya.parse import SourceLocation
from puya.program_refs import ContractReference
from puya.utils import attrs_extend

if typing.TYPE_CHECKING:
    from puya.ir.builder.main import FunctionIRBuilder

TMP_VAR_INDICATOR = "%"

@attrs.frozen(kw_only=True)
class IRBuildContext(CompileContext):
    awst: awst_nodes.AWST
    subroutines: dict[awst_nodes.Function, Subroutine]
    embedded_funcs_lookup: Mapping[str, Subroutine]
    root: awst_nodes.Contract | awst_nodes.LogicSignature | None = None
    routers: dict[ContractReference, Subroutine] = attrs.field(factory=dict)

    @cached_property
    def _awst_lookup(self) -> Mapping[str, awst_nodes.RootNode]:
        return {node.id: node for node in self.awst}

    def for_root(self, root: awst_nodes.Contract | awst_nodes.LogicSignature) -> typing.Self:
        return attrs.evolve(
            self,
            root=root,
            # copy subroutines so that contract specific subroutines do not pollute other passes
            subroutines=self.subroutines.copy(),
        )

    def for_function(
        self, function: awst_nodes.Function, subroutine: Subroutine, visitor: "FunctionIRBuilder"
    ) -> "IRFunctionBuildContext":
        return attrs_extend(
            IRFunctionBuildContext, self, visitor=visitor, function=function, subroutine=subroutine
        )

    def resolve_function_reference(
        self,
        target: awst_nodes.SubroutineTarget,
        source_location: SourceLocation,
        caller: awst_nodes.Function,
    ) -> awst_nodes.Function:
        if isinstance(target, awst_nodes.SubroutineID):
            func: awst_nodes.Node | None = self._awst_lookup.get(target.target)
        else:
            contract = self.root
            if not (
                isinstance(contract, awst_nodes.Contract)
                and isinstance(caller, awst_nodes.ContractMethod)
            ):
                raise CodeError(
                    "call to contract method from outside of contract class",
                    source_location,
                )
            match target:
                case awst_nodes.InstanceMethodTarget(member_name=member_name):
                    func = contract.resolve_contract_method(member_name)
                case awst_nodes.ContractMethodTarget(cref=start_at, member_name=member_name):
                    func = contract.resolve_contract_method(
                        member_name, source_location, start=start_at
                    )
                case awst_nodes.InstanceSuperMethodTarget(member_name=member_name):
                    func = contract.resolve_contract_method(
                        member_name, source_location, start=caller.cref, skip=True
                    )
                case unexpected:
                    typing.assert_never(unexpected)
        if func is None:
            raise CodeError(f"unable to resolve function reference {target}", source_location)
        if not isinstance(func, awst_nodes.Function):
            raise CodeError(
                f"function reference {target} resolved to non-function {func}", source_location
            )
        return func

    @property
    def default_fallback(self) -> SourceLocation | None:
        if self.root:
            return self.root.source_location
        return None

    @contextlib.contextmanager
    def log_exceptions(self, fallback_location: SourceLocation | None = None) -> Iterator[None]:
        with log_exceptions(fallback_location or self.default_fallback):
            yield

@attrs.frozen(kw_only=True)
class IRFunctionBuildContext(IRBuildContext):
    """Context when building from an awst Function node"""

    function: awst_nodes.Function
    subroutine: Subroutine
    visitor: "FunctionIRBuilder"
    block_builder: BlocksBuilder = attrs.field()
    _tmp_counters: defaultdict[str, Iterator[int]] = attrs.field(
        factory=lambda: defaultdict(itertools.count)
    )

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

    @block_builder.default
    def _block_builder_factory(self) -> BlocksBuilder:
        return BlocksBuilder(self.subroutine.parameters, self.function.source_location)

    def next_tmp_name(self, description: str) -> str:
        counter_value = next(self._tmp_counters[description])
        return f"{description}{TMP_VAR_INDICATOR}{counter_value}"

    @property
    def default_fallback(self) -> SourceLocation | None:
        return self.function.source_location

    def resolve_subroutine(
        self,
        target: awst_nodes.SubroutineTarget,
        source_location: SourceLocation,
        *,
        caller: awst_nodes.Function | None = None,
    ) -> Subroutine:
        func = self.resolve_function_reference(
            target=target, source_location=source_location, caller=caller or self.function
        )
        return self.subroutines[func]


import itertools
import typing
import typing as t
from collections.abc import Iterable

import attrs

from puya import log
from puya.avm import AVMType
from puya.ir import models
from puya.ir.types_ import IRType
from puya.ir.visitor_mem_replacer import MemoryReplacer
from puya.ir.vla import VariableLifetimeAnalysis
from puya.options import LocalsCoalescingStrategy
from puya.utils import StableSet

logger = log.get_logger(__name__)

@attrs.define
class MemoryReplacerWithRedundantAssignmentRemoval(MemoryReplacer):
    def visit_assignment(self, op: models.Assignment) -> models.Assignment | None:
        ass = super().visit_assignment(op)
        if ass is None or ass.targets == (ass.source,):
            return None
        return ass

class CoalesceGroupStrategy(t.Protocol):
    def get_group_key(self, reg: models.Register) -> object: ...

    def determine_group_replacement(self, regs: Iterable[models.Register]) -> models.Register: ...

def coalesce_registers(group_strategy: CoalesceGroupStrategy, sub: models.Subroutine) -> int:
    """
    A local can be merged with another local if they are never live at the same time.

    For each local that is being defined, check to see what the live-out locals are.

    It can be merged with another local set if:
     - This local is not in the "live-out" of any local in set
     - The "live-out" of this local does not intersect the local set
    """
    vla = VariableLifetimeAnalysis.analyze(sub)

    # TODO: this uses a basic definition of interference by looking at live-ranges,
    #       a better option is to continue with https://inria.hal.science/inria-00349925v1/document
    #       which has already been partially implemented (phase 1 + 4 have been, anyway)

    variables_live_at_definition = dict[models.Register, StableSet[models.Register]]()
    for param in sub.parameters:
        variables_live_at_definition[param] = StableSet.from_iter(sub.parameters)
    for block in sub.body:
        for op in block.ops:
            match op:
                case models.Assignment(targets=targets):
                    op_live_out = vla.get_live_out_variables(op)
                    for defined_reg in targets:
                        live_set = variables_live_at_definition.setdefault(
                            defined_reg, StableSet()
                        )
                        live_set |= op_live_out

    coalescable_groups_by_key = dict[
        object, list[tuple[StableSet[models.Register], StableSet[models.Register]]]
    ]()
    for defined_reg, live_set in variables_live_at_definition.items():
        coalescable_groups = coalescable_groups_by_key.setdefault(
            group_strategy.get_group_key(defined_reg), []
        )
        for coalescable_register_set, combined_live_out in coalescable_groups:
            # conditions:
            # 1) this register/variable must not be "alive" _after_ the
            #    definition of any other variable in this set
            # 2) no register already in the set should be live out at the
            #    definition of this register
            # regardless of the order the definitions are processed in, this guarantees that:
            # for all A and B in coalescable_register_set such that A != B:
            #       A is not live-out whenever B is assigned
            #   AND B is not live-out whenever A is assigned
            if defined_reg not in combined_live_out and live_set.isdisjoint(
                coalescable_register_set
            ):
                coalescable_register_set.add(defined_reg)
                combined_live_out |= live_set
                break
        else:
            coalescable_groups.append((StableSet(defined_reg), StableSet.from_iter(live_set)))

    replacements = dict[models.Register, models.Register]()
    for group in coalescable_groups_by_key.values():
        for coalescable_register_set, _ in group:
            if len(coalescable_register_set) < 2:
                continue
            replacement = group_strategy.determine_group_replacement(coalescable_register_set)
            find = coalescable_register_set - {replacement}

            logger.debug(f"Coalescing {replacement} with [{', '.join(map(str, find))}]")
            replacements.update({to_find: replacement for to_find in find})
    total_replacements = MemoryReplacerWithRedundantAssignmentRemoval.apply(
        sub.body, replacements=replacements
    )
    return total_replacements

class RootOperandGrouping(CoalesceGroupStrategy):
    def __init__(self, isolate: frozenset[models.Register] | None = None) -> None:
        self._isolate = isolate or frozenset()

    def get_group_key(self, reg: models.Register) -> object:
        if reg in self._isolate:
            return reg
        return reg.name, reg.ir_type

    def determine_group_replacement(self, regs: Iterable[models.Register]) -> models.Register:
        return min(regs, key=lambda r: r.version)

class AggressiveGrouping(CoalesceGroupStrategy):
    def __init__(self, sub: models.Subroutine) -> None:
        self._params = frozenset(sub.parameters)
        self._counter = itertools.count()

    def get_group_key(self, reg: models.Register) -> object:
        if reg in self._params:
            return reg
        else:
            return reg.atype

    def determine_group_replacement(self, regs: Iterable[models.Register]) -> models.Register:
        next_id = next(self._counter)
        (atype,) = {r.atype for r in regs}
        match atype:
            case AVMType.uint64:
                ir_type = IRType.uint64
            case AVMType.bytes:
                ir_type = IRType.bytes
            case _:
                typing.assert_never(atype)
        return models.Register(
            name="",
            version=next_id,
            ir_type=ir_type,
            source_location=None,
        )

def coalesce_locals(subroutine: models.Subroutine, strategy: LocalsCoalescingStrategy) -> None:
    match strategy:
        case LocalsCoalescingStrategy.root_operand:
            group_strategy: CoalesceGroupStrategy = RootOperandGrouping()
        case LocalsCoalescingStrategy.root_operand_excluding_args:
            group_strategy = RootOperandGrouping(isolate=frozenset(subroutine.parameters))
        case LocalsCoalescingStrategy.aggressive:
            group_strategy = AggressiveGrouping(subroutine)
    logger.debug(
        f"Coalescing local variables in {subroutine.id}"
        f" using strategy {type(group_strategy).__name__}"
    )
    replacements = coalesce_registers(group_strategy, subroutine)
    logger.debug(f"Coalescing resulted in {replacements} replacement/s")


import attrs

from puya import log
from puya.context import CompileContext
from puya.ir import models
from puya.ir.destructure.coalesce_locals import coalesce_locals
from puya.ir.destructure.optimize import post_ssa_optimizer
from puya.ir.destructure.parcopy import sequentialize_parallel_copies
from puya.ir.destructure.remove_phi import convert_to_cssa, destructure_cssa

logger = log.get_logger(__name__)

def destructure_ssa(context: CompileContext, program: models.Program) -> None:
    for subroutine in program.all_subroutines:
        logger.debug(f"Performing SSA IR destructuring for {subroutine.id}")
        convert_to_cssa(subroutine)
        subroutine.validate_with_ssa()
        destructure_cssa(subroutine)
        coalesce_locals(subroutine, context.options.locals_coalescing_strategy)
        sequentialize_parallel_copies(subroutine)
        post_ssa_optimizer(subroutine, context.options.optimization_level)
        attrs.validate(subroutine)


import contextlib
import itertools

from puya import log
from puya.ir import models
from puya.ir.optimize.collapse_blocks import BlockReferenceReplacer
from puya.utils import unique

logger = log.get_logger(__name__)

def post_ssa_optimizer(sub: models.Subroutine, optimization_level: int) -> None:
    logger.debug(f"Performing post-SSA optimizations at level {optimization_level}")
    if optimization_level >= 1:
        _remove_linear_jumps(sub)
    if optimization_level >= 2:
        _block_deduplication(sub)

def _remove_linear_jumps(subroutine: models.Subroutine) -> None:
    #  P = {p0, p1, ..., pn} -> {j} -> {t}
    # remove {j} from subroutine
    # point P at t:
    #  update references within P from j to t
    #  ensure P are all in predecessors of t
    # This exists here and not in main IR optimization loop because we only want to do it for
    # blocks that are _truly_ empty, not ones that contain phi-node magic that results in copies
    # build a map of any blocks that are just an unconditional branch to their targets
    jumps = dict[models.BasicBlock, models.BasicBlock]()
    for block in subroutine.body.copy():
        match block:
            case models.BasicBlock(
                ops=[], terminator=models.Goto(target=target)
            ) if target is not block:
                jumps[block] = target
                logger.debug(f"Removing jump block {block}")
                with contextlib.suppress(ValueError):
                    target.predecessors.remove(block)
                subroutine.body.remove(block)

    # now back-propagate any chains
    replacements = dict[models.BasicBlock, models.BasicBlock]()
    for src, target in jumps.items():
        while True:
            try:
                target = jumps[target]
            except KeyError:
                break
        logger.debug(f"branching to {src} will be replaced with {target}")
        replacements[src] = target
        BlockReferenceReplacer.apply(find=src, replacement=target, blocks=subroutine.body)
        for pred in src.predecessors:
            if pred not in target.predecessors:
                target.predecessors.append(pred)

def _block_deduplication(subroutine: models.Subroutine) -> None:
    seen = dict[tuple[object, ...], models.BasicBlock]()
    for block in subroutine.body.copy():
        all_ops = tuple(op.freeze() for op in block.all_ops)
        if existing := seen.get(all_ops):
            logger.debug(
                f"Removing duplicated block {block} and updating references to {existing}"
            )
            BlockReferenceReplacer.apply(find=block, replacement=existing, blocks=subroutine.body)
            subroutine.body.remove(block)
            existing.predecessors = unique(
                itertools.chain(existing.predecessors, block.predecessors)
            )
        else:
            seen[all_ops] = block


import itertools
from collections.abc import Callable, Iterable

from puya import log
from puya.ir import models
from puya.ir.context import TMP_VAR_INDICATOR

logger = log.get_logger(__name__)

def sequentialize_parallel_copies(sub: models.Subroutine) -> None:
    logger.debug(f"Sequentializing parallel copies in {sub.id}")
    our_tmp_prefix = f"parcopy{TMP_VAR_INDICATOR}"
    max_tmp_id = max(
        (
            int(r.name.split(TMP_VAR_INDICATOR)[1])
            for r in sub.get_assigned_registers()
            if r.name.startswith(our_tmp_prefix)
        ),
        default=-1,
    )
    next_tmp_id = itertools.count(max_tmp_id + 1)

    def make_temp(x: models.Value | models.Register) -> models.Register:
        return models.Register(
            ir_type=x.ir_type,
            name=f"{our_tmp_prefix}{next(next_tmp_id)}",
            version=0,
            source_location=x.source_location,
        )

    for block in sub.body:
        ops = list[models.Op]()
        for op in block.ops:
            match op:
                case models.Assignment(targets=targets, source=models.ValueTuple(values=sources)):
                    seqd = _sequentialize(zip(targets, sources, strict=True), mktmp=make_temp)
                    for dst, src in seqd:
                        assert isinstance(dst, models.Register)  # TODO: this is bad
                        ops.append(
                            models.Assignment(
                                targets=[dst],
                                source=src,
                                source_location=op.source_location,
                            )
                        )
                case _:
                    ops.append(op)
        block.ops = ops

def _sequentialize[T](
    copies: Iterable[tuple[T, T]],
    mktmp: Callable[[T], T],
    *,
    filter_dup_dests: bool = True,
    allow_fan_out: bool = True,
) -> list[tuple[T, T]]:
    # If filter_dup_dests is True, consider pairs ordered, and if multiple
    # pairs have the same dest var, the last one takes effect. Otherwise,
    # such duplicate dest vars is an error.
    if filter_dup_dests:
        # If there're multiple assignments to the same var, keep only the latest
        copies = list(dict(copies).items())

    ready = []
    to_do = []
    pred = {}
    loc = dict[T, T | None]()
    res = []

    for b, _ in copies:
        loc[b] = None

    for b, a in copies:
        loc[a] = a
        pred[b] = a

        # Extra check
        if not filter_dup_dests and b in to_do:
            raise ValueError(f"Conflicting assignments to destination {b}, latest: {(b, a)}")

        to_do.append(b)

    for b, _ in copies:
        if loc[b] is None:
            ready.append(b)

    logger.debug("loc: %s", "{" + ", ".join(f"{k}={v}" for k, v in loc.items()) + "}")
    logger.debug("pred: %s", "{" + ", ".join(f"{k}={v}" for k, v in pred.items()) + "}")
    logger.debug("ready: %s", ", ".join(map(str, ready)))
    logger.debug("to_do: %s", ", ".join(map(str, to_do)))

    while to_do:
        while ready:
            b = ready.pop()
            logger.debug("* avail %s", b)
            if b not in pred:
                continue
            a = pred[b]
            c = loc[a]
            assert c is not None
            # print("%s <- %s" % (b, a))
            res.append((b, c))

            # Addition by Paul Sokolovsky to handle fan-out case (when same
            # source is assigned to multiple destinations).
            if allow_fan_out and c in to_do:
                to_do.remove(c)

            loc[a] = b
            if a == c:
                ready.append(a)

        # Addition to handle fan-out.
        if allow_fan_out and not to_do:
            break

        b = to_do.pop()
        logger.debug("* to_do %s", b)
        if b != loc[pred[b]]:
            tmp = mktmp(b)
            # print("%s <- %s" % (b, a))
            res.append((tmp, b))
            loc[b] = tmp
            ready.append(b)

    return res


from puya import log
from puya.ir import models
from puya.ir.visitor_mem_replacer import MemoryReplacer

logger = log.get_logger(__name__)

def convert_to_cssa(sub: models.Subroutine) -> None:
    logger.debug("Converting to CSSA")
    max_versions = dict[str, int]()
    for reg in sub.get_assigned_registers():
        max_versions[reg.name] = max(max_versions.get(reg.name, 0), reg.version)

    def make_prime(mem: models.Register) -> models.Register:
        max_versions[mem.name] += 1
        return models.Register(
            source_location=mem.source_location,
            ir_type=mem.ir_type,
            name=mem.name,
            version=max_versions[mem.name],
        )

    block_exit_copies = dict[models.BasicBlock, list[tuple[models.Register, models.Register]]]()
    for phi_block in sub.body:
        prime_phis = list[models.Phi]()
        prime_copies = list[tuple[models.Register, models.Register]]()
        for phi in phi_block.phis:
            if not phi.args:
                phi_block.ops.insert(
                    0,
                    models.Assignment(
                        targets=[phi.register],
                        source=models.Undefined(
                            ir_type=phi.ir_type, source_location=phi.source_location
                        ),
                        source_location=phi.source_location,
                    ),
                )
            else:
                prime_args = list[models.PhiArgument]()
                for phi_arg in phi.args:
                    prime_arg_reg = make_prime(phi_arg.value)
                    prime_arg = models.PhiArgument(
                        value=prime_arg_reg,
                        through=phi_arg.through,
                    )
                    prime_args.append(prime_arg)
                    # insert copy to prime arg at end of predecessor
                    block_exit_copies.setdefault(phi_arg.through, []).append(
                        (prime_arg_reg, phi_arg.value)
                    )
                phi_prime = models.Phi(register=make_prime(phi.register), args=prime_args)
                prime_phis.append(phi_prime)
                prime_copies.append((phi.register, phi_prime.register))
        phi_block.phis = prime_phis
        if prime_copies:
            phi_block.ops.insert(0, _make_copy_assignment(prime_copies))

    for block, copies in block_exit_copies.items():
        block.ops.append(_make_copy_assignment(copies))

def _make_copy_assignment(
    copies: list[tuple[models.Register, models.Register]],
) -> models.Assignment:
    if len(copies) == 1:
        ((dst, src),) = copies
        targets = [dst]
        source: models.ValueProvider = src
    else:
        targets_tup, sources_tup = zip(*copies, strict=True)
        targets = list(targets_tup)
        sources_list = list(sources_tup)
        source = models.ValueTuple(source_location=None, values=sources_list)
    return models.Assignment(
        targets=targets,
        source=source,
        source_location=None,
    )

def destructure_cssa(sub: models.Subroutine) -> None:
    logger.debug("Removing Phi nodes")
    # Once the sub is in CSSA, destructuring is trivial.
    # All variables involved in a Phi can be given the same name, and the Phi can be removed.
    replacements = dict[models.Register, models.Register]()
    for block in sub.body:
        phis = block.phis
        block.phis = []
        replacements.update({arg.value: phi.register for phi in phis for arg in phi.args})
    MemoryReplacer.apply(sub.body, replacements=replacements)


from puya.ir.avm_ops import AVMOp
from puya.ir.models import Intrinsic, Value
from puya.ir.types_ import IRType
from puya.parse import SourceLocation

def select(
    *, condition: Value, false: Value, true: Value, type_: IRType, source_location: SourceLocation
) -> Intrinsic:
    return Intrinsic(
        op=AVMOp.select,
        args=[
            false,
            true,
            condition,
        ],
        types=[type_],
        source_location=source_location,
    )


import contextlib
import copy
import typing
from collections import defaultdict
from collections.abc import Collection, Iterator
from pathlib import Path

from immutabledict import immutabledict

from puya import artifact_metadata, log
from puya.avm import AVMType
from puya.awst import nodes as awst_nodes
from puya.awst.awst_traverser import AWSTTraverser
from puya.awst.function_traverser import FunctionTraverser
from puya.awst.serialize import awst_from_json
from puya.context import CompileContext
from puya.ir import arc4_router
from puya.ir._contract_metadata import build_contract_metadata
from puya.ir._utils import make_subroutine
from puya.ir.arc4_router import AWSTContractMethodSignature
from puya.ir.builder.main import FunctionIRBuilder
from puya.ir.context import IRBuildContext
from puya.ir.models import Contract, LogicSignature, ModuleArtifact, Program, Subroutine
from puya.ir.optimize.dead_code_elimination import remove_unused_subroutines
from puya.ir.types_ import wtype_to_ir_type
from puya.program_refs import ProgramKind
from puya.utils import StableSet, attrs_extend, coalesce, set_add, set_remove

logger = log.get_logger(__name__)

CalleesLookup: typing.TypeAlias = defaultdict[awst_nodes.Function, set[awst_nodes.Function]]
_EMBEDDED_LIB = Path(__file__).parent / "_puya_lib.awst.json"

def awst_to_ir(context: CompileContext, awst: awst_nodes.AWST) -> Iterator[ModuleArtifact]:
    compilation_set = _CompilationSetCollector.collect(context, awst)

    ir_ctx = _build_subroutines(context, awst)

    for node in compilation_set:
        artifact_ctx = ir_ctx.for_root(node)
        with artifact_ctx.log_exceptions():
            if isinstance(node, awst_nodes.LogicSignature):
                yield _build_logic_sig_ir(artifact_ctx, node)
            else:
                yield _build_contract_ir(artifact_ctx, node)

def _build_subroutines(ctx: CompileContext, awst: awst_nodes.AWST) -> IRBuildContext:
    embedded_awst = awst_from_json(_EMBEDDED_LIB.read_text())
    embedded_subroutines = {
        node: make_subroutine(node, allow_implicits=False)
        for node in embedded_awst
        if isinstance(node, awst_nodes.Subroutine)
    }
    user_subroutines = {
        node: make_subroutine(node, allow_implicits=True)
        for node in awst
        if isinstance(node, awst_nodes.Subroutine)
    }
    ir_ctx = attrs_extend(
        IRBuildContext,
        ctx,
        subroutines={**embedded_subroutines, **user_subroutines},
        awst=[*embedded_awst, *awst],
        embedded_funcs_lookup={func.id: sub for func, sub in embedded_subroutines.items()},
    )
    for func, sub in ir_ctx.subroutines.items():
        FunctionIRBuilder.build_body(ir_ctx, function=func, subroutine=sub)
    return ir_ctx

class _CompilationSetCollector(AWSTTraverser):
    def __init__(self, awst: awst_nodes.AWST, *, explicit_compilation_set: StableSet[str]) -> None:
        super().__init__()
        self._remaining_explicit_set: typing.Final = explicit_compilation_set
        self.compilation_set: typing.Final = dict[
            str, awst_nodes.Contract | awst_nodes.LogicSignature
        ]()
        self._nodes_by_id: typing.Final = immutabledict[str, awst_nodes.RootNode](
            {n.id: n for n in awst}
        )

    @property
    def unresolved_explict_ids(self) -> Collection[str]:
        return self._remaining_explicit_set

    @classmethod
    def collect(
        cls, context: CompileContext, awst: awst_nodes.AWST
    ) -> Collection[awst_nodes.Contract | awst_nodes.LogicSignature]:
        collector = cls(
            awst, explicit_compilation_set=StableSet.from_iter(context.compilation_set)
        )
        for node in awst:
            node.accept(collector)
        for unresolved_id in collector.unresolved_explict_ids:
            logger.error(f"unable to resolve compilation artifact '{unresolved_id}'")
        return collector.compilation_set.values()

    def visit_compiled_contract(self, compiled_contract: awst_nodes.CompiledContract) -> None:
        super().visit_compiled_contract(compiled_contract)
        node = self._nodes_by_id.get(compiled_contract.contract)
        match node:
            case awst_nodes.Contract() as contract:
                self._visit_contract_or_lsig(contract, reference=True)
            case None:
                logger.error(
                    "unable to resolve contract reference",
                    location=compiled_contract.source_location,
                )
            case _:
                logger.error(
                    "reference is not a contract", location=compiled_contract.source_location
                )

    def visit_compiled_logicsig(self, compiled_lsig: awst_nodes.CompiledLogicSig) -> None:
        super().visit_compiled_logicsig(compiled_lsig)
        node = self._nodes_by_id.get(compiled_lsig.logic_sig)
        match node:
            case awst_nodes.LogicSignature() as lsig:
                self._visit_contract_or_lsig(lsig, reference=True)
            case None:
                logger.error(
                    "unable to resolve logic signature reference",
                    location=compiled_lsig.source_location,
                )
            case _:
                logger.error(
                    "reference is not a logic signature", location=compiled_lsig.source_location
                )

    def visit_contract(self, contract: awst_nodes.Contract) -> None:
        return self._visit_contract_or_lsig(contract)

    def visit_logic_signature(self, lsig: awst_nodes.LogicSignature) -> None:
        return self._visit_contract_or_lsig(lsig)

    def _visit_contract_or_lsig(
        self,
        node: awst_nodes.Contract | awst_nodes.LogicSignature,
        *,
        reference: bool = False,
    ) -> None:
        if node.id in self.compilation_set:
            return  # already visited
        direct = set_remove(self._remaining_explicit_set, node.id)
        if direct or reference:
            self.compilation_set[node.id] = node
            match node:
                case awst_nodes.Contract():
                    super().visit_contract(node)
                case awst_nodes.LogicSignature():
                    super().visit_logic_signature(node)
                case unexpected:
                    typing.assert_never(unexpected)

def _build_logic_sig_ir(
    ctx: IRBuildContext, logic_sig: awst_nodes.LogicSignature
) -> LogicSignature:
    metadata = artifact_metadata.LogicSignatureMetaData(
        ref=logic_sig.id,
        description=logic_sig.docstring,
        name=logic_sig.short_name,
    )

    avm_version = coalesce(logic_sig.avm_version, ctx.options.target_avm_version)
    sig_ir = _make_program(
        ctx, logic_sig.program, kind=ProgramKind.logic_signature, avm_version=avm_version
    )
    return LogicSignature(
        program=sig_ir, metadata=metadata, source_location=logic_sig.source_location
    )

def _build_contract_ir(ctx: IRBuildContext, contract: awst_nodes.Contract) -> Contract:
    metadata, arc4_methods = build_contract_metadata(ctx, contract)
    routing_data = {
        md: AWSTContractMethodSignature(
            target=awst_nodes.ContractMethodTarget(cref=cm.cref, member_name=cm.member_name),
            return_type=cm.return_type,
            parameter_types=[a.wtype for a in cm.args],
        )
        for cm, md in arc4_methods.items()
    }
    arc4_router_awst = arc4_router.create_abi_router(contract, routing_data)
    ctx.routers[contract.id] = ctx.subroutines[arc4_router_awst] = make_subroutine(
        arc4_router_awst, allow_implicits=False
    )

    # Build callees list, excluding calls from router.
    # Used to if a function should implicitly return mutable reference parameters.
    callees = SubroutineCollector.collect(
        ctx, contract.approval_program, contract.clear_program, *arc4_methods
    )

    # construct unique Subroutine objects for each function
    # that was referenced through either entry point
    for method in contract.methods:
        # ABI methods which are only called by the ABI router in the approval_program do not need
        # to implicitly return anything as we know our router is not interested in anything but the
        # explicit return value.
        allow_implicits = bool(callees[method])
        # make the emtpy subroutine, because functions reference other functions
        func_ir = make_subroutine(method, allow_implicits=allow_implicits)
        ctx.subroutines[method] = func_ir

    # now construct the subroutine IR
    for func, sub in ctx.subroutines.items():
        if not sub.body:  # in case something is pre-built (ie from embedded lib)
            FunctionIRBuilder.build_body(ctx, function=func, subroutine=sub)

    avm_version = coalesce(contract.avm_version, ctx.options.target_avm_version)
    approval_ir = _make_program(
        ctx, contract.approval_program, kind=ProgramKind.approval, avm_version=avm_version
    )
    clear_state_ir = _make_program(
        ctx, contract.clear_program, kind=ProgramKind.clear_state, avm_version=avm_version
    )
    return Contract(
        approval_program=approval_ir,
        clear_program=clear_state_ir,
        metadata=metadata,
        source_location=contract.source_location,
    )

def _make_program(
    ctx: IRBuildContext,
    main: awst_nodes.Function,
    *,
    kind: ProgramKind,
    avm_version: int,
) -> Program:
    assert not main.args, "main method should not have args"
    return_type = wtype_to_ir_type(main.return_type)
    assert return_type.avm_type == AVMType.uint64, "main method should return uint64 backed type"
    main_sub = Subroutine(
        id=main.full_name,
        short_name=main.short_name,
        parameters=[],
        returns=[return_type],
        body=[],
        inline=False,
        source_location=main.source_location,
    )
    FunctionIRBuilder.build_body(ctx, function=main, subroutine=main_sub)
    program = Program(
        kind=kind,
        main=main_sub,
        subroutines=tuple(ctx.subroutines.values()),
        avm_version=avm_version,
    )
    remove_unused_subroutines(program)
    program = copy.deepcopy(program)
    return program

class SubroutineCollector(FunctionTraverser):
    def __init__(self, context: IRBuildContext, callees: CalleesLookup) -> None:
        self.context = context
        self._seen = StableSet[awst_nodes.Function]()
        self.callees = callees
        self._func_stack = list[awst_nodes.Function]()

    @classmethod
    def collect(cls, context: IRBuildContext, *entry_points: awst_nodes.Function) -> CalleesLookup:
        callees = CalleesLookup(set)
        collector = cls(context, callees)
        for start in entry_points:
            with collector._enter_func(start):  # noqa: SLF001
                start.body.accept(collector)
        return callees

    @typing.override
    def visit_subroutine_call_expression(self, expr: awst_nodes.SubroutineCallExpression) -> None:
        super().visit_subroutine_call_expression(expr)
        callee = self._func_stack[-1]
        func = self.context.resolve_function_reference(
            expr.target, expr.source_location, caller=callee
        )
        self.callees[func].add(callee)
        if set_add(self._seen, func):
            with self._enter_func(func):
                func.body.accept(self)

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