from collections.abc import Iterator

import attrs

from puya import log
from puya.awst import (
    nodes as awst_nodes,
    wtypes,
)
from puya.awst.awst_traverser import AWSTTraverser

logger = log.get_logger(__name__)

class ARC4CopyValidator(AWSTTraverser):
    @classmethod
    def validate(cls, module: awst_nodes.AWST) -> None:
        validator = cls()
        for module_statement in module:
            module_statement.accept(validator)

    def __init__(self) -> None:
        super().__init__()
        self._for_items: awst_nodes.Lvalue | None = None

    # for nodes that can't modify the input don't need to check for copies unless an assignment
    # expression is being used
    def visit_submit_inner_transaction(self, call: awst_nodes.SubmitInnerTransaction) -> None:
        if _HasAssignmentVisitor.check(call):
            super().visit_submit_inner_transaction(call)

    def visit_intrinsic_call(self, call: awst_nodes.IntrinsicCall) -> None:
        if _HasAssignmentVisitor.check(call):
            super().visit_intrinsic_call(call)

    def visit_emit(self, emit: awst_nodes.Emit) -> None:
        if _HasAssignmentVisitor.check(emit):
            super().visit_emit(emit)

    def visit_assert_expression(self, expr: awst_nodes.AssertExpression) -> None:
        if expr.condition and _HasAssignmentVisitor.check(expr.condition):
            super().visit_assert_expression(expr)

    def visit_assignment_statement(self, statement: awst_nodes.AssignmentStatement) -> None:
        _check_assignment(statement.target, statement.value)
        statement.value.accept(self)

    def visit_tuple_expression(self, expr: awst_nodes.TupleExpression) -> None:
        super().visit_tuple_expression(expr)
        if expr is not self._for_items:
            for item in expr.items:
                _check_for_arc4_copy(item, "being passed to a tuple expression")

    def visit_for_in_loop(self, statement: awst_nodes.ForInLoop) -> None:
        # statement.items is immediately checked before entering the for body
        # so don't need to worry about preserving _for_items through multiple loops
        self._for_items = statement.items
        super().visit_for_in_loop(statement)
        self._for_items = None

        # looping is essentially assigning so check sequence
        sequence = statement.sequence
        while isinstance(sequence, awst_nodes.Enumeration | awst_nodes.Reversed):
            sequence = sequence.expr
        if (  # mutable tuples cannot be iterated in a semantically correct way
            isinstance(sequence.wtype, wtypes.WTuple)
            and _is_referable_expression(sequence)
            and _is_arc4_mutable(sequence.wtype)
        ):
            logger.error(
                "tuple of mutable ARC4 values cannot be iterated",
                location=sequence.source_location,
            )
        elif (  # arrays of mutable types, must be modified and iterated by index
            isinstance(sequence.wtype, wtypes.ARC4Array)
            and _is_referable_expression(sequence)
            and _is_arc4_mutable(sequence.wtype.element_type)
        ):
            logger.error(
                "cannot directly iterate an ARC4 array of mutable objects,"
                " construct a for-loop over the indexes instead",
                location=sequence.source_location,
            )

    def visit_assignment_expression(self, expr: awst_nodes.AssignmentExpression) -> None:
        _check_assignment(expr.target, expr.value)
        expr.value.accept(self)

    def visit_subroutine_call_expression(self, expr: awst_nodes.SubroutineCallExpression) -> None:
        for arg_ in expr.args:
            for arg in _expand_tuple_items(arg_.value):
                match arg:
                    case awst_nodes.VarExpression():
                        # Var expressions don't need copy as we implicitly return the latest value
                        # and update the var
                        continue
                    case awst_nodes.AppStateExpression() | awst_nodes.AppAccountStateExpression():
                        message = "being passed to a subroutine from state"
                    case _:
                        message = "being passed to a subroutine"
                _check_for_arc4_copy(arg, message)

    def visit_new_array(self, expr: awst_nodes.NewArray) -> None:
        super().visit_new_array(expr)
        if isinstance(expr.wtype, wtypes.ARC4Array):
            for v in expr.values:
                _check_for_arc4_copy(v, "being passed to an array constructor")

    def visit_new_struct(self, expr: awst_nodes.NewStruct) -> None:
        super().visit_new_struct(expr)
        if isinstance(expr.wtype, wtypes.ARC4Struct):
            for v in expr.values.values():
                _check_for_arc4_copy(v, "being passed to a struct constructor")

    def visit_arc4_encode(self, expr: awst_nodes.ARC4Encode) -> None:
        super().visit_arc4_encode(expr)

        for item in _expand_tuple_items(expr.value):
            _check_for_arc4_copy(item, "being passed to a constructor")

def _is_referable_expression(expr: awst_nodes.Expression) -> bool:
    """
    Returns True if expr represents something that can be referenced multiple times.
    """
    match expr:
        case (
            awst_nodes.VarExpression()
            | awst_nodes.AppStateExpression()
            | awst_nodes.AppAccountStateExpression()
            | awst_nodes.StateGet()
            | awst_nodes.StateGetEx()
            | awst_nodes.BoxValueExpression()
        ):
            return True
        case (
            awst_nodes.IndexExpression(base=base_expr)
            | awst_nodes.TupleItemExpression(base=base_expr)
            | awst_nodes.FieldExpression(base=base_expr)
        ):
            return _is_referable_expression(base_expr)
    return False

def _check_assignment(target: awst_nodes.Expression, value: awst_nodes.Expression) -> None:
    if not isinstance(target, awst_nodes.TupleExpression):
        _check_for_arc4_copy(value, "being assigned to another variable")
    elif _is_referable_expression(value) and any(_is_arc4_mutable(i) for i in target.wtype.types):
        logger.error(
            "tuples containing a mutable reference to an ARC4-encoded value cannot be unpacked,"
            " use index access instead",
            location=value.source_location,
        )

def _check_for_arc4_copy(expr: awst_nodes.Expression, context_desc: str) -> None:
    if _is_arc4_mutable(expr.wtype) and _is_referable_expression(expr):
        logger.error(
            "mutable reference to ARC4-encoded value"
            f" must be copied using .copy() when {context_desc}",
            location=expr.source_location,
        )

def _expand_tuple_items(expr: awst_nodes.Expression) -> Iterator[awst_nodes.Expression]:
    match expr:
        case awst_nodes.TupleExpression(items=items):
            for item in items:
                yield from _expand_tuple_items(item)
        case _:
            yield expr

def _is_arc4_mutable(wtype: wtypes.WType) -> bool:
    """
    Returns True if expr represents an arc4 type that is mutable
    """
    match wtype:
        case wtypes.ARC4Type(immutable=False):
            return True
        case wtypes.WTuple(types=types):
            return any(_is_arc4_mutable(t) for t in types)
    return False

@attrs.define
class _HasAssignmentVisitor(AWSTTraverser):
    has_assignment: bool = False

    @classmethod
    def check(cls, expr: awst_nodes.Expression) -> bool:
        visitor = _HasAssignmentVisitor()
        expr.accept(visitor)
        return visitor.has_assignment

    def visit_assignment_expression(self, _: awst_nodes.AssignmentExpression) -> None:
        self.has_assignment = True


import contextlib
import typing
from collections.abc import Iterator

from puya import log
from puya.awst import nodes as awst_nodes
from puya.awst.awst_traverser import AWSTTraverser
from puya.awst.nodes import (
    ContractMethodTarget,
    InstanceMethodTarget,
    InstanceSuperMethodTarget,
    SubroutineID,
)

logger = log.get_logger(__name__)

class BaseInvokerValidator(AWSTTraverser):
    def __init__(self) -> None:
        self._contract: awst_nodes.Contract | None = None

    @property
    def contract(self) -> awst_nodes.Contract | None:
        return self._contract

    @contextlib.contextmanager
    def _enter_contract(self, contract: awst_nodes.Contract) -> Iterator[None]:
        assert self._contract is None
        self._contract = contract
        try:
            yield
        finally:
            self._contract = None

    @classmethod
    def validate(cls, module: awst_nodes.AWST) -> None:
        validator = cls()
        for module_statement in module:
            module_statement.accept(validator)

    @typing.override
    def visit_contract(self, statement: awst_nodes.Contract) -> None:
        with self._enter_contract(statement):
            super().visit_contract(statement)

    @typing.override
    def visit_subroutine_call_expression(self, expr: awst_nodes.SubroutineCallExpression) -> None:
        super().visit_subroutine_call_expression(expr)
        match expr.target:
            case SubroutineID():
                # always okay
                pass
            case InstanceMethodTarget() | InstanceSuperMethodTarget():
                if self.contract is None:
                    logger.error(
                        "invocation of instance method outside of a contract method",
                        location=expr.source_location,
                    )
            case ContractMethodTarget(cref=target_class):
                caller_class = self.contract
                if caller_class is None:
                    logger.error(
                        "invocation of contract method outside of a contract method",
                        location=expr.source_location,
                    )
                else:
                    caller_ref = caller_class.id
                    if (
                        target_class != caller_ref
                        and target_class not in caller_class.method_resolution_order
                    ):
                        logger.error(
                            "invocation of a contract method outside of current hierarchy",
                            location=expr.source_location,
                        )
            case invalid:
                typing.assert_never(invalid)


from puya import log
from puya.awst import nodes as awst_nodes
from puya.awst.awst_traverser import AWSTTraverser

logger = log.get_logger(__name__)

class ImmutableValidator(AWSTTraverser):
    @classmethod
    def validate(cls, module: awst_nodes.AWST) -> None:
        validator = cls()
        for module_statement in module:
            module_statement.accept(validator)

    def visit_assignment_expression(self, expr: awst_nodes.AssignmentExpression) -> None:
        super().visit_assignment_expression(expr)
        _validate_lvalue(expr.target)

    def visit_assignment_statement(self, statement: awst_nodes.AssignmentStatement) -> None:
        super().visit_assignment_statement(statement)
        _validate_lvalue(statement.target)

    def visit_array_pop(self, expr: awst_nodes.ArrayPop) -> None:
        super().visit_array_pop(expr)
        if expr.base.wtype.immutable:
            logger.error(
                "cannot modify - object is immutable",
                location=expr.source_location,
            )

    def visit_array_extend(self, expr: awst_nodes.ArrayExtend) -> None:
        super().visit_array_extend(expr)
        if expr.base.wtype.immutable:
            logger.error(
                "cannot modify - object is immutable",
                location=expr.source_location,
            )

def _validate_lvalue(lvalue: awst_nodes.Expression) -> None:
    if isinstance(lvalue, awst_nodes.FieldExpression | awst_nodes.IndexExpression):
        if lvalue.base.wtype.immutable:
            logger.error(
                "expression is not valid as an assignment target - object is immutable",
                location=lvalue.source_location,
            )
    elif isinstance(lvalue, awst_nodes.TupleExpression):
        for item in lvalue.items:
            _validate_lvalue(item)


import contextlib
from collections.abc import Iterator

from puya import log
from puya.awst import (
    nodes as awst_nodes,
    wtypes,
)
from puya.awst.awst_traverser import AWSTTraverser
from puya.awst.wtypes import WInnerTransaction, WInnerTransactionFields
from puya.parse import SourceLocation

logger = log.get_logger(__name__)

INNER_TRANSACTION_ASSIGNMENT_EXPRESSION_ERROR = (
    "inner transactions cannot be used in assignment expressions"
)
INNER_TRANSACTION_COPY_REQUIRED_ERROR = (
    "inner transaction fields must be copied using .copy() when assigning to a new local"
)
INNER_TRANSACTION_LOOP_MODIFICATION_ERROR = (
    "inner transaction fields cannot be modified after submission while in a loop"
)
INNER_TRANSACTION_MAYBE_STALE_ERROR = (
    "inner transaction array field can not be reliably accessed due to other inner transaction "
    " submissions or subroutine calls, move array field access closer to {stale_var!r} definition"
)
INNER_TRANSACTION_MAYBE_STALE_WARNING = (
    "inner transaction {stale_var!r} potentially becomes stale here"
)
INNER_TRANSACTION_SOURCE_ERROR = "inner transactions can not be used like this"
INNER_TRANSACTION_SUBROUTINE_ERROR = (
    "inner transactions cannot be used as a subroutine argument or return value"
)

class InnerTransactionsValidator(AWSTTraverser):
    """
    Validates that expressions of type WInnerTransaction and WInnerTransactionFields are only
    used in the ways currently supported. Emits errors for:

    Reassigning expressions of type WInnerTransaction
    Reassigning expressions of type WInnerTransactionFields without a copy()
    Using WInnerTransaction or WInnerTransactionFields in a subroutine or assignment expression
    """

    @classmethod
    def validate(cls, module: awst_nodes.AWST) -> None:
        for module_statement in module:
            validator = cls()
            module_statement.accept(validator)

    def visit_contract_method(self, statement: awst_nodes.ContractMethod) -> None:
        _check_method_types(statement)
        super().visit_contract_method(statement)

    def visit_subroutine(self, statement: awst_nodes.Subroutine) -> None:
        _check_method_types(statement)
        super().visit_subroutine(statement)

    def visit_assignment_statement(self, statement: awst_nodes.AssignmentStatement) -> None:
        self._check_inner_transaction_assignment(statement.value)
        self._check_inner_transaction_fields_assignment(statement.value)
        super().visit_assignment_statement(statement)

    def _check_inner_transaction_assignment(self, value: awst_nodes.Expression) -> None:
        if _is_itxn_wtype(value.wtype) and not _is_assignable_itxn_expr(value):
            logger.error(INNER_TRANSACTION_SOURCE_ERROR, location=value.source_location)

    def _check_inner_transaction_fields_assignment(self, value: awst_nodes.Expression) -> None:
        match value:
            case awst_nodes.CreateInnerTransaction() | awst_nodes.Copy():
                pass  # ok
            case (
                awst_nodes.VarExpression(wtype=wtype)
                | awst_nodes.TupleItemExpression(wtype=wtype)
            ) if isinstance(wtype, WInnerTransactionFields):
                logger.error(
                    INNER_TRANSACTION_COPY_REQUIRED_ERROR,
                    location=value.source_location,
                )
            case awst_nodes.Expression(wtype=wtype) if (
                isinstance(wtype, WInnerTransactionFields)
            ):
                logger.error(
                    INNER_TRANSACTION_SOURCE_ERROR,
                    location=value.source_location,
                )

    def visit_assignment_expression(self, expr: awst_nodes.AssignmentExpression) -> None:
        super().visit_assignment_expression(expr)
        if _is_either_itxn_wtype(expr.wtype):
            logger.error(
                INNER_TRANSACTION_ASSIGNMENT_EXPRESSION_ERROR,
                location=expr.source_location,
            )

    def visit_subroutine_call_expression(self, expr: awst_nodes.SubroutineCallExpression) -> None:
        super().visit_subroutine_call_expression(expr)
        for arg in expr.args:
            if _is_either_itxn_wtype(arg.value.wtype):
                logger.error(
                    INNER_TRANSACTION_SUBROUTINE_ERROR,
                    location=expr.source_location,
                )

class InnerTransactionUsedInALoopValidator(AWSTTraverser):
    """
    Validates that expressions of type WInnerTransactionFields are not modified after submission
    while in a loop
    Modifying after submission while in a loop
    """

    def __init__(self) -> None:
        super().__init__()
        self._current_itxn_var_stack = list[list[str]]()

    @classmethod
    def validate(cls, module: awst_nodes.AWST) -> None:
        for module_statement in module:
            validator = cls()
            module_statement.accept(validator)

    @property
    def _current_loop_itxn_vars(self) -> list[str] | None:
        return self._current_itxn_var_stack[-1] if self._current_itxn_var_stack else None

    @contextlib.contextmanager
    def _enter_loop(self) -> Iterator[None]:
        self._current_itxn_var_stack.append(
            self._current_loop_itxn_vars.copy() if self._current_loop_itxn_vars else []
        )
        try:
            yield
        finally:
            self._current_itxn_var_stack.pop()

    def visit_for_in_loop(self, statement: awst_nodes.ForInLoop) -> None:
        with self._enter_loop():
            super().visit_for_in_loop(statement)

    def visit_while_loop(self, statement: awst_nodes.WhileLoop) -> None:
        with self._enter_loop():
            super().visit_while_loop(statement)

    def visit_assignment_statement(self, statement: awst_nodes.AssignmentStatement) -> None:
        value = statement.value
        match value:
            case awst_nodes.CreateInnerTransaction() | awst_nodes.Copy():
                self._check_itxn_params_not_submitted_in_loop(statement.target)
        super().visit_assignment_statement(statement)

    def visit_submit_inner_transaction(self, call: awst_nodes.SubmitInnerTransaction) -> None:
        if self._current_loop_itxn_vars is not None:
            for itxn_params in call.itxns:
                match itxn_params:
                    case awst_nodes.TupleItemExpression(
                        base=awst_nodes.VarExpression(name=var_name)
                    ):
                        self._current_loop_itxn_vars.append(var_name)
                    case awst_nodes.VarExpression(name=var_name):
                        self._current_loop_itxn_vars.append(var_name)
        super().visit_submit_inner_transaction(call)

    def visit_update_inner_transaction(self, call: awst_nodes.UpdateInnerTransaction) -> None:
        super().visit_update_inner_transaction(call)
        self._check_itxn_params_not_submitted_in_loop(call.itxn)

    def _check_itxn_params_not_submitted_in_loop(self, expr: awst_nodes.Expression) -> None:
        if (
            self._current_loop_itxn_vars
            and isinstance(expr, awst_nodes.VarExpression)
            and expr.name in self._current_loop_itxn_vars
        ):
            logger.error(
                INNER_TRANSACTION_LOOP_MODIFICATION_ERROR,
                location=expr.source_location,
            )

def _check_method_types(stmt: awst_nodes.Function) -> None:
    for arg in stmt.args:
        if _is_either_itxn_wtype(arg.wtype):
            logger.error(
                INNER_TRANSACTION_SUBROUTINE_ERROR,
                location=arg.source_location,
            )
    if _is_either_itxn_wtype(stmt.return_type):
        logger.error(
            INNER_TRANSACTION_SUBROUTINE_ERROR,
            location=stmt.source_location,
        )

def _is_assignable_itxn_expr(expr: awst_nodes.Expression) -> bool:
    if not _is_itxn_wtype(expr.wtype):  # non itxn expressions are assignable
        return True
    match expr:
        case awst_nodes.VarExpression():  # local itxn expressions can be copied
            return True
        case awst_nodes.SubmitInnerTransaction():  # submit expressions are assignable
            return True
        case awst_nodes.TupleExpression(
            items=items
        ):  # tuple expressions composed of assignable expressions are assignable
            return all(map(_is_assignable_itxn_expr, items))
        case awst_nodes.TupleItemExpression(
            base=base
        ):  # tuple items are assignable if their base is assignable
            return _is_assignable_itxn_expr(base)
        case awst_nodes.SingleEvaluation(source=source):
            return _is_assignable_itxn_expr(source)
        case awst_nodes.SliceExpression(
            base=base, wtype=wtypes.WTuple()
        ):  # tuple slices can be assigned if base can be
            return _is_assignable_itxn_expr(base)
    # anything else is not considered assignable
    return False

def _is_either_itxn_wtype(wtype: wtypes.WType) -> bool:
    return _is_itxn_wtype(wtype) or _is_itxn_fields_wtype(wtype)

def _is_itxn_wtype(wtype: wtypes.WType) -> bool:
    return isinstance(wtype, WInnerTransaction) or (
        isinstance(wtype, wtypes.WTuple) and any(map(_is_itxn_wtype, wtype.types))
    )

def _is_itxn_fields_wtype(wtype: wtypes.WType) -> bool:
    return isinstance(wtype, WInnerTransactionFields) or (
        isinstance(wtype, wtypes.WTuple) and any(map(_is_itxn_fields_wtype, wtype.types))
    )

class StaleInnerTransactionsValidator(AWSTTraverser):
    """Validates that inner transaction array access are not stale"""

    @classmethod
    def validate(cls, module: awst_nodes.AWST) -> None:
        for module_statement in module:
            validator = cls()
            module_statement.accept(validator)

    def __init__(self) -> None:
        super().__init__()
        self._maybe_stale_itxn_vars = dict[str, SourceLocation]()
        self._active_itxn_vars = list[str]()

    def visit_assignment_statement(self, stmt: awst_nodes.AssignmentStatement) -> None:
        super().visit_assignment_statement(stmt)
        match stmt.value:
            case awst_nodes.SubmitInnerTransaction():
                new_itxn_var_names = self._get_var_names(stmt.target)
                self._update_active_var_names(new_itxn_var_names)
            case awst_nodes.TupleExpression(items=items) if any(
                isinstance(item.wtype, WInnerTransaction) for item in items
            ):
                var_names = self._get_var_names(stmt.target)
                self._update_active_var_names(var_names)

    def visit_intrinsic_call(self, call: awst_nodes.IntrinsicCall) -> None:
        super().visit_intrinsic_call(call)
        if call.op_code == "itxn_submit":
            self._update_maybe_stale_itxn_vars(call.source_location)

    def visit_submit_inner_transaction(self, call: awst_nodes.SubmitInnerTransaction) -> None:
        super().visit_submit_inner_transaction(call)
        self._update_maybe_stale_itxn_vars(call.source_location)

    def visit_inner_transaction_field(self, itxn_field: awst_nodes.InnerTransactionField) -> None:
        super().visit_inner_transaction_field(itxn_field)
        match itxn_field.itxn:
            case awst_nodes.VarExpression(name=stale_var):
                pass
            case awst_nodes.TupleItemExpression(base=awst_nodes.VarExpression(name=stale_var)):
                pass
            case _:
                return
        if itxn_field.field.is_array:
            try:
                stale_var_loc = self._maybe_stale_itxn_vars[stale_var]
            except KeyError:
                return
            logger.error(
                INNER_TRANSACTION_MAYBE_STALE_ERROR.format(stale_var=stale_var),
                location=itxn_field.itxn.source_location,
            )
            logger.warning(
                INNER_TRANSACTION_MAYBE_STALE_WARNING.format(stale_var=stale_var),
                location=stale_var_loc,
            )

    def visit_subroutine_call_expression(self, expr: awst_nodes.SubroutineCallExpression) -> None:
        super().visit_subroutine_call_expression(expr)
        self._update_maybe_stale_itxn_vars(expr.source_location)

    def _get_var_names(self, expr: awst_nodes.Expression) -> list[str]:
        match expr:
            case awst_nodes.VarExpression(name=name):
                return [name]
            case awst_nodes.TupleExpression(items=items):
                return [self._get_var_names(item)[0] for item in items]
            case _:
                return []

    def _update_active_var_names(self, var_names: list[str]) -> None:
        self._active_itxn_vars = var_names
        # if a var_name is reassigned then it is not considered stale
        for var_name in self._active_itxn_vars:
            with contextlib.suppress(KeyError):
                del self._maybe_stale_itxn_vars[var_name]

    def _update_maybe_stale_itxn_vars(self, staling_expr_loc: SourceLocation) -> None:
        for stale_var_name in self._active_itxn_vars:
            self._maybe_stale_itxn_vars[stale_var_name] = staling_expr_loc
        self._active_itxn_vars = []


from collections import defaultdict

from puya import log
from puya.awst import nodes as awst_nodes
from puya.awst.function_traverser import FunctionTraverser

logger = log.get_logger(__name__)

class LabelsValidator(FunctionTraverser):
    @classmethod
    def validate(cls, module: awst_nodes.AWST) -> None:
        for module_statement in module:
            if isinstance(module_statement, awst_nodes.Subroutine):
                cls(module_statement)
            elif isinstance(module_statement, awst_nodes.Contract):
                for method in module_statement.all_methods:
                    cls(method)

    def __init__(self, function: awst_nodes.Function) -> None:
        self._labelled_blocks = dict[awst_nodes.Label, awst_nodes.Block]()
        self._seen_targets = defaultdict[awst_nodes.Label, list[awst_nodes.Goto]](list)
        function.body.accept(self)
        for target, goto_list in self._seen_targets.items():
            if target not in self._labelled_blocks:
                for goto in goto_list:
                    logger.error(
                        f"label target {target} does not exist", location=goto.source_location
                    )

    def visit_goto(self, goto: awst_nodes.Goto) -> None:
        self._seen_targets[goto.target].append(goto)

    def visit_block(self, block: awst_nodes.Block) -> None:
        if block.label is not None:
            first_seen = self._labelled_blocks.setdefault(block.label, block)
            if block is not first_seen:
                logger.error(
                    f"block has duplicate label {block.label}", location=block.source_location
                )
                logger.info("label first seen here", location=first_seen.source_location)


from puya.awst import nodes as awst_nodes
from puya.awst.validation.arc4_copy import ARC4CopyValidator
from puya.awst.validation.base_invoker import BaseInvokerValidator
from puya.awst.validation.immutable import ImmutableValidator
from puya.awst.validation.inner_transactions import (
    InnerTransactionsValidator,
    InnerTransactionUsedInALoopValidator,
    StaleInnerTransactionsValidator,
)
from puya.awst.validation.labels import LabelsValidator
from puya.awst.validation.scratch_slots import ScratchSlotReservationValidator
from puya.awst.validation.storage import StorageTypesValidator

def validate_awst(module: awst_nodes.AWST) -> None:
    ARC4CopyValidator.validate(module)
    ScratchSlotReservationValidator.validate(module)
    InnerTransactionsValidator.validate(module)
    InnerTransactionUsedInALoopValidator.validate(module)
    StaleInnerTransactionsValidator.validate(module)
    BaseInvokerValidator.validate(module)
    StorageTypesValidator.validate(module)
    LabelsValidator.validate(module)
    ImmutableValidator.validate(module)


from collections.abc import Collection, Iterator

from puya import log
from puya.awst import nodes as awst_nodes
from puya.awst.awst_traverser import AWSTTraverser
from puya.parse import SourceLocation

logger = log.get_logger(__name__)

class ScratchSlotReservationValidator(AWSTTraverser):
    @classmethod
    def validate(cls, module: awst_nodes.AWST) -> None:
        for module_statement in module:
            validator = cls()
            module_statement.accept(validator)
            for slot, loc in validator.invalid_slot_usages:
                logger.error(
                    f"Scratch slot {slot} has not been reserved.",
                    location=loc,
                )

    def __init__(self) -> None:
        super().__init__()
        self._reserved_slots: Collection[int] = ()
        self._used_slots = list[tuple[int, SourceLocation]]()

    @property
    def invalid_slot_usages(self) -> Iterator[tuple[int, SourceLocation]]:
        for slot, loc in self._used_slots:
            if slot not in self._reserved_slots:
                yield slot, loc

    def visit_contract(self, statement: awst_nodes.Contract) -> None:
        super().visit_contract(statement)
        self._reserved_slots = statement.reserved_scratch_space

    def visit_intrinsic_call(self, call: awst_nodes.IntrinsicCall) -> None:
        super().visit_intrinsic_call(call)
        match call.op_code, call.stack_args:
            case "loads", [awst_nodes.IntegerConstant(value=slot, source_location=loc)]:
                self._used_slots.append((slot, loc))
            case "stores", [awst_nodes.IntegerConstant(value=slot, source_location=loc), *_]:
                self._used_slots.append((slot, loc))
        match call.op_code, call.immediates:
            case "load", [int(slot)]:
                self._used_slots.append((slot, call.source_location))
            case "store", [int(slot)]:
                self._used_slots.append((slot, call.source_location))


import typing
from collections import defaultdict

from puya import log
from puya.awst import (
    nodes as awst_nodes,
    wtypes,
)
from puya.awst.awst_traverser import AWSTTraverser
from puya.awst.nodes import AppStorageKind
from puya.utils import set_add

logger = log.get_logger(__name__)

class StorageTypesValidator(AWSTTraverser):
    @classmethod
    def validate(cls, module: awst_nodes.AWST) -> None:
        for module_statement in module:
            # create a new instance for each top level construct,
            # either subroutine or contract class, so that we can de-dupe
            # messages (where possible) there
            validator = cls()
            module_statement.accept(validator)

    def __init__(self) -> None:
        super().__init__()
        self._seen_keys = defaultdict[AppStorageKind, set[bytes]](set)

    @typing.override
    def visit_app_storage_definition(self, defn: awst_nodes.AppStorageDefinition) -> None:
        super().visit_app_storage_definition(defn)
        wtypes.validate_persistable(defn.storage_wtype, defn.source_location)
        if defn.key_wtype is not None:
            wtypes.validate_persistable(defn.key_wtype, defn.source_location)

    @typing.override
    def visit_app_state_expression(self, expr: awst_nodes.AppStateExpression) -> None:
        super().visit_app_state_expression(expr)
        self._validate_usage(expr, AppStorageKind.app_global)

    @typing.override
    def visit_app_account_state_expression(
        self, expr: awst_nodes.AppAccountStateExpression
    ) -> None:
        super().visit_app_account_state_expression(expr)
        self._validate_usage(expr, AppStorageKind.account_local)

    @typing.override
    def visit_box_value_expression(self, expr: awst_nodes.BoxValueExpression) -> None:
        super().visit_box_value_expression(expr)
        self._validate_usage(expr, AppStorageKind.box)

    def _validate_usage(self, expr: awst_nodes.StorageExpression, kind: AppStorageKind) -> None:
        if isinstance(expr.key, awst_nodes.BytesConstant) and not set_add(
            self._seen_keys[kind], expr.key.value
        ):
            return
        wtypes.validate_persistable(expr.wtype, expr.source_location)


from __future__ import annotations  # needed to break import cycle

import typing as t
from abc import ABC, abstractmethod

if t.TYPE_CHECKING:
    import puya.awst.nodes

class StatementVisitor[T](ABC):
    @abstractmethod
    def visit_block(self, statement: puya.awst.nodes.Block) -> T: ...

    @abstractmethod
    def visit_if_else(self, statement: puya.awst.nodes.IfElse) -> T: ...

    @abstractmethod
    def visit_switch(self, statement: puya.awst.nodes.Switch) -> T: ...

    @abstractmethod
    def visit_while_loop(self, statement: puya.awst.nodes.WhileLoop) -> T: ...

    @abstractmethod
    def visit_loop_exit(self, statement: puya.awst.nodes.LoopExit) -> T: ...

    @abstractmethod
    def visit_return_statement(self, statement: puya.awst.nodes.ReturnStatement) -> T: ...

    @abstractmethod
    def visit_loop_continue(self, statement: puya.awst.nodes.LoopContinue) -> T: ...

    @abstractmethod
    def visit_expression_statement(self, statement: puya.awst.nodes.ExpressionStatement) -> T: ...

    @abstractmethod
    def visit_uint64_augmented_assignment(
        self, statement: puya.awst.nodes.UInt64AugmentedAssignment
    ) -> T: ...

    @abstractmethod
    def visit_biguint_augmented_assignment(
        self, statement: puya.awst.nodes.BigUIntAugmentedAssignment
    ) -> T: ...

    @abstractmethod
    def visit_bytes_augmented_assignment(
        self, statement: puya.awst.nodes.BytesAugmentedAssignment
    ) -> T: ...

    @abstractmethod
    def visit_for_in_loop(self, statement: puya.awst.nodes.ForInLoop) -> T: ...

    @abstractmethod
    def visit_assignment_statement(self, statement: puya.awst.nodes.AssignmentStatement) -> T: ...

    @abstractmethod
    def visit_goto(self, statement: puya.awst.nodes.Goto) -> T: ...

class RootNodeVisitor[T](ABC):
    @abstractmethod
    def visit_subroutine(self, statement: puya.awst.nodes.Subroutine) -> T: ...

    @abstractmethod
    def visit_contract(self, statement: puya.awst.nodes.Contract) -> T: ...

    @abstractmethod
    def visit_logic_signature(self, statement: puya.awst.nodes.LogicSignature) -> T: ...

class ContractMemberVisitor[T](ABC):
    @abstractmethod
    def visit_contract_method(self, statement: puya.awst.nodes.ContractMethod) -> T: ...

    @abstractmethod
    def visit_app_storage_definition(
        self, statement: puya.awst.nodes.AppStorageDefinition
    ) -> T: ...

class ExpressionVisitor[T](ABC):
    @abstractmethod
    def visit_state_delete(self, expr: puya.awst.nodes.StateDelete) -> T: ...

    @abstractmethod
    def visit_assignment_expression(self, expr: puya.awst.nodes.AssignmentExpression) -> T: ...

    @abstractmethod
    def visit_uint64_binary_operation(self, expr: puya.awst.nodes.UInt64BinaryOperation) -> T: ...

    @abstractmethod
    def visit_biguint_binary_operation(
        self, expr: puya.awst.nodes.BigUIntBinaryOperation
    ) -> T: ...

    @abstractmethod
    def visit_integer_constant(self, expr: puya.awst.nodes.IntegerConstant) -> T: ...

    @abstractmethod
    def visit_decimal_constant(self, expr: puya.awst.nodes.DecimalConstant) -> T: ...

    @abstractmethod
    def visit_bool_constant(self, expr: puya.awst.nodes.BoolConstant) -> T: ...

    @abstractmethod
    def visit_bytes_constant(self, expr: puya.awst.nodes.BytesConstant) -> T: ...

    @abstractmethod
    def visit_string_constant(self, expr: puya.awst.nodes.StringConstant) -> T: ...

    @abstractmethod
    def visit_void_constant(self, expr: puya.awst.nodes.VoidConstant) -> T: ...

    @abstractmethod
    def visit_address_constant(self, expr: puya.awst.nodes.AddressConstant) -> T: ...

    @abstractmethod
    def visit_compiled_contract(self, expr: puya.awst.nodes.CompiledContract) -> T: ...

    @abstractmethod
    def visit_compiled_logicsig(self, expr: puya.awst.nodes.CompiledLogicSig) -> T: ...

    @abstractmethod
    def visit_numeric_comparison_expression(
        self, expr: puya.awst.nodes.NumericComparisonExpression
    ) -> T: ...

    @abstractmethod
    def visit_var_expression(self, expr: puya.awst.nodes.VarExpression) -> T: ...

    @abstractmethod
    def visit_intrinsic_call(self, call: puya.awst.nodes.IntrinsicCall) -> T: ...

    @abstractmethod
    def visit_puya_lib_call(self, call: puya.awst.nodes.PuyaLibCall) -> T: ...

    @abstractmethod
    def visit_group_transaction_reference(
        self, ref: puya.awst.nodes.GroupTransactionReference
    ) -> T: ...

    @abstractmethod
    def visit_create_inner_transaction(
        self, create_itxn: puya.awst.nodes.CreateInnerTransaction
    ) -> T: ...

    @abstractmethod
    def visit_update_inner_transaction(
        self, update_itxn: puya.awst.nodes.UpdateInnerTransaction
    ) -> T: ...

    @abstractmethod
    def visit_submit_inner_transaction(
        self, submit: puya.awst.nodes.SubmitInnerTransaction
    ) -> T: ...

    @abstractmethod
    def visit_inner_transaction_field(
        self, itxn_field: puya.awst.nodes.InnerTransactionField
    ) -> T: ...
    @abstractmethod
    def visit_assert_expression(self, statement: puya.awst.nodes.AssertExpression) -> T: ...

    @abstractmethod
    def visit_checked_maybe(self, call: puya.awst.nodes.CheckedMaybe) -> T: ...

    @abstractmethod
    def visit_arc4_decode(self, expr: puya.awst.nodes.ARC4Decode) -> T: ...

    @abstractmethod
    def visit_arc4_encode(self, expr: puya.awst.nodes.ARC4Encode) -> T: ...

    @abstractmethod
    def visit_array_concat(self, expr: puya.awst.nodes.ArrayConcat) -> T: ...

    @abstractmethod
    def visit_array_extend(self, expr: puya.awst.nodes.ArrayExtend) -> T: ...

    @abstractmethod
    def visit_tuple_expression(self, expr: puya.awst.nodes.TupleExpression) -> T: ...

    @abstractmethod
    def visit_tuple_item_expression(self, expr: puya.awst.nodes.TupleItemExpression) -> T: ...

    @abstractmethod
    def visit_field_expression(self, expr: puya.awst.nodes.FieldExpression) -> T: ...

    @abstractmethod
    def visit_index_expression(self, expr: puya.awst.nodes.IndexExpression) -> T: ...

    @abstractmethod
    def visit_slice_expression(self, expr: puya.awst.nodes.SliceExpression) -> T: ...

    @abstractmethod
    def visit_conditional_expression(self, expr: puya.awst.nodes.ConditionalExpression) -> T: ...

    @abstractmethod
    def visit_single_evaluation(self, expr: puya.awst.nodes.SingleEvaluation) -> T: ...

    @abstractmethod
    def visit_app_state_expression(self, expr: puya.awst.nodes.AppStateExpression) -> T: ...

    @abstractmethod
    def visit_app_account_state_expression(
        self, expr: puya.awst.nodes.AppAccountStateExpression
    ) -> T: ...

    @abstractmethod
    def visit_new_array(self, expr: puya.awst.nodes.NewArray) -> T: ...

    @abstractmethod
    def visit_new_struct(self, expr: puya.awst.nodes.NewStruct) -> T: ...

    @abstractmethod
    def visit_bytes_comparison_expression(
        self, expr: puya.awst.nodes.BytesComparisonExpression
    ) -> T: ...

    @abstractmethod
    def visit_subroutine_call_expression(
        self, expr: puya.awst.nodes.SubroutineCallExpression
    ) -> T: ...

    @abstractmethod
    def visit_bytes_binary_operation(self, expr: puya.awst.nodes.BytesBinaryOperation) -> T: ...

    @abstractmethod
    def visit_boolean_binary_operation(
        self, expr: puya.awst.nodes.BooleanBinaryOperation
    ) -> T: ...

    @abstractmethod
    def visit_uint64_unary_operation(self, expr: puya.awst.nodes.UInt64UnaryOperation) -> T: ...

    @abstractmethod
    def visit_bytes_unary_operation(self, expr: puya.awst.nodes.BytesUnaryOperation) -> T: ...

    @abstractmethod
    def visit_not_expression(self, expr: puya.awst.nodes.Not) -> T: ...

    @abstractmethod
    def visit_reinterpret_cast(self, expr: puya.awst.nodes.ReinterpretCast) -> T: ...

    @abstractmethod
    def visit_enumeration(self, expr: puya.awst.nodes.Enumeration) -> T: ...

    @abstractmethod
    def visit_method_constant(self, expr: puya.awst.nodes.MethodConstant) -> T: ...

    @abstractmethod
    def visit_array_pop(self, expr: puya.awst.nodes.ArrayPop) -> T: ...

    @abstractmethod
    def visit_copy(self, expr: puya.awst.nodes.Copy) -> T: ...

    @abstractmethod
    def visit_reversed(self, expr: puya.awst.nodes.Reversed) -> T: ...

    @abstractmethod
    def visit_state_get(self, expr: puya.awst.nodes.StateGet) -> T: ...

    @abstractmethod
    def visit_state_get_ex(self, expr: puya.awst.nodes.StateGetEx) -> T: ...

    @abstractmethod
    def visit_state_exists(self, expr: puya.awst.nodes.StateExists) -> T: ...

    @abstractmethod
    def visit_template_var(self, expr: puya.awst.nodes.TemplateVar) -> T: ...

    @abstractmethod
    def visit_intersection_slice_expression(
        self, expr: puya.awst.nodes.IntersectionSliceExpression
    ) -> T: ...

    @abstractmethod
    def visit_box_value_expression(self, expr: puya.awst.nodes.BoxValueExpression) -> T: ...

    @abstractmethod
    def visit_uint64_postfix_unary_operation(
        self, expr: puya.awst.nodes.UInt64PostfixUnaryOperation
    ) -> T: ...

    @abstractmethod
    def visit_biguint_postfix_unary_operation(
        self, expr: puya.awst.nodes.BigUIntPostfixUnaryOperation
    ) -> T: ...

    @abstractmethod
    def visit_arc4_router(self, expr: puya.awst.nodes.ARC4Router) -> T: ...

    @abstractmethod
    def visit_range(self, node: puya.awst.nodes.Range) -> T: ...

    @abstractmethod
    def visit_emit(self, emit: puya.awst.nodes.Emit) -> T: ...


import typing
from collections.abc import Iterable, Mapping
from functools import cached_property

import attrs
from immutabledict import immutabledict

from puya import log
from puya.avm import AVMType, TransactionType
from puya.errors import CodeError, InternalError
from puya.parse import SourceLocation
from puya.utils import unique

logger = log.get_logger(__name__)

@attrs.frozen(kw_only=True)
class WType:
    name: str
    scalar_type: typing.Literal[AVMType.uint64, AVMType.bytes, None]
    "the (unbound) AVM stack type, if any"
    ephemeral: bool = False
    """ephemeral types are not suitable for naive storage / persistence,
     even if their underlying type is a simple stack value"""
    immutable: bool

    def __str__(self) -> str:
        return self.name

void_wtype: typing.Final = WType(
    name="void",
    scalar_type=None,
    immutable=True,
)

bool_wtype: typing.Final = WType(
    name="bool",
    scalar_type=AVMType.uint64,
    immutable=True,
)

uint64_wtype: typing.Final = WType(
    name="uint64",
    scalar_type=AVMType.uint64,
    immutable=True,
)

biguint_wtype: typing.Final = WType(
    name="biguint",
    scalar_type=AVMType.bytes,
    immutable=True,
)

bytes_wtype: typing.Final = WType(
    name="bytes",
    scalar_type=AVMType.bytes,
    immutable=True,
)
string_wtype: typing.Final = WType(
    name="string",
    scalar_type=AVMType.bytes,
    immutable=True,
)
asset_wtype: typing.Final = WType(
    name="asset",
    scalar_type=AVMType.uint64,
    immutable=True,
)

account_wtype: typing.Final = WType(
    name="account",
    scalar_type=AVMType.bytes,
    immutable=True,
)

application_wtype: typing.Final = WType(
    name="application",
    scalar_type=AVMType.uint64,
    immutable=True,
)

state_key: typing.Final = WType(
    name="state_key",
    scalar_type=AVMType.bytes,
    immutable=True,
)
box_key: typing.Final = WType(
    name="box_key",
    scalar_type=AVMType.bytes,
    immutable=True,
)

uint64_range_wtype: typing.Final = WType(
    name="uint64_range",
    scalar_type=None,
    immutable=True,
)

@attrs.frozen
class WEnumeration(WType):
    sequence_type: WType
    name: str = attrs.field(init=False)
    immutable: bool = attrs.field(default=True, init=False)
    scalar_type: None = attrs.field(default=None, init=False)
    ephemeral: bool = attrs.field(default=False, init=False)

    @name.default
    def _name_factory(self) -> str:
        return f"enumerate_{self.sequence_type.name}"

@attrs.frozen
class _TransactionRelatedWType(WType):
    transaction_type: TransactionType | None
    ephemeral: bool = attrs.field(default=True, init=False)
    immutable: bool = attrs.field(default=True, init=False)

@typing.final
@attrs.frozen
class WGroupTransaction(_TransactionRelatedWType):
    scalar_type: typing.Literal[AVMType.uint64] = attrs.field(default=AVMType.uint64, init=False)

    @classmethod
    def from_type(cls, transaction_type: TransactionType | None) -> "WGroupTransaction":
        name = "group_transaction"
        if transaction_type:
            name = f"{name}_{transaction_type.name}"
        return cls(name=name, transaction_type=transaction_type)

@typing.final
@attrs.frozen
class WInnerTransactionFields(_TransactionRelatedWType):
    scalar_type: None = attrs.field(default=None, init=False)

    @classmethod
    def from_type(cls, transaction_type: TransactionType | None) -> "WInnerTransactionFields":
        name = "inner_transaction_fields"
        if transaction_type:
            name = f"{name}_{transaction_type.name}"
        return cls(name=name, transaction_type=transaction_type)

@typing.final
@attrs.frozen
class WInnerTransaction(_TransactionRelatedWType):
    scalar_type: None = attrs.field(default=None, init=False)

    @classmethod
    def from_type(cls, transaction_type: TransactionType | None) -> "WInnerTransaction":
        name = "inner_transaction"
        if transaction_type:
            name = f"{name}_{transaction_type.name}"
        return cls(name=name, transaction_type=transaction_type)

@typing.final
@attrs.frozen
class WStructType(WType):
    fields: immutabledict[str, WType] = attrs.field(converter=immutabledict)
    frozen: bool
    immutable: bool = attrs.field(init=False)
    scalar_type: None = attrs.field(default=None, init=False)
    source_location: SourceLocation | None = attrs.field(eq=False)
    desc: str | None = None

    @immutable.default
    def _immutable(self) -> bool:
        # TODO: determine correct behaviour when implementing native structs
        raise NotImplementedError

    @fields.validator
    def _fields_validator(self, _: object, fields: immutabledict[str, WType]) -> None:
        if not fields:
            raise CodeError("struct needs fields", self.source_location)
        if void_wtype in fields.values():
            raise CodeError("struct should not contain void types", self.source_location)

@typing.final
@attrs.frozen
class WArray(WType):
    element_type: WType = attrs.field()
    name: str = attrs.field(init=False)
    scalar_type: None = attrs.field(default=None, init=False)
    source_location: SourceLocation | None = attrs.field(eq=False)
    immutable: bool = attrs.field(default=False, init=False)

    @element_type.validator
    def _element_type_validator(self, _: object, element_type: WType) -> None:
        if element_type == void_wtype:
            raise CodeError("array element type cannot be void", self.source_location)

    @name.default
    def _name(self) -> str:
        return f"array<{self.element_type.name}>"

@typing.final
@attrs.frozen(eq=False)
class WTuple(WType):
    types: tuple[WType, ...] = attrs.field(converter=tuple[WType, ...])
    source_location: SourceLocation | None = attrs.field(default=None)
    scalar_type: None = attrs.field(default=None, init=False)
    immutable: bool = attrs.field(default=True, init=False)
    name: str = attrs.field(kw_only=True)
    names: tuple[str, ...] | None = attrs.field(default=None)
    desc: str | None = None

    def __eq__(self, other: object) -> bool:
        # this custom equality check ensures that
        # tuple field names are only considered when both sides
        # have defined names
        if not isinstance(other, WTuple):
            return False
        return self.types == other.types and (
            self.names == other.names or None in (self.names, other.names)
        )

    def __hash__(self) -> int:
        return hash(self.types)

    @types.validator
    def _types_validator(self, _attribute: object, types: tuple[WType, ...]) -> None:
        if void_wtype in types:
            raise CodeError("tuple should not contain void types", self.source_location)

    @name.default
    def _name(self) -> str:
        return f"tuple<{','.join([t.name for t in self.types])}>"

    @names.validator
    def _names_validator(self, _attribute: object, names: tuple[str, ...] | None) -> None:
        if names is None:
            return
        if len(names) != len(self.types):
            raise InternalError("mismatch between tuple item names length and types")
        if len(names) != len(unique(names)):
            raise CodeError("tuple item names are not unique", self.source_location)

    @cached_property
    def fields(self) -> Mapping[str, WType]:
        """Mapping of item names to types if `names` is defined, otherwise empty."""
        if self.names is None:
            return {}
        return dict(zip(self.names, self.types, strict=True))

    def name_to_index(self, name: str, source_location: SourceLocation) -> int:
        if self.names is None:
            raise CodeError(
                "cannot access tuple item by name of an unnamed tuple", source_location
            )
        try:
            return self.names.index(name)
        except ValueError:
            raise CodeError(f"{name} is not a member of {self.name}") from None

@attrs.frozen(kw_only=True)
class ARC4Type(WType):
    scalar_type: typing.Literal[AVMType.bytes] = attrs.field(default=AVMType.bytes, init=False)
    arc4_name: str = attrs.field(eq=False)  # exclude from equality in case of aliasing
    native_type: WType | None

    def can_encode_type(self, wtype: WType) -> bool:
        return wtype == self.native_type

arc4_bool_wtype: typing.Final = ARC4Type(
    name="arc4.bool",
    arc4_name="bool",
    immutable=True,
    native_type=bool_wtype,
)

@typing.final
@attrs.frozen(kw_only=True)
class ARC4UIntN(ARC4Type):
    immutable: bool = attrs.field(default=True, init=False)
    native_type: WType = attrs.field(default=biguint_wtype, init=False)
    n: int = attrs.field()
    arc4_name: str = attrs.field(eq=False)
    name: str = attrs.field(init=False)
    source_location: SourceLocation | None = attrs.field(default=None, eq=False)

    @n.validator
    def _n_validator(self, _attribute: object, n: int) -> None:
        if not (n % 8 == 0):
            raise CodeError("Bit size must be multiple of 8", self.source_location)
        if not (8 <= n <= 512):
            raise CodeError("Bit size must be between 8 and 512 inclusive", self.source_location)

    @arc4_name.default
    def _arc4_name(self) -> str:
        return f"uint{self.n}"

    @name.default
    def _name(self) -> str:
        return f"arc4.{self._arc4_name()}"

    def can_encode_type(self, wtype: WType) -> bool:
        return wtype in (bool_wtype, uint64_wtype, biguint_wtype)

@typing.final
@attrs.frozen(kw_only=True)
class ARC4UFixedNxM(ARC4Type):
    n: int = attrs.field()
    m: int = attrs.field()
    immutable: bool = attrs.field(default=True, init=False)
    arc4_name: str = attrs.field(init=False, eq=False)
    name: str = attrs.field(init=False)
    source_location: SourceLocation | None = attrs.field(default=None, eq=False)
    native_type: None = attrs.field(default=None, init=False)

    @arc4_name.default
    def _arc4_name(self) -> str:
        return f"ufixed{self.n}x{self.m}"

    @name.default
    def _name(self) -> str:
        return f"arc4.{self.arc4_name}"

    @n.validator
    def _n_validator(self, _attribute: object, n: int) -> None:
        if not (n % 8 == 0):
            raise CodeError("Bit size must be multiple of 8", self.source_location)
        if not (8 <= n <= 512):
            raise CodeError("Bit size must be between 8 and 512 inclusive", self.source_location)

    @m.validator
    def _m_validator(self, _attribute: object, m: int) -> None:
        if not (1 <= m <= 160):
            raise CodeError("Precision must be between 1 and 160 inclusive", self.source_location)

def _required_arc4_wtypes(wtypes: Iterable[WType]) -> tuple[ARC4Type, ...]:
    result = []
    for wtype in wtypes:
        if not isinstance(wtype, ARC4Type):
            raise CodeError(f"expected ARC4 type: {wtype}")
        result.append(wtype)
    return tuple(result)

@typing.final
@attrs.frozen(kw_only=True)
class ARC4Tuple(ARC4Type):
    source_location: SourceLocation | None = attrs.field(default=None, eq=False)
    types: tuple[ARC4Type, ...] = attrs.field(converter=_required_arc4_wtypes)
    name: str = attrs.field(init=False)
    arc4_name: str = attrs.field(init=False, eq=False)
    immutable: bool = attrs.field(init=False)
    native_type: WTuple = attrs.field(init=False)

    @name.default
    def _name(self) -> str:
        return f"arc4.tuple<{','.join(t.name for t in self.types)}>"

    @arc4_name.default
    def _arc4_name(self) -> str:
        return f"({','.join(item.arc4_name for item in self.types)})"

    @immutable.default
    def _immutable(self) -> bool:
        return all(typ.immutable for typ in self.types)

    @native_type.default
    def _native_type(self) -> WTuple:
        return WTuple(self.types, self.source_location)

    def can_encode_type(self, wtype: WType) -> bool:
        return super().can_encode_type(wtype) or _is_arc4_encodeable_tuple(wtype, self.types)

def _is_arc4_encodeable_tuple(
    wtype: WType, target_types: tuple[ARC4Type, ...]
) -> typing.TypeGuard[WTuple]:
    return (
        isinstance(wtype, WTuple)
        and len(wtype.types) == len(target_types)
        and all(
            arc4_wtype == encode_wtype or arc4_wtype.can_encode_type(encode_wtype)
            for arc4_wtype, encode_wtype in zip(target_types, wtype.types, strict=True)
        )
    )

def _expect_arc4_type(wtype: WType) -> ARC4Type:
    if not isinstance(wtype, ARC4Type):
        raise CodeError(f"expected ARC4 type: {wtype}")
    return wtype

@attrs.frozen(kw_only=True)
class ARC4Array(ARC4Type):
    element_type: ARC4Type = attrs.field(converter=_expect_arc4_type)
    native_type: WType | None = None
    immutable: bool = False

@typing.final
@attrs.frozen(kw_only=True)
class ARC4DynamicArray(ARC4Array):
    name: str = attrs.field(init=False)
    arc4_name: str = attrs.field(eq=False)
    source_location: SourceLocation | None = attrs.field(default=None, eq=False)

    @name.default
    def _name(self) -> str:
        return f"arc4.dynamic_array<{self.element_type.name}>"

    @arc4_name.default
    def _arc4_name(self) -> str:
        return f"{self.element_type.arc4_name}[]"

@typing.final
@attrs.frozen(kw_only=True)
class ARC4StaticArray(ARC4Array):
    array_size: int = attrs.field(validator=attrs.validators.ge(0))
    name: str = attrs.field(init=False)
    arc4_name: str = attrs.field(eq=False)
    source_location: SourceLocation | None = attrs.field(default=None, eq=False)

    @name.default
    def _name(self) -> str:
        return f"arc4.static_array<{self.element_type.name}, {self.array_size}>"

    @arc4_name.default
    def _arc4_name(self) -> str:
        return f"{self.element_type.arc4_name}[{self.array_size}]"

def _require_arc4_fields(fields: Mapping[str, WType]) -> immutabledict[str, ARC4Type]:
    if not fields:
        raise CodeError("arc4.Struct needs at least one element")
    non_arc4_fields = [
        field_name
        for field_name, field_type in fields.items()
        if not isinstance(field_type, ARC4Type)
    ]
    if non_arc4_fields:
        raise CodeError(
            "invalid ARC4 Struct declaration,"
            f" the following fields are not ARC4 encoded types: {', '.join(non_arc4_fields)}",
        )
    return immutabledict(fields)

@typing.final
@attrs.frozen(kw_only=True)
class ARC4Struct(ARC4Type):
    fields: immutabledict[str, ARC4Type] = attrs.field(converter=_require_arc4_fields)
    frozen: bool
    immutable: bool = attrs.field(init=False)
    source_location: SourceLocation | None = attrs.field(default=None, eq=False)
    arc4_name: str = attrs.field(init=False, eq=False)
    native_type: None = attrs.field(default=None, init=False)
    desc: str | None = None

    @immutable.default
    def _immutable(self) -> bool:
        return self.frozen and all(typ.immutable for typ in self.fields.values())

    @arc4_name.default
    def _arc4_name(self) -> str:
        return f"({','.join(item.arc4_name for item in self.types)})"

    @cached_property
    def names(self) -> tuple[str, ...]:
        return tuple(self.fields.keys())

    @cached_property
    def types(self) -> tuple[ARC4Type, ...]:
        return tuple(self.fields.values())

    def can_encode_type(self, wtype: WType) -> bool:
        return super().can_encode_type(wtype) or (
            _is_arc4_encodeable_tuple(wtype, self.types)
            and (wtype.names is None or wtype.names == self.names)
        )

arc4_byte_alias: typing.Final = ARC4UIntN(
    n=8,
    arc4_name="byte",
    source_location=None,
)

arc4_string_alias: typing.Final = ARC4DynamicArray(
    arc4_name="string",
    element_type=arc4_byte_alias,
    native_type=string_wtype,
    immutable=True,
    source_location=None,
)

arc4_address_alias: typing.Final = ARC4StaticArray(
    arc4_name="address",
    element_type=arc4_byte_alias,
    native_type=account_wtype,
    array_size=32,
    immutable=True,
    source_location=None,
)

def persistable_stack_type(
    wtype: WType, location: SourceLocation
) -> typing.Literal[AVMType.uint64, AVMType.bytes]:
    match _storage_type_or_error(wtype):
        case str(error):
            raise CodeError(error, location=location)
        case result:
            return result

def validate_persistable(wtype: WType, location: SourceLocation) -> bool:
    match _storage_type_or_error(wtype):
        case str(error):
            logger.error(error, location=location)
            return False
        case _:
            return True

def _storage_type_or_error(wtype: WType) -> str | typing.Literal[AVMType.uint64, AVMType.bytes]:
    if wtype.ephemeral:
        return "ephemeral types (such as transaction related types) are not suitable for storage"
    if wtype.scalar_type is None:
        return "type is not suitable for storage"
    return wtype.scalar_type


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

import attrs

from puya.artifact_metadata import ContractMetaData, LogicSignatureMetaData
from puya.parse import SourceLocation
from puya.program_refs import ContractReference, LogicSigReference

TemplateValue = tuple[int | bytes, SourceLocation | None]

class DebugEvent(typing.TypedDict, total=False):
    """Describes various attributes for a particular PC location"""

    subroutine: str
    """Subroutine name"""
    params: Mapping[str, str]
    """Describes a subroutines parameters and their types"""
    block: str
    """Name of a block"""
    stack_in: Sequence[str]
    """Variable names on the stack BEFORE the next op executes"""
    op: str
    """Op description"""
    callsub: str
    """The subroutine that is about to be called"""
    retsub: bool
    """Returns from current subroutine"""
    stack_out: Sequence[str]
    """Variable names on the stack AFTER the next op executes"""
    defined_out: Sequence[str]
    """Variable names that are defined AFTER the next op executes"""
    error: str
    """Error message if failure occurs at this op"""

@attrs.frozen
class DebugInfo:
    version: int
    sources: list[str]
    mappings: str
    op_pc_offset: int
    pc_events: Mapping[int, DebugEvent]

class CompiledProgram(abc.ABC):
    @property
    @abc.abstractmethod
    def teal_src(self) -> str: ...

    @property
    @abc.abstractmethod
    def bytecode(self) -> bytes | None:
        """
        bytecode can only be produced if no template variables are used OR template values are
        provided, for this reason bytecode is only provided if output_bytecode is True
        """

    @property
    @abc.abstractmethod
    def debug_info(self) -> DebugInfo | None: ...

    @property
    @abc.abstractmethod
    def template_variables(self) -> Mapping[str, int | bytes | None]: ...

class CompiledContract(abc.ABC):
    @property
    @abc.abstractmethod
    def source_location(self) -> SourceLocation | None: ...

    @property
    @abc.abstractmethod
    def approval_program(self) -> CompiledProgram: ...

    @property
    @abc.abstractmethod
    def clear_program(self) -> CompiledProgram: ...

    @property
    @abc.abstractmethod
    def metadata(self) -> ContractMetaData: ...

    @typing.final
    @property
    def id(self) -> ContractReference:
        return self.metadata.ref

class CompiledLogicSig(abc.ABC):
    @property
    @abc.abstractmethod
    def source_location(self) -> SourceLocation | None: ...
    @property
    @abc.abstractmethod
    def program(self) -> CompiledProgram: ...

    @property
    @abc.abstractmethod
    def metadata(self) -> LogicSignatureMetaData: ...

    @typing.final
    @property
    def id(self) -> LogicSigReference:
        return self.metadata.ref

CompilationArtifact: typing.TypeAlias = CompiledContract | CompiledLogicSig


import itertools
import shutil
import typing
from collections import defaultdict
from collections.abc import Iterator, Mapping, Sequence, Set
from pathlib import Path

import attrs
from cattrs.preconf.json import make_converter
from immutabledict import immutabledict

from puya import log
from puya.arc32 import create_arc32_json
from puya.arc56 import create_arc56_json
from puya.artifact_metadata import ContractMetaData, LogicSignatureMetaData, StateTotals
from puya.artifact_sorter import Artifact, ArtifactCompilationSorter
from puya.awst.nodes import AWST
from puya.awst.validation.main import validate_awst
from puya.compilation_artifacts import (
    CompilationArtifact,
    CompiledContract,
    CompiledLogicSig,
    CompiledProgram,
    DebugInfo,
    TemplateValue,
)
from puya.context import (
    ArtifactCompileContext,
    CompileContext,
    CompiledProgramProvider,
    OutputPathProvider,
)
from puya.errors import CodeError, InternalError
from puya.ir import models as ir_models
from puya.ir.destructure.main import destructure_ssa
from puya.ir.main import awst_to_ir
from puya.ir.optimize.main import optimize_program_ir
from puya.ir.to_text_visitor import render_program
from puya.ir.validation.main import validate_module_artifact
from puya.log import LoggingContext
from puya.mir.main import program_ir_to_mir
from puya.options import PuyaOptions
from puya.parse import SourceLocation
from puya.program_refs import ContractReference, LogicSigReference, ProgramKind
from puya.teal.main import mir_to_teal
from puya.teal.models import TealProgram, TealSubroutine
from puya.teal.output import emit_teal
from puya.ussemble.main import assemble_program
from puya.utils import attrs_extend, make_path_relative_to, make_path_relative_to_cwd

logger = log.get_logger(__name__)

def awst_to_teal(
    log_ctx: LoggingContext,
    options: PuyaOptions,
    compilation_set: Mapping[ContractReference | LogicSigReference, Path],
    sources_by_path: Mapping[Path, Sequence[str] | None],
    awst: AWST,
    *,
    write: bool = True,
) -> list[CompilationArtifact]:
    validate_awst(awst)
    log_ctx.exit_if_errors()
    context = CompileContext(
        options=options,
        compilation_set=compilation_set,
        sources_by_path=sources_by_path,
    )
    log_ctx.exit_if_errors()
    ir = list(awst_to_ir(context, awst))
    log_ctx.exit_if_errors()
    teal = list(_ir_to_teal(log_ctx, context, ir))
    log_ctx.exit_if_errors()
    if write:
        _write_artifacts(context, teal)
    return teal

def _ir_to_teal(
    log_ctx: LoggingContext, context: CompileContext, all_ir: Sequence[ir_models.ModuleArtifact]
) -> Iterator[CompilationArtifact]:
    compiled_program_provider = _CompiledProgramProviderImpl(
        compile_context=context,
        state_totals={
            ir.metadata.ref: ir.metadata.state_totals
            for ir in all_ir
            if isinstance(ir, ir_models.Contract)
        },
    )

    # used to check for conflicts that would occur on output
    artifacts_by_output_base = dict[Path, Artifact]()
    for artifact in ArtifactCompilationSorter.sort(all_ir):
        output_path_provider = None
        if out_dir_setting := context.compilation_set.get(artifact.id):
            name = artifact.ir.metadata.name
            maybe_out_dir = out_dir_setting / f"{name}.ir"
            first_seen = artifacts_by_output_base.setdefault(maybe_out_dir, artifact)
            if artifact is not first_seen:
                logger.error(f"duplicate contract name {name}", location=artifact.source_location)
                logger.info(
                    f"contract name {name} first seen here", location=first_seen.source_location
                )
            else:
                out_dir = maybe_out_dir
                shutil.rmtree(out_dir, ignore_errors=True)
                output_path_provider = _SequentialOutputPathProvider(
                    metadata=artifact.ir.metadata, out_dir=out_dir
                )

        artifact_context = attrs_extend(
            ArtifactCompileContext,
            context,
            output_path_provider=output_path_provider,
            compiled_program_provider=compiled_program_provider,
        )

        num_errors_before_optimization = log_ctx.num_errors
        artifact_ir = _optimize_and_destructure_ir(artifact_context, artifact.ir)
        # IR validation that occurs at the end of optimize_and_destructure_ir may have revealed
        # further errors, add dummy artifacts and continue so other artifacts can still be lowered
        # and report any errors they encounter
        errors_in_optimization = log_ctx.num_errors > num_errors_before_optimization
        if not errors_in_optimization:
            compiled: _CompiledContract | _CompiledLogicSig
            if isinstance(artifact_ir, ir_models.Contract):
                compiled = _contract_ir_to_teal(artifact_context, artifact_ir)
            else:
                compiled = _logic_sig_to_teal(artifact_context, artifact_ir)
            yield compiled
            compiled_program_provider.add_compiled_result(artifact, compiled)

@attrs.define
class _CompiledProgramProviderImpl(CompiledProgramProvider):
    compile_context: CompileContext
    state_totals: Mapping[ContractReference, StateTotals]
    _compiled_artifacts: dict[
        ContractReference | LogicSigReference, "_CompiledContract | _CompiledLogicSig"
    ] = attrs.field(factory=dict, init=False)
    _bytecode_cache: dict[
        tuple[
            ContractReference | LogicSigReference, ProgramKind, immutabledict[str, TemplateValue]
        ],
        bytes,
    ] = attrs.field(factory=dict, init=False)

    def add_compiled_result(
        self, artifact: Artifact, result: "_CompiledContract | _CompiledLogicSig"
    ) -> None:
        self._compiled_artifacts[artifact.id] = result

    @typing.override
    def build_program_bytecode(
        self,
        ref: ContractReference | LogicSigReference,
        kind: ProgramKind,
        *,
        template_constants: immutabledict[str, TemplateValue],
    ) -> bytes:
        cache_key = (ref, kind, template_constants)
        try:
            return self._bytecode_cache[cache_key]
        except KeyError:
            pass

        try:
            comp_ref = self._compiled_artifacts[ref]
        except KeyError:
            raise CodeError(f"invalid reference: {ref}") from None
        match kind, comp_ref:
            case ProgramKind.logic_signature, _CompiledLogicSig(program=program):
                pass
            case ProgramKind.approval, _CompiledContract(approval_program=program):
                pass
            case ProgramKind.clear_state, _CompiledContract(clear_program=program):
                pass
            case _:
                raise InternalError(f"invalid kind: {kind}, {type(comp_ref)}")
        assembled = assemble_program(
            self.compile_context,
            ref,
            program.teal,
            template_constants=template_constants,
            is_reference=True,
        )
        result = assembled.bytecode
        self._bytecode_cache[cache_key] = result
        return result

    @typing.override
    def get_state_totals(self, ref: ContractReference) -> StateTotals:
        return self.state_totals[ref]

@attrs.frozen
class _SequentialOutputPathProvider(OutputPathProvider):
    _metadata: ContractMetaData | LogicSignatureMetaData
    _out_dir: Path
    _output_seq: defaultdict[str, Iterator[int]] = attrs.field(
        factory=lambda: defaultdict(itertools.count), init=False
    )

    @typing.override
    def __call__(self, *, kind: str, qualifier: str, suffix: str) -> Path:
        out_dir = self._out_dir
        out_dir.mkdir(exist_ok=True)
        if qualifier:
            qualifier = f".{qualifier}"
        qualifier = f"{next(self._output_seq[kind])}{qualifier}"
        if kind is not ProgramKind.logic_signature:
            qualifier = f"{kind}.{qualifier}"
        return out_dir / f"{self._metadata.name}.{qualifier}.{suffix}"

def _optimize_and_destructure_ir(
    context: ArtifactCompileContext, artifact_ir: ir_models.ModuleArtifact
) -> ir_models.ModuleArtifact:
    if isinstance(artifact_ir, ir_models.LogicSignature):
        routable_method_ids = None
    else:
        routable_method_ids = {a4m.id for a4m in artifact_ir.metadata.arc4_methods}
    for program in artifact_ir.all_programs():
        _optimize_and_destructure_program_ir(
            context, artifact_ir.metadata.ref, program, routable_method_ids=routable_method_ids
        )
    # validation is run as the last step, in case we've accidentally inserted something,
    # and in particular post subroutine removal, because some things that are "linked"
    # are not necessarily used from the current artifact
    validate_module_artifact(context, artifact_ir)
    return artifact_ir

def _optimize_and_destructure_program_ir(
    context: ArtifactCompileContext,
    ref: ContractReference | LogicSigReference,
    program: ir_models.Program,
    routable_method_ids: Set[str] | None = None,
) -> None:
    if context.options.output_ssa_ir:
        render_program(context, program, qualifier="ssa")
    logger.info(
        f"optimizing {program.kind} program of {ref} at level {context.options.optimization_level}"
    )
    optimize_program_ir(context, program, routable_method_ids=routable_method_ids)
    destructure_ssa(context, program)
    if context.options.output_destructured_ir:
        render_program(context, program, qualifier="destructured")

@attrs.frozen
class _CompiledProgram(CompiledProgram):
    teal: TealProgram
    teal_src: str
    template_variables: Mapping[str, int | bytes | None]
    debug_info: DebugInfo | None = None
    bytecode: bytes | None = None

@attrs.frozen
class _CompiledContract(CompiledContract):
    source_location: SourceLocation | None
    approval_program: _CompiledProgram
    clear_program: _CompiledProgram
    metadata: ContractMetaData

@attrs.frozen
class _CompiledLogicSig(CompiledLogicSig):
    source_location: SourceLocation | None
    program: _CompiledProgram
    metadata: LogicSignatureMetaData

def _dummy_program() -> _CompiledProgram:
    from puya.mir.models import Signature

    return _CompiledProgram(
        teal=TealProgram(
            avm_version=0,
            main=TealSubroutine(
                is_main=True,
                signature=Signature(
                    name="",
                    parameters=(),
                    returns=(),
                ),
                blocks=[],
                source_location=None,
            ),
            subroutines=[],
        ),
        teal_src="",
        template_variables={},
    )

def _contract_ir_to_teal(
    context: ArtifactCompileContext, contract_ir: ir_models.Contract
) -> _CompiledContract:
    approval_mir = program_ir_to_mir(context, contract_ir.approval_program)
    clear_state_mir = program_ir_to_mir(context, contract_ir.clear_program)
    approval_teal = mir_to_teal(context, approval_mir)
    clear_state_teal = mir_to_teal(context, clear_state_mir)
    ref = contract_ir.metadata.ref
    approval_program = _compile_program(context, ref, approval_teal)
    clear_program = _compile_program(context, ref, clear_state_teal)
    return _CompiledContract(
        approval_program=approval_program,
        clear_program=clear_program,
        metadata=contract_ir.metadata,
        source_location=contract_ir.source_location,
    )

def _logic_sig_to_teal(
    context: ArtifactCompileContext, logic_sig_ir: ir_models.LogicSignature
) -> _CompiledLogicSig:
    program_mir = program_ir_to_mir(context, logic_sig_ir.program)
    teal_program = mir_to_teal(context, program_mir)
    program = _compile_program(context, logic_sig_ir.metadata.ref, teal_program)
    return _CompiledLogicSig(
        program=program,
        metadata=logic_sig_ir.metadata,
        source_location=logic_sig_ir.source_location,
    )

def _compile_program(
    context: ArtifactCompileContext,
    ref: ContractReference | LogicSigReference,
    program: TealProgram,
) -> _CompiledProgram:
    assembled = assemble_program(context, ref, program)
    teal_src = emit_teal(context, program)
    return _CompiledProgram(
        teal=program,
        teal_src=teal_src,
        bytecode=assembled.bytecode,
        debug_info=assembled.debug_info,
        template_variables=assembled.template_variables,
    )

def _write_artifacts(
    context: CompileContext, compiled_artifacts: list[CompilationArtifact]
) -> None:
    if not compiled_artifacts:
        logger.warning("No contracts or logic signatures discovered in any source files")
        return
    for artifact in compiled_artifacts:
        out_dir = context.compilation_set.get(artifact.id)
        if out_dir is None:
            continue
        teal_file_stem = artifact.metadata.name
        artifact_base_path = out_dir / teal_file_stem
        match artifact:
            case CompiledLogicSig(program=program):
                programs = {"": program}
            case CompiledContract(approval_program=approval, clear_program=clear) as contract:
                programs = {
                    ".approval": approval,
                    ".clear": clear,
                }
                if contract.metadata.is_arc4:
                    if context.options.output_arc32:
                        app_spec_json = create_arc32_json(
                            approval.teal_src,
                            clear.teal_src,
                            contract.metadata,
                        )
                        _write_output(
                            artifact_base_path,
                            {".arc32.json": app_spec_json.encode("utf8")},
                        )
                    if context.options.output_arc56:
                        app_spec_json = create_arc56_json(
                            metadata=contract.metadata,
                            approval_program=approval,
                            clear_program=clear,
                            template_prefix=context.options.template_vars_prefix,
                        )
                        _write_output(
                            artifact_base_path,
                            {".arc56.json": app_spec_json.encode("utf8")},
                        )
            case _:
                typing.assert_never(artifact)
        if context.options.output_teal:
            _write_output(
                artifact_base_path,
                {
                    f"{suffix}.teal": program.teal_src.encode("utf8")
                    for suffix, program in programs.items()
                },
            )
        if context.options.output_bytecode:
            _write_output(
                artifact_base_path,
                {f"{suffix}.bin": program.bytecode for suffix, program in programs.items()},
            )
        if context.options.output_source_map:
            _write_output(
                artifact_base_path,
                {
                    f"{suffix}.puya.map": (
                        _debug_info_as_json(program.debug_info, out_dir)
                        if program.debug_info
                        else None
                    )
                    for suffix, program in programs.items()
                },
            )

_debug_info_converter = make_converter(omit_if_default=True)

def _debug_info_as_json(info: DebugInfo, base_path: Path) -> bytes:
    # make sources relative to output
    info = attrs.evolve(
        info,
        sources=[
            make_path_relative_to(path=Path(s), to=base_path, walk_up=True) for s in info.sources
        ],
    )
    json = _debug_info_converter.dumps(info, DebugInfo, indent=2)
    return json.encode("utf-8")

def _write_output(base_path: Path, programs: dict[str, bytes | None]) -> None:
    for suffix, program in programs.items():
        output_path = base_path.with_suffix(suffix)
        if program is None:
            logger.critical(f"Unable to output {make_path_relative_to_cwd(output_path)}")
        else:
            logger.info(f"Writing {make_path_relative_to_cwd(output_path)}")
            output_path.write_bytes(program)


import typing
from collections.abc import Mapping, Sequence
from pathlib import Path

import attrs
from immutabledict import immutabledict

from puya import log
from puya.artifact_metadata import StateTotals
from puya.compilation_artifacts import TemplateValue
from puya.options import PuyaOptions
from puya.parse import SourceLocation
from puya.program_refs import ContractReference, LogicSigReference, ProgramKind

logger = log.get_logger(__name__)

@attrs.define(kw_only=True)
class CompileContext:
    options: PuyaOptions
    compilation_set: Mapping[ContractReference | LogicSigReference, Path]
    sources_by_path: Mapping[Path, Sequence[str] | None]

    def try_get_source(self, location: SourceLocation | None) -> Sequence[str] | None:
        return try_get_source(self.sources_by_path, location)

def try_get_source(
    sources_by_path: Mapping[Path, Sequence[str] | None], location: SourceLocation | None
) -> Sequence[str] | None:
    if not location or not location.file:
        return None
    source_lines = sources_by_path.get(location.file)
    if source_lines is None:
        return None

    src_content = list(source_lines[location.line - 1 : location.end_line])
    if not src_content:
        logger.warning(f"could not locate source: {location}", location=None)
    else:
        end_column = location.end_column
        if end_column is not None:
            src_content[-1] = src_content[-1][:end_column]
        start_column = location.column
        if start_column is not None:
            src_content[0] = src_content[0][start_column:]
    return src_content

class CompiledProgramProvider(typing.Protocol):
    def build_program_bytecode(
        self,
        ref: ContractReference | LogicSigReference,
        kind: ProgramKind,
        *,
        template_constants: immutabledict[str, TemplateValue],
    ) -> bytes: ...

    def get_state_totals(self, ref: ContractReference) -> StateTotals: ...

class OutputPathProvider(typing.Protocol):
    def __call__(self, *, kind: str, qualifier: str, suffix: str) -> Path: ...

@attrs.define(kw_only=True)
class ArtifactCompileContext(CompileContext):
    _compiled_program_provider: CompiledProgramProvider = attrs.field(
        on_setattr=attrs.setters.frozen
    )
    _output_path_provider: OutputPathProvider | None = attrs.field(on_setattr=attrs.setters.frozen)

    def build_output_path(self, kind: str, qualifier: str, suffix: str) -> Path | None:
        if self._output_path_provider is None:
            return None
        return self._output_path_provider(kind=kind, qualifier=qualifier, suffix=suffix)

    @typing.overload
    def build_program_bytecode(
        self,
        ref: ContractReference,
        kind: typing.Literal[ProgramKind.approval, ProgramKind.clear_state],
        *,
        template_constants: immutabledict[str, TemplateValue],
    ) -> bytes: ...

    @typing.overload
    def build_program_bytecode(
        self,
        ref: LogicSigReference,
        kind: typing.Literal[ProgramKind.logic_signature],
        *,
        template_constants: immutabledict[str, TemplateValue],
    ) -> bytes: ...

    def build_program_bytecode(
        self,
        ref: ContractReference | LogicSigReference,
        kind: ProgramKind,
        *,
        template_constants: immutabledict[str, TemplateValue],
    ) -> bytes:
        return self._compiled_program_provider.build_program_bytecode(
            ref, kind, template_constants=template_constants
        )

    def get_state_totals(self, ref: ContractReference) -> StateTotals:
        return self._compiled_program_provider.get_state_totals(ref)


import contextlib
import enum
import sys
import traceback
from collections.abc import Iterator

from puya import log
from puya.parse import SourceLocation

logger = log.get_logger(__name__)

class ErrorExitCode(enum.IntEnum):
    code = 1
    internal = 2

class PuyaError(Exception):
    def __init__(self, msg: str, location: SourceLocation | None = None):
        super().__init__(msg)
        self.msg = msg
        self.location = location

class InternalError(PuyaError):
    """Base class for all exceptions that indicate a fault in the compiler."""

class CodeError(PuyaError):
    """Base class for all exceptions that indicate a fault in the code being compiled."""

@contextlib.contextmanager
def log_exceptions(fallback_location: SourceLocation | None = None) -> Iterator[None]:
    try:
        yield
    except CodeError as ex:
        logger.error(ex.msg, location=ex.location or fallback_location)  # noqa: TRY400
    except InternalError as ex:
        _log_traceback()
        logger.critical(ex.msg, location=ex.location or fallback_location)
        sys.exit(ErrorExitCode.internal)
    except Exception as ex:
        _log_traceback()
        logger.critical(f"{type(ex).__name__}: {ex}", location=fallback_location)
        sys.exit(ErrorExitCode.internal)

def _log_traceback() -> None:
    traceback_lines = traceback.format_exc()
    logger.debug(traceback_lines.rstrip("\n"))


IR_SSA = "ssa"
IR_OPTIMIZATION = "pass"
IR_FINAL = "final"

IR_ALL = (IR_SSA, IR_OPTIMIZATION, IR_FINAL)


import typing
from collections.abc import Mapping

import attrs

import puya.ir.models as ir
from puya import (
    artifact_metadata as models,
    log,
)
from puya.avm import AVMType, OnCompletionAction
from puya.awst import (
    nodes as awst_nodes,
    wtypes,
)
from puya.awst.arc4_types import maybe_avm_to_arc4_equivalent_type, wtype_to_arc4
from puya.context import CompiledProgramProvider
from puya.errors import CodeError
from puya.ir._utils import make_subroutine
from puya.ir.builder.main import FunctionIRBuilder
from puya.ir.context import IRBuildContext
from puya.ir.optimize.context import IROptimizationContext
from puya.ir.optimize.main import get_subroutine_optimizations
from puya.options import PuyaOptions
from puya.parse import SourceLocation

logger = log.get_logger(__name__)

def convert_default_args(
    ctx: IRBuildContext,
    contract: awst_nodes.Contract,
    m: awst_nodes.ContractMethod,
    config: awst_nodes.ARC4ABIMethodConfig,
) -> Mapping[awst_nodes.SubroutineArgument, models.MethodArgDefault | None]:
    state_by_name = {s.member_name: s for s in contract.app_state}
    return {
        a: _convert_default_arg(
            ctx, contract, state_by_name, a, config.default_args.get(a.name), method_id=m.full_name
        )
        for a in m.args
    }

def _convert_default_arg(
    ctx: IRBuildContext,
    contract: awst_nodes.Contract,
    state_by_name: Mapping[str, awst_nodes.AppStorageDefinition],
    param: awst_nodes.SubroutineArgument,
    default: awst_nodes.ABIMethodArgDefault | None,
    *,
    method_id: str,
) -> models.MethodArgDefault | None:
    match default:
        case None:
            return None
        case awst_nodes.ABIMethodArgMemberDefault(name=member_name):
            result_or_error = _convert_member_arg_default(
                contract, state_by_name, param, member_name
            )
            match result_or_error:
                case str(error_message):
                    logger.error(error_message, location=param.source_location)
                    return None
                case result:
                    return result
        case awst_nodes.ABIMethodArgConstantDefault(value=expr):
            return _compile_arc4_default_constant(ctx, method_id, param, expr)

def _convert_member_arg_default(
    contract: awst_nodes.Contract,
    state_by_name: Mapping[str, awst_nodes.AppStorageDefinition],
    param: awst_nodes.SubroutineArgument,
    member_name: str,
) -> models.MethodArgDefaultFromState | models.MethodArgDefaultFromMethod | str:
    param_arc4_type = wtype_to_arc4(param.wtype)

    # special handling for reference types
    match param_arc4_type:
        case "asset" | "application":
            param_arc4_type = "uint64"
        case "account":
            param_arc4_type = "address"

    if (state_source := state_by_name.get(member_name)) is not None:
        if state_source.key_wtype is not None:
            return "state member is a map"
        storage_type = wtypes.persistable_stack_type(
            state_source.storage_wtype, param.source_location
        )
        if (
            storage_type is AVMType.uint64
            # storage can provide an int to types <= uint64
            # TODO: check what ATC does with ufixed, see if it can be added
            and not (param_arc4_type == "byte" or param_arc4_type.startswith("uint"))
        ) or (
            storage_type is AVMType.bytes
            # storage can provide fixed byte arrays
            and not (
                param_arc4_type == "address"
                or (param_arc4_type.startswith("byte[") and param_arc4_type != "byte[]")
            )
        ):
            return f"{member_name!r} cannot provide {param_arc4_type!r} type"
        return models.MethodArgDefaultFromState(
            kind=state_source.kind,
            key=state_source.key.value,
            key_type=(
                "AVMString"
                if state_source.key.encoding == awst_nodes.BytesEncoding.utf8
                else "AVMBytes"
            ),
        )
    elif (method_source := contract.resolve_contract_method(member_name)) is not None:
        abi_method_config = method_source.arc4_method_config
        if not isinstance(abi_method_config, awst_nodes.ARC4ABIMethodConfig):
            return "only ARC-4 ABI methods can be used as default values"
        if OnCompletionAction.NoOp not in abi_method_config.allowed_completion_types:
            return f"{member_name!r} does not allow no_op on completion calls"
        if abi_method_config.create == awst_nodes.ARC4CreateOption.require:
            return f"{member_name!r} can only be used for create calls"
        if not abi_method_config.readonly:
            return f"{member_name!r} is not readonly"
        if method_source.args:
            return f"{member_name!r} does not take zero arguments"
        if method_source.return_type is wtypes.void_wtype:
            return f"{member_name!r} does not provide a value"
        return_type_arc4 = wtype_to_arc4(method_source.return_type)
        if return_type_arc4 != param_arc4_type:
            return f"{member_name!r} does not provide {param_arc4_type!r} type"
        return models.MethodArgDefaultFromMethod(
            name=abi_method_config.name,
            return_type=return_type_arc4,
            readonly=abi_method_config.readonly,
        )
    else:
        return f"{member_name!r} is not a known state or method attribute"

def _compile_arc4_default_constant(
    ctx: IRBuildContext,
    method_id: str,
    param: awst_nodes.SubroutineArgument,
    expr: awst_nodes.Expression,
) -> models.MethodArgDefaultConstant | None:
    location = expr.source_location
    logger.debug("Building IR for ARC4 method argument default constant", location=location)

    if param.wtype != expr.wtype:
        logger.error("mismatch between parameter type and default value type", location=location)
        return None

    if isinstance(expr.wtype, wtypes.ARC4Type):
        arc4_type_name = expr.wtype.arc4_name
    else:
        arc4_type = maybe_avm_to_arc4_equivalent_type(expr.wtype)
        if arc4_type is None:
            logger.error("unsupported type for argument default", location=location)
            return None
        expr = awst_nodes.ARC4Encode(value=expr, wtype=arc4_type, source_location=location)
        arc4_type_name = arc4_type.arc4_name

    fake_name = f"#default:{param.name}"
    awst_subroutine = awst_nodes.Subroutine(
        id=method_id + fake_name,
        name=fake_name,
        source_location=expr.source_location,
        args=[],
        return_type=expr.wtype,
        body=awst_nodes.Block(
            body=[awst_nodes.ReturnStatement(value=expr, source_location=expr.source_location)],
            source_location=expr.source_location,
        ),
        documentation=awst_nodes.MethodDocumentation(),
        inline=False,
    )
    ir_subroutine = make_subroutine(awst_subroutine, allow_implicits=False)
    FunctionIRBuilder.build_body(ctx, awst_subroutine, ir_subroutine)

    bytes_result = _try_extract_byte_constant(ir_subroutine)
    if bytes_result is None:
        _optimize_subroutine(ctx, ir_subroutine, location)
        bytes_result = _try_extract_byte_constant(ir_subroutine)
    if bytes_result is None:
        logger.error("could not determine constant value", location=location)
        return None
    return models.MethodArgDefaultConstant(data=bytes_result, type_=arc4_type_name)

def _optimize_subroutine(
    ctx: IRBuildContext, subroutine: ir.Subroutine, location: SourceLocation
) -> None:
    optimization_level = max(2, ctx.options.optimization_level)
    logger.debug(
        f"Running optimizer at level {optimization_level}"
        f" to encode compile time constant to bytes",
        location=location,
    )
    options = PuyaOptions(
        optimization_level=optimization_level,
        target_avm_version=ctx.options.target_avm_version,
    )
    dummy_program_provider = _NoCompiledProgramProvider(location)
    pass_context = IROptimizationContext(
        compilation_set=ctx.compilation_set,
        sources_by_path=ctx.sources_by_path,
        options=options,
        compiled_program_provider=dummy_program_provider,
        output_path_provider=None,
        expand_all_bytes=True,
    )
    pipeline = get_subroutine_optimizations(optimization_level=optimization_level)
    while True:
        modified = False
        for optimizer in pipeline:
            if optimizer.optimize(pass_context, subroutine):
                modified = True
        if not modified:
            return

def _try_extract_byte_constant(subroutine: ir.Subroutine) -> bytes | None:
    match subroutine.body:
        case [
            ir.BasicBlock(
                phis=[],
                ops=[],
                terminator=ir.SubroutineReturn(result=[ir.BytesConstant(value=result)]),
            )
        ]:
            return result
    return None

@attrs.frozen
class _NoCompiledProgramProvider(CompiledProgramProvider):
    source_location: SourceLocation

    @typing.override
    def build_program_bytecode(self, *args: object, **kwargs: object) -> typing.Never:
        raise CodeError(
            "compilation references are not supported as method default constants",
            self.source_location,
        )

    @typing.override
    def get_state_totals(self, *args: object, **kwargs: object) -> typing.Never:
        raise CodeError(
            "compilation references are not supported as method default constants",
            self.source_location,
        )