import abc
import typing
import typing as t
from collections.abc import Iterable, Iterator, Mapping, Sequence, Set

import attrs
from immutabledict import immutabledict

from puya import log
from puya.artifact_metadata import ContractMetaData, LogicSignatureMetaData
from puya.avm import AVMType
from puya.awst.txn_fields import TxnField
from puya.errors import CodeError, InternalError
from puya.ir.avm_ops import AVMOp
from puya.ir.avm_ops_models import ImmediateKind, OpSignature, StackType, Variant
from puya.ir.types_ import AVMBytesEncoding, IRType, stack_type_to_avm_type, stack_type_to_ir_type
from puya.ir.visitor import IRVisitor
from puya.parse import SourceLocation
from puya.program_refs import (
    ContractReference,
    LogicSigReference,
    ProgramKind,
)
from puya.utils import unique

logger = log.get_logger(__name__)

T = t.TypeVar("T")

class Context(t.Protocol):
    source_location: SourceLocation | None

class IRVisitable(Context, abc.ABC):
    @abc.abstractmethod
    def accept(self, visitor: IRVisitor[T]) -> T: ...

    def __str__(self) -> str:
        from puya.ir.to_text_visitor import ToTextVisitor

        return self.accept(ToTextVisitor())

class _Freezable(abc.ABC):
    def freeze(self) -> object:
        data = self._frozen_data()
        hash(data)  # check we can hash
        if data is self:
            return data
        return self.__class__, data

    @abc.abstractmethod
    def _frozen_data(self) -> object: ...

# NOTE! we don't want structural equality in the IR, everything needs to have eq=False
#       the workaround to do this (trivial in Python to extend a decorator) AND have mypy
#       not complain is ... well, see for yourself:
#       https://www.attrs.org/en/stable/extending.html#wrapping-the-decorator
@attrs.define(eq=False)
class ValueProvider(IRVisitable, _Freezable, abc.ABC):
    """A node that provides/produces a value"""

    source_location: SourceLocation | None = attrs.field(eq=False)

    @property
    @abc.abstractmethod
    def types(self) -> Sequence[IRType]: ...

    @property
    @t.final
    def atypes(self) -> Sequence[AVMType]:
        return tuple(t.avm_type for t in self.types)

@attrs.frozen
class Value(ValueProvider, abc.ABC):
    """Base class for value types.

    This is anything that *is* a value, so excludes
    value *producers* such as subroutine invocations
    """

    ir_type: IRType = attrs.field(repr=lambda x: x.name)

    @property
    @t.final
    def atype(self) -> typing.Literal[AVMType.uint64, AVMType.bytes]:
        return self.ir_type.avm_type

    @property
    def types(self) -> Sequence[IRType]:
        return (self.ir_type,)

    def _frozen_data(self) -> object:
        return self

@attrs.frozen
class Undefined(Value):
    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_undefined(self)

class Constant(Value, abc.ABC):
    """Base class for value constants - any value that is known at compile time"""

@attrs.define(eq=False)
class Op(IRVisitable, _Freezable, abc.ABC):
    """Base class for non-control-flow, non-phi operations

    This is anything other than a Phi can appear inside a BasicBlock before the terminal ControlOp.
    """

@attrs.define(eq=False)
class ControlOp(IRVisitable, _Freezable, abc.ABC):
    """Base class for control-flow operations

    These appear in a BasicBlock as the terminal node.
    """

    source_location: SourceLocation | None

    @abc.abstractmethod
    def targets(self) -> Sequence["BasicBlock"]:
        """For graph traversal - result could be empty if it's the end of the current graph"""

    @property
    def unique_targets(self) -> list["BasicBlock"]:
        return unique(self.targets())

@attrs.frozen
class Register(Value):
    """A register is an abstraction of "local variable storage".

    This could be mapped to either value on the stack or a scratch slot during code generation.
    """

    name: str
    version: int

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_register(self)

    @property
    def local_id(self) -> str:
        return f"{self.name}#{self.version}"

@attrs.define(eq=False)
class PhiArgument(IRVisitable):
    value: Register
    through: "BasicBlock"
    source_location: None = attrs.field(default=None, init=False)

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_phi_argument(self)

@attrs.define(eq=False)
class Phi(IRVisitable, _Freezable):
    """Phi nodes are oracles that, given a list of other variables, always produce the
    one that has been defined in the control flow thus far.

    The term phi node comes from the literature on Static Single Assignment
    """

    register: Register
    args: list[PhiArgument] = attrs.field(factory=list)
    source_location: None = attrs.field(default=None, init=False)

    @property
    def ir_type(self) -> IRType:
        return self.register.ir_type

    @property
    def atype(self) -> AVMType:
        return self.ir_type.avm_type

    @property
    def non_self_args(self) -> Sequence[PhiArgument]:
        return tuple(a for a in self.args if a.value != self.register)

    def _frozen_data(self) -> object:
        return (
            self.register.freeze(),
            tuple((arg.through.id, arg.value.freeze()) for arg in self.args),
        )

    @args.validator
    def check_args(self, _attribute: object, args: Sequence[PhiArgument]) -> None:
        bad_args = [
            arg for arg in args if arg.value.ir_type.maybe_avm_type != self.ir_type.maybe_avm_type
        ]
        if bad_args:
            raise InternalError(
                f"Phi node (register={self.register}) received arguments with unexpected type(s):"
                f" {', '.join(map(str, bad_args))}, "
            )
        seen_blocks = set[BasicBlock]()
        for arg in args:
            if arg.through in seen_blocks:
                raise InternalError(f"Duplicate source to phi node: {arg.through}")
            seen_blocks.add(arg.through)

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_phi(self)

@attrs.frozen(kw_only=True)
class UInt64Constant(Constant):
    value: int
    ir_type: IRType = attrs.field(default=IRType.uint64)
    teal_alias: str | None = None

    @ir_type.validator
    def _validate_ir_type(self, _attribute: object, ir_type: IRType) -> None:
        if ir_type.maybe_avm_type is not AVMType.uint64:
            raise InternalError(
                f"Invalid type for UInt64Constant: {ir_type}", self.source_location
            )

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_uint64_constant(self)

@attrs.frozen(kw_only=True)
class ITxnConstant(Constant):
    value: int
    ir_type: IRType = attrs.field()

    @ir_type.validator
    def _validate_ir_type(self, _attribute: object, ir_type: IRType) -> None:
        if ir_type not in (IRType.itxn_group_idx, IRType.itxn_field_set):
            raise InternalError(f"invalid type for ITxnConstant: {ir_type}", self.source_location)

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_itxn_constant(self)

@attrs.frozen
class BigUIntConstant(Constant):
    value: int
    ir_type: IRType = attrs.field(default=IRType.biguint, init=False)

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_biguint_constant(self)

@attrs.frozen
class TemplateVar(Value):
    name: str
    ir_type: IRType

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_template_var(self)

@attrs.frozen(kw_only=True)
class BytesConstant(Constant):
    """Constant for types that are logically bytes"""

    ir_type: IRType = attrs.field(default=IRType.bytes)

    encoding: AVMBytesEncoding
    value: bytes

    @ir_type.validator
    def _validate_ir_type(self, _attribute: object, ir_type: IRType) -> None:
        if ir_type.maybe_avm_type is not AVMType.bytes:
            raise InternalError(f"invalid type for BytesConstant: {ir_type}", self.source_location)

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_bytes_constant(self)

@attrs.define
class CompiledContractReference(Value):
    artifact: ContractReference
    field: TxnField
    template_variables: Mapping[str, Value] = attrs.field(converter=immutabledict)
    """
    template variable keys here are fully qualified with their appropriate prefix
    """
    source_location: SourceLocation | None = attrs.field(eq=False)
    program_page: int | None = None  # used for approval and clear_state fields

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_compiled_contract_reference(self)

@attrs.define
class CompiledLogicSigReference(Value):
    artifact: LogicSigReference
    template_variables: Mapping[str, Value] = attrs.field(converter=immutabledict)
    source_location: SourceLocation | None = attrs.field(eq=False)

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_compiled_logicsig_reference(self)

@attrs.frozen
class AddressConstant(Constant):
    """Constant for address literals"""

    ir_type: IRType = attrs.field(default=IRType.bytes, init=False)
    value: str = attrs.field()

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_address_constant(self)

@attrs.frozen
class MethodConstant(Constant):
    """Constant for method literals"""

    ir_type: IRType = attrs.field(default=IRType.bytes, init=False)
    value: str

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_method_constant(self)

@attrs.define(eq=False)
class InnerTransactionField(ValueProvider):
    field: str
    group_index: Value
    is_last_in_group: Value
    array_index: Value | None
    type: IRType

    def _frozen_data(self) -> object:
        return self.field, self.group_index, self.is_last_in_group, self.array_index, self.type

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_inner_transaction_field(self)

    @property
    def types(self) -> Sequence[IRType]:
        return (self.type,)

@attrs.define(eq=False)
class Intrinsic(Op, ValueProvider):
    """Any TEAL op (or pseudo-op) that doesn't interrupt control flow, in the "basic block" sense.

    refs:
    - https://developer.algorand.org/docs/get-details/dapps/avm/teal/opcodes/
    - https://developer.algorand.org/docs/get-details/dapps/avm/teal/specification/#assembler-syntax
    """

    op: AVMOp
    # TODO: consider treating ops with no args (only immediates) as Value types
    #       e.g. `txn NumAppArgs` or `txna ApplicationArgs 0`
    immediates: list[str | int] = attrs.field(factory=list)
    args: list[Value] = attrs.field(factory=list)
    error_message: str | None = None
    """If the program fails at this op, error_message will be displayed as the reason"""
    _types: Sequence[IRType] = attrs.field(converter=tuple[IRType, ...])

    @_types.default
    def _default_types(self) -> tuple[IRType, ...]:
        types = list[IRType]()
        for stack_type in self.op_signature.returns:
            ir_type = stack_type_to_ir_type(stack_type)
            if ir_type is None:
                raise InternalError(
                    f"Intrinsic op {self.op.name} requires return type information",
                    self.source_location,
                )
            types.append(ir_type)
        return tuple(types)

    def _frozen_data(self) -> object:
        return self.op, tuple(self.immediates), tuple(self.args), self.error_message

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_intrinsic_op(self)

    @property
    def types(self) -> Sequence[IRType]:
        return self._types

    @property
    def op_signature(self) -> OpSignature:
        return self.op_variant.signature

    @property
    def op_variant(self) -> Variant:
        return self.op.get_variant(self.immediates)

    @_types.validator
    def _validate_types(self, _attribute: object, types: Sequence[IRType]) -> None:
        self._check_stack_types("return", self.op_signature.returns, types)

    @args.validator
    def _validate_args(self, _attribute: object, args: list[Value]) -> None:
        arg_types = [a.ir_type for a in args]
        self._check_stack_types("argument", self.op_signature.args, arg_types)

    @immediates.validator
    def _validate_immediates(self, _attribute: object, immediates: list[int | str]) -> None:
        if len(self.op.immediate_types) != len(immediates):
            logger.error("Incorrect number of immediates", location=self.source_location)
            return
        for imm_type, imm in zip(self.op.immediate_types, immediates, strict=True):
            match imm_type:
                case ImmediateKind.uint8:
                    if not isinstance(imm, int) or not (0 <= imm <= 255):
                        logger.critical(
                            "Invalid immediate, expected value between 0 and 255",
                            location=self.source_location,
                        )
                case ImmediateKind.arg_enum:
                    if not isinstance(imm, str):
                        logger.critical(
                            "Invalid immediate, expected enum value",
                            location=self.source_location,
                        )
                case _:
                    typing.assert_never(imm_type)

    def _check_stack_types(
        self,
        context: str,
        expected_types: Sequence[StackType],
        source_types: Sequence[IRType],
    ) -> None:
        target_types = [stack_type_to_avm_type(a) for a in expected_types]
        if len(target_types) != len(source_types) or not all(
            tt & st.avm_type for tt, st in zip(target_types, source_types, strict=True)
        ):
            logger.error(
                (
                    f"Incompatible {context} types on Intrinsic"
                    f"({self.op} {' '.join(map(str, self.immediates))}): "
                    f" received = ({', '.join(map(str, source_types))}),"
                    f" expected = ({', '.join(map(str, target_types))})"
                ),
                location=self.source_location,
            )

@attrs.define(eq=False)
class InvokeSubroutine(Op, ValueProvider):
    """Subroutine invocation

    opcode: callsub"""

    target: "Subroutine"
    # TODO: validation for args
    args: list[Value]

    def _frozen_data(self) -> object:
        return self.target.id, tuple(self.args)

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_invoke_subroutine(self)

    @property
    def types(self) -> Sequence[IRType]:
        return self.target.returns

@attrs.define(eq=False)
class ValueTuple(ValueProvider):
    values: Sequence[Value]

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_value_tuple(self)

    @property
    def types(self) -> Sequence[IRType]:
        return [val.ir_type for val in self.values]

    def _frozen_data(self) -> object:
        return tuple(self.values)

@attrs.define(eq=False)
class Assignment(Op):
    """
    Assignment of either a value or the result of something that produces a value to register(s)
    """

    source_location: SourceLocation | None
    targets: Sequence[Register] = attrs.field(
        validator=[attrs.validators.min_len(1)], converter=tuple[Register, ...]
    )
    source: ValueProvider = attrs.field()

    def _frozen_data(self) -> object:
        return tuple(self.targets), self.source.freeze()

    @source.validator
    def _check_types(self, _attribute: object, source: ValueProvider) -> None:
        target_ir_types = [target.ir_type for target in self.targets]
        source_ir_types = list(source.types)
        if target_ir_types != source_ir_types:
            # TODO: need to update some optimiser code and/or
            #       introduce ReinterpretCast ValueProvider
            #       here before we can remove this fallback to AVMType here
            target_atypes = [tt.maybe_avm_type for tt in target_ir_types]
            source_atypes = [st.maybe_avm_type for st in source_ir_types]
            if target_atypes != source_atypes:
                raise CodeError(
                    f"Incompatible types on assignment:"
                    f" source = ({', '.join(map(str, source_ir_types))}),"
                    f" target = ({', '.join(map(str, target_ir_types))})",
                    self.source_location,
                )

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_assignment(self)

@attrs.define(eq=False, str=False, kw_only=True)
class BasicBlock(Context):
    """IR basic block.

    Contains a sequence of operations and ends with a terminating ControlOp.
    Only the last op can be a ControlOp.

    All generated Ops live in basic blocks. Basic blocks determine the
    order of evaluation and control flow within a function. A basic
    block is always associated with a single Subroutine or program main.

    Ops that may terminate the program due to a panic or similar aren't treated as exits.
    """

    source_location: SourceLocation  # the location that caused the block to be constructed
    phis: list[Phi] = attrs.field(factory=list)
    ops: list[Op] = attrs.field(factory=list)
    terminator: ControlOp | None = attrs.field(default=None)
    predecessors: "list[BasicBlock]" = attrs.field(factory=list)
    id: int | None = None
    label: str | None = None
    comment: str | None = None

    @phis.validator
    def _check_phis(self, _attribute: object, phis: list[Phi]) -> None:
        for phi in phis:
            attrs.validate(phi)

    @ops.validator
    def _check_ops(self, _attribute: object, ops: list[Op]) -> None:
        for op in ops:
            attrs.validate(op)

    @terminator.validator
    def _check_terminator(self, _attribute: object, terminator: ControlOp | None) -> None:
        if terminator is not None:
            attrs.validate(terminator)

    @property
    def terminated(self) -> bool:
        return self.terminator is not None

    @property
    def successors(self) -> "Sequence[BasicBlock]":
        if self.terminator is None:
            return ()
        return self.terminator.unique_targets

    @property
    def is_empty(self) -> bool:
        return not (self.phis or self.ops or self.terminator)

    @property
    def all_ops(self) -> Iterator[Phi | Op | ControlOp]:
        yield from self.phis.copy()
        yield from self.ops.copy()  # copy in case ops is modified
        if self.terminator is not None:
            yield self.terminator

    def get_assigned_registers(self) -> Iterator[Register]:
        for phi in self.phis:
            yield phi.register
        for op in self.ops:
            if isinstance(op, Assignment):
                yield from op.targets

    def __str__(self) -> str:
        return f"block@{self.id}"

@attrs.define(eq=False, kw_only=True)
class ConditionalBranch(ControlOp):
    """Branch based on zero-ness

    opcode: bz+b or bnz+b
    """

    condition: Value = attrs.field()
    non_zero: BasicBlock
    zero: BasicBlock

    def _frozen_data(self) -> object:
        return self.condition.freeze(), self.non_zero.id, self.zero.id

    @condition.validator
    def check(self, _attribute: object, result: Value) -> None:
        if result.atype != AVMType.uint64:
            raise CodeError(
                "Branch condition can only be uint64 backed value", self.source_location
            )

    def targets(self) -> Sequence[BasicBlock]:
        return self.zero, self.non_zero

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_conditional_branch(self)

@attrs.define(eq=False)
class Goto(ControlOp):
    """Unconditional jump

    opcode: b
    """

    target: BasicBlock

    def _frozen_data(self) -> object:
        return self.target.id

    def targets(self) -> Sequence[BasicBlock]:
        return (self.target,)

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_goto(self)

@attrs.define(eq=False)
class GotoNth(ControlOp):
    """Jump to the nth block in a list where n is defined by a UInt16 Value.
    Jumps to the default if n is larger than the number of blocks

    opcode: switch+b"""

    value: Value
    blocks: list[BasicBlock] = attrs.field()
    default: BasicBlock

    def _frozen_data(self) -> object:
        return self.value, tuple(b.id for b in self.blocks), self.default.id

    def targets(self) -> Sequence[BasicBlock]:
        return *self.blocks, self.default

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_goto_nth(self)

@attrs.define(eq=False)
class Switch(ControlOp):
    """Jump based on comparison between a value and a fixed list of other values.
    If no match is found, jumps to "default". This default option doesn't exist in
    the underlying op code, but we need it to make this truly terminate the basic block
    it's in, otherwise it'd violate certain CFG constraints.

    opcode: match+b
    """

    value: Value
    cases: dict[Value, BasicBlock] = attrs.field()
    default: BasicBlock

    @cases.validator
    def _check_cases(self, _attribute: object, cases: dict[Value, BasicBlock]) -> None:
        if any(case.atype != self.value.atype for case in cases):
            raise CodeError(
                "Switch cases types mismatch with value to match", self.source_location
            )

    def _frozen_data(self) -> object:
        return (
            self.value,
            tuple((v, b.id) for v, b in self.cases.items()),
            self.default.id,
        )

    def targets(self) -> Sequence[BasicBlock]:
        return *self.cases.values(), self.default

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_switch(self)

@attrs.define(eq=False)
class SubroutineReturn(ControlOp):
    """Return from within a subroutine

    opcode: retsub
    """

    result: list[Value]

    def _frozen_data(self) -> object:
        return tuple(self.result)

    def targets(self) -> Sequence[BasicBlock]:
        return ()

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_subroutine_return(self)

@attrs.define(eq=False)
class ProgramExit(ControlOp):
    """Return from and exit the program immediately

    opcode: return
    """

    result: Value = attrs.field()

    @result.validator
    def check(self, _attribute: object, result: Value) -> None:
        if result.atype != AVMType.uint64:
            raise CodeError("Can only exit with uint64 backed value", self.source_location)

    def _frozen_data(self) -> object:
        return self.result

    def targets(self) -> Sequence[BasicBlock]:
        return ()

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_program_exit(self)

@attrs.define(eq=False)
class Fail(ControlOp):
    """Exits immediately with a failure condition.

    assert statements with a compile time constant that is false
    should be translated to this node type in order to become ControlOp

    opcode: err
    """

    error_message: str | None

    def targets(self) -> Sequence[BasicBlock]:
        return ()

    def accept(self, visitor: IRVisitor[T]) -> T:
        return visitor.visit_fail(self)

    def _frozen_data(self) -> object:
        return self.error_message

@attrs.frozen
class Parameter(Register):
    implicit_return: bool

@attrs.define(eq=False, kw_only=True)
class Subroutine(Context):
    id: str
    short_name: str
    # source_location might be None if it was synthesized e.g. ARC4 approval method
    source_location: SourceLocation | None
    parameters: Sequence[Parameter]
    _returns: Sequence[IRType]
    body: list[BasicBlock] = attrs.field()
    inline: bool | None

    @property
    def returns(self) -> list[IRType]:
        return [*self._returns, *(p.ir_type for p in self.parameters if p.implicit_return)]

    @body.validator
    def _check_blocks(self, _attribute: object, body: list[BasicBlock]) -> None:
        blocks = frozenset(body)
        for block in body:
            attrs.validate(block)
            if block.terminator is None:
                raise InternalError(
                    f"Unterminated block {block} assigned to subroutine {self.id}",
                    block.source_location,
                )
            for successor in block.successors:
                if block not in successor.predecessors:
                    # Note: this check is here rather than on BasicBlock only because of
                    # circular validation issues where you're trying to update the CFG by
                    # replacing a terminator
                    raise InternalError(
                        f"{block} does not appear in all {block.terminator}"
                        f" target's predecessor lists - missing from {successor.id} at least",
                        block.terminator.source_location
                        or block.source_location
                        or self.source_location,
                    )
            if not blocks.issuperset(block.predecessors):
                raise InternalError(
                    f"{block} of subroutine {self.id} has predecessor block(s) outside of list",
                    block.source_location,
                )
            if not blocks.issuperset(block.successors):
                raise InternalError(
                    f"{block} of subroutine {self.id} has predecessor block(s) outside of list",
                    block.source_location,
                )
            block_predecessors = dict.fromkeys(block.predecessors)
            for phi in block.phis:
                phi_blocks = dict.fromkeys(a.through for a in phi.args)
                if block_predecessors.keys() != phi_blocks.keys():
                    phi_block_labels = list(map(str, phi_blocks.keys()))
                    pred_block_labels = list(map(str, block_predecessors.keys()))
                    raise InternalError(
                        f"{self.id}: mismatch between phi predecessors ({phi_block_labels})"
                        f" and {block} predecessors ({pred_block_labels})"
                        f" for phi node {phi}",
                        self.source_location,
                    )
        used_registers = _get_used_registers(body)
        defined_registers = frozenset(self.parameters) | frozenset(_get_assigned_registers(body))
        bad_reads = used_registers - defined_registers
        if bad_reads:
            raise InternalError(
                f"The following variables are used but never defined:"
                f" {', '.join(map(str, bad_reads))}",
                self.source_location,
            )

    @property
    def entry(self) -> BasicBlock:
        return self.body[0]

    def get_assigned_registers(self) -> Iterator[Register]:
        yield from self.parameters
        yield from _get_assigned_registers(self.body)

    def get_used_registers(self) -> Iterator[Register]:
        yield from _get_used_registers(self.body)

    def validate_with_ssa(self) -> None:
        all_assigned = set[Register]()
        for block in self.body:
            for register in block.get_assigned_registers():
                if register in all_assigned:
                    raise InternalError(
                        f"SSA constraint violated:"
                        f" {register.local_id} is assigned multiple times"
                    )
                all_assigned.add(register)
        attrs.validate(self)

def _get_assigned_registers(blocks: Sequence[BasicBlock]) -> Iterator[Register]:
    # TODO: replace with visitor
    for block in blocks:
        for phi in block.phis:
            yield phi.register
        for op in block.ops:
            if isinstance(op, Assignment):
                yield from op.targets

def _get_used_registers(blocks: Sequence[BasicBlock]) -> Set[Register]:
    from puya.ir.register_read_collector import RegisterReadCollector

    collector = RegisterReadCollector()
    for block in blocks:
        for op in block.all_ops:
            op.accept(collector)
    return collector.used_registers

@attrs.define(kw_only=True, eq=False)
class Program(Context):
    """An individual TEAL output unit - e.g. an approval program, clear program, lsig"""

    # note: main is represented as a subroutine for simplified handling,
    # but should be "inlined" as main contract body during codegen.
    # ie, it could be generated as subroutine "main" with proto 0 1,
    # and then start of contract becomes:
    #
    # callsub main
    # return
    # main:
    #   proto 0 1
    #   ...
    #   retsub
    #
    # but, to save program size + op codes, this should be simplified to:
    # ...
    # return
    #
    # ie, just omit the subroutine header, and replace any&all retsub ops with a return instead
    kind: ProgramKind
    main: Subroutine
    subroutines: Sequence[Subroutine]
    avm_version: int
    source_location: SourceLocation | None = None

    def __attrs_post_init__(self) -> None:
        if self.source_location is None:
            self.source_location = self.main.source_location

    @property
    def all_subroutines(self) -> Iterable[Subroutine]:
        yield self.main
        yield from self.subroutines

@attrs.define(eq=False)
class Contract(Context):
    source_location: SourceLocation
    approval_program: Program
    clear_program: Program
    metadata: ContractMetaData

    def all_subroutines(self) -> Iterable[Subroutine]:
        from itertools import chain

        yield from unique(
            chain(
                self.approval_program.all_subroutines,
                self.clear_program.all_subroutines,
            )
        )

    def all_programs(self) -> Iterable[Program]:
        return [self.approval_program, self.clear_program]

@attrs.define(eq=False)
class LogicSignature(Context):
    source_location: SourceLocation
    program: Program
    metadata: LogicSignatureMetaData

    def all_subroutines(self) -> Iterable[Subroutine]:
        return self.program.all_subroutines

    def all_programs(self) -> Iterable[Program]:
        return [self.program]

ModuleArtifact: t.TypeAlias = Contract | LogicSignature


from collections import Counter
from functools import cached_property

import networkx as nx  # type: ignore[import-untyped]

from puya.ir import models

class CallGraph:
    def __init__(self, program: models.Program) -> None:
        self._program = program

    @cached_property
    def _graph(self) -> nx.MultiDiGraph:
        graph = nx.MultiDiGraph()
        for sub in self._program.all_subroutines:
            graph.add_node(sub.id, ref=sub)
        for sub in self._program.all_subroutines:
            for block in sub.body:
                for op in block.ops:
                    match op:
                        case (
                            models.InvokeSubroutine(target=target)
                            | models.Assignment(source=models.InvokeSubroutine(target=target))
                        ):
                            graph.add_edge(sub.id, target.id)
        return graph

    @cached_property
    def _paths(self) -> dict[str, dict[str, object]]:
        return dict(nx.all_pairs_shortest_path(self._graph))

    def callees(self, sub: models.Subroutine) -> list[tuple[str, int]]:
        return list(Counter(callee_id for callee_id, _ in self._graph.in_edges(sub.id)).items())

    def has_path(self, from_: str, to: str) -> bool:
        try:
            self._paths[from_][to]
        except KeyError:
            return False
        else:
            return True

    def is_auto_recursive(self, sub: models.Subroutine) -> bool:
        return bool(self._graph.has_predecessor(sub.id, sub.id))


from puya.errors import InternalError
from puya.ir import models

def get_definition(
    subroutine: models.Subroutine, register: models.Register, *, should_exist: bool = True
) -> models.Assignment | models.Phi | None:
    if register in subroutine.parameters:
        return None
    for block in subroutine.body:
        for phi in block.phis:
            if phi.register == register:
                return phi
        for op in block.ops:
            if isinstance(op, models.Assignment) and register in op.targets:
                return op
    if should_exist:
        raise InternalError(f"Register is not defined: {register}", subroutine.source_location)
    return None


from puya import log
from puya.context import CompileContext
from puya.ir import models
from puya.ir.context import TMP_VAR_INDICATOR
from puya.ir.visitor_mem_replacer import MemoryReplacer

logger = log.get_logger(__name__)

def copy_propagation(_context: CompileContext, subroutine: models.Subroutine) -> bool:
    set_lookup = dict[models.Register, list[models.Register]]()
    all_equivalence_sets = list[list[models.Register]]()

    modified = False
    for block in subroutine.body:
        for op in block.ops.copy():
            match op:
                case models.Assignment(targets=[target], source=models.Register() as source):
                    try:
                        equiv_set = set_lookup[source]
                        assert source in equiv_set
                    except KeyError:
                        set_lookup[source] = equiv_set = [source]
                        all_equivalence_sets.append(equiv_set)
                    equiv_set.append(target)
                    set_lookup[target] = equiv_set
                    block.ops.remove(op)
                    modified = True

    replacements = dict[models.Register, models.Register]()
    for equivalence_set in all_equivalence_sets:
        assert len(equivalence_set) >= 2
        equiv_set_ids = ", ".join(x.local_id for x in equivalence_set)
        logger.debug(f"Found equivalence set: {equiv_set_ids}")
        for reg in equivalence_set:
            if TMP_VAR_INDICATOR not in reg.name:
                replacement = reg
                break
        else:
            replacement = equivalence_set[0]
        for r in equivalence_set:
            if r is not replacement:
                replacements[r] = replacement

    for block in subroutine.body:
        for phi in block.phis.copy():
            try:
                (single_register,) = {replacements.get(arg.value, arg.value) for arg in phi.args}
            except ValueError:
                continue
            else:
                replacements[phi.register] = single_register
                block.phis.remove(phi)
                modified = True
    replaced = MemoryReplacer.apply(subroutine.body, replacements=replacements)
    if replaced:
        logger.debug(f"Copy propagation made {replaced} modifications")
        modified = True

    return modified


import contextlib
import itertools
from collections.abc import Iterable

import attrs

from puya import log
from puya.context import CompileContext
from puya.ir import models
from puya.ir.visitor_mutator import IRMutator
from puya.utils import unique

logger = log.get_logger(__name__)

@attrs.define
class BlockReferenceReplacer(IRMutator):
    find: models.BasicBlock
    replacement: models.BasicBlock

    @classmethod
    def apply(
        cls,
        find: models.BasicBlock,
        replacement: models.BasicBlock,
        blocks: Iterable[models.BasicBlock],
    ) -> None:
        replacer = cls(find=find, replacement=replacement)
        for block in blocks:
            replacer.visit_block(block)

    def visit_block(self, block: models.BasicBlock) -> None:
        super().visit_block(block)
        if self.find in block.predecessors:
            block.predecessors = [
                self.replacement if b is self.find else b for b in block.predecessors
            ]
            logger.debug(f"Replaced predecessor {self.find} with {self.replacement} in {block}")

    def visit_phi_argument(self, arg: models.PhiArgument) -> models.PhiArgument:
        if arg.through == self.find:
            arg.through = self.replacement
        return arg

    def visit_conditional_branch(self, branch: models.ConditionalBranch) -> models.ControlOp:
        if branch.zero == self.find:
            branch.zero = self.replacement
        if branch.non_zero == self.find:
            branch.non_zero = self.replacement
        return _replace_single_target_with_goto(branch)

    def visit_goto(self, goto: models.Goto) -> models.Goto:
        if goto.target == self.find:
            goto.target = self.replacement
        return goto

    def visit_goto_nth(self, goto_nth: models.GotoNth) -> models.ControlOp:
        for index, block in enumerate(goto_nth.blocks):
            if block == self.find:
                goto_nth.blocks[index] = self.replacement
        if goto_nth.default == self.find:
            goto_nth.default = self.replacement
        return _replace_single_target_with_goto(goto_nth)

    def visit_switch(self, switch: models.Switch) -> models.ControlOp:
        for case, target in switch.cases.items():
            if target == self.find:
                switch.cases[case] = self.replacement
        if switch.default == self.find:
            switch.default = self.replacement
        return _replace_single_target_with_goto(switch)

def _replace_single_target_with_goto(terminator: models.ControlOp) -> models.ControlOp:
    """
    If a ControlOp has a single target, replace it with a Goto, otherwise return the original op.
    """
    match terminator:
        case models.ControlOp(unique_targets=[single_target]):
            replacement = models.Goto(
                source_location=terminator.source_location,
                target=single_target,
            )
            logger.debug(f"replaced {terminator} with {replacement}")
            return replacement
        case _:
            return terminator

def remove_linear_jump(_context: CompileContext, subroutine: models.Subroutine) -> bool:
    changes = False
    for block in subroutine.body[1:]:
        match block.predecessors:
            case [models.BasicBlock(terminator=models.Goto(target=successor)) as predecessor]:
                assert successor is block
                # can merge blocks when there is an unconditional jump between them
                predecessor.phis.extend(block.phis)
                predecessor.ops.extend(block.ops)

                # this will update the predecessors of all block.successors to
                # now point back to predecessor e.g.
                # predecessor <-> block <-> [ss1, ...]
                # predecessor <-> [ss1, ...]
                BlockReferenceReplacer.apply(
                    find=block, replacement=predecessor, blocks=block.successors
                )

                predecessor.terminator = block.terminator

                # update block to reflect modifications
                subroutine.body.remove(block)
                changes = True
                logger.debug(f"Merged linear {block} into {predecessor}")
    return changes

def remove_empty_blocks(_context: CompileContext, subroutine: models.Subroutine) -> bool:
    changes = False
    for block in subroutine.body.copy():
        if not block.phis and not block.ops and isinstance(block.terminator, models.Goto):
            empty_block = block
            target = block.terminator.target

            if target.phis:
                logger.debug(
                    f"Not removing empty block {empty_block} because it's used by phi nodes"
                )
                continue

            # this will replace any ops that pointed to block
            BlockReferenceReplacer.apply(
                find=empty_block,
                replacement=target,
                blocks=empty_block.predecessors,
            )

            # remove the empty block from the targets predecessors, and add and of the empty block
            # predecessors that aren't already present
            target.predecessors = unique(
                itertools.chain(empty_block.predecessors, target.predecessors)
            )
            # might have already been replaced by BlockReferenceReplacer
            with contextlib.suppress(ValueError):
                target.predecessors.remove(empty_block)

            if empty_block is subroutine.body[0]:
                # place target at start of body so it's now the new entry block
                subroutine.body.remove(target)
                subroutine.body.insert(0, target)
            subroutine.body.remove(empty_block)
            changes = True
            logger.debug(f"Removed empty block: {empty_block}")
    return changes


"""
The compiled reference replacement is part of the optimizer pipeline for two reasons:

1.) It relies on any template variables provided being optimized into constant values
2.) Once compiled references are replaced there are additional optimizations that can occur
"""

import typing
from collections.abc import Mapping

import attrs
from immutabledict import immutabledict

from puya import log
from puya.algo_constants import HASH_PREFIX_PROGRAM, MAX_BYTES_LENGTH
from puya.awst.txn_fields import TxnField
from puya.compilation_artifacts import TemplateValue
from puya.context import ArtifactCompileContext
from puya.errors import CodeError, InternalError
from puya.ir import models as ir
from puya.ir.types_ import AVMBytesEncoding
from puya.ir.visitor_mutator import IRMutator
from puya.program_refs import ProgramKind
from puya.utils import (
    Address,
    biguint_bytes_eval,
    calculate_extra_program_pages,
    method_selector_hash,
    sha512_256_hash,
)

logger = log.get_logger(__name__)

def replace_compiled_references(
    context: ArtifactCompileContext, subroutine: ir.Subroutine
) -> bool:
    replacer = CompiledReferenceReplacer(context)
    for block in subroutine.body:
        replacer.visit_block(block)
    return replacer.modified

@attrs.define
class CompiledReferenceReplacer(IRMutator):
    context: ArtifactCompileContext
    modified: bool = False

    def visit_compiled_logicsig_reference(  # type: ignore[override]
        self,
        const: ir.CompiledLogicSigReference,
    ) -> ir.CompiledLogicSigReference | ir.Constant:
        if not _is_constant(const.template_variables):
            return const
        template_constants = _get_template_constants(const.template_variables)
        program_bytecode = self.context.build_program_bytecode(
            const.artifact, ProgramKind.logic_signature, template_constants=template_constants
        )
        address_public_key = sha512_256_hash(HASH_PREFIX_PROGRAM + program_bytecode)
        return ir.AddressConstant(
            value=Address.from_public_key(address_public_key).address,
            source_location=const.source_location,
        )

    def visit_compiled_contract_reference(  # type: ignore[override]
        self,
        const: ir.CompiledContractReference,
    ) -> ir.CompiledContractReference | ir.Constant:
        field = const.field
        if field in (
            TxnField.GlobalNumUint,
            TxnField.GlobalNumByteSlice,
            TxnField.LocalNumUint,
            TxnField.LocalNumByteSlice,
        ):
            state_total = self.context.get_state_totals(const.artifact)
            match field:
                case TxnField.GlobalNumUint:
                    total = state_total.global_uints
                case TxnField.GlobalNumByteSlice:
                    total = state_total.global_bytes
                case TxnField.LocalNumUint:
                    total = state_total.local_uints
                case TxnField.LocalNumByteSlice:
                    total = state_total.local_bytes
                case _:
                    raise InternalError(
                        f"Invalid state total field: {field.name}", const.source_location
                    )
            return ir.UInt64Constant(
                value=total,
                source_location=const.source_location,
            )

        if not _is_constant(const.template_variables):
            return const
        template_constants = _get_template_constants(const.template_variables)
        match field:
            case TxnField.ApprovalProgramPages | TxnField.ClearStateProgramPages:
                page = const.program_page
                if page is None:
                    raise InternalError("expected non-none value for page", const.source_location)
                program_bytecode = self.context.build_program_bytecode(
                    const.artifact,
                    (
                        ProgramKind.approval
                        if field == TxnField.ApprovalProgramPages
                        else ProgramKind.clear_state
                    ),
                    template_constants=template_constants,
                )
                program_page = program_bytecode[
                    page * MAX_BYTES_LENGTH : (page + 1) * MAX_BYTES_LENGTH
                ]
                return ir.BytesConstant(
                    value=program_page,
                    encoding=AVMBytesEncoding.base64,
                    source_location=const.source_location,
                )
            case TxnField.ExtraProgramPages:
                approval_bytecode = self.context.build_program_bytecode(
                    const.artifact, ProgramKind.approval, template_constants=template_constants
                )
                clear_bytecode = self.context.build_program_bytecode(
                    const.artifact, ProgramKind.clear_state, template_constants=template_constants
                )
                return ir.UInt64Constant(
                    value=calculate_extra_program_pages(
                        len(approval_bytecode), len(clear_bytecode)
                    ),
                    source_location=const.source_location,
                )
        raise InternalError(
            f"Unhandled compiled reference field: {field.name}", const.source_location
        )

def _is_constant(
    template_variables: Mapping[str, ir.Value],
) -> typing.TypeGuard[Mapping[str, ir.Constant]]:
    return all(isinstance(var, ir.Constant) for var in template_variables.values())

def _get_template_constants(
    template_variables: Mapping[str, ir.Constant],
) -> immutabledict[str, TemplateValue]:
    result = {
        var: (_extract_constant_value(value), value.source_location)
        for var, value in template_variables.items()
    }
    return immutabledict(result)

def _extract_constant_value(value: ir.Constant) -> int | bytes:
    match value:
        case ir.UInt64Constant(value=int_value):
            return int_value
        case ir.BytesConstant(value=bytes_value):
            return bytes_value
        case ir.BigUIntConstant(value=biguint):
            return biguint_bytes_eval(biguint)
        case ir.AddressConstant(value=addr):
            address = Address.parse(addr)
            return address.public_key
        case ir.MethodConstant(value=method):
            return method_selector_hash(method)
        case ir.ITxnConstant():
            raise CodeError(
                "inner transactions cannot be used as a template variable",
                value.source_location,
            )
        case _:
            raise InternalError(
                f"unhandled constant type: {type(value).__name__}",
                location=value.source_location,
            )


from collections import defaultdict

import attrs

from puya import log
from puya.context import CompileContext
from puya.ir import models
from puya.ir.register_read_collector import RegisterReadCollector
from puya.ir.visitor_mutator import IRMutator

logger = log.get_logger(__name__)

_AnyOp = models.Op | models.ControlOp

def constant_replacer(_context: CompileContext, subroutine: models.Subroutine) -> bool:
    constants = dict[models.Register, models.Constant]()
    ssa_reads = defaultdict[models.Register, list[_AnyOp]](list)
    for block in subroutine.body:
        assert block.terminator is not None
        ops: tuple[_AnyOp, ...] = (*block.ops, block.terminator)
        for op in ops:
            match op:
                case models.Assignment(targets=[register], source=models.Constant() as constant):
                    constants[register] = constant
                case _:
                    collector = RegisterReadCollector()
                    op.accept(collector)
                    for read_reg in collector.used_registers:
                        ssa_reads[read_reg].append(op)

    modified = 0
    work_list = constants.copy()
    while work_list:
        const_reg, const_val = work_list.popitem()
        replacer = ConstantRegisterReplacer({const_reg: const_val})
        for const_read in ssa_reads[const_reg]:
            if isinstance(const_read, models.Assignment) and const_read.source == const_reg:
                (register,) = const_read.targets
                constants[register] = work_list[register] = const_val
            const_read.accept(replacer)
        modified += replacer.modified

    phi_replace = {
        phi.register: phi_constant
        for block in subroutine.body
        for phi in block.phis
        if (phi_constant := _get_singular_phi_constant(phi, constants)) is not None
    }
    if phi_replace:
        modified += ConstantRegisterReplacer.apply(phi_replace, to=subroutine)
    return modified > 0

def _get_singular_phi_constant(
    phi: models.Phi, constants: dict[models.Register, models.Constant]
) -> models.Constant | None:
    try:
        (constant,) = {constants[phi_arg.value] for phi_arg in phi.args}
    except (KeyError, ValueError):
        return None
    else:
        return constant

@attrs.define
class ConstantRegisterReplacer(IRMutator):
    constants: dict[models.Register, models.Constant]
    modified: int = 0

    @classmethod
    def apply(
        cls, constants: dict[models.Register, models.Constant], to: models.Subroutine
    ) -> int:
        replacer = cls(constants)
        for block in to.body:
            replacer.visit_block(block)
        return replacer.modified

    def visit_assignment(self, ass: models.Assignment) -> models.Assignment:
        # don't visit target(s), needs to stay as Register
        ass.source = ass.source.accept(self)
        return ass

    def visit_phi(self, phi: models.Phi) -> models.Phi:
        # don't visit phi nodes, needs to stay as Register
        return phi

    def visit_register(self, reg: models.Register) -> models.Register:
        try:
            const = self.constants[reg]
        except KeyError:
            return reg
        self.modified += 1
        return const  # type: ignore[return-value]


import attrs

from puya.context import ArtifactCompileContext

@attrs.define(kw_only=True)
class IROptimizationContext(ArtifactCompileContext):
    expand_all_bytes: bool

    inlineable_calls: set[tuple[str, str]] = attrs.field(factory=set)
    """src -> dst pairs that can/should be inlined"""
    constant_with_constant_args: dict[str, bool] = attrs.field(factory=dict)


import attrs

from puya import log
from puya.avm import AVMType
from puya.context import CompileContext
from puya.ir import models
from puya.ir.avm_ops import AVMOp
from puya.ir.optimize._utils import get_definition
from puya.ir.ssa import TrivialPhiRemover
from puya.ir.types_ import IRType
from puya.utils import unique

logger = log.get_logger(__name__)

# the ratio of default cases to all cases when a match of constant values
# can be simplified to a goto-nth
_SWITCH_SPARSENESS_SIMPLIFICATION_RATIO = 0.5

def simplify_control_ops(_context: CompileContext, subroutine: models.Subroutine) -> bool:
    changes = False
    modified_phis = []

    def remove_target(parent: models.BasicBlock, to_remove: models.BasicBlock) -> None:
        to_remove.predecessors.remove(parent)
        for other_phi in to_remove.phis:
            other_phi.args = [arg for arg in other_phi.args if arg.through is not parent]
            modified_phis.append(other_phi)

    for block in subroutine.body:
        terminator = block.terminator
        match terminator:
            case models.ProgramExit(result=models.UInt64Constant(value=0)):
                logger.debug("simplifying exit 0 to err")
                block.terminator = models.Fail(
                    source_location=terminator.source_location, error_message=None
                )
            case models.ConditionalBranch(
                condition=models.UInt64Constant(value=value), zero=zero, non_zero=non_zero
            ):
                logger.debug("simplifying conditional branch with a constant into a goto")
                if value == 0:
                    goto, other = zero, non_zero
                else:
                    goto, other = non_zero, zero
                block.terminator = models.Goto(
                    source_location=terminator.source_location, target=goto
                )
                if other is not goto:
                    remove_target(block, other)
            case models.ConditionalBranch(
                condition=condition,
                zero=models.BasicBlock(
                    phis=[], ops=[], terminator=models.Fail(error_message=fail_comment)
                ) as err_block,
                non_zero=non_zero,
                source_location=source_location,
            ):
                logger.debug("inlining condition branch to err block into an assert true")
                block.ops.append(
                    models.Intrinsic(
                        op=AVMOp.assert_,
                        args=[condition],
                        error_message=fail_comment,
                        source_location=source_location,
                    )
                )
                block.terminator = models.Goto(target=non_zero, source_location=source_location)
                if err_block not in block.successors:
                    remove_target(block, err_block)
            case models.ConditionalBranch(
                condition=models.Register() as condition,
                non_zero=models.BasicBlock(
                    phis=[], ops=[], terminator=models.Fail(error_message=fail_comment)
                ) as err_block,
                zero=zero,
                source_location=source_location,
            ):
                logger.debug("inlining condition branch to err block into an assert false")
                not_condition = models.Register(
                    name=f"not%{condition.name}",
                    ir_type=IRType.bool,
                    version=condition.version,
                    source_location=source_location,
                )
                if get_definition(subroutine, not_condition, should_exist=False) is None:
                    block.ops.append(
                        models.Assignment(
                            targets=[not_condition],
                            source=models.Intrinsic(
                                op=AVMOp.not_,
                                args=[condition],
                                source_location=source_location,
                            ),
                            source_location=source_location,
                        )
                    )
                block.ops.append(
                    models.Intrinsic(
                        op=AVMOp.assert_,
                        args=[not_condition],
                        error_message=fail_comment,
                        source_location=source_location,
                    )
                )
                block.terminator = models.Goto(target=zero, source_location=source_location)
                if err_block not in block.successors:
                    remove_target(block, err_block)
            case (
                models.ConditionalBranch(
                    condition=models.Register() as condition,
                ) as branch
            ) if (
                isinstance(defn := get_definition(subroutine, condition), models.Assignment)
                and isinstance(defn.source, models.Intrinsic)
                and defn.source.op is AVMOp.not_
            ):
                logger.debug(
                    f"simplified branch on !{condition} by swapping zero and non-zero targets"
                )
                block.terminator = attrs.evolve(
                    branch,
                    zero=branch.non_zero,
                    non_zero=branch.zero,
                    condition=defn.source.args[0],
                )
            case (
                models.Switch(
                    value=models.Value(atype=AVMType.uint64) as value,
                    cases=cases,
                    default=default_block,
                    source_location=source_location,
                ) as switch
            ) if _can_simplify_switch(switch):
                logger.debug("simplifying a switch with constants into goto nth")
                # reduce to GotoNth
                block_map = dict[int, models.BasicBlock]()
                for case, case_block in cases.items():
                    assert isinstance(case, models.UInt64Constant)
                    block_map[case.value] = case_block
                max_value = max(block_map)
                block.terminator = models.GotoNth(
                    value=value,
                    blocks=[block_map.get(i, default_block) for i in range(max_value + 1)],
                    source_location=source_location,
                    default=default_block,
                )
            case models.GotoNth(
                value=models.UInt64Constant(value=value),
                blocks=blocks,
                default=default_block,
            ):
                logger.debug("simplifying a goto nth with a constant into a goto")
                goto = blocks[value] if value < len(blocks) else default_block
                block.terminator = models.Goto(
                    source_location=terminator.source_location, target=goto
                )
                for target in unique(terminator.targets()):
                    if target is not goto:
                        remove_target(block, target)
            case models.GotoNth(
                value=value,
                blocks=[zero],  # the constant here is the size of blocks
                default=non_zero,
            ):  # reduces to ConditionalBranch
                logger.debug("simplifying a goto nth with two targets into a conditional branch")
                block.terminator = models.ConditionalBranch(
                    condition=value,
                    zero=zero,
                    non_zero=non_zero,
                    source_location=terminator.source_location,
                )
            case _:
                continue
        changes = True
        logger.debug(f"simplified terminator of {block} from {terminator} to {block.terminator}")
    for phi in modified_phis:
        TrivialPhiRemover.try_remove(phi, subroutine.body)
    return changes

def _can_simplify_switch(switch: models.Switch) -> bool:
    total_targets = 0
    for case in switch.cases:
        if not isinstance(case, models.UInt64Constant):
            return False
        total_targets = max(total_targets, case.value)
    default_targets = total_targets - len(switch.cases)
    return default_targets < (total_targets * _SWITCH_SPARSENESS_SIMPLIFICATION_RATIO)


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

import attrs

from puya import log
from puya.context import CompileContext
from puya.errors import InternalError
from puya.ir import models, visitor
from puya.ir._utils import bfs_block_order
from puya.ir.ssa import TrivialPhiRemover
from puya.utils import StableSet

logger = log.get_logger(__name__)

PURE_AVM_OPS = frozenset(
    [
        # group: ops that can't fail at runtime
        # `txn FirstValidTime` technically could fail, but shouldn't happen on mainnet?
        "txn",
        "sha256",
        "keccak256",
        "sha3_256",
        "sha512_256",
        "bitlen",
        # group: could only fail on a type error
        "!",
        "!=",
        "&",
        "&&",
        "<",
        "<=",
        "==",
        ">",
        ">=",
        "|",
        "||",
        "~",
        "addw",
        "mulw",
        "itob",
        "len",
        "select",
        "sqrt",
        "shl",
        "shr",
        # group: fail if an input is zero
        "%",
        "/",
        "expw",
        "divmodw",
        "divw",
        # group: fail on over/underflow
        "*",
        "+",
        "-",
        "^",
        "exp",
        # group: fail on index out of bounds
        "arg",
        "arg_0",
        "arg_1",
        "arg_2",
        "arg_3",
        "args",
        "extract",
        "extract3",
        "extract_uint16",
        "extract_uint32",
        "extract_uint64",
        "replace2",
        "replace3",
        "setbit",
        "setbyte",
        "getbit",
        "getbyte",
        "gaid",
        "gaids",
        "gload",
        "gloads",
        "gloadss",
        "substring",
        "substring3",
        "txna",
        "txnas",
        "gtxn",
        "gtxna",
        "gtxnas",
        "gtxns",
        "gtxnsa",
        "gtxnsas",
        "block",
        # group: fail on input too large
        "b%",
        "b*",
        "b+",
        "b-",
        "b/",
        "b^",
        "btoi",
        # group: might fail on input too large? TODO: verify these
        "b!=",
        "b<",
        "b<=",
        "b==",
        "b>",
        "b>=",
        "b&",
        "b|",
        "b~",
        "bsqrt",
        # group: fail on output too large
        "concat",
        "bzero",
        # group: fail on input format / byte lengths
        "base64_decode",
        "json_ref",
        "ecdsa_pk_decompress",
        "ecdsa_pk_recover",
        "ec_add",
        "ec_pairing_check",
        "ec_scalar_mul",
        "ec_subgroup_check",
        "ec_multi_scalar_mul",
        "ec_map_to",
        "ecdsa_verify",
        "ed25519verify",
        "ed25519verify_bare",
        "vrf_verify",
        # group: v11
        "falcon_verify",
        "sumhash512",
    ]
)

# ops that have no observable side effects outside the function
# note: originally generated basd on all ops that:
#       - return a stack value (this, as of v10, yields no false negatives)
#       - AND isn't in the generate_avm_ops.py list of exclusions (which are all control flow
#             or pure stack manipulations)
#       - AND isn't box_create or box_del, they were the only remaining false positives
IMPURE_SIDE_EFFECT_FREE_AVM_OPS = frozenset(
    [
        # group: ops that can't fail at runtime
        "global",  # OpcodeBudget is non-const, otherwise this could be pure
        # group: could only fail on a type error
        "app_global_get",
        "app_global_get_ex",
        "load",
        # group: fail on resource not "available"
        # TODO: determine if any of this group is pure
        "acct_params_get",
        "app_opted_in",
        "app_params_get",
        "asset_holding_get",
        "asset_params_get",
        "app_local_get",
        "app_local_get_ex",
        "balance",
        "min_balance",
        "box_extract",
        "box_get",
        "box_len",
        # group: fail on index out of bounds
        "loads",
        # group: might fail depending on state
        "itxn",
        "itxna",
        "itxnas",
        "gitxn",
        "gitxna",
        "gitxnas",
    ]
)

_should_be_empty = PURE_AVM_OPS & IMPURE_SIDE_EFFECT_FREE_AVM_OPS
assert not _should_be_empty, _should_be_empty
SIDE_EFFECT_FREE_AVM_OPS = frozenset([*PURE_AVM_OPS, *IMPURE_SIDE_EFFECT_FREE_AVM_OPS])

@attrs.define
class SubroutineCollector(visitor.IRTraverser):
    subroutines: StableSet[models.Subroutine] = attrs.field(factory=StableSet)

    @classmethod
    def collect(cls, program: models.Program) -> StableSet[models.Subroutine]:
        collector = cls()
        collector.visit_subroutine(program.main)
        return collector.subroutines

    def visit_subroutine(self, subroutine: models.Subroutine) -> None:
        if subroutine not in self.subroutines:
            self.subroutines.add(subroutine)
            self.visit_all_blocks(subroutine.body)

    def visit_invoke_subroutine(self, callsub: models.InvokeSubroutine) -> None:
        self.visit_subroutine(callsub.target)

def remove_unused_subroutines(program: models.Program) -> bool:
    subroutines = SubroutineCollector.collect(program)
    if modified := (len(subroutines) != (1 + len(program.subroutines))):
        to_keep = [p for p in program.subroutines if p in subroutines]
        for p in program.subroutines:
            if p not in subroutines:
                logger.debug(f"removing unused subroutine {p.id}")
        program.subroutines = to_keep
    return modified

def remove_unused_variables(_context: CompileContext, subroutine: models.Subroutine) -> bool:
    modified = 0
    assignments = dict[tuple[models.BasicBlock, models.Assignment], set[models.Register]]()
    for block, op, register in UnusedRegisterCollector.collect(subroutine):
        if isinstance(op, models.Assignment):
            assignments.setdefault((block, op), set()).add(register)
        else:
            assert register == op.register
            block.phis.remove(op)
            logger.debug(f"Removing unused variable {register.local_id}")
            modified += 1

    for (block, ass), registers in assignments.items():
        if registers.symmetric_difference(ass.targets):
            pass  # some registers still used
        elif isinstance(ass.source, models.Value | models.InnerTransactionField) or (
            isinstance(ass.source, models.Intrinsic)
            and ass.source.op.code in SIDE_EFFECT_FREE_AVM_OPS
        ):
            for reg in sorted(registers, key=lambda r: r.local_id):
                logger.debug(f"Removing unused variable {reg.local_id}")
            block.ops.remove(ass)
            modified += 1
        else:
            logger.debug(
                f"Not removing unused assignment since source is not marked as pure: {ass}"
            )
    return modified > 0

@attrs.define(kw_only=True)
class UnusedRegisterCollector(visitor.IRTraverser):
    used: set[models.Register] = attrs.field(factory=set)
    assigned: dict[models.Register, tuple[models.BasicBlock, models.Assignment | models.Phi]] = (
        attrs.field(factory=dict)
    )
    active_block: models.BasicBlock

    @classmethod
    def collect(
        cls, sub: models.Subroutine
    ) -> Iterable[tuple[models.BasicBlock, models.Assignment | models.Phi, models.Register]]:
        collector = cls(active_block=sub.entry)
        collector.visit_all_blocks(sub.body)
        for reg, (block, ass) in collector.assigned.items():
            if reg not in collector.used:
                yield block, ass, reg

    @typing.override
    def visit_block(self, block: models.BasicBlock) -> None:
        self.active_block = block
        super().visit_block(block)

    @typing.override
    def visit_assignment(self, ass: models.Assignment) -> None:
        for target in ass.targets:
            self.assigned[target] = (self.active_block, ass)
        ass.source.accept(self)

    @typing.override
    def visit_phi(self, phi: models.Phi) -> None:
        # don't visit phi.register, as this would mean the phi can never be considered unused
        for arg in phi.args:
            arg.accept(self)
        self.assigned[phi.register] = (self.active_block, phi)

    @typing.override
    def visit_register(self, reg: models.Register) -> None:
        self.used.add(reg)

def remove_unreachable_blocks(_context: CompileContext, subroutine: models.Subroutine) -> bool:
    reachable_set = frozenset(bfs_block_order(subroutine.entry))
    unreachable_blocks = [b for b in subroutine.body if b not in reachable_set]
    if not unreachable_blocks:
        return False

    logger.debug(f"Removing unreachable blocks: {', '.join(map(str, unreachable_blocks))}")

    reachable_blocks = []
    for block in subroutine.body:
        if block in reachable_set:
            reachable_blocks.append(block)
            if not reachable_set.issuperset(block.successors):
                raise InternalError(
                    f"Block {block} has unreachable successor(s),"
                    f" but was not marked as unreachable itself"
                )
            if not reachable_set.issuperset(block.predecessors):
                block.predecessors = [b for b in block.predecessors if b in reachable_set]
                logger.debug(f"Removed unreachable predecessors from {block}")

    UnreachablePhiArgsRemover.apply(unreachable_blocks, reachable_blocks)
    subroutine.body = reachable_blocks
    return True

@attrs.define
class UnreachablePhiArgsRemover(visitor.IRTraverser):
    _unreachable_blocks: Set[models.BasicBlock]
    _reachable_blocks: Sequence[models.BasicBlock]

    @classmethod
    def apply(
        cls,
        unreachable_blocks: Sequence[models.BasicBlock],
        reachable_blocks: Sequence[models.BasicBlock],
    ) -> None:
        collector = cls(frozenset(unreachable_blocks), reachable_blocks)
        collector.visit_all_blocks(reachable_blocks)

    def visit_phi(self, phi: models.Phi) -> None:
        args_to_remove = [a for a in phi.args if a.through in self._unreachable_blocks]
        if not args_to_remove:
            return
        logger.debug(
            "Removing unreachable phi arguments: " + ", ".join(sorted(map(str, args_to_remove)))
        )
        phi.args = [a for a in phi.args if a not in args_to_remove]
        if not phi.non_self_args:
            raise InternalError(
                f"undefined phi created when removing args through "
                f"{', '.join(map(str, self._unreachable_blocks))}"
            )
        TrivialPhiRemover.try_remove(phi, self._reachable_blocks)


import copy
import itertools
import typing
from collections import defaultdict
from collections.abc import Collection, Iterable, Iterator, Mapping, Sequence

import attrs
import networkx as nx  # type: ignore[import-untyped]

from puya import log
from puya.ir import models
from puya.ir.context import TMP_VAR_INDICATOR
from puya.ir.optimize._call_graph import CallGraph
from puya.ir.optimize.context import IROptimizationContext
from puya.ir.optimize.intrinsic_simplification import COMPILE_TIME_CONSTANT_OPS
from puya.ir.visitor import IRTraverser
from puya.ir.visitor_mutator import IRMutator
from puya.utils import lazy_setdefault

logger = log.get_logger(__name__)

def analyse_subroutines_for_inlining(
    context: IROptimizationContext,
    program: models.Program,
    routable_method_ids: Collection[str] | None,
) -> None:
    context.inlineable_calls.clear()
    context.constant_with_constant_args.clear()

    call_graph = CallGraph(program)
    for sub in program.subroutines:
        if sub.inline is False:
            pass  # nothing to do
        elif any(phi.args for phi in sub.entry.phis):
            logger.debug(
                f"function has phi node(s) with arguments in entry block: {sub.id}",
                location=sub.source_location,
            )
            if sub.inline is True:
                logger.warning(
                    "function not suitable for inlining due to complex control flow",
                    location=sub.source_location,
                )
            sub.inline = False
        elif call_graph.is_auto_recursive(sub):
            logger.debug(f"function is auto-recursive: {sub.id}", location=sub.source_location)
            if sub.inline is True:
                logger.warning("unable to inline recursive function", location=sub.source_location)
            sub.inline = False
        elif sub.inline is True:
            for callee_id, _ in call_graph.callees(sub):
                if not call_graph.has_path(sub.id, callee_id):
                    context.inlineable_calls.add((callee_id, sub.id))
                else:
                    logger.warning(
                        f"not inlining call from {callee_id} to {sub.id}"
                        f" because call may be re-entrant",
                        location=sub.source_location,
                    )
        else:
            typing.assert_type(sub.inline, None)
    if context.options.optimization_level == 0:
        return

    # for optimization levels below 2, skip auto-inlining routable methods into the approval
    # program. mostly this can be beneficial in terms of #ops, but not always.
    # also, it impacts the debugging experience
    skip_routable_ids = frozenset[str]()
    if context.options.optimization_level < 2:
        skip_routable_ids = frozenset(routable_method_ids or ())

    for sub in program.subroutines:
        if sub.inline is None:
            if _is_trivial(sub):
                # special case for trivial methods, even when routable
                logger.debug(f"marking trivial method {sub.id} as inlineable")
                for callee_id, _ in call_graph.callees(sub):
                    # don't need to check re-entrancy here since there's no callsub
                    context.inlineable_calls.add((callee_id, sub.id))
            elif sub.id not in skip_routable_ids:
                match call_graph.callees(sub):
                    case [(callee_id, 1)]:
                        assert (
                            callee_id != sub.id
                        ), f"function {sub.id} is auto-recursive and disconnected"
                        logger.debug(f"marking single-use function {sub.id} for inlining")
                        context.inlineable_calls.add((callee_id, sub.id))
    if context.inlineable_calls:
        return
    for sub in program.subroutines:
        if sub.inline is None and sub.id not in skip_routable_ids:
            complexity = sum(
                len(b.phis) + len(b.ops) + len(_not_none(b.terminator).targets()) for b in sub.body
            )
            threshold = max(3, 1 + len(sub._returns) + len(sub.parameters))  # noqa: SLF001
            if complexity <= threshold:
                logger.debug(
                    f"marking simple function {sub.id} for inlining"
                    f" ({complexity=} <= {threshold=})"
                )
                for callee_id, _ in call_graph.callees(sub):
                    if not call_graph.has_path(sub.id, callee_id):
                        context.inlineable_calls.add((callee_id, sub.id))

def _is_trivial(sub: models.Subroutine) -> bool:
    match sub.body:
        case [models.BasicBlock(phis=[], ops=[])]:
            return True
    return False

def perform_subroutine_inlining(
    context: IROptimizationContext, subroutine: models.Subroutine
) -> bool:
    inline_calls_to = {
        to_id for from_id, to_id in context.inlineable_calls if from_id == subroutine.id
    }
    if not (inline_calls_to or context.options.optimization_level >= 2):
        return False
    modified = False
    blocks_to_visit = subroutine.body.copy()
    max_block_id = max(_not_none(block.id) for block in blocks_to_visit)
    next_id = itertools.count(max_block_id + 1)
    while blocks_to_visit:
        block = blocks_to_visit.pop()
        for op_index, op in enumerate(block.ops):
            match op:
                case models.Assignment(
                    targets=return_targets, source=models.InvokeSubroutine() as call
                ):
                    pass
                case models.InvokeSubroutine() as call:
                    return_targets = []
                case _:
                    continue
            if call.target.id in inline_calls_to:
                logger.debug(
                    f"inlining call to {call.target.id} in {subroutine.id}",
                    location=op.source_location,
                )
            elif (
                context.options.optimization_level >= 2
                and call.target.inline is None
                and all(isinstance(arg, models.Constant) for arg in call.args)
                and lazy_setdefault(
                    context.constant_with_constant_args,
                    call.target.id,
                    lambda _: _ConstantFunctionDetector.is_constant_with_constant_args(
                        call.target
                    ),
                )
            ):
                logger.debug(
                    f"constant function call to {call.target.id} in {subroutine.id}",
                    location=op.source_location,
                )
            else:
                continue

            remainder, created_blocks = _inline_call(
                block,
                call,
                next_id,
                op_index,
                return_targets,
                list(subroutine.get_assigned_registers()),
            )
            # only visit the remainder of the block, the newly created blocks
            # come from a different subroutine, so shouldn't be tested for inlining
            # within the curren subroutine
            blocks_to_visit.append(remainder)
            created_blocks.append(remainder)
            idx_after_block = subroutine.body.index(block) + 1
            subroutine.body[idx_after_block:idx_after_block] = created_blocks
            modified = True
            break  # we're done with this block, since it's been split
    return modified

def _inline_call(
    block: models.BasicBlock,
    call: models.InvokeSubroutine,
    next_id: Iterator[int],
    op_index: int,
    return_targets: Sequence[models.Register],
    host_assigned_registers: list[models.Register],
) -> tuple[models.BasicBlock, list[models.BasicBlock]]:
    # make a copy of the entire block graph, and adjust register versions
    # to avoid any collisions, as well as updating block IDs
    register_offsets = defaultdict[str, int](int)
    for host_reg in host_assigned_registers:
        register_offsets[host_reg.name] = max(
            register_offsets[host_reg.name], host_reg.version + 1
        )
    new_blocks = _inlined_blocks(call.target, register_offsets)
    for new_block in new_blocks:
        new_block.id = next(next_id)
    # split the block after the callsub instruction
    ops_after = block.ops[op_index + 1 :]
    del block.ops[op_index:]
    terminator_after = block.terminator
    assert terminator_after is not None
    for succ in terminator_after.unique_targets:
        succ.predecessors.remove(block)
    inlined_entry = new_blocks[0]
    block.terminator = models.Goto(target=inlined_entry, source_location=call.source_location)
    inlined_entry.predecessors.append(block)
    for phi in inlined_entry.phis:
        # TODO: assign undefined and use that as phi arg so we can inline phi nodes with args.
        #       this requires finding a new register though, and is quite hard to write code
        #       that would trigger this, so leaving for now
        assert not phi.args, "entry phi with arguments should have been prevented from inlining"
        inlined_entry.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,
            ),
        )
    inlined_entry.phis.clear()
    # assign parameters before call
    for arg, param in zip(call.args, call.target.parameters, strict=True):
        updated_param = models.Register(
            name=param.name,
            ir_type=param.ir_type,
            version=param.version + register_offsets[param.name],
            source_location=param.source_location,
        )
        block.ops.append(
            models.Assignment(
                targets=[updated_param], source=arg, source_location=arg.source_location
            )
        )

    returning_blocks = [
        (new_block, new_block.terminator.result)
        for new_block in new_blocks
        if isinstance(new_block.terminator, models.SubroutineReturn)
    ]
    # create the return block and insert as a predecessor of the original blocks target(s)
    remainder = models.BasicBlock(
        id=next(next_id),
        phis=[],
        ops=ops_after,
        terminator=terminator_after,
        predecessors=[],
        comment=f"after_inlined_{call.target.id}",
        source_location=block.source_location,
    )
    for succ in remainder.successors:
        for phi in succ.phis:
            for phi_arg in phi.args:
                if phi_arg.through == block:
                    phi_arg.through = remainder
        succ.predecessors.append(remainder)
    # replace inlined retsubs with unconditional branches to the second block half
    for new_block, _ in returning_blocks:
        new_block.terminator = models.Goto(target=remainder, source_location=call.source_location)
        remainder.predecessors.append(new_block)

    num_returns = len(returning_blocks)
    if num_returns == 1:
        # if there is a single retsub, we can assign to the return variables in that block
        # directly without violating SSA
        ((new_block, return_values),) = returning_blocks
        if return_targets:
            new_block.ops.append(
                models.Assignment(
                    targets=return_targets,
                    source=models.ValueTuple(values=return_values, source_location=None),
                    source_location=None,
                )
            )
    elif num_returns > 1:
        # otherwise when there's more than on restsub block,
        # return value(s) become phi node(s) in the second block half
        return_phis = [models.Phi(register=ret_target) for ret_target in return_targets]
        for new_block_idx, (new_block, return_values) in enumerate(returning_blocks):
            for ret_idx, ret_phi in enumerate(return_phis):
                ret_value = return_values[ret_idx]
                if not isinstance(ret_value, models.Register):
                    tmp_value = ret_value
                    tmp_reg_name = f"{call.target.id}{TMP_VAR_INDICATOR}{ret_idx}"
                    ret_value = models.Register(
                        ir_type=ret_value.ir_type,
                        source_location=ret_value.source_location,
                        name=tmp_reg_name,
                        version=new_block_idx + register_offsets[tmp_reg_name],
                    )
                    new_block.ops.append(
                        models.Assignment(
                            targets=[ret_value],
                            source=tmp_value,
                            source_location=tmp_value.source_location,
                        )
                    )
                ret_phi.args.append(models.PhiArgument(value=ret_value, through=new_block))

        remainder.phis = return_phis
    return remainder, new_blocks

def _not_none[T](x: T | None) -> T:
    assert x is not None
    return x

def _inlined_blocks(
    sub: models.Subroutine, register_offsets: Mapping[str, int]
) -> list[models.BasicBlock]:
    ref_collector = _SubroutineReferenceCollector()
    ref_collector.visit_all_blocks(sub.body)
    memo = {id(s): s for s in ref_collector.subroutines}
    blocks = copy.deepcopy(sub.body, memo=memo)
    _OffsetRegisterVersions.apply(blocks, register_offsets=register_offsets)
    return blocks

@attrs.define
class _OffsetRegisterVersions(IRMutator):
    register_offsets: Mapping[str, int]

    @classmethod
    def apply(
        cls, blocks: Iterable[models.BasicBlock], *, register_offsets: Mapping[str, int]
    ) -> None:
        replacer = cls(register_offsets=register_offsets)
        for block in blocks:
            replacer.visit_block(block)

    def visit_register(self, reg: models.Register) -> models.Register:
        return models.Register(
            name=reg.name,
            ir_type=reg.ir_type,
            version=reg.version + self.register_offsets[reg.name],
            source_location=reg.source_location,
        )

@attrs.define
class _SubroutineReferenceCollector(IRTraverser):
    subroutines: set[models.Subroutine] = attrs.field(factory=set)

    @typing.override
    def visit_invoke_subroutine(self, callsub: models.InvokeSubroutine) -> None:
        self.subroutines.add(callsub.target)
        super().visit_invoke_subroutine(callsub)

class _NonConstantFunctionError(Exception):
    pass

class _ConstantFunctionDetector(IRTraverser):
    def __init__(self) -> None:
        self._block_graph = nx.DiGraph()

    def has_cycle(self, *, start_from: int | None = None) -> bool:
        try:
            nx.find_cycle(self._block_graph, source=start_from)
        except nx.NetworkXNoCycle:
            return False
        else:
            return True

    @classmethod
    def is_constant_with_constant_args(cls, sub: models.Subroutine) -> bool:
        """Detect if a function is constant assuming all argument are constants"""
        visitor = cls()
        try:
            visitor.visit_all_blocks(sub.body)
        except _NonConstantFunctionError:
            return False
        return not visitor.has_cycle(start_from=sub.entry.id)

    @typing.override
    def visit_block(self, block: models.BasicBlock) -> None:
        super().visit_block(block)
        self._block_graph.add_node(block.id)
        for target in block.successors:
            self._block_graph.add_edge(block.id, target.id)

    @typing.override
    def visit_invoke_subroutine(self, callsub: models.InvokeSubroutine) -> None:
        raise _NonConstantFunctionError

    @typing.override
    def visit_template_var(self, deploy_var: models.TemplateVar) -> None:
        raise _NonConstantFunctionError

    @typing.override
    def visit_compiled_contract_reference(self, const: models.CompiledContractReference) -> None:
        raise _NonConstantFunctionError

    @typing.override
    def visit_compiled_logicsig_reference(self, const: models.CompiledLogicSigReference) -> None:
        raise _NonConstantFunctionError

    @typing.override
    def visit_intrinsic_op(self, intrinsic: models.Intrinsic) -> None:
        if intrinsic.op.code not in COMPILE_TIME_CONSTANT_OPS:
            raise _NonConstantFunctionError
        super().visit_intrinsic_op(intrinsic)

    @typing.override
    def visit_itxn_constant(self, const: models.ITxnConstant) -> None:
        raise _NonConstantFunctionError

    @typing.override
    def visit_inner_transaction_field(self, field: models.InnerTransactionField) -> None:
        raise _NonConstantFunctionError


import attrs

from puya.context import CompileContext
from puya.ir import models
from puya.ir.avm_ops import AVMOp
from puya.ir.visitor_mutator import IRMutator

def inner_txn_field_replacer(_context: CompileContext, subroutine: models.Subroutine) -> bool:
    modified = IntrinsicFieldReplacer.apply(to=subroutine)
    return modified > 0

@attrs.define
class IntrinsicFieldReplacer(IRMutator):
    modified: int = 0

    @classmethod
    def apply(cls, to: models.Subroutine) -> int:
        replacer = cls()
        for block in to.body:
            replacer.visit_block(block)
        return replacer.modified

    def visit_inner_transaction_field(
        self, field: models.InnerTransactionField
    ) -> models.InnerTransactionField | models.Intrinsic:
        match field.group_index:
            case models.ITxnConstant(value=group_index):
                pass
            case _:
                return field
        match field.is_last_in_group:
            case models.UInt64Constant(value=is_last_in_group):
                pass
            case _:
                is_last_in_group = 0
        self.modified += 1
        return (
            models.Intrinsic(
                op=AVMOp.itxnas if field.array_index else AVMOp.itxn,
                immediates=[field.field],
                args=[field.array_index] if field.array_index else [],
                source_location=field.source_location,
            )
            if is_last_in_group
            else models.Intrinsic(
                op=AVMOp.gitxnas if field.array_index else AVMOp.gitxn,
                immediates=[group_index, field.field],
                args=[field.array_index] if field.array_index else [],
                source_location=field.source_location,
            )
        )