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

import attrs

from puya import log
from puya.errors import InternalError
from puya.mir import models as mir
from puya.mir.context import SubroutineCodeGenContext
from puya.mir.stack import Stack

logger = log.get_logger(__name__)

@attrs.define(eq=False, repr=False)
class BlockRecord:
    block: mir.MemoryBasicBlock
    local_references: list[mir.AbstractStore | mir.AbstractLoad]
    live_in: Set[str]
    live_out: Set[str]
    children: "list[BlockRecord]" = attrs.field(factory=list)
    parents: "list[BlockRecord]" = attrs.field(factory=list)
    co_parents: "list[BlockRecord]" = attrs.field(factory=list)
    siblings: "list[BlockRecord]" = attrs.field(factory=list)

    x_stack_in: Sequence[str] | None = None
    x_stack_out: Sequence[str] | None = None

    def __repr__(self) -> str:
        # due to recursive nature of BlockRecord, provide str implementation to
        # simplify output
        return f"BlockRecord({self.block})"

    @staticmethod
    def by_index(block: "BlockRecord") -> int:
        return block.block.id

@attrs.frozen
class EdgeSet:
    out_blocks: Sequence[BlockRecord] = attrs.field(converter=tuple[BlockRecord, ...])
    in_blocks: Sequence[BlockRecord] = attrs.field(converter=tuple[BlockRecord, ...])

def sort_by_appearance(
    variables: Set[str], block: mir.MemoryBasicBlock, *, load: bool = True
) -> Sequence[str]:
    appearance = list[str]()
    block_ops = block.ops if load else reversed(block.ops)
    if load:
        virtual_ops = (o.local_id for o in block_ops if isinstance(o, mir.AbstractLoad))
    else:
        virtual_ops = (o.local_id for o in block_ops if isinstance(o, mir.AbstractStore))
    for local_id in virtual_ops:
        if local_id in variables and local_id not in appearance:
            appearance.append(local_id)
            # don't keep searching once we are done
            if len(appearance) == len(variables):
                break
    return appearance

def len_and_value(value: tuple[str, ...]) -> tuple[int, tuple[str, ...]]:
    return len(value), value

def find_shared_x_stack(x_stack_candidates: Sequence[Sequence[str]]) -> Sequence[str]:
    """Find a common subsequence that is shared by all x-stacks"""
    cache = dict[tuple[tuple[str, ...], tuple[str, ...]], tuple[str, ...]]()

    def lcs(s1: tuple[str, ...], s2: tuple[str, ...]) -> tuple[str, ...]:
        key = (s1, s2)
        result = cache.get(key)
        if result is None:
            i = len(s1)
            j = len(s2)
            if i == 0 or j == 0:
                result = ()
            elif s1[-1] == s2[-1]:
                result = (*lcs(s1[:-1], s2[:-1]), s1[-1])
            else:
                result = max(lcs(s1[:-1], s2), lcs(s1, s2[:-1]), key=len_and_value)
            cache[key] = result
        return result

    shared, *others = sorted({tuple(s) for s in x_stack_candidates}, key=len_and_value)
    for other in others:
        shared = lcs(shared, other)
    return shared

def get_x_stack_load_ops(record: BlockRecord) -> set[mir.AbstractLoad]:
    block = record.block
    assert block.x_stack_in is not None

    remaining = set(block.x_stack_in)
    load_ops = []
    for ref in record.local_references:
        if isinstance(ref, mir.AbstractLoad) and ref.local_id in remaining:
            remaining.remove(ref.local_id)
            load_ops.append(ref)

    if remaining:
        raise InternalError(
            f"Failed to move {', '.join(remaining)} from the x-stack",
            location=block.source_location,
        )

    return set(load_ops)

def get_x_stack_store_ops(record: BlockRecord) -> set[mir.AbstractStore]:
    block = record.block
    assert block.x_stack_out is not None

    remaining = set(block.x_stack_out)
    store_ops = []
    for ref in reversed(record.local_references):
        if isinstance(ref, mir.AbstractStore) and ref.local_id in remaining:
            remaining.remove(ref.local_id)
            store_ops.append(ref)

    if remaining:
        raise InternalError(
            f"Failed to copy {', '.join(remaining)} to the x-stack",
            location=block.source_location,
        )

    return set(store_ops)

def add_x_stack_ops(ctx: SubroutineCodeGenContext, record: BlockRecord) -> None:
    block = record.block
    # determine ops to replace
    load_ops = get_x_stack_load_ops(record)
    store_ops = get_x_stack_store_ops(record)

    stack = Stack.begin_block(ctx.subroutine, block)
    for index, op in enumerate(block.mem_ops):
        if op in store_ops:
            assert isinstance(op, mir.AbstractStore)
            # can replace virtual store op because only variables that could be fully
            # scheduled are on the x-stack
            block.mem_ops[index] = op = mir.StoreXStack(
                local_id=op.local_id,
                depth=stack.xl_height - 1,  # store to bottom
                atype=op.atype,
                source_location=op.source_location,
            )
        elif op in load_ops:
            assert isinstance(op, mir.AbstractLoad)
            block.mem_ops[index] = op = mir.LoadXStack(
                local_id=op.local_id,
                depth=stack.xl_height - stack.x_stack.index(op.local_id) - 1,
                atype=op.atype,
                source_location=op.source_location,
            )
        op.accept(stack)

def add_x_stack_ops_to_edge_sets(
    ctx: SubroutineCodeGenContext, edge_sets: Sequence[EdgeSet]
) -> None:
    records = dict.fromkeys(
        b
        for edge_set in edge_sets
        for b in itertools.chain(edge_set.out_blocks, edge_set.in_blocks)
        if b.x_stack_in or b.x_stack_out
    )
    for record in records:
        assert record.x_stack_in is not None
        assert record.x_stack_out is not None
        record.block.x_stack_in = record.x_stack_in
        record.block.x_stack_out = record.x_stack_out
        add_x_stack_ops(ctx, record)

def _unique_ordered_blocks(blocks: Iterable[BlockRecord]) -> list[BlockRecord]:
    return sorted(set(blocks), key=BlockRecord.by_index)

def get_edge_set(block: BlockRecord) -> EdgeSet | None:
    out_blocks = _unique_ordered_blocks(itertools.chain((block,), block.co_parents))
    # keep expanding out_blocks (and consequently in_blocks) until out_blocks stabilize
    while True:
        in_blocks = _unique_ordered_blocks(s for p in out_blocks for s in p.children)
        new_out_blocks = _unique_ordered_blocks(p for s in in_blocks for p in s.parents)
        if new_out_blocks == out_blocks:
            break
        out_blocks = new_out_blocks

    return EdgeSet(out_blocks, in_blocks) if in_blocks else None

def get_edge_sets(ctx: SubroutineCodeGenContext) -> Sequence[EdgeSet]:
    subroutine = ctx.subroutine
    vla = ctx.vla
    records = {
        block: BlockRecord(
            block=block,
            local_references=[
                op for op in block.ops if isinstance(op, mir.AbstractStore | mir.AbstractLoad)
            ],
            live_in=vla.get_live_in_variables(block.ops[0]),
            live_out=vla.get_live_out_variables(block.ops[-1]),
        )
        for block in subroutine.body
    }

    # given blocks 1,2,3,4,5,6 and 7
    # edges: 1->5, 2->4, 2->5, 2->6, 3->5, 7->6, 7->8
    #
    # e.g  1  2  3   7
    #       \/|\/   / \
    #      / \|/ \ /   \
    #     4   5   6     8
    #
    #   consider 2
    #   4, 5 & 6 are children
    #   1 & 3 are co-parents of 5, 7 is a co-parent of 6
    # 1, 2, 3 and 7 form the out_blocks of an edge set
    # 4, 5 & 6 are the in_blocks of the same edge set

    # 1. first pass
    # populate children and parents
    blocks = [records[b] for b in subroutine.body]
    for block in blocks:
        block.children = [records[subroutine.get_block(c)] for c in block.block.successors]
        for child in block.children:
            child.parents.append(block)

    # 2. second pass - boundary mapping
    for block in blocks:
        # determine siblings
        for parent in block.parents:
            for child in parent.children:
                if child is not block and child not in block.siblings:
                    block.siblings.append(child)

        # determine co-parents
        for child in block.children:
            for parent in child.parents:
                if parent is not block and parent not in block.co_parents:
                    block.co_parents.append(parent)

    edge_sets = dict[EdgeSet, None]()
    for block in blocks:
        # keep expanding edge set until it stabilizes
        edge_set = get_edge_set(block)
        if edge_set:
            edge_sets[edge_set] = None
        else:
            block.x_stack_out = ()

        if not block.parents:
            block.x_stack_in = ()

    return list(edge_sets.keys())

def schedule_sets(ctx: SubroutineCodeGenContext, edge_sets: Sequence[EdgeSet]) -> None:
    # determine all blocks referencing variables, so we can track if all references to a
    # variable are scheduled to x-stack
    stores = dict[str, set[mir.MemoryBasicBlock]]()
    loads = dict[str, set[mir.MemoryBasicBlock]]()
    vla = ctx.vla
    for variable in vla.all_variables:
        stores[variable] = vla.get_store_blocks(variable)
        loads[variable] = vla.get_load_blocks(variable)

    for edge_set in edge_sets:
        in_blocks = edge_set.in_blocks
        out_blocks = edge_set.out_blocks

        # get potential l-stacks (unordered)
        l_stacks = [
            *(b.live_out for b in out_blocks),
            *(b.live_in for b in in_blocks),
        ]

        # determine shared l-stack variables for this edge_set
        # determine all valid x-stacks (ordered)
        first, *others = l_stacks
        common_locals = frozenset(first).intersection(*others)

        # TODO: better results might be possible if we allow reordering of x-stack
        x_stack_candidates = [
            *(sort_by_appearance(common_locals, b.block, load=False) for b in out_blocks),
            *(sort_by_appearance(common_locals, b.block, load=True) for b in in_blocks),
        ]

        # find an x_stack for this EdgeSet
        x_stack = find_shared_x_stack(x_stack_candidates)

        for block in out_blocks:
            assert block.x_stack_out is None
            block.x_stack_out = x_stack
            for x in x_stack:
                stores[x].remove(block.block)

        for block in in_blocks:
            assert block.x_stack_in is None
            block.x_stack_in = x_stack
            for x in x_stack:
                loads[x].remove(block.block)

    # adjust final x-stacks based on what could be fully scheduled
    variables_not_fully_scheduled = {
        var for var, blocks in itertools.chain(stores.items(), loads.items()) if len(blocks) > 0
    }
    variables_successfully_scheduled = sorted(stores.keys() - variables_not_fully_scheduled)
    for block in {b for es in edge_sets for b in itertools.chain((*es.out_blocks, *es.in_blocks))}:
        assert block.x_stack_out is not None
        assert block.x_stack_in is not None
        block.x_stack_out = tuple(
            x for x in block.x_stack_out if x in variables_successfully_scheduled
        )
        block.x_stack_in = tuple(
            x for x in block.x_stack_in if x in variables_successfully_scheduled
        )

    if variables_successfully_scheduled:
        ctx.invalidate_vla()
        logger.debug(
            f"Allocated {len(variables_successfully_scheduled)} "
            f"variable/s to x-stack: {', '.join(variables_successfully_scheduled)}"
        )

def validate_pair(parent: BlockRecord, child: BlockRecord) -> bool:
    parent_x = parent.x_stack_out
    child_x = child.x_stack_in
    assert parent_x is not None
    assert child_x is not None
    if parent_x != child_x:
        logger.error(
            f"x-stacks do not match for {parent.block} -> {child.block}: "
            f"{', '.join(parent_x)} -> {', '.join(child_x)}"
        )
        return False
    if parent_x:
        logger.debug(f"shared x-stack for {parent.block} -> {child.block}: {', '.join(parent_x)}")
    return True

def validate_x_stacks(edge_sets: Sequence[EdgeSet]) -> bool:
    ok = True
    for edge_set in edge_sets:
        for parent in edge_set.out_blocks:
            for child in edge_set.in_blocks:
                ok = validate_pair(parent, child) and ok
    return ok

def x_stack_allocation(ctx: SubroutineCodeGenContext) -> None:
    # this is basically baileys algorithm
    edge_sets = get_edge_sets(ctx)
    if not edge_sets:
        # nothing to do
        return

    logger.debug(f"Found {len(edge_sets)} edge set/s for {ctx.subroutine.signature.name}")
    schedule_sets(ctx, edge_sets)

    if not validate_x_stacks(edge_sets):
        raise InternalError("Could not schedule x-stack")

    add_x_stack_ops_to_edge_sets(ctx, edge_sets)


from __future__ import annotations

import abc
import typing

if typing.TYPE_CHECKING:
    from puya.mir import models

class MIRVisitor[T](abc.ABC):
    @abc.abstractmethod
    def visit_int(self, push: models.Int) -> T: ...

    @abc.abstractmethod
    def visit_byte(self, push: models.Byte) -> T: ...

    @abc.abstractmethod
    def visit_undefined(self, push: models.Undefined) -> T: ...

    @abc.abstractmethod
    def visit_comment(self, comment: models.Comment) -> T: ...

    @abc.abstractmethod
    def visit_store_l_stack(self, store: models.StoreLStack) -> T: ...

    @abc.abstractmethod
    def visit_load_l_stack(self, load: models.LoadLStack) -> T: ...

    @abc.abstractmethod
    def visit_store_x_stack(self, store: models.StoreXStack) -> T: ...

    @abc.abstractmethod
    def visit_load_x_stack(self, load: models.LoadXStack) -> T: ...

    @abc.abstractmethod
    def visit_store_f_stack(self, store: models.StoreFStack) -> T: ...

    @abc.abstractmethod
    def visit_load_f_stack(self, load: models.LoadFStack) -> T: ...

    @abc.abstractmethod
    def visit_load_param(self, load: models.LoadParam) -> T: ...

    @abc.abstractmethod
    def visit_store_param(self, load: models.StoreParam) -> T: ...

    @abc.abstractmethod
    def visit_abstract_store(self, store: models.AbstractStore) -> T: ...

    @abc.abstractmethod
    def visit_abstract_load(self, load: models.AbstractLoad) -> T: ...

    @abc.abstractmethod
    def visit_allocate(self, allocate: models.Allocate) -> T: ...

    @abc.abstractmethod
    def visit_pop(self, pop: models.Pop) -> T: ...

    @abc.abstractmethod
    def visit_callsub(self, callsub: models.CallSub) -> T: ...

    @abc.abstractmethod
    def visit_intrinsic(self, intrinsic: models.IntrinsicOp) -> T: ...

    @abc.abstractmethod
    def visit_retsub(self, retsub: models.RetSub) -> T: ...

    @abc.abstractmethod
    def visit_program_exit(self, op: models.ProgramExit) -> T: ...

    @abc.abstractmethod
    def visit_err(self, op: models.Err) -> T: ...

    @abc.abstractmethod
    def visit_goto(self, op: models.Goto) -> T: ...

    @abc.abstractmethod
    def visit_conditional_branch(self, op: models.ConditionalBranch) -> T: ...

    @abc.abstractmethod
    def visit_switch(self, op: models.Switch) -> T: ...

    @abc.abstractmethod
    def visit_match(self, op: models.Match) -> T: ...

    @abc.abstractmethod
    def visit_address(self, addr: models.Address) -> T: ...

    @abc.abstractmethod
    def visit_method(self, method: models.Method) -> T: ...

    @abc.abstractmethod
    def visit_template_var(self, deploy_var: models.TemplateVar) -> T: ...


import itertools
import typing
from collections.abc import Sequence, Set
from functools import cached_property

import attrs

from puya.mir import models
from puya.utils import StableSet

@attrs.define(kw_only=True)
class _OpLifetime:
    block: models.MemoryBasicBlock
    used: StableSet[str] = attrs.field(on_setattr=attrs.setters.frozen)
    defined: StableSet[str] = attrs.field(on_setattr=attrs.setters.frozen)
    successors: list[typing.Self] = attrs.field(factory=list)
    predecessors: list[typing.Self] = attrs.field(factory=list)

    live_in: StableSet[str] = attrs.field(factory=StableSet)
    live_out: StableSet[str] = attrs.field(factory=StableSet)

@attrs.define
class VariableLifetimeAnalysis:
    """Performs VLA analysis for a subroutine, providing a mapping of ops to sets of live local_ids
    see https://www.classes.cs.uchicago.edu/archive/2004/spring/22620-1/docs/liveness.pdf"""

    subroutine: models.MemorySubroutine
    _op_lifetimes: dict[models.BaseOp, _OpLifetime] = attrs.field(init=False)

    @cached_property
    def all_variables(self) -> Sequence[str]:
        return sorted(
            {v for live in self._op_lifetimes.values() for v in (*live.defined, *live.used)}
        )

    @_op_lifetimes.default
    def _op_lifetimes_factory(self) -> dict[models.BaseOp, _OpLifetime]:
        result = dict[models.BaseOp, _OpLifetime]()
        block_map = {b.block_name: b.ops[0] for b in self.subroutine.body}
        for block in self.subroutine.body:
            for op in block.ops:
                used = StableSet[str]()
                defined = StableSet[str]()
                if isinstance(op, models.AbstractStore):
                    defined.add(op.local_id)
                elif isinstance(op, models.AbstractLoad):
                    used.add(op.local_id)
                result[op] = _OpLifetime(
                    block=block,
                    used=used,
                    defined=defined,
                )
        for block in self.subroutine.body:
            for op, next_op in itertools.zip_longest(block.ops, block.ops[1:]):
                if isinstance(op, models.ControlOp):
                    assert next_op is None
                    # note: control ops that end the current subroutine don't have any logical
                    # successors
                    successors = [result[block_map[s]] for s in op.targets()]
                else:
                    assert next_op is not None
                    successors = [result[next_op]]
                op_lifetime = result[op]
                op_lifetime.successors = successors
                for s in successors:
                    s.predecessors.append(op_lifetime)
        return result

    def get_live_out_variables(self, op: models.BaseOp) -> Set[str]:
        return self._op_lifetimes[op].live_out

    def get_live_in_variables(self, op: models.BaseOp) -> Set[str]:
        return self._op_lifetimes[op].live_in

    def get_store_blocks(self, variable: str) -> set[models.MemoryBasicBlock]:
        return {op.block for op in self._op_lifetimes.values() if variable in op.defined}

    def get_load_blocks(self, variable: str) -> set[models.MemoryBasicBlock]:
        return {op.block for op in self._op_lifetimes.values() if variable in op.used}

    @classmethod
    def analyze(cls, subroutine: models.MemorySubroutine) -> typing.Self:
        analysis = cls(subroutine)
        analysis._analyze()  # noqa: SLF001
        return analysis

    def _analyze(self) -> None:
        changed = list(self._op_lifetimes.values())
        while changed:
            orig_changed = changed
            changed = []
            for n in orig_changed:
                # For OUT, find out the union of previous variables
                # in the IN set for each succeeding node of n.

                # out[n] = U s ∈ succ[n] in[s]
                live_out = StableSet[str]()
                for s in n.successors:
                    live_out |= s.live_in

                # in[n] = use[n] U (out[n] - def [n])
                live_in = n.used | (live_out - n.defined)

                if live_out != n.live_out or live_in != n.live_in:
                    n.live_in = live_in
                    n.live_out = live_out
                    changed.extend(n.predecessors)


import enum
from collections.abc import Mapping
from functools import cached_property

import attrs

from puya.algo_constants import MAINNET_AVM_VERSION

class LocalsCoalescingStrategy(enum.StrEnum):
    root_operand = enum.auto()
    root_operand_excluding_args = enum.auto()
    aggressive = enum.auto()

@attrs.frozen(kw_only=True)
class PuyaOptions:
    output_teal: bool = False
    output_source_map: bool = False
    output_arc32: bool = False
    output_arc56: bool = False
    output_ssa_ir: bool = False
    output_optimization_ir: bool = False
    output_destructured_ir: bool = False
    output_memory_ir: bool = False
    output_bytecode: bool = False
    debug_level: int = 1
    optimization_level: int = 1
    target_avm_version: int = MAINNET_AVM_VERSION
    cli_template_definitions: Mapping[str, int | bytes] = attrs.field(factory=dict)
    template_vars_prefix: str = "TMPL_"
    # TODO: the below is probably not scalable as a set of optimisation on/off flags,
    #       but it'll do for now
    locals_coalescing_strategy: LocalsCoalescingStrategy = LocalsCoalescingStrategy.root_operand

    @cached_property
    def template_variables(self) -> Mapping[str, int | bytes]:
        return {self.template_vars_prefix + k: v for k, v in self.cli_template_definitions.items()}


import functools
import typing
from collections.abc import Iterable
from pathlib import Path

import attrs

from puya.utils import make_path_relative_to_cwd

@attrs.frozen(kw_only=True, repr=False, str=False)
class SourceLocation:
    file: Path | None = attrs.field()
    line: int = attrs.field(validator=attrs.validators.ge(1))
    end_line: int = attrs.field()
    comment_lines: int = attrs.field(default=0, validator=attrs.validators.ge(0))
    """the number of lines preceding `line` to take as a comment"""
    column: int | None = attrs.field(
        default=None, validator=attrs.validators.optional(attrs.validators.ge(0))
    )
    end_column: int | None = attrs.field(
        default=None, validator=attrs.validators.optional(attrs.validators.ge(1))
    )

    @file.validator
    def _file_validator(self, _attribute: object, value: Path) -> None:
        # this check is simply to make sure relative paths aren't accidentally passed in.
        # so we use root rather than is_absolute(), because that requires a drive on Windows,
        # which we naturally don't supply for synthetic paths such as embedded lib.
        if value is not None and not value.root:
            raise ValueError(f"source file locations cannot be relative, got {value}")

    @end_line.default
    def _end_line_default(self) -> int:
        return self.line

    @end_line.validator
    def _end_line_validator(self, _attribute: object, value: int) -> None:
        if value < self.line:
            raise ValueError(f"end_line = {value} is before start line = {self.line}")

    @end_column.validator
    def _end_column_validator(self, _attribute: object, value: int | None) -> None:
        if (
            self.end_line == self.line
            and value is not None
            and self.column is not None
            and value <= self.column
        ):
            raise ValueError(
                f"source location end column = {value} is before start column = {self.column}"
            )

    @property
    def line_count(self) -> int:
        return self.end_line - self.line + 1

    def __str__(self) -> str:
        relative_path = make_path_relative_to_cwd(self.file) if self.file else "INTERNAL"
        result = f"{relative_path}:{self.line}"
        if self.end_line != self.line:
            result += f"-{self.end_line}"
        return result

    def __repr__(self) -> str:
        result = str(self)
        if self.column is not None:
            result += f":{self.column}"
            if self.end_column is not None:
                result += f"-{self.end_column}"
        return result

    def with_comments(self) -> "SourceLocation":
        if self.comment_lines == 0:
            return self
        return attrs.evolve(
            self,
            line=self.line - self.comment_lines,
            column=None,
            comment_lines=0,
        )

    def try_merge(self, other: "SourceLocation | None") -> "SourceLocation":
        """Attempt to merge this source location with another, if they are either adjacent
        or overlapping in lines. If not, the source location is returned unmodified."""
        if other is None or other.file != self.file:
            return self
        file = self.file
        # if they both start at the same line, not only is there overlap,
        # but things are also much simpler
        if self.line == other.line:
            line = self.line
            # expand to the largest end_line
            end_line = max(self.end_line, other.end_line)
            # in theory this should be the same value, but just in case, we can take the max
            comment_lines = max(self.comment_lines, other.comment_lines)
            # if either location is not column-bounded, then the result shouldn't be either
            # otherwise take the minimum of the columns, since the line numbers are the same
            if self.column is None or other.column is None:
                column = None
            else:
                column = min(self.column, other.column)
        else:
            # if they don't start on the same line, one must start first
            first, second = (self, other) if self.line < other.line else (other, self)
            line_after_first = first.end_line + 1
            # TODO: maybe consider fetching the source to exclude blank lines?
            if line_after_first < second.line:
                return self
            # first starts first, so... that's where we start
            line = first.line
            # whilst we know first starts before second,
            # it's also possible that first ends after second
            end_line = max(second.end_line, first.end_line)
            # naturally, comment line count needs to come from the first location
            comment_lines = first.comment_lines
            # same first starting column
            column = first.column
        # the logic for computing the end_column is the same regardless of whether
        # they start on the same line or not
        if self.end_line == other.end_line:
            # if either location is not end_column-bounded, then the result shouldn't be either
            # otherwise take the maximum of the end_columns, since the line numbers are the same
            if self.end_column is None or other.end_column is None:
                end_column = None
            else:
                end_column = max(self.end_column, other.end_column)
        elif self.end_line > other.end_line:
            # if self ends last, take it's end column
            end_column = self.end_column
        else:
            # otherwise other ends last, so take it's end column
            end_column = other.end_column

        return SourceLocation(
            file=file,
            line=line,
            end_line=end_line,
            comment_lines=comment_lines,
            column=column,
            end_column=end_column,
        )

@typing.overload
def sequential_source_locations_merge(sources: Iterable[SourceLocation]) -> SourceLocation: ...

@typing.overload
def sequential_source_locations_merge(
    sources: Iterable[SourceLocation | None],
) -> SourceLocation | None: ...

def sequential_source_locations_merge(
    sources: Iterable[SourceLocation | None],
) -> SourceLocation | None:
    """Given a sequence of SourceLocations, try merging them one at a one in order.

    If all sources are None, then None is returned.

    If there are no sources, then a TypeError will be raised.
    """
    return functools.reduce(_try_merge_source_locations, sources)

def _try_merge_source_locations(
    source: SourceLocation | None, merge: SourceLocation | None
) -> SourceLocation | None:
    if source is None:
        return merge
    return source.try_merge(merge)


import abc
import enum
import typing

import attrs

class ContractReference(str):  # can't use typing.NewType with pattern matching
    __slots__ = ()

class LogicSigReference(str):  # can't use typing.NewType with pattern matching
    __slots__ = ()

class ProgramKind(enum.StrEnum):
    approval = "approval"
    clear_state = "clear"
    logic_signature = "lsig"

class ProgramReference(abc.ABC):
    @property
    @abc.abstractmethod
    def kind(self) -> ProgramKind: ...

    @property
    @abc.abstractmethod
    def reference(self) -> ContractReference | LogicSigReference: ...

    @property
    @abc.abstractmethod
    def id(self) -> str: ...

@attrs.frozen
class LogicSigProgramReference(ProgramReference):
    kind: typing.Literal[ProgramKind.logic_signature] = attrs.field(
        default=ProgramKind.logic_signature, init=False
    )
    reference: LogicSigReference

    @property
    def id(self) -> str:
        return self.reference

@attrs.frozen
class ContractProgramReference(ProgramReference):
    kind: typing.Literal[ProgramKind.approval, ProgramKind.clear_state]
    reference: ContractReference
    program_name: str

    @property
    def id(self) -> str:
        return f"{self.reference}.{self.program_name}"


import typing
from collections.abc import Sequence

import attrs

from puya.errors import InternalError
from puya.teal import models

_T = typing.TypeVar("_T", bound=models.TealOp)

def preserve_stack_manipulations(
    ops: list[_T],
    window: slice,
    new: Sequence[_T],
) -> None:
    """Replaces window of ops with new, preserving any stack manipulations from original window

    new is not empty: added to last op in new
    new is empty, window starts after first index: appended to prior op
    new is empty, window ends before last index: prepended to subsequent op
    """
    if not new:
        # expand window to include at least 1 op
        if window.start > 0:
            # expand window to prior op
            window = slice(window.start - 1, window.stop)
            new = [ops[window.start]]
        elif window.stop < len(ops):  # must be start of a block
            # expand window to subsequent op
            new = [ops[window.stop]]
            window = slice(window.start, window.stop + 1)
        else:
            # can this even happen? if it does, maybe attach to block instead?
            raise InternalError("could not preserve stack manipulations")

    # clear existing stack_manipulations on new sequence
    new = [attrs.evolve(op, stack_manipulations=()) for op in new]
    # add original stack manipulations to last op in new sequence
    new[-1] = attrs.evolve(
        new[-1],
        stack_manipulations=[sm for op in ops[window] for sm in op.stack_manipulations],
    )
    # replace original ops with new ops
    ops[window] = new


import typing
from itertools import zip_longest

import attrs

from puya.avm import AVMType
from puya.errors import InternalError
from puya.ir.types_ import AVMBytesEncoding
from puya.mir import models as mir
from puya.mir.visitor import MIRVisitor
from puya.teal import models as teal

@attrs.frozen
class TealBuilder(MIRVisitor[None]):
    next_block_label: str | None
    use_frame: bool
    label_stack: list[str]
    ops: list[teal.TealOp] = attrs.field(factory=list)

    @classmethod
    def build_subroutine(cls, mir_sub: mir.MemorySubroutine) -> teal.TealSubroutine:
        result = teal.TealSubroutine(
            is_main=mir_sub.is_main,
            signature=mir_sub.signature,
            blocks=[],
            source_location=mir_sub.source_location,
        )
        entry_block = mir_sub.body[0]
        label_stack = [entry_block.block_name]
        blocks_by_label = {
            b.block_name: (b, None if next_b is None else next_b.block_name)
            for b, next_b in zip_longest(mir_sub.body, mir_sub.body[1:])
        }
        while label_stack:
            label = label_stack.pop()
            mir_block, next_block_label = blocks_by_label.pop(label, (None, None))
            if mir_block is None:
                continue
            builder = cls(
                next_block_label=next_block_label,
                use_frame=not mir_sub.is_main,
                label_stack=label_stack,
            )
            if mir_block is entry_block and not mir_sub.is_main:
                builder.ops.append(
                    teal.Proto(
                        parameters=len(mir_sub.signature.parameters),
                        returns=len(mir_sub.signature.returns),
                        source_location=mir_sub.source_location,
                    )
                )
            for op in mir_block.ops:
                op.accept(builder)
            teal_block = teal.TealBlock(
                label=mir_block.block_name,
                ops=builder.ops,
                x_stack_in=mir_block.x_stack_in or (),
                entry_stack_height=mir_block.entry_stack_height,
                exit_stack_height=mir_block.exit_stack_height,
            )
            teal_block.validate_stack_height()
            result.blocks.append(teal_block)
        return result

    def _add_op(self, op: teal.TealOp) -> None:
        self.ops.append(op)

    def visit_int(self, const: mir.Int) -> None:
        self._add_op(
            teal.Int(
                const.value,
                stack_manipulations=_lstack_manipulations(const),
                source_location=const.source_location,
            )
        )

    def visit_byte(self, const: mir.Byte) -> None:
        self._add_op(
            teal.Byte(
                const.value,
                const.encoding,
                stack_manipulations=_lstack_manipulations(const),
                source_location=const.source_location,
            )
        )

    def visit_undefined(self, push: mir.Undefined) -> None:
        match push.atype:
            case AVMType.uint64:
                self._add_op(
                    teal.Byte(
                        b"",
                        AVMBytesEncoding.utf8,
                        stack_manipulations=_lstack_manipulations(push),
                        source_location=push.source_location,
                    )
                )
            case AVMType.bytes:
                self._add_op(
                    teal.Int(
                        0,
                        stack_manipulations=_lstack_manipulations(push),
                        source_location=push.source_location,
                    )
                )
            case unexpected:
                typing.assert_never(unexpected)

    def visit_template_var(self, const: mir.TemplateVar) -> None:
        self._add_op(
            teal.TemplateVar(
                name=const.name,
                op_code=const.op_code,
                stack_manipulations=_lstack_manipulations(const),
                source_location=const.source_location,
            )
        )

    def visit_address(self, const: mir.Address) -> None:
        self._add_op(
            teal.Address(
                const.value,
                stack_manipulations=_lstack_manipulations(const),
                source_location=const.source_location,
            )
        )

    def visit_method(self, const: mir.Method) -> None:
        self._add_op(
            teal.Method(
                const.value,
                stack_manipulations=_lstack_manipulations(const),
                source_location=const.source_location,
            )
        )

    def visit_comment(self, _comment: mir.Comment) -> None:
        pass

    def visit_abstract_store(self, store: mir.AbstractStore) -> typing.Never:
        raise InternalError(
            "AbstractStore op encountered during TEAL generation", store.source_location
        )

    def visit_abstract_load(self, load: mir.AbstractLoad) -> typing.Never:
        raise InternalError(
            "AbstractLoad op encountered during TEAL generation", load.source_location
        )

    def _store_f_stack(self, store: mir.StoreFStack) -> None:
        local_id = store.local_id
        source_location = store.source_location
        define = teal.StackDefine(local_id)
        if self.use_frame:
            op: teal.TealOp = teal.FrameBury(
                store.frame_index,
                stack_manipulations=[*_lstack_manipulations(store), define],
                source_location=source_location,
            )
        else:
            op = teal.Bury(
                store.depth,
                stack_manipulations=[*_lstack_manipulations(store), define],
                source_location=source_location,
            )
        self._add_op(op)

    def _insert_f_stack(self, store: mir.StoreFStack) -> None:
        local_id = store.local_id
        source_location = store.source_location
        self._add_op(
            teal.Cover(
                store.depth,
                stack_manipulations=[
                    *_lstack_manipulations(store),
                    teal.StackInsert(store.depth, local_id),
                    teal.StackDefine(local_id),
                ],
                source_location=source_location,
            )
        )

    def visit_store_f_stack(self, store: mir.StoreFStack) -> None:
        if store.insert:
            self._insert_f_stack(store)
        else:
            self._store_f_stack(store)

    def visit_load_f_stack(self, load: mir.LoadFStack) -> None:
        sm = _lstack_manipulations(load)
        loc = load.source_location
        if self.use_frame:  # and load.depth:
            op: teal.TealOp = teal.FrameDig(
                n=load.frame_index, stack_manipulations=sm, source_location=loc
            )
        else:
            op = teal.Dig(n=load.depth, stack_manipulations=sm, source_location=loc)
        self._add_op(op)

    def visit_store_x_stack(self, store: mir.StoreXStack) -> None:
        self._add_op(
            teal.Cover(
                store.depth,
                stack_manipulations=[
                    *_lstack_manipulations(store),
                    teal.StackInsert(store.depth, store.local_id),
                    teal.StackDefine(store.local_id),
                ],
                source_location=store.source_location,
            )
        )

    def visit_load_x_stack(self, load: mir.LoadXStack) -> None:
        self._add_op(
            teal.Uncover(
                load.depth,
                stack_manipulations=[
                    teal.StackPop(load.depth),
                    *_lstack_manipulations(load),
                ],
                source_location=load.source_location,
            )
        )

    def visit_store_l_stack(self, store: mir.StoreLStack) -> None:
        if store.copy:
            self._add_op(
                teal.Dup(
                    stack_manipulations=[
                        # re-alias top of stack
                        teal.StackConsume(1),
                        *_lstack_manipulations(store),
                        # actual dup
                        teal.StackExtend([f"{store.local_id} (copy)"]),
                    ],
                    source_location=store.source_location,
                ),
            )
        cover = store.depth
        if store.copy:
            cover += 1
        self._add_op(
            teal.Cover(
                cover,
                stack_manipulations=[
                    teal.StackConsume(1),
                    # store
                    teal.StackInsert(cover, store.local_id),
                    teal.StackDefine([store.local_id]),
                ],
                source_location=store.source_location,
            )
        )

    def visit_load_l_stack(self, load: mir.LoadLStack) -> None:
        uncover = load.depth
        assert uncover is not None, "expected l-stack depths to be assigned"
        if load.copy:
            self._add_op(
                teal.Dig(
                    uncover,
                    stack_manipulations=_lstack_manipulations(load),
                    source_location=load.source_location,
                )
            )
        else:
            self._add_op(
                teal.Uncover(
                    uncover,
                    stack_manipulations=[teal.StackPop(uncover), *_lstack_manipulations(load)],
                    source_location=load.source_location,
                )
            )

    def visit_load_param(self, load: mir.LoadParam) -> None:
        self._add_op(
            teal.FrameDig(
                load.index,
                source_location=load.source_location,
                stack_manipulations=_lstack_manipulations(load),
            )
        )

    def visit_store_param(self, store: mir.StoreParam) -> None:
        self._add_op(
            teal.FrameBury(
                store.index,
                stack_manipulations=_lstack_manipulations(store),
                source_location=store.source_location,
            )
        )

    def visit_allocate(self, allocate: mir.Allocate) -> None:
        bad_bytes_value = teal.Int(
            0,
            source_location=allocate.source_location,
        )
        bad_uint_value = teal.Byte(
            value=b"",
            encoding=AVMBytesEncoding.utf8,
            source_location=allocate.source_location,
        )

        for idx, local_id in enumerate(allocate.allocate_on_entry):
            bad_value = bad_bytes_value if idx < allocate.num_bytes else bad_uint_value
            self._add_op(
                attrs.evolve(
                    bad_value,
                    stack_manipulations=[teal.StackExtend([local_id])],
                )
            )

    def visit_pop(self, pop: mir.Pop) -> None:
        self._add_op(
            teal.PopN(
                n=pop.n,
                stack_manipulations=_lstack_manipulations(pop),
                source_location=pop.source_location,
            )
        )

    def visit_callsub(self, callsub: mir.CallSub) -> None:
        self._add_op(
            teal.CallSub(
                target=callsub.target,
                consumes=callsub.parameters,
                produces=callsub.returns,
                stack_manipulations=_lstack_manipulations(callsub),
                source_location=callsub.source_location,
            )
        )

    def visit_retsub(self, retsub: mir.RetSub) -> None:
        fx_height = retsub.fx_height
        if retsub.returns < fx_height:
            # move returns to base of frame in order
            for n in reversed(range(retsub.returns)):
                self._add_op(teal.FrameBury(n, source_location=retsub.source_location))
        else:
            # f-stack + x-stack is smaller than number of returns, so move it out of the way
            n = retsub.returns + fx_height - 1
            for _ in range(fx_height):
                self._add_op(teal.Uncover(n, source_location=retsub.source_location))
        self._add_op(
            teal.RetSub(
                consumes=retsub.returns,
                source_location=retsub.source_location,
            )
        )

    def visit_program_exit(self, op: mir.ProgramExit) -> None:
        self._add_op(
            teal.Return(
                error_message=op.error_message,
                stack_manipulations=_lstack_manipulations(op),
                source_location=op.source_location,
            )
        )

    def visit_err(self, op: mir.Err) -> None:
        self._add_op(
            teal.Err(
                error_message=op.error_message,
                stack_manipulations=_lstack_manipulations(op),
                source_location=op.source_location,
            )
        )

    def visit_goto(self, op: mir.Goto) -> None:
        self._add_op(
            teal.Branch(
                target=op.target,
                error_message=op.error_message,
                stack_manipulations=_lstack_manipulations(op),
                source_location=op.source_location,
            )
        )
        self.label_stack.append(op.target)

    def visit_conditional_branch(self, op: mir.ConditionalBranch) -> None:
        condition_op: type[teal.BranchNonZero | teal.BranchZero]
        if op.nonzero_target == self.next_block_label:
            condition_op = teal.BranchZero
            condition_op_target = op.zero_target
            other_target = op.nonzero_target
        else:
            condition_op = teal.BranchNonZero
            condition_op_target = op.nonzero_target
            other_target = op.zero_target

        self._add_op(
            condition_op(
                target=condition_op_target,
                error_message=op.error_message,
                stack_manipulations=_lstack_manipulations(op),
                source_location=op.source_location,
            )
        )
        self.label_stack.append(condition_op_target)
        self._add_op(
            teal.Branch(
                target=other_target,
                error_message="",
                stack_manipulations=[],
                source_location=op.source_location,
            )
        )
        self.label_stack.append(other_target)

    def visit_switch(self, op: mir.Switch) -> None:
        self._add_op(
            teal.Switch(
                targets=op.switch_targets,
                error_message=op.error_message,
                stack_manipulations=_lstack_manipulations(op),
                source_location=op.source_location,
            )
        )
        self.label_stack.extend(op.switch_targets)
        self._add_op(
            teal.Branch(
                target=op.default_target,
                error_message="",
                stack_manipulations=[],
                source_location=op.source_location,
            )
        )
        self.label_stack.append(op.default_target)

    def visit_match(self, op: mir.Match) -> None:
        self._add_op(
            teal.Match(
                targets=op.match_targets,
                error_message=op.error_message,
                stack_manipulations=_lstack_manipulations(op),
                source_location=op.source_location,
            )
        )
        self.label_stack.extend(op.match_targets)
        self._add_op(
            teal.Branch(
                target=op.default_target,
                error_message="",
                stack_manipulations=[],
                source_location=op.source_location,
            )
        )
        self.label_stack.append(op.default_target)

    def visit_intrinsic(self, intrinsic: mir.IntrinsicOp) -> None:
        self._add_op(
            teal.Intrinsic(
                op_code=intrinsic.op_code,
                immediates=intrinsic.immediates,
                error_message=intrinsic.error_message,
                consumes=intrinsic.consumes,
                produces=len(intrinsic.produces),
                stack_manipulations=_lstack_manipulations(intrinsic),
                source_location=intrinsic.source_location,
            )
        )

def _lstack_manipulations(op: mir.BaseOp) -> list[teal.StackManipulation]:
    result = list[teal.StackManipulation]()
    if op.consumes:
        result.append(teal.StackConsume(op.consumes))
    if op.produces:
        result.append(teal.StackExtend(op.produces))
        result.append(teal.StackDefine(op.produces))
    return result


from puya import log
from puya.context import CompileContext
from puya.mir import models as mir
from puya.teal import models as teal_models
from puya.teal.builder import TealBuilder
from puya.teal.optimize.main import optimize_teal_program

logger = log.get_logger(__name__)

def mir_to_teal(context: CompileContext, program_mir: mir.Program) -> teal_models.TealProgram:
    main = TealBuilder.build_subroutine(program_mir.main)
    subroutines = [TealBuilder.build_subroutine(mir_sub) for mir_sub in program_mir.subroutines]
    teal = teal_models.TealProgram(
        avm_version=program_mir.avm_version,
        main=main,
        subroutines=subroutines,
    )
    optimize_teal_program(context, teal)
    return teal


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

import attrs

from puya.avm import OnCompletionAction, TransactionType
from puya.errors import InternalError
from puya.ir.types_ import AVMBytesEncoding
from puya.ir.utils import format_bytes, format_error_comment
from puya.mir import models as mir
from puya.parse import SourceLocation
from puya.utils import valid_bytes, valid_int64

MAX_NUMBER_CONSTANTS = 256
TEAL_ALIASES = {
    **{e.name: e.value for e in OnCompletionAction},
    **{e.name: e.value for e in TransactionType},
}

@attrs.frozen
class StackConsume:
    n: int

@attrs.frozen
class StackExtend:
    local_ids: Sequence[str]

@attrs.frozen
class StackInsert:
    depth: int
    local_id: str

@attrs.frozen
class StackPop:
    depth: int

@attrs.frozen
class StackDefine:
    local_ids: Sequence[str]

StackManipulation = StackConsume | StackExtend | StackDefine | StackInsert | StackPop

@attrs.frozen(kw_only=True)
class TealOp:
    op_code: str
    consumes: int
    produces: int
    source_location: SourceLocation | None = attrs.field(eq=False)
    comment: str | None = None
    """A comment that is always emitted after the op in TEAL"""
    error_message: str | None = None
    """Error message to display if program fails at this op"""
    stack_manipulations: Sequence[StackManipulation] = attrs.field(
        default=(),
        converter=tuple[StackManipulation, ...],
        eq=False,
    )

    @property
    def immediates(self) -> Sequence[int | str]:
        return ()

    def teal(self) -> str:
        teal_args = [self.op_code, *map(str, self.immediates)]
        if self.comment or self.error_message:
            error_message = (
                format_error_comment(self.op_code, self.error_message)
                if self.error_message
                else ""
            )
            comment_lines = error_message.splitlines()
            comment_lines += (self.comment or "").splitlines()
            comment = "\n//".join(comment_lines)
            teal_args.append(f"// {comment}")
        return " ".join(teal_args)

    @property
    def stack_height_delta(self) -> int:
        return self.produces - self.consumes

@attrs.frozen
class Dup(TealOp):
    op_code: str = attrs.field(default="dup", init=False)
    consumes: int = attrs.field(default=1, init=False)
    produces: int = attrs.field(default=2, init=False)

@attrs.frozen
class Dup2(TealOp):
    op_code: str = attrs.field(default="dup2", init=False)
    consumes: int = attrs.field(default=2, init=False)
    produces: int = attrs.field(default=4, init=False)

@attrs.frozen
class Pop(TealOp):
    op_code: str = attrs.field(default="pop", init=False)
    consumes: int = attrs.field(default=1, init=False)
    produces: int = attrs.field(default=0, init=False)

@attrs.frozen
class TealOpN(TealOp):
    n: int = attrs.field(validator=[attrs.validators.ge(0), attrs.validators.le(255)])

    @property
    def immediates(self) -> Sequence[int | str]:
        return (self.n,)

@attrs.frozen
class TealOpUInt8(TealOpN):
    n: int = attrs.field(validator=[attrs.validators.ge(0), attrs.validators.le(255)])

@attrs.frozen
class TealOpInt8(TealOpN):
    n: int = attrs.field(validator=[attrs.validators.ge(-128), attrs.validators.le(127)])

@attrs.frozen
class Cover(TealOpUInt8):
    op_code: str = attrs.field(default="cover", init=False)
    consumes: int = attrs.field(init=False)
    produces: int = attrs.field(init=False)

    @consumes.default
    def _consumes(self) -> int:
        return self.n + 1

    @produces.default
    def _produces(self) -> int:
        return self.n + 1

@attrs.frozen
class Uncover(TealOpUInt8):
    op_code: str = attrs.field(default="uncover", init=False)
    consumes: int = attrs.field(init=False)
    produces: int = attrs.field(init=False)

    @consumes.default
    def _consumes(self) -> int:
        return self.n + 1

    @produces.default
    def _produces(self) -> int:
        return self.n + 1

@attrs.frozen
class Swap(TealOp):
    op_code: str = attrs.field(default="swap", init=False)
    consumes: int = 2
    produces: int = 2

@attrs.frozen
class Dig(TealOpUInt8):
    op_code: str = attrs.field(default="dig", init=False)
    consumes: int = attrs.field(init=False)
    produces: int = attrs.field(init=False)

    @consumes.default
    def _consumes(self) -> int:
        return self.n + 1

    @produces.default
    def _produces(self) -> int:
        return self.n + 2

@attrs.frozen
class Bury(TealOpUInt8):
    op_code: str = attrs.field(default="bury", init=False)
    consumes: int = attrs.field(init=False)
    produces: int = attrs.field(init=False)

    @consumes.default
    def _consumes(self) -> int:
        return self.n

    @produces.default
    def _produces(self) -> int:
        return self.n - 1

def _valid_uint64(node: TealOp, _attribute: object, value: int) -> None:
    if not valid_int64(value):
        raise InternalError(
            "Invalid UInt64 value",
            node.source_location,
        )

def _valid_bytes(node: TealOp, _attribute: object, value: bytes) -> None:
    if not valid_bytes(value):
        raise InternalError("Invalid Bytes value", node.source_location)

def _valid_ref(node: TealOp, _attribute: object, value: int) -> None:
    if value < 0 or value >= MAX_NUMBER_CONSTANTS:
        raise InternalError(
            "Invalid constant reference",
            node.source_location,
        )

@attrs.frozen
class IntBlock(TealOp):
    op_code: str = attrs.field(default="intcblock", init=False)
    constants: Mapping[int | str, SourceLocation | None]
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(default=0, init=False)

    @property
    def immediates(self) -> Sequence[int | str]:
        return tuple(self.constants)

@attrs.frozen
class IntC(TealOp):
    index: int = attrs.field(validator=_valid_ref)
    op_code: str = attrs.field(init=False)
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(default=1, init=False)

    @op_code.default
    def _op_code(self) -> str:
        if self.index < 4:
            return f"intc_{self.index}"
        else:
            return "intc"

    @property
    def immediates(self) -> Sequence[int]:
        if self.index < 4:
            return ()
        else:
            return (self.index,)

@attrs.frozen
class PushInt(TealOp):
    op_code: str = attrs.field(default="pushint", init=False)
    value: int = attrs.field(validator=_valid_uint64)
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(default=1, init=False)

    @property
    def immediates(self) -> Sequence[int]:
        return (self.value,)

@attrs.frozen
class PushInts(TealOp):
    op_code: str = attrs.field(default="pushints", init=False)
    values: list[int] = attrs.field(validator=attrs.validators.deep_iterable(_valid_uint64))
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(init=False)

    @produces.default
    def _produces(self) -> int:
        return len(self.values)

    @property
    def immediates(self) -> Sequence[int]:
        return self.values

@attrs.frozen
class BytesBlock(TealOp):
    op_code: str = attrs.field(default="bytecblock", init=False)
    constants: Mapping[bytes | str, tuple[AVMBytesEncoding, SourceLocation | None]]
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(default=0, init=False)

    @property
    def immediates(self) -> Sequence[str]:
        return tuple(
            _encoded_bytes(c, es[0]) if isinstance(c, bytes) else c
            for c, es in self.constants.items()
        )

@attrs.frozen
class BytesC(TealOp):
    index: int = attrs.field(validator=_valid_ref)
    op_code: str = attrs.field(init=False)
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(default=1, init=False)

    @op_code.default
    def _op_code(self) -> str:
        if self.index < 4:
            return f"bytec_{self.index}"
        else:
            return "bytec"

    @property
    def immediates(self) -> Sequence[int]:
        if self.index < 4:
            return ()
        else:
            return (self.index,)

@attrs.frozen
class PushBytes(TealOp):
    op_code: str = attrs.field(default="pushbytes", init=False)
    value: bytes = attrs.field(validator=_valid_bytes)
    # exclude encoding from equality so for example 0x and "" can be combined
    encoding: AVMBytesEncoding = attrs.field(eq=False)
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(default=1, init=False)

    @property
    def immediates(self) -> Sequence[str]:
        return (_encoded_bytes(self.value, self.encoding),)

@attrs.frozen
class PushBytess(TealOp):
    op_code: str = attrs.field(default="pushbytess", init=False)
    values: Sequence[tuple[bytes, AVMBytesEncoding]] = attrs.field()
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(init=False)

    @produces.default
    def _produces(self) -> int:
        return len(self.values)

    @values.validator
    def _values_validator(
        self, _: object, value: Sequence[tuple[bytes, AVMBytesEncoding]]
    ) -> None:
        if not all(valid_bytes(b) for b, _ in value):
            raise InternalError("invalid bytes value", self.source_location)

    @property
    def immediates(self) -> Sequence[str]:
        return tuple(_encoded_bytes(c, e) for c, e in self.values)

@attrs.frozen
class FrameDig(TealOpInt8):
    op_code: str = attrs.field(default="frame_dig", init=False)
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(default=1, init=False)

@attrs.frozen
class FrameBury(TealOpInt8):
    op_code: str = attrs.field(default="frame_bury", init=False)
    consumes: int = attrs.field(default=1, init=False)
    produces: int = attrs.field(default=0, init=False)

@attrs.frozen
class Int(TealOp):
    value: int | str
    op_code: str = attrs.field(default="int", init=False)
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(default=1, init=False)

    @property
    def immediates(self) -> Sequence[int | str]:
        return (self.value,)

@attrs.frozen
class PopN(TealOpUInt8):
    op_code: str = attrs.field(default="popn", init=False)
    consumes: int = attrs.field(init=False)
    produces: int = attrs.field(default=0, init=False)

    @consumes.default
    def _consumes(self) -> int:
        return self.n

@attrs.frozen
class DupN(TealOpUInt8):
    op_code: str = attrs.field(default="dupn", init=False)
    consumes: int = attrs.field(default=1, init=False)
    produces: int = attrs.field(init=False)

    @produces.default
    def _produces(self) -> int:
        return self.n + 1

@attrs.frozen
class Proto(TealOp):
    parameters: int
    returns: int
    op_code: str = attrs.field(default="proto", init=False)
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(default=0, init=False)

    @property
    def immediates(self) -> Sequence[int | str]:
        return self.parameters, self.returns

@attrs.frozen
class Byte(TealOp):
    value: bytes
    # exclude encoding from equality so for example 0x and "" can be combined
    encoding: AVMBytesEncoding = attrs.field(eq=False)
    op_code: str = attrs.field(default="byte", init=False)
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(default=1, init=False)

    @property
    def immediates(self) -> Sequence[int | str]:
        return (_encoded_bytes(self.value, self.encoding),)

@attrs.frozen
class TemplateVar(TealOp):
    name: str
    op_code: typing.Literal["int", "byte"]
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(default=1, init=False)

    @property
    def immediates(self) -> Sequence[int | str]:
        return (self.name,)

@attrs.frozen
class Address(TealOp):
    value: str
    op_code: str = attrs.field(default="addr", init=False)
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(default=1, init=False)

    @property
    def immediates(self) -> Sequence[int | str]:
        return (self.value,)

@attrs.frozen
class Method(TealOp):
    value: str
    op_code: str = attrs.field(default="method", init=False)
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(default=1, init=False)

    @property
    def immediates(self) -> Sequence[int | str]:
        return (f'"{self.value}"',)

@attrs.frozen
class Intrinsic(TealOp):
    immediates: Sequence[int | str]

@attrs.frozen
class ControlOp(TealOp, abc.ABC):
    @property
    @abc.abstractmethod
    def targets(self) -> Sequence[str]: ...

@attrs.frozen
class RetSub(ControlOp):
    op_code: str = attrs.field(default="retsub", init=False)
    produces: int = attrs.field(default=0, init=False)

    @property
    def targets(self) -> Sequence[str]:
        return ()

@attrs.frozen
class Return(ControlOp):
    op_code: str = attrs.field(default="return", init=False)
    consumes: int = attrs.field(default=1, init=False)
    produces: int = attrs.field(default=0, init=False)

    @property
    def targets(self) -> Sequence[str]:
        return ()

@attrs.frozen
class Err(ControlOp):
    op_code: str = attrs.field(default="err", init=False)
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(default=0, init=False)

    @property
    def targets(self) -> Sequence[str]:
        return ()

@attrs.frozen
class Branch(ControlOp):
    target: str
    op_code: str = attrs.field(default="b", init=False)
    consumes: int = attrs.field(default=0, init=False)
    produces: int = attrs.field(default=0, init=False)

    @property
    def immediates(self) -> Sequence[int | str]:
        return (self.target,)

    @property
    def targets(self) -> Sequence[str]:
        return (self.target,)

@attrs.frozen
class BranchNonZero(ControlOp):
    target: str
    op_code: str = attrs.field(default="bnz", init=False)
    consumes: int = attrs.field(default=1, init=False)
    produces: int = attrs.field(default=0, init=False)

    @property
    def immediates(self) -> Sequence[int | str]:
        return (self.target,)

    @property
    def targets(self) -> Sequence[str]:
        return (self.target,)

@attrs.frozen
class BranchZero(ControlOp):
    target: str
    op_code: str = attrs.field(default="bz", init=False)
    consumes: int = attrs.field(default=1, init=False)
    produces: int = attrs.field(default=0, init=False)

    @property
    def immediates(self) -> Sequence[int | str]:
        return (self.target,)

    @property
    def targets(self) -> Sequence[str]:
        return (self.target,)

@attrs.frozen
class Switch(ControlOp):
    targets: Sequence[str] = attrs.field(converter=tuple[str, ...])
    op_code: str = attrs.field(default="switch", init=False)
    consumes: int = attrs.field(default=1, init=False)
    produces: int = attrs.field(default=0, init=False)

    @property
    def immediates(self) -> Sequence[int | str]:
        return self.targets

@attrs.frozen
class Match(ControlOp):
    targets: Sequence[str] = attrs.field(converter=tuple[str, ...])
    op_code: str = attrs.field(default="match", init=False)
    consumes: int = attrs.field(init=False)
    produces: int = attrs.field(default=0, init=False)

    @consumes.default
    def _consumes(self) -> int:
        return len(self.targets) + 1

    @property
    def immediates(self) -> Sequence[int | str]:
        return self.targets

@attrs.frozen
class CallSub(TealOp):
    target: str
    op_code: str = attrs.field(default="callsub", init=False)

    @property
    def immediates(self) -> Sequence[int | str]:
        return (self.target,)

@attrs.frozen
class TealBlock:
    label: str
    ops: list[TealOp]
    x_stack_in: Sequence[str]
    entry_stack_height: int = attrs.field(validator=attrs.validators.ge(0))
    exit_stack_height: int = attrs.field(validator=attrs.validators.ge(0))

    def validate_stack_height(self) -> None:
        stack_height = self.entry_stack_height
        for op in self.ops:
            stack_height -= op.consumes
            if stack_height < 0:
                raise InternalError("Access below stack height", op.source_location)
            stack_height += op.produces
        expected_exit_height = self.exit_stack_height
        if stack_height != expected_exit_height and not (
            self.ops and self.ops[-1].op_code in ("return", "retsub", "err")
        ):
            raise InternalError(
                f"Stack size at block {self.label} exit is {stack_height},"
                f" expected {expected_exit_height}",
                self.ops[-1].source_location,
            )

@attrs.frozen
class TealSubroutine:
    is_main: bool
    signature: mir.Signature
    blocks: list[TealBlock]
    source_location: SourceLocation | None

@attrs.define
class TealProgram:
    avm_version: int
    main: TealSubroutine
    subroutines: list[TealSubroutine]

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

def _encoded_bytes(value: bytes, encoding: AVMBytesEncoding) -> str:
    # not all encodings can handle an empty bytes, so use base16 if bytes is empty
    if not value and encoding in (AVMBytesEncoding.base32, AVMBytesEncoding.base64):
        encoding = AVMBytesEncoding.base16
    bytes_str = format_bytes(value, encoding)
    if encoding in (
        AVMBytesEncoding.utf8,
        AVMBytesEncoding.base16,
        AVMBytesEncoding.unknown,
    ):
        return bytes_str
    hint = encoding.name
    return f"{hint}({bytes_str})"


LOAD_OP_CODES = frozenset(
    [
        "addr",
        "arg",
        *(f"arg_{i}" for i in range(4)),
        "byte",
        *(f"bytec_{i}" for i in range(4)),  # valid as long as we don't push through a bytecblock
        "gaid",
        "gload",
        # global OpcodeBudget could be affected by shuffling,
        # but we don't guarantee ordering of stack manipulations
        "global",
        "gtxn",
        "gtxna",
        "int",
        *(f"intc_{i}" for i in range(4)),  # valid as long as we don't push through an intcblock
        "load",  # valid as long as we don't push through a store/stores
        "method",
        "pushbytes",
        "pushint",
        "txn",
        "txna",
        # below are valid as long as we don't push through an itxn_submit
        "gitxn",
        "gitxna",
        "itxn",
        "itxna",
    ]
)
LOAD_OP_CODES_INCL_OFFSET = frozenset(
    [
        *LOAD_OP_CODES,
        "dig",
        "frame_dig",
    ]
)
STORE_OPS_INCL_OFFSET = frozenset(
    [
        "app_global_del",
        # "bury", TODO: this one is very hard, and doesn't show up in any of our examples
        # "assert", TODO: enable this but only for O2 or higher
        "frame_bury",
        "itxn_field",
        "pop",
        "store",
    ]
)

COMMUTATIVE_OPS = frozenset(
    [
        "+",
        "*",
        "&",
        "&&",
        "|",
        "||",
        "^",
        "==",
        "!=",
        "b*",
        "b+",
        "b&",
        "b|",
        "b^",
        "b==",
        "b!=",
        "addw",
        "mulw",
    ]
)

ORDERING_OPS = frozenset(["<", "<=", ">", ">=", "b<", "b<=", "b>", "b>="])


import itertools
import typing
from collections.abc import Sequence

from puya.parse import sequential_source_locations_merge
from puya.teal import models

def combine_pushes(program: models.TealProgram) -> None:
    for block in itertools.chain.from_iterable(sub.blocks for sub in program.all_subroutines):
        pushes = list[models.PushInt | models.PushBytes]()
        result = list[models.TealOp]()
        for op in block.ops:
            if _is_different_push_type(pushes, op):
                result.append(_combine_ops(pushes))
                pushes = []
            if isinstance(op, models.PushInt | models.PushBytes):
                pushes.append(op)
            else:
                result.append(op)
        if pushes:
            result.append(_combine_ops(pushes))
        block.ops[:] = result

def _is_different_push_type(
    consecutive: list[models.PushInt | models.PushBytes], next_op: models.TealOp
) -> bool:
    return bool(consecutive) and type(consecutive[-1]) is not type(next_op)

def _combine_ops(consecutive: Sequence[models.PushInt | models.PushBytes]) -> models.TealOp:
    if len(consecutive) == 1:
        return consecutive[0]
    loc = sequential_source_locations_merge(op.source_location for op in consecutive)
    stack_manipulations = list(
        itertools.chain.from_iterable(op.stack_manipulations for op in consecutive)
    )
    if isinstance(consecutive[0], models.PushInt):
        consecutive = typing.cast(Sequence[models.PushInt], consecutive)
        return models.PushInts(
            values=[v.value for v in consecutive],
            stack_manipulations=stack_manipulations,
            source_location=loc,
            comment=_comment_ops(consecutive),
        )
    else:
        consecutive = typing.cast(Sequence[models.PushBytes], consecutive)
        return models.PushBytess(
            values=[(v.value, v.encoding) for v in consecutive],
            stack_manipulations=stack_manipulations,
            source_location=loc,
            comment=_comment_ops(consecutive),
        )

def _comment_ops(consecutive: Sequence[models.PushInt | models.PushBytes]) -> str:
    return ", ".join(
        map(str, ((v.comment or " ".join(map(str, v.immediates))) for v in consecutive))
    )


import itertools
import typing
from collections import Counter

from puya.errors import CodeError, InternalError
from puya.ir.types_ import AVMBytesEncoding
from puya.parse import SourceLocation
from puya.teal import models
from puya.utils import Address, coalesce, method_selector_hash, unique

_T = typing.TypeVar("_T")

def gather_program_constants(program: models.TealProgram) -> None:
    # collect constants & template vars
    all_ints = list[int | str]()
    all_bytes = list[bytes | str]()
    bytes_encodings = dict[bytes | str, AVMBytesEncoding]()
    tmpl_locs = dict[bytes | int | str, SourceLocation | None]()

    # collect constants
    for block in itertools.chain.from_iterable(sub.blocks for sub in program.all_subroutines):
        for idx, op in enumerate(block.ops):
            # replace Method & Address constants with Byte before gathering
            match op:
                case models.Method(value=method_value):
                    op = block.ops[idx] = models.Byte(
                        value=method_selector_hash(method_value),
                        encoding=AVMBytesEncoding.base16,
                        comment=op.teal(),
                        stack_manipulations=op.stack_manipulations,
                        source_location=op.source_location,
                    )
                case models.Address(value=address_value, source_location=loc):
                    address = Address.parse(address_value)
                    if not address.is_valid:
                        raise InternalError(f"Invalid address literal: {address_value}", loc)
                    op = block.ops[idx] = models.Byte(
                        value=address.public_key,
                        encoding=AVMBytesEncoding.base32,
                        comment=op.teal(),
                        stack_manipulations=op.stack_manipulations,
                        source_location=op.source_location,
                    )
            match op:
                case models.Int(value=int_or_alias):
                    all_ints.append(_resolve_teal_alias(int_or_alias))
                case models.Byte(value=bytes_value) as byte:
                    all_bytes.append(bytes_value)
                    # preserve bytes encoding if it matches
                    if bytes_encodings.setdefault(bytes_value, byte.encoding) != byte.encoding:
                        bytes_encodings[bytes_value] = AVMBytesEncoding.base16
                # put template vars in constant blocks regardless of optimization level
                case models.TemplateVar(name=name, op_code=op_code):
                    # capture first defined tmpl location
                    tmpl_locs[name] = tmpl_locs.get(name) or op.source_location
                    match op_code:
                        case "int":
                            all_ints.append(name)
                        case "byte":
                            all_bytes.append(name)
                            bytes_encodings[name] = AVMBytesEncoding.base16
                        case _:
                            typing.assert_never(op_code)

    all_templates = unique(val for val in (*all_ints, *all_bytes) if isinstance(val, str))
    int_block = _sort_and_filter_constants(all_ints)
    byte_block = _sort_and_filter_constants(all_bytes)

    assert all(
        t in int_block or t in byte_block for t in all_templates
    ), "expected all template variables to be in constant block"
    # insert constant blocks
    entry_block = program.main.blocks[0]
    if byte_block:
        entry_block.ops.insert(
            0,
            models.BytesBlock(
                constants={b: (bytes_encodings[b], tmpl_locs.get(b)) for b in byte_block},
                source_location=None,
            ),
        )
    if int_block:
        entry_block.ops.insert(
            0,
            models.IntBlock(
                constants={i: tmpl_locs.get(i) for i in int_block}, source_location=None
            ),
        )

    # replace constants and template vars with constant load ops
    for block in itertools.chain.from_iterable(sub.blocks for sub in program.all_subroutines):
        for idx, op in enumerate(block.ops):
            match op:
                case models.Int(value=int_value):
                    comment = coalesce(
                        op.comment,
                        int_value if isinstance(int_value, str) else None,
                        str(int_value),
                    )
                    int_value = _resolve_teal_alias(int_value)
                    try:
                        const_index = int_block[int_value]
                    except KeyError:
                        block.ops[idx] = models.PushInt(
                            value=int_value,
                            stack_manipulations=op.stack_manipulations,
                            source_location=op.source_location,
                            comment=comment,
                        )
                    else:
                        block.ops[idx] = models.IntC(
                            index=const_index,
                            stack_manipulations=op.stack_manipulations,
                            source_location=op.source_location,
                            comment=comment,
                        )
                case models.Byte(value=bytes_value) as byte_op:
                    try:
                        const_index = byte_block[bytes_value]
                    except KeyError:
                        block.ops[idx] = models.PushBytes(
                            value=bytes_value,
                            encoding=byte_op.encoding,
                            stack_manipulations=op.stack_manipulations,
                            source_location=op.source_location,
                            comment=op.comment,
                        )
                    else:
                        block.ops[idx] = models.BytesC(
                            index=const_index,
                            stack_manipulations=op.stack_manipulations,
                            source_location=op.source_location,
                            comment=op.comment or " ".join(map(str, op.immediates)),
                        )
                case models.TemplateVar(name=name, op_code=op_code):
                    match op_code:
                        case "int":
                            block.ops[idx] = models.IntC(
                                index=int_block[name],
                                stack_manipulations=op.stack_manipulations,
                                source_location=op.source_location,
                                comment=op.comment or name,
                            )
                        case "byte":
                            block.ops[idx] = models.BytesC(
                                index=byte_block[name],
                                stack_manipulations=op.stack_manipulations,
                                source_location=op.source_location,
                                comment=op.comment or name,
                            )
                        case _:
                            typing.assert_never(op_code)

def _sort_and_filter_constants(values: list[_T | str]) -> dict[_T | str, int]:
    value_frequencies = {
        value: freq
        for value, freq in Counter(values).most_common()
        if (freq > 1 or isinstance(value, str))
    }
    # filter constants based on their size * frequency when used as a constant vs inline
    freq_idx = 0
    ordered_values = list(value_frequencies)
    while freq_idx < len(ordered_values):
        value = ordered_values[freq_idx]
        if not isinstance(value, str):
            encoded_size = _encoded_size(value)
            const_usage_size = 1 if freq_idx < 4 else 2
            inline_usage_size = 1 + encoded_size
            freq = value_frequencies[value]
            # include constant block size in consideration if only one value
            block_size = 2 if len(ordered_values) == 1 else 0
            inline = inline_usage_size * freq
            const = const_usage_size * freq + encoded_size + block_size
            if inline <= const:
                ordered_values.pop(freq_idx)
                continue
        freq_idx += 1

    # ensure any template variables that sit beyond MAX_NUMBER_CONSTANTS are always included
    overflow_template_vars = [
        value for value in ordered_values[models.MAX_NUMBER_CONSTANTS :] if isinstance(value, str)
    ]
    if len(overflow_template_vars) > models.MAX_NUMBER_CONSTANTS:
        # this doesn't have to be an error as the only consequence of not having all template vars
        # in the constant block is that it breaks the debug mapping assumptions
        # however for now it is easier to just fail, as this many template variables is unlikely
        raise CodeError(f"cannot exceed {models.MAX_NUMBER_CONSTANTS} template values")
    constants = (
        *ordered_values[: models.MAX_NUMBER_CONSTANTS - len(overflow_template_vars)],
        *overflow_template_vars,
    )
    assert len(constants) <= models.MAX_NUMBER_CONSTANTS, "constant block size exceeded"
    return {value: index for index, value in enumerate(constants)}

def _encoded_size(value: object) -> int:
    # varuint encodes 7 bits per byte
    if isinstance(value, int):
        # max accounts for 0, which still requires at least 1 byte
        return max(1, (value.bit_length() + 6) // 7)
    elif isinstance(value, bytes):
        return _encoded_size(len(value)) + len(value)
    else:
        raise TypeError(f"unencodable type: {value!r}")

def _resolve_teal_alias(value: int | str) -> int:
    return models.TEAL_ALIASES[value] if isinstance(value, str) else value


import attrs

from puya.parse import sequential_source_locations_merge
from puya.teal import models
from puya.teal._util import preserve_stack_manipulations
from puya.teal.optimize._data import LOAD_OP_CODES

def perform_constant_stack_shuffling(block: models.TealBlock) -> bool:
    result = block.ops.copy()
    loads = list[models.TealOp]()
    loads_modified = modified = False
    start_idx = idx = 0
    while idx < len(result):
        op = result[idx]
        if _is_constant_load(op):
            if not loads:
                start_idx = idx
                loads_modified = False
            loads.append(op)
        elif loads and op.op_code in ("dup", "dupn"):
            modified = loads_modified = True
            (n,) = op.immediates or (1,)
            assert isinstance(n, int)
            # extend loads with n copies of the last load
            loads.extend([attrs.evolve(loads[-1], source_location=op.source_location)] * n)
        elif loads:
            match op:
                case models.Uncover(n=n) if n < len(loads):
                    modified = loads_modified = True
                    uncovered = loads.pop(-(n + 1))
                    loads.append(uncovered)
                case models.Cover(n=n) if n < len(loads):
                    modified = loads_modified = True
                    covered = loads.pop()
                    loads.insert(len(loads) - n, covered)
                case _:
                    if loads_modified:
                        window = slice(start_idx, idx)
                        preserve_stack_manipulations(result, window, loads)
                        idx = start_idx + len(loads)
                    loads = []
        idx += 1

    if loads_modified and loads:
        window = slice(start_idx, len(result))
        preserve_stack_manipulations(result, window, loads)
    block.ops[:] = result
    return modified

def constant_dupn_insertion(block: models.TealBlock) -> bool:
    result = list[models.TealOp]()
    loads = list[models.TealOp]()
    modified = False
    for op in block.ops:
        if loads and op == loads[0]:
            loads.append(op)
        else:
            if loads:
                modified = _collapse_loads(loads) or modified
                result.extend(loads)
                loads = []
            if _is_constant_load(op):
                loads.append(op)
            else:
                result.append(op)
    if loads:
        modified = _collapse_loads(loads) or modified
        result.extend(loads)
    block.ops[:] = result
    return modified

def constant_dup2_insertion(block: models.TealBlock) -> bool:
    result = block.ops
    modified = False
    idx = 0
    while (idx + 4) <= len(block.ops):
        load_a, load_b, load_a2, load_b2 = block.ops[idx : idx + 4]
        if (
            _is_constant_load(load_a)
            and _is_constant_load(load_b)
            and (load_a, load_b) == (load_a2, load_b2)
        ):
            loc = sequential_source_locations_merge(
                (load_a2.source_location, load_b2.source_location)
            )
            dup2 = models.Dup2(source_location=loc)
            preserve_stack_manipulations(result, slice(idx + 2, idx + 4), [dup2])
            modified = True
            idx += 3
        else:
            idx += 1
    block.ops[:] = result
    return modified

def _collapse_loads(loads: list[models.TealOp]) -> bool:
    n = len(loads) - 1
    if n < 1:
        return False

    dup_source_location = sequential_source_locations_merge(op.source_location for op in loads[1:])
    if n == 1:
        dup_op: models.TealOp = models.Dup(source_location=dup_source_location)
    else:
        dup_op = models.DupN(n=n, source_location=dup_source_location)

    preserve_stack_manipulations(loads, slice(1, None), [dup_op])
    return True

def _is_constant_load(op: models.TealOp) -> bool:
    return op.op_code in LOAD_OP_CODES or (isinstance(op, models.FrameDig) and op.n < 0)


import itertools
from collections import defaultdict

import attrs

from puya import log
from puya.context import CompileContext
from puya.teal import models
from puya.teal._util import preserve_stack_manipulations
from puya.teal.optimize.combine_pushes import combine_pushes
from puya.teal.optimize.constant_block import gather_program_constants
from puya.teal.optimize.constant_stack_shuffling import (
    constant_dup2_insertion,
    constant_dupn_insertion,
    perform_constant_stack_shuffling,
)
from puya.teal.optimize.peephole import peephole
from puya.teal.optimize.repeated_rotations import simplify_repeated_rotation_ops, simplify_swap_ops
from puya.teal.optimize.repeated_rotations_search import repeated_rotation_ops_search

logger = log.get_logger(__name__)

def optimize_teal_program(context: CompileContext, teal_program: models.TealProgram) -> None:
    for teal_sub in teal_program.all_subroutines:
        _optimize_subroutine(context, teal_sub)
    before = [_get_all_stack_manipulations(sub) for sub in teal_program.subroutines]
    gather_program_constants(teal_program)
    if context.options.optimization_level > 0:
        combine_pushes(teal_program)
    after = [_get_all_stack_manipulations(sub) for sub in teal_program.subroutines]
    assert before == after, "expected stack manipulations to be preserved after optimization"

def _get_all_stack_manipulations(sub: models.TealSubroutine) -> list[models.StackManipulation]:
    return [sm for block in sub.blocks for op in block.ops for sm in op.stack_manipulations]

def _optimize_subroutine(context: CompileContext, teal_sub: models.TealSubroutine) -> None:
    logger.debug(
        f"optimizing TEAL subroutine {teal_sub.signature}", location=teal_sub.source_location
    )
    before = _get_all_stack_manipulations(teal_sub)
    for teal_block in teal_sub.blocks:
        _optimize_block(teal_block, level=context.options.optimization_level)
        teal_block.validate_stack_height()
    after = _get_all_stack_manipulations(teal_sub)
    assert before == after, "expected stack manipulations to be preserved after optimization"
    if context.options.optimization_level > 0:
        # at this point, blocks should still be almost "basic"
        # - control flow should only enter at the start of a block.
        # - control flow should only leave at the end of the block (although this might be
        #   spread over multiple ops in the case of say a switch or match)
        # - the final op must be an unconditional control flow op (e.g. retusb, b, return, err)
        _inline_jump_chains(teal_sub)
        # now all blocks that are just a b should have been inlined/removed.
        # any single-op blocks left must be non-branching control ops, ie ops
        # that terminate either exit the program or the subroutine.
        # inlining these are only possible when they are unconditionally branched to,
        # thus this still maintains the "almost basic" structure as outlined above.
        _inline_single_op_blocks(teal_sub)
        _inline_singly_referenced_blocks(teal_sub)
    _remove_jump_fallthroughs(teal_sub)

def _optimize_block(block: models.TealBlock, *, level: int) -> None:
    modified = True
    while modified:
        modified = False
        if level > 0:
            modified = perform_constant_stack_shuffling(block) or modified
            modified = simplify_repeated_rotation_ops(block) or modified
        modified = peephole(block, level) or modified

    # we don't do dup/dupn collapse in the above loop, but after it.
    # it's easier to deal with expanded dup/dupn instructions above when looking at
    # stack shuffling etc, but once it's done we save ops / program size by collapsing them
    constant_dupn_insertion(block)
    constant_dup2_insertion(block)
    if level >= 2:
        # this is a brute-force search which can be slow at times,
        # so it's only done once and only at higher optimisation levels
        block.ops[:] = repeated_rotation_ops_search(block.ops)

    # simplifying uncover/cover 1 to swap is easier to do after other rotation optimizations
    simplify_swap_ops(block)

def _inline_jump_chains(teal_sub: models.TealSubroutine) -> None:
    # build a map of any blocks that are just an unconditional branch to their targets
    jumps = dict[str, str]()
    for block_idx, block in enumerate(teal_sub.blocks.copy()):
        match block.ops:
            case [models.Branch(target=target_label) as b]:
                if b.stack_manipulations:
                    logger.debug(
                        "not inlining jump-block due to stack manipulations",
                        location=b.source_location,
                    )
                else:
                    jumps[block.label] = target_label
                    logger.debug(f"removing jump-chain block {block.label}")
                    teal_sub.blocks.pop(block_idx)
    # now back-propagate any chains
    replacements = dict[str, str]()
    for src, target in jumps.items():
        while True:
            try:
                target = jumps[target]
            except KeyError:
                break
        replacements[src] = target
        logger.debug(f"branching to {src} will be replaced with {target}")
    for block in teal_sub.blocks:
        for op_idx, op in enumerate(block.ops):
            if isinstance(op, models.Branch | models.BranchNonZero | models.BranchZero):
                if op.target in replacements:
                    block.ops[op_idx] = attrs.evolve(op, target=replacements[op.target])
            elif isinstance(op, models.Switch | models.Match):
                block.ops[op_idx] = attrs.evolve(
                    op, targets=[replacements.get(t, t) for t in op.targets]
                )

def _inline_single_op_blocks(teal_sub: models.TealSubroutine) -> None:
    # TODO: this should only encounter exiting ops, so we don't need a traversal to find unused,
    #       just keep track of predecessors??
    single_op_blocks = {b.label: b.ops for b in teal_sub.blocks if len(b.ops) == 1}
    modified = False
    for teal_block, next_block in itertools.zip_longest(teal_sub.blocks, teal_sub.blocks[1:]):
        match teal_block.ops[-1]:
            case models.Branch(target=target_label) as branch_op if (
                (replace_ops := single_op_blocks.get(target_label))
                and (next_block is None or target_label != next_block.label)
            ):
                modified = True
                (replace_op,) = replace_ops
                # we shouldn't encounter any branching ops, since any block that
                # is just an unconditional branch has already been inlined, and
                # at this point blocks should still have an unconditional exit as the final op,
                # which rules out bz/bnz/match/switch, leaving only exiting ops
                # like retsub, return, or err.
                # this also means we can keep track of which blocks to eliminate without having
                # to do a traversal, thus the assertion
                assert isinstance(replace_op, models.ControlOp) and not replace_op.targets
                logger.debug(
                    f"replacing `{branch_op.teal()}` with `{replace_op.teal()}`",
                    location=branch_op.source_location,
                )
                preserve_stack_manipulations(teal_block.ops, slice(-1, None), replace_ops)
    if modified:
        # if any were inlined, they may no longer be referenced and thus removable
        _remove_unreachable_blocks(teal_sub)

def _remove_unreachable_blocks(teal_sub: models.TealSubroutine) -> None:
    entry = teal_sub.blocks[0]
    to_visit = [entry]
    unreachable_blocks_by_label = {b.label: b for b in teal_sub.blocks[1:]}
    while to_visit:
        current_block = to_visit.pop()
        for op in current_block.ops:
            if isinstance(op, models.ControlOp):
                to_visit.extend(
                    target
                    for target_label in op.targets
                    if (target := unreachable_blocks_by_label.pop(target_label, None))
                )
    if unreachable_blocks_by_label:
        logger.debug(f"removing unreachable blocks: {list(map(str, unreachable_blocks_by_label))}")
        teal_sub.blocks[:] = [
            b for b in teal_sub.blocks if b.label not in unreachable_blocks_by_label
        ]

def _inline_singly_referenced_blocks(teal_sub: models.TealSubroutine) -> None:
    predecessors = defaultdict[str, list[str]](list)
    predecessors[teal_sub.blocks[0].label].append("<entrypoint>")
    for block in teal_sub.blocks:
        for op in block.ops:
            if isinstance(op, models.ControlOp):
                for target_label in op.targets:
                    predecessors[target_label].append(block.label)

    pairs = dict[str, str]()
    inlineable = set[str]()
    for block in teal_sub.blocks:
        match block.ops[-1]:
            case models.Branch(target=target_label):
                target_predecessors = predecessors[target_label]
                if len(target_predecessors) == 1:
                    assert target_predecessors == [block.label]
                    pairs[block.label] = target_label
                    inlineable.add(target_label)

    blocks_by_label = {b.label: b for b in teal_sub.blocks}
    result = list[models.TealBlock]()
    for block in teal_sub.blocks:
        this_label = block.label
        if this_label not in inlineable:
            result.append(block)
            while (next_label := pairs.get(this_label)) is not None:
                logger.debug(f"inlining single reference block {next_label} into {block.label}")
                next_block = blocks_by_label[next_label]
                preserve_stack_manipulations(block.ops, slice(-1, None), next_block.ops)
                this_label = next_label
    teal_sub.blocks[:] = result

def _remove_jump_fallthroughs(teal_sub: models.TealSubroutine) -> None:
    for block, next_block in zip(teal_sub.blocks, teal_sub.blocks[1:], strict=False):
        match block.ops[-1]:
            # we guard against having stack manipulations but only one op, thus nowhere to put
            # them, but this should only occur in O0 as in higher levels, blocks with just a b
            # will already be inlined
            case models.Branch(
                target=target_label, stack_manipulations=sm
            ) if target_label == next_block.label and (len(block.ops) > 1 or not sm):
                logger.debug(f"removing explicit jump to fall-through block {next_block.label}")
                block.ops.pop()
                if block.ops:  # guard is only for O0 case
                    block.ops[-1] = attrs.evolve(
                        block.ops[-1],
                        stack_manipulations=(*block.ops[-1].stack_manipulations, *sm),
                    )