#!/usr/bin/env python3
import subprocess
import sys
import typing
from collections.abc import Callable
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent / "src" / "puyapy" / "_vendor"))

import attrs
import mypy.build
import mypy.nodes
from mypy.visitor import NodeVisitor

from puyapy.compile import get_mypy_options
from puyapy.parse import parse_and_typecheck

SCRIPTS_DIR = Path(__file__).parent
VCS_ROOT = SCRIPTS_DIR.parent
SRC_DIR = VCS_ROOT / "src"
DOCS_DIR = VCS_ROOT / "docs"
STUBS_DIR = VCS_ROOT / "stubs" / "algopy-stubs"
STUBS_DOC_DIR = DOCS_DIR / "algopy-stubs"

@attrs.define
class ModuleImports:
    from_imports: dict[str, str | None] = attrs.field(factory=dict)
    import_all: bool = False
    import_module: bool = False

def main() -> None:
    manager, _ = parse_and_typecheck([STUBS_DIR], get_mypy_options())
    output_doc_stubs(manager)
    run_sphinx()

def output_doc_stubs(manager: mypy.build.BuildManager) -> None:
    # parse and output reformatted __init__.pyi
    stub = DocStub.process_module(manager, "algopy")
    algopy_direct_imports = stub.collected_imports["algopy"]
    # remove any algopy imports that are now defined in __init__.py itself
    output_combined_stub(stub, STUBS_DOC_DIR / "__init__.pyi")

    # remaining imports from algopy are other public modules
    # parse and output them too
    for other_stub_name in algopy_direct_imports.from_imports:
        stub = DocStub.process_module(manager, f"algopy.{other_stub_name}")
        output_combined_stub(stub, STUBS_DOC_DIR / f"{other_stub_name}.pyi")

def output_combined_stub(stubs: "DocStub", output: Path) -> None:
    # remove algopy imports that have been inlined
    lines = ["# ruff: noqa: A001, E501, F403, PYI021, PYI034, W291"]
    rexported = list[str]()
    for module, imports in stubs.collected_imports.items():
        if imports.import_module:
            lines.append(f"import {module}")
        if imports.from_imports:
            rexported.extend(filter(None, imports.from_imports.values()))
            from_imports = ", ".join(_name_as(k, v) for k, v in imports.from_imports.items())
            lines.append(f"from {module} import {from_imports}")
    lines.extend(["", ""])

    # assemble __all__
    lines.append("__all__ = [")
    for symbol in (*rexported, *stubs.collected_symbols):
        if symbol.startswith("_"):
            continue
        lines.append(f'    "{symbol}",')
    lines.append("]")

    # assemble symbols
    lines.extend(stubs.collected_symbols.values())

    # output and linting
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text("\n".join(lines))

    subprocess.run(["ruff", "format", str(output)], check=True, cwd=VCS_ROOT)
    subprocess.run(["ruff", "check", "--fix", str(output)], check=True, cwd=VCS_ROOT)

def run_sphinx() -> None:
    subprocess.run(
        ["sphinx-build", ".", "_build", "-W", "--keep-going", "-n", "-E"], check=True, cwd=DOCS_DIR
    )

@attrs.define(kw_only=True)
class ClassBases:
    klass: mypy.nodes.ClassDef
    bases: list[mypy.nodes.Expression]
    protocol_bases: list[tuple[mypy.nodes.MypyFile, mypy.nodes.ClassDef]]

@attrs.define
class SymbolCollector(NodeVisitor[None]):
    file: mypy.nodes.MypyFile
    read_source: Callable[[str], list[str] | None]
    all_classes: dict[str, tuple[mypy.nodes.MypyFile, mypy.nodes.ClassDef]]
    inlined_protocols: dict[str, set[str]]
    symbols: dict[str, str] = attrs.field(factory=dict)
    last_stmt: mypy.nodes.Statement | None = None

    def get_src(
        self, node: mypy.nodes.Context, *, path: str | None = None, entire_lines: bool = True
    ) -> str:
        columns: tuple[int, int] | None = None
        if node.end_column and not entire_lines:
            columns = (node.column, node.end_column)
        return self.get_src_from_lines(node.line, node.end_line or node.line, path, columns)

    def get_src_from_lines(
        self,
        line: int,
        end_line: int,
        path: str | None = None,
        columns: tuple[int, int] | None = None,
    ) -> str:
        src = self.read_source(path or self.file.path)
        if not src:
            raise Exception("Could not get src")
        lines = src[line - 1 : end_line]
        if columns:
            lines[-1] = lines[-1][: columns[1]]
            lines[0] = lines[0][columns[0] :]
        return "\n".join(lines)

    def visit_mypy_file(self, o: mypy.nodes.MypyFile) -> None:
        for stmt in o.defs:
            stmt.accept(self)
            self.last_stmt = stmt

    def _get_bases(self, klass: mypy.nodes.ClassDef) -> ClassBases:
        bases = list[mypy.nodes.Expression]()
        inline = list[tuple[mypy.nodes.MypyFile, mypy.nodes.ClassDef]]()
        for base in klass.base_type_exprs:
            if (
                isinstance(base, mypy.nodes.NameExpr)
                and _should_inline_module(base.fullname)
                and self._is_protocol(base.fullname)
            ):
                inline.append(self.all_classes[base.fullname])
            else:
                bases.append(base)
        return ClassBases(klass=klass, bases=bases, protocol_bases=inline)

    def _get_inlined_class(self, klass: ClassBases) -> str:
        # TODO: what about class keywords
        klass_str = f"class {klass.klass.name}"
        if klass.bases:
            klass_str += f"({', '.join(self.get_src(b, entire_lines=False) for b in klass.bases)})"
        src = [f"{klass_str}:"]
        src.extend(self.get_src(member) for member in klass.klass.defs.body)
        for base_class_file, base_class in klass.protocol_bases:
            self.inlined_protocols.setdefault(base_class_file.fullname, set()).add(base_class.name)
            src.extend(
                self.get_src(member, path=base_class_file.path) for member in base_class.defs.body
            )
        return "\n".join(src)

    def visit_class_def(self, o: mypy.nodes.ClassDef) -> None:
        self.all_classes[o.fullname] = self.file, o
        class_bases = self._get_bases(o)
        if class_bases.protocol_bases:
            self.symbols[o.name] = self._get_inlined_class(class_bases)
        else:
            self.symbols[o.name] = self.get_src(o)

    def visit_func_def(self, o: mypy.nodes.FuncDef) -> None:
        self.symbols[o.name] = self.get_src(o)

    def visit_overloaded_func_def(self, o: mypy.nodes.OverloadedFuncDef) -> None:
        line = o.line
        end_line = o.end_line or o.line
        for item in o.items:
            end_line = max(end_line, item.end_line or item.line)
        overloaded_src = self.get_src_from_lines(line, end_line)
        best_sig = _get_documented_overload(o)

        if not best_sig:
            src = overloaded_src
        else:
            best_sig_src = self.get_src(best_sig)
            src = f"{overloaded_src}\n{best_sig_src}"

        self.symbols[o.name] = src

    def visit_assignment_stmt(self, o: mypy.nodes.AssignmentStmt) -> None:
        try:
            (lvalue,) = o.lvalues
        except ValueError as ex:
            raise ValueError(f"Multi assignments are not supported: {o}") from ex
        if not isinstance(lvalue, mypy.nodes.NameExpr):
            raise TypeError(f"Multi assignments are not supported: {lvalue}")
        # find actual rvalue src location by taking the entire statement and subtracting the lvalue
        loc = mypy.nodes.Context()
        loc.set_line(o)
        if lvalue.end_column:
            loc.column = lvalue.end_column
        self.symbols[lvalue.name] = self.get_src(loc)

    def visit_expression_stmt(self, o: mypy.nodes.ExpressionStmt) -> None:
        if isinstance(o.expr, mypy.nodes.StrExpr) and isinstance(
            self.last_stmt, mypy.nodes.AssignmentStmt
        ):
            (lvalue,) = self.last_stmt.lvalues
            if isinstance(lvalue, mypy.nodes.NameExpr):
                self.symbols[lvalue.name] += "\n" + self.get_src(o.expr)

    def _is_protocol(self, fullname: str) -> bool:
        try:
            klass = self.all_classes[fullname]
        except KeyError:
            return False
        info: mypy.nodes.TypeInfo = klass[1].info
        return info.is_protocol

def _get_documented_overload(o: mypy.nodes.OverloadedFuncDef) -> mypy.nodes.FuncDef | None:
    best_overload: mypy.nodes.FuncDef | None = None
    for overload in o.items:
        match overload:
            case mypy.nodes.Decorator(func=func_def):
                pass
            case mypy.nodes.FuncDef() as func_def:
                pass
            case _:
                raise Exception("Only function overloads supported")

        docstring = func_def.docstring

        # this is good enough until a more complex case arises
        if docstring and (
            not best_overload or len(func_def.arguments) > len(best_overload.arguments)
        ):
            best_overload = func_def
    return best_overload

@attrs.define
class ImportCollector(NodeVisitor[None]):
    collected_imports: dict[str, ModuleImports]

    def get_imports(self, module_id: str) -> ModuleImports:
        try:
            imports = self.collected_imports[module_id]
        except KeyError:
            imports = self.collected_imports[module_id] = ModuleImports()
        return imports

    def visit_mypy_file(self, o: mypy.nodes.MypyFile) -> None:
        for stmt in o.defs:
            stmt.accept(self)

    def visit_import_from(self, o: mypy.nodes.ImportFrom) -> None:
        imports = self.get_imports(o.id)
        for name, name_as in o.names:
            imports.from_imports[name] = name_as

    def visit_import(self, o: mypy.nodes.Import) -> None:
        for name, name_as in o.ids:
            if name != (name_as or name):
                raise Exception("Aliasing symbols in stubs is not supported")

            imports = self.get_imports(name)
            imports.import_module = True

@attrs.define
class DocStub(NodeVisitor[None]):
    read_source: Callable[[str], list[str] | None]
    file: mypy.nodes.MypyFile
    modules: dict[str, mypy.nodes.MypyFile]
    parsed_modules: dict[str, SymbolCollector] = attrs.field(factory=dict)
    all_classes: dict[str, tuple[mypy.nodes.MypyFile, mypy.nodes.ClassDef]] = attrs.field(
        factory=dict
    )
    collected_imports: dict[str, ModuleImports] = attrs.field(factory=dict)
    inlined_protocols: dict[str, set[str]] = attrs.field(factory=dict)
    collected_symbols: dict[str, str] = attrs.field(factory=dict)

    @classmethod
    def process_module(cls, manager: mypy.build.BuildManager, module_id: str) -> typing.Self:
        read_source = manager.errors.read_source
        assert read_source
        modules = manager.modules
        module: mypy.nodes.MypyFile = modules[module_id]
        stub = cls(read_source=read_source, file=module, modules=modules)
        module.accept(stub)
        stub._remove_inlined_symbols()  # noqa: SLF001
        return stub

    def _get_module(self, module_id: str) -> SymbolCollector:
        try:
            return self.parsed_modules[module_id]
        except KeyError:
            file = self.modules[module_id]
            self.parsed_modules[module_id] = collector = SymbolCollector(
                file=file,
                read_source=self.read_source,
                all_classes=self.all_classes,
                inlined_protocols=self.inlined_protocols,
            )
            file.accept(collector)
            self._collect_imports(file)
            return collector

    def _collect_imports(self, o: mypy.nodes.Node) -> None:
        o.accept(ImportCollector(self.collected_imports))
        self._remove_inlined_symbols()

    def _remove_inlined_symbols(self) -> None:
        for module, imports in self.collected_imports.items():
            inlined_protocols = self.inlined_protocols.get(module, ())
            if imports.import_module and module in self.collected_symbols:
                raise Exception(f"Symbol/import collision: {module}")
            for name, name_as in list(imports.from_imports.items()):
                if name in inlined_protocols:
                    print(f"Removed inlined protocol: {name}")
                    del imports.from_imports[name]
                    del self.collected_symbols[name]
                elif name in self.collected_symbols:
                    if name_as is None:
                        del imports.from_imports[name]
                    else:
                        print(f"Symbol/import collision: from {module} import {name} as {name_as}")

    def visit_mypy_file(self, o: mypy.nodes.MypyFile) -> None:
        for stmt in o.defs:
            stmt.accept(self)
        self._add_all_symbols(o.fullname)

    def visit_import_from(self, o: mypy.nodes.ImportFrom) -> None:
        if not _should_inline_module(o.id):
            self._collect_imports(o)
            return
        module = self._get_module(o.id)
        name_mapping = dict(o.names)
        for name in module.symbols:
            try:
                name_as = name_mapping[name]
            except KeyError:
                continue
            if name != (name_as or name):
                raise Exception("Aliasing symbols in stubs is not supported")
            self.add_symbol(module, name)

    def visit_import_all(self, o: mypy.nodes.ImportAll) -> None:
        if _should_inline_module(o.id):
            self._add_all_symbols(o.id)
        else:
            self._collect_imports(o)

    def _add_all_symbols(self, module_id: str) -> None:
        module = self._get_module(module_id)
        for sym in module.symbols:
            self.add_symbol(module, sym)

    def visit_import(self, o: mypy.nodes.Import) -> None:
        self._collect_imports(o)

    def add_symbol(self, module: SymbolCollector, name: str) -> None:
        lines = module.symbols[name]
        existing = self.collected_symbols.get(name)
        if existing is not None and existing != lines:
            raise Exception(f"Duplicate definitions are not supported: {name}\n{lines}")
        self.collected_symbols[name] = lines

def _name_as(name: str, name_as: str | None) -> str:
    if name_as is None:
        return name
    return f"{name} as {name_as}"

def _should_inline_module(module_id: str) -> bool:
    return module_id.startswith("algopy._")

if __name__ == "__main__":
    main()


#!/usr/bin/env python3

import subprocess
from pathlib import Path

SCRIPTS_DIR = Path(__file__).parent
VCS_ROOT = SCRIPTS_DIR.parent
LIB_NAME = "_puya_lib"

def main() -> None:
    # compile puya lib
    # normalize source_location.path
    # save
    subprocess.run(["puyapy", "--output-awst-json", f"src/{LIB_NAME}"], check=True, cwd=VCS_ROOT)
    awst_path = VCS_ROOT / "module.awst.json"
    puya_lib_path = VCS_ROOT / "src" / LIB_NAME
    output_path = VCS_ROOT / "src" / "puya" / "ir" / "_puya_lib.awst.json"
    replace_awst = awst_path.read_text()
    for lib_path in puya_lib_path.glob("*.py"):
        path_as_str = str(lib_path).replace("\\", "\\\\")
        find_str = f'"file": "{path_as_str}",'
        replace_str = '"file": null,'
        replace_awst = replace_awst.replace(find_str, replace_str)
    output_path.write_text(replace_awst)
    awst_path.unlink(missing_ok=True)

if __name__ == "__main__":
    main()


#!/usr/bin/env python3
import builtins
import copy
import json
import keyword
import subprocess
import textwrap
import typing
from collections.abc import Iterable, Iterator, Sequence
from pathlib import Path

import attrs

from puya import log
from puya.algo_constants import SUPPORTED_AVM_VERSIONS
from puyapy.awst_build import pytypes
from puyapy.awst_build.intrinsic_models import FunctionOpMapping, OpMappingWithOverloads
from puyapy.awst_build.utils import snake_case
from scripts.transform_lang_spec import (
    ArgEnum,
    Immediate,
    ImmediateKind,
    LanguageSpec,
    Op,
    StackType,
    StackValue,
)

logger = log.get_logger(__name__)
INDENT = " " * 4
VCS_ROOT = Path(__file__).parent.parent
MIN_SUPPORTED_VERSION = min(SUPPORTED_AVM_VERSIONS)

PYTHON_ENUM_CLASS = {
    "Mimc Configurations": "MiMCConfigurations",
}
PYTYPE_TO_LITERAL: dict[pytypes.PyType, pytypes.LiteralOnlyType | None] = {
    pytypes.BytesType: pytypes.BytesLiteralType,
    pytypes.UInt64Type: pytypes.IntLiteralType,
    pytypes.AccountType: None,  # pytypes.StrLiteralType, # TODO: should we enable this?
    pytypes.BigUIntType: pytypes.IntLiteralType,
    pytypes.BoolType: None,  # already a Python type
    pytypes.ApplicationType: pytypes.IntLiteralType,
    pytypes.AssetType: pytypes.IntLiteralType,
    pytypes.TransactionTypeType: None,
    pytypes.OnCompleteActionType: None,
}
PYTYPE_REPR = {
    value: f"pytypes.{key}"
    for key, value in pytypes.__dict__.items()
    if isinstance(value, pytypes.PyType)
}
STACK_TYPE_MAPPING: dict[StackType, Sequence[pytypes.PyType]] = {
    StackType.address_or_index: [pytypes.AccountType, pytypes.UInt64Type],
    StackType.application: [pytypes.ApplicationType, pytypes.UInt64Type],
    StackType.asset: [pytypes.AssetType, pytypes.UInt64Type],
    StackType.bytes: [pytypes.BytesType],
    StackType.bytes_8: [pytypes.BytesType],
    StackType.bytes_32: [pytypes.BytesType],
    StackType.bytes_33: [pytypes.BytesType],
    StackType.bytes_64: [pytypes.BytesType],
    StackType.bytes_80: [pytypes.BytesType],
    StackType.bytes_1232: [pytypes.BytesType],
    StackType.bytes_1793: [pytypes.BytesType],
    StackType.bool: [pytypes.BoolType, pytypes.UInt64Type],
    StackType.uint64: [pytypes.UInt64Type],
    StackType.any: [pytypes.BytesType, pytypes.UInt64Type],
    StackType.box_name: [pytypes.BytesType],  # TODO: should this be another type..?
    StackType.address: [pytypes.AccountType],
    StackType.bigint: [pytypes.BigUIntType],
    StackType.state_key: [pytypes.BytesType],  # TODO: should this be another type..?
}

BYTES_LITERAL = "bytes"
UINT64_LITERAL = "int"
STUB_NAMESPACE = "op"
ALGORAND_OP_URL = "https://developer.algorand.org/docs/get-details/dapps/avm/teal/opcodes/v10/"

class OpCodeGroup(typing.Protocol):
    def handled_ops(self) -> Iterator[str]: ...

@attrs.define(kw_only=True)
class RenamedOpCode(OpCodeGroup):
    name: str
    stack_aliases: dict[str, list[str]] = attrs.field(factory=dict)
    """ops that are aliases for other ops that take stack values instead of immediates"""
    op: str

    def handled_ops(self) -> Iterator[str]:
        yield self.op
        yield from self.stack_aliases.keys()

@attrs.define(kw_only=True)
class MergedOpCodes(OpCodeGroup):
    name: str
    doc: str
    ops: dict[str, dict[str, list[str]]]

    def handled_ops(self) -> Iterator[str]:
        for op, aliases in self.ops.items():
            yield op
            yield from aliases.keys()

@attrs.define(kw_only=True)
class GroupedOpCodes(OpCodeGroup):
    name: str
    """ops that are aliases for other ops that take stack values instead of immediates"""
    doc: str
    ops: dict[str, str] = attrs.field(factory=dict)
    """ops to include in group, mapped to their new name"""

    def handled_ops(self) -> Iterator[str]:
        yield from self.ops.keys()

GROUPED_OP_CODES = [
    GroupedOpCodes(
        name="AppGlobal",
        doc="Get or modify Global app state",
        ops={
            "app_global_get": "get",
            "app_global_get_ex": "get_ex",
            "app_global_del": "delete",
            "app_global_put": "put",
        },
    ),
    GroupedOpCodes(
        name="Scratch",
        doc="Load or store scratch values",
        ops={"loads": "load", "stores": "store"},
    ),
    GroupedOpCodes(
        name="AppLocal",
        doc="Get or modify Local app state",
        ops={
            "app_local_get": "get",
            "app_local_get_ex": "get_ex",
            "app_local_del": "delete",
            "app_local_put": "put",
        },
    ),
    GroupedOpCodes(
        name="Box",
        doc="Get or modify box state",
        ops={
            "box_create": "create",
            "box_del": "delete",
            "box_extract": "extract",
            "box_get": "get",
            "box_len": "length",
            "box_put": "put",
            "box_replace": "replace",
            "box_resize": "resize",
            "box_splice": "splice",
        },
    ),
    GroupedOpCodes(
        name="EllipticCurve",
        doc="Elliptic Curve functions",
        ops={
            "ec_add": "add",
            "ec_map_to": "map_to",
            "ec_multi_scalar_mul": "scalar_mul_multi",
            "ec_pairing_check": "pairing_check",
            "ec_scalar_mul": "scalar_mul",
            "ec_subgroup_check": "subgroup_check",
        },
    ),
    GroupedOpCodes(
        name="ITxnCreate",
        doc="Create inner transactions",
        ops={
            "itxn_begin": "begin",
            "itxn_next": "next",
            "itxn_submit": "submit",
            "itxn_field": "set",
        },
    ),
]
MERGED_OP_CODES = [
    MergedOpCodes(
        name="Txn",
        doc="Get values for the current executing transaction",
        ops={
            "txn": {},
            "txnas": {
                "txna": ["F", "I"],
            },
        },
    ),
    MergedOpCodes(
        name="GTxn",
        doc="Get values for transactions in the current group",
        ops={
            "gtxns": {
                "gtxn": ["F", "T"],
            },
            # field is immediate, first stack arg is txn index, second stack arg is array index
            "gtxnsas": {
                "gtxnsa": ["F", "A", "I"],  # group index on stack
                "gtxna": ["F", "T", "I"],  # no stack args
                "gtxnas": ["F", "T", "A"],  # array index on stack
            },
        },
    ),
    MergedOpCodes(
        name="ITxn",
        doc="Get values for the last inner transaction",
        ops={
            "itxn": {},
            "itxnas": {
                "itxna": ["F", "I"],
            },
        },
    ),
    MergedOpCodes(
        name="GITxn",
        doc="Get values for inner transaction in the last group submitted",
        ops={
            "gitxn": {},
            "gitxnas": {
                "gitxna": ["T", "F", "I"],
            },
        },
    ),
    MergedOpCodes(
        name="Global",
        doc="Get Global values",
        ops={"global": {}},
    ),
]
RENAMED_OP_CODES = [
    RenamedOpCode(
        name="arg",
        op="args",
        stack_aliases={"arg": ["N"]},
    ),
    RenamedOpCode(
        name="extract",
        op="extract3",
        stack_aliases={
            "extract": ["A", "S", "L"],
        },
    ),
    RenamedOpCode(
        name="replace",
        op="replace3",
        stack_aliases={
            "replace2": ["A", "S", "B"],
        },
    ),
    RenamedOpCode(
        name="substring",
        op="substring3",
        stack_aliases={
            "substring": ["A", "S", "E"],
        },
    ),
    RenamedOpCode(
        name="gload",
        op="gloadss",
        stack_aliases={
            "gload": ["T", "I"],
            "gloads": ["A", "I"],
        },
    ),
    RenamedOpCode(
        name="gaid",
        op="gaids",
        stack_aliases={"gaid": ["T"]},
    ),
    RenamedOpCode(
        name="exit",
        op="return",
    ),
]

EXCLUDED_OPCODES = {
    # low level flow control
    "bnz",
    "bz",
    "b",
    "callsub",
    "retsub",
    "proto",
    "switch",
    "match",
    # low level stack manipulation
    "intc",
    *[f"intc_{i}" for i in range(4)],
    "bytec",
    *[f"bytec_{i}" for i in range(4)],
    "pushbytes",
    "pushbytess",
    "pushint",
    "pushints",
    "frame_dig",
    "frame_bury",
    "bury",
    "cover",
    "dig",
    "dup",
    "dup2",
    "dupn",
    "pop",
    "popn",
    "swap",
    "uncover",
    # program scratch slot read/modification (for current program)
    "load",
    "loads",
    "store",
    "stores",
    # maninuplates what other low level ops point to
    "intcblock",
    "bytecblock",
    # implicit immediates, covered by optimiser and/or assembler
    "arg_0",
    "arg_1",
    "arg_2",
    "arg_3",
    # have a higher level abstraction that supersedes it
    "log",
}

# which ops to treat as properties in the generated stubs
PROPERTY_OPS = {
    "global": {"exclude": ["opcode_budget"]},
    "txn": {"exclude": list[str]()},
}

@attrs.define
class TypedName:
    name: str
    type: StackType | ImmediateKind | str
    doc: str | None

@attrs.define(kw_only=True)
class FunctionDef:
    name: str
    doc: list[str]
    is_property: bool
    args: list[TypedName] = attrs.field(factory=list)
    return_docs: list[str] = attrs.field(factory=list)
    op_mapping: OpMappingWithOverloads
    min_avm_version: int

@attrs.define
class ClassDef:
    name: str
    doc: str
    methods: list[FunctionDef] = attrs.field()
    ops: list[str]

def main() -> None:
    spec_path = VCS_ROOT / "langspec.puya.json"

    lang_spec_json = json.loads(spec_path.read_text(encoding="utf-8"))
    lang_spec = LanguageSpec.from_json(lang_spec_json)

    non_simple_ops = {
        *EXCLUDED_OPCODES,
        *dir(builtins),
        *keyword.kwlist,  # TODO: maybe consider softkwlist too?
    }
    function_defs = list[FunctionDef]()
    class_defs = list[ClassDef]()
    enums_to_build = dict[str, bool]()

    for merged in MERGED_OP_CODES:
        non_simple_ops.update(merged.handled_ops())
        class_defs.append(build_merged_ops(lang_spec, merged))
    for grouped in GROUPED_OP_CODES:
        non_simple_ops.update(grouped.handled_ops())
        class_defs.append(build_grouped_ops(lang_spec, grouped, enums_to_build))
    for aliased in RENAMED_OP_CODES:
        function_defs.extend(build_aliased_ops(lang_spec, aliased))
        non_simple_ops.update(aliased.handled_ops())

    for op in lang_spec.ops.values():
        if op.name in non_simple_ops or not op.name.isidentifier():
            logger.info(f"Ignoring: {op.name}")
            continue
        overriding_immediate = get_overriding_immediate(op)
        if overriding_immediate:
            class_defs.append(
                build_class_from_overriding_immediate(
                    lang_spec,
                    op,
                    class_name=get_python_enum_class(op.name),
                    class_doc=" ".join(op.doc),
                    immediate=overriding_immediate,
                    aliases=[],
                )
            )
        else:
            for immediate in op.immediate_args:
                if immediate.immediate_type == ImmediateKind.arg_enum and (
                    immediate.modifies_stack_input is None
                    and immediate.modifies_stack_output is None
                ):
                    assert immediate.arg_enum is not None
                    enums_to_build[immediate.arg_enum] = True
            function_defs.extend(build_operation_methods(op, op.name, []))

    function_defs.sort(key=lambda x: x.name)
    class_defs.sort(key=lambda x: x.name)

    enum_names = list(enums_to_build.keys())
    output_stub(lang_spec, enum_names, function_defs, class_defs)
    output_awst_data(lang_spec, enum_names, function_defs, class_defs)

def sub_types(type_name: StackType, *, covariant: bool) -> Sequence[pytypes.PyType]:
    try:
        typs = STACK_TYPE_MAPPING[type_name]
    except KeyError as ex:
        raise NotImplementedError(
            f"Could not map stack type {type_name} to an algopy type"
        ) from ex
    else:
        last_index = None if covariant else 1
        return typs[:last_index]

def immediate_kind_to_type(kind: ImmediateKind) -> type[int | str]:
    match kind:
        case ImmediateKind.uint8 | ImmediateKind.int8 | ImmediateKind.varuint:
            return int
        case ImmediateKind.arg_enum:
            return str
        case _:
            raise ValueError(f"Unexpected ImmediateKind: {kind}")

def get_python_type(
    typ: StackType | ImmediateKind | str, *, covariant: bool, any_as: str | None
) -> str:
    match typ:
        case StackType() as stack_type:
            if any_as and stack_type == StackType.any:
                return any_as
            ptypes_ = sub_types(stack_type, covariant=covariant)
            names = [str(wt).removeprefix("algopy.") for wt in ptypes_]
            if covariant:
                for pt in ptypes_:
                    lit_t = PYTYPE_TO_LITERAL[pt]
                    if lit_t is not None:
                        lit_name = str(lit_t)
                        if lit_name not in names:
                            names.append(lit_name)
            return " | ".join(names)
        case ImmediateKind() as immediate_kind:
            return immediate_kind_to_type(immediate_kind).__name__
        case _:
            return typ

def build_method_stub(function: FunctionDef, prefix: str = "") -> Iterable[str]:
    signature = list[str]()
    doc = function.doc[:]
    signature.append(f"def {function.name}(")
    args = list[str]()
    for arg in function.args:
        python_type = get_python_type(arg.type, covariant=True, any_as=None)
        args.append(f"{arg.name}: {python_type}")
        if arg.doc:
            doc.append(f":param {python_type} {arg.name}: {arg.doc}")
    if function.args:
        args.append("/")  # TODO: remove once we support kwargs
    signature.append(", ".join(args))

    return_docs = function.return_docs
    returns = pytype_stub_repr(function.op_mapping.result)
    if return_docs:
        if doc:
            doc.append(f":returns {returns}: {return_docs[0]}")
            doc.extend(return_docs[1:])
        else:
            doc = return_docs
    signature.append(f") -> {returns}:")
    teal_ops = sorted({op.op_code for op in function.op_mapping.overloads})
    teal_op_desc = ", ".join(_get_algorand_doc(teal_op) for teal_op in teal_ops)
    doc.append("")
    doc.append(f"Native TEAL opcode: {teal_op_desc}")
    body = list[str]()
    if doc:
        body.append('"""')
        body.extend(doc)
        body.append('"""')
    else:
        body.append("...")

    yield prefix + "".join(signature)
    yield from [textwrap.indent(line, prefix=prefix + INDENT) for line in body]

def build_stub_class(klass: ClassDef) -> Iterable[str]:
    ops = [f"{_get_algorand_doc(op)}" for op in klass.ops]
    docstring = "\n".join(
        [
            INDENT + '"""',
            INDENT + klass.doc,
            INDENT + f"Native TEAL op{'s' if len(ops) > 1 else ''}: {', '.join(ops)}",
            INDENT + '"""',
        ]
    )
    method_preamble = f"{INDENT}@staticmethod"
    yield f"class {klass.name}:"
    yield docstring
    for method in klass.methods:
        if method.is_property:
            yield from build_class_var_stub(method, INDENT)
        else:
            yield method_preamble
            yield from build_method_stub(method, prefix=INDENT)
        yield ""

def build_class_var_stub(function: FunctionDef, indent: str) -> Iterable[str]:
    returns = pytype_stub_repr(function.op_mapping.result)
    return_docs = function.return_docs
    doc = return_docs if return_docs else function.doc[:]
    _maybe_add_min_version_doc(doc, function.min_avm_version)
    yield f"{indent}{function.name}: typing.Final[{returns}] = ..."
    yield f'{indent}"""'
    for doc_line in doc:
        yield f"{indent}{doc_line}"
    yield f'{indent}"""'

def _get_modified_stack_value(alias: Op) -> StackValue:
    immediate = get_overriding_immediate(alias)
    assert immediate
    if immediate.modifies_stack_input is not None:
        return alias.stack_inputs[immediate.modifies_stack_input]
    else:
        assert immediate.modifies_stack_output is not None
        return alias.stack_outputs[immediate.modifies_stack_output]

AliasT: typing.TypeAlias = tuple[Op, list[str]]

def build_class_from_overriding_immediate(
    spec: LanguageSpec,
    op: Op,
    class_name: str,
    class_doc: str,
    immediate: Immediate,
    aliases: list[AliasT],
) -> ClassDef:
    assert immediate.arg_enum
    logger.info(f"Using overriding immediate for {op.name}")

    arg_enum_values = spec.arg_enums[immediate.arg_enum]

    # copy inputs so they can be mutated safely
    op = copy.deepcopy(op)
    aliases = copy.deepcopy(aliases)

    # obtain a list of stack values that will be modified for each enum
    stacks_to_modify = [_get_modified_stack_value(o) for o, _ in [(op, None), *aliases]]

    # build a method for each arg enum value
    methods = list[FunctionDef]()
    class_ops = {op.name}
    for value in arg_enum_values:
        stack_type = value.stack_type
        assert stack_type

        for stack_to_modify in stacks_to_modify:
            stack_to_modify.stack_type = stack_type
            stack_to_modify.doc = value.doc

        method = build_operation_method(
            op, snake_case(value.name), aliases, const_immediate_value=(immediate, value)
        )
        # some enums are reused across ops, so need to take the max minimum of op and enum version
        method.min_avm_version = max(op.min_avm_version, value.min_avm_version)
        _maybe_add_min_version_doc(method.doc, method.min_avm_version)

        for op_mapping in method.op_mapping.overloads:
            class_ops.add(op_mapping.op_code)

        methods.append(method)

    return ClassDef(name=class_name, doc=class_doc, methods=methods, ops=sorted(class_ops))

def get_op_doc(op: Op) -> list[str]:
    doc = [d.replace("\\", "\\\\") for d in op.doc]
    _maybe_add_min_version_doc(doc, op.min_avm_version)
    return doc

def get_python_enum_class(arg_enum: str) -> str:
    try:
        return PYTHON_ENUM_CLASS[arg_enum]
    except KeyError:
        pass
    # don't change acronyms
    if arg_enum.isupper():
        return arg_enum
    return snake_case(arg_enum).replace("_", " ").title().replace(" ", "")

def get_overriding_immediate(op: Op) -> Immediate | None:
    return next(
        (
            immediate
            for immediate in op.immediate_args
            if immediate.modifies_stack_input is not None
            or immediate.modifies_stack_output is not None
        ),
        None,
    )

def build_enum(spec: LanguageSpec, arg_enum: str) -> Iterable[str]:
    values = spec.arg_enums[arg_enum]
    enum_name = get_python_enum_class(arg_enum)
    yield f"class {enum_name}(str):"
    yield f'{INDENT}"""Available values for the `{arg_enum}` enum"""'
    for value in values:
        yield f"{INDENT}{value.name}: {enum_name} = ..."
        enum_doc = []
        if value.doc:
            enum_doc.append(value.doc)
        _maybe_add_min_version_doc(enum_doc, value.min_avm_version)
        if enum_doc:
            yield f'{INDENT}"""'
            for doc_line in enum_doc:
                yield f"{INDENT}{doc_line}"
            yield f'{INDENT}"""'
    yield ""

def _maybe_add_min_version_doc(doc: list[str], version: int) -> None:
    # only output min AVM version if it is greater than our min supported version
    if version > MIN_SUPPORTED_VERSION:
        doc.append(f"Min AVM version: {version}")

def build_operation_method(
    op: Op,
    op_function_name: str,
    aliases: list[AliasT],
    const_immediate_value: tuple[Immediate, ArgEnum] | None = None,
) -> FunctionDef:
    args = []
    # python stub args can be different to mapping args, due to immediate args
    # that are inferred based on the method/property used
    function_args = []
    doc = get_op_doc(op)
    for immediate in op.immediate_args:
        arg_type: ImmediateKind | str
        if immediate.immediate_type == ImmediateKind.arg_enum:
            assert immediate.arg_enum, "Arg enum expected"
            arg_type = get_python_enum_class(immediate.arg_enum)
        else:
            arg_type = immediate.immediate_type
        im_arg = TypedName(name=immediate.name.lower(), type=arg_type, doc=immediate.doc)
        args.append(im_arg)
        if const_immediate_value and const_immediate_value[0] == immediate:
            # omit immediate arg from signature
            doc = []
        else:
            function_args.append(im_arg)
    for si in op.stack_inputs:
        stack_arg = TypedName(name=si.name.lower(), type=si.stack_type, doc=si.doc)
        args.append(stack_arg)
        function_args.append(stack_arg)

    if op.halts:
        return_docs = ["Halts program"]
    else:
        return_docs = [so.doc for so in op.stack_outputs if so.doc]

    try:
        property_op = PROPERTY_OPS[op.name]
    except KeyError:
        is_property = False
    else:
        is_property = op_function_name not in property_op["exclude"]

    if op.halts:
        result_typ = pytypes.NeverType
    else:
        # replace immediate reference to arg enum with a constant enum value
        result_ptypes = [sub_types(o.stack_type, covariant=False)[0] for o in op.stack_outputs]
        if not result_ptypes:
            result_typ = pytypes.NoneType
        elif len(op.stack_outputs) == 1:
            (result_typ,) = result_ptypes
        else:
            result_typ = pytypes.GenericTupleType.parameterise(result_ptypes, source_location=None)
        if result_typ == pytypes.UInt64Type:
            if op_function_name == "on_completion":
                result_typ = pytypes.OnCompleteActionType
            elif op_function_name == "type_enum":
                result_typ = pytypes.TransactionTypeType
    op_mappings = []
    ops_with_aliases = [(op, list[str]()), *aliases]
    for map_op, alias_args in ops_with_aliases:
        assert map_op.stack_outputs == op.stack_outputs
        if alias_args:
            # map the stack or immediate input name to the function signature position
            name_to_sig_idx = {n: idx2 for idx2, n in enumerate(alias_args)}
        else:
            name_to_sig_idx = {tn.name.upper(): idx2 for idx2, tn in enumerate(args)}
        map_immediates = list[str | int | type[str | int]]()
        map_args_map = dict[int, Sequence[pytypes.PyType] | int]()
        for idx, i_arg in enumerate(map_op.immediate_args):
            if const_immediate_value and const_immediate_value[0] == i_arg:
                map_immediates.append(const_immediate_value[1].name)
            else:
                im_typ = immediate_kind_to_type(i_arg.immediate_type)
                map_immediates.append(im_typ)
                sig_idx = name_to_sig_idx[i_arg.name]
                map_args_map[sig_idx] = idx

        for s_arg in map_op.stack_inputs:
            allowed_types = tuple(sub_types(s_arg.stack_type, covariant=True))
            sig_idx = name_to_sig_idx[s_arg.name]
            map_args_map[sig_idx] = allowed_types

        op_mappings.append(
            FunctionOpMapping(
                op_code=map_op.name,
                immediates=map_immediates,
                args=[map_args_map[k] for k in sorted(map_args_map)],
            )
        )
    proto_function = FunctionDef(
        name=op_function_name,
        doc=doc,
        is_property=is_property,
        args=function_args,
        return_docs=return_docs,
        op_mapping=OpMappingWithOverloads(
            arity=len(function_args),
            result=result_typ,
            overloads=op_mappings,
        ),
        min_avm_version=op.min_avm_version,
    )

    return proto_function

def build_operation_methods(
    op: Op, op_function_name: str, aliases: list[AliasT]
) -> Iterable[FunctionDef]:
    logger.info(f"Mapping {op.name} to {op_function_name}")

    if StackType.any in (s.stack_type for s in op.stack_outputs):
        logger.info(f"Found any output for {op.name}")
        for replace_any_with in (StackType.bytes, StackType.uint64):
            new_op = op_any_replaced(op, replace_any_with)
            new_name = f"{op_function_name}_{replace_any_with.name}"
            new_aliases = [
                (op_any_replaced(alias_op, replace_any_with), names) for alias_op, names in aliases
            ]
            yield build_operation_method(new_op, new_name, new_aliases)
    else:
        yield build_operation_method(op, op_function_name, aliases)

def op_any_replaced(op: Op, replace_any_with: StackType) -> Op:
    stack_inputs = []
    input_replaced = 0
    for si in op.stack_inputs:
        if si.stack_type != StackType.any:
            stack_inputs.append(si)
        else:
            input_replaced += 1
            stack_inputs.append(attrs.evolve(si, stack_type=replace_any_with))

    stack_outputs = []
    outputs_replaced = 0
    for so in op.stack_outputs:
        if so.stack_type != StackType.any:
            stack_outputs.append(so)
        else:
            outputs_replaced += 1
            stack_outputs.append(attrs.evolve(so, stack_type=replace_any_with))
    assert outputs_replaced == 1
    return attrs.evolve(op, stack_inputs=stack_inputs, stack_outputs=stack_outputs)

def build_aliased_ops(spec: LanguageSpec, group: RenamedOpCode) -> Iterable[FunctionDef]:
    op = spec.ops[group.op]
    aliases = [
        (spec.ops[stack_alias], arg_map) for stack_alias, arg_map in group.stack_aliases.items()
    ]
    methods = build_operation_methods(op, group.name, aliases)
    return methods

def build_merged_ops(spec: LanguageSpec, group: MergedOpCodes) -> ClassDef:
    merge_methods = dict[str, FunctionDef]()

    for other_op_name, alias_dict in group.ops.items():
        aliases = [(spec.ops[alias_op], arg_map) for alias_op, arg_map in alias_dict.items()]
        other_op = spec.ops[other_op_name]
        overriding_immediate = get_overriding_immediate(other_op)
        assert overriding_immediate
        other_class = build_class_from_overriding_immediate(
            spec,
            other_op,
            class_name=group.name,
            class_doc=group.doc,
            immediate=overriding_immediate,
            aliases=aliases,
        )
        for method in other_class.methods:
            merge_methods[method.name] = method

    methods = list(merge_methods.values())
    return ClassDef(name=group.name, doc=group.doc, methods=methods, ops=sorted(group.ops))

def build_grouped_ops(
    spec: LanguageSpec, group: GroupedOpCodes, enums_to_build: dict[str, bool]
) -> ClassDef:
    methods = list[FunctionDef]()
    for rename_op_name, python_name in group.ops.items():
        rename_op = spec.ops[rename_op_name]
        rename_immediate = get_overriding_immediate(rename_op)
        if rename_immediate:
            rename_class = build_class_from_overriding_immediate(
                spec,
                rename_op,
                class_name=group.name,
                class_doc=group.doc,
                immediate=rename_immediate,
                aliases=[],
            )
            # when grouping an op with immediate overrides, treat python_name as a prefix
            for method in rename_class.methods:
                method.name = f"{python_name}_{method.name}"
            methods.extend(rename_class.methods)

        else:
            methods.extend(build_operation_methods(rename_op, python_name, aliases=[]))
        for arg in rename_op.immediate_args:
            if arg.immediate_type == ImmediateKind.arg_enum and (
                arg.modifies_stack_input is None and arg.modifies_stack_output is None
            ):
                assert arg.arg_enum is not None
                enums_to_build[arg.arg_enum] = True

    class_def = ClassDef(
        name=group.name,
        doc=group.doc,
        methods=methods,
        ops=sorted(group.ops),
    )
    return class_def

def pytype_repr(typ: pytypes.PyType) -> str:
    try:
        return PYTYPE_REPR[typ]
    except KeyError:
        pass
    match typ:
        case pytypes.TupleType(items=tuple_items) if len(tuple_items) > 1:
            item_strs = [pytype_repr(item) for item in tuple_items]
            return (
                f"pytypes.GenericTupleType.parameterise("
                f"({', '.join(item_strs)}), source_location=None)"
            )
    raise ValueError(f"Unexpected pytype: {typ}")

def build_op_specification_body(function: FunctionDef) -> Iterable[str]:
    if function.is_property:
        (op_mapping,) = function.op_mapping.overloads
        (immediate,) = op_mapping.immediates
        yield (
            f"{function.name}=PropertyOpMapping("
            f"{op_mapping.op_code!r}, {immediate!r}, {pytype_repr(function.op_mapping.result)},"
            f"),"
        )
    else:
        yield f"{function.name}=OpMappingWithOverloads("
        if function.op_mapping.result is not pytypes.NoneType:
            yield f" result={pytype_repr(function.op_mapping.result)},"
        yield f" arity={function.op_mapping.arity}, "
        yield " overloads=["
        for op_mapping in function.op_mapping.overloads:
            yield f"FunctionOpMapping({op_mapping.op_code!r},"
            if op_mapping.immediates:
                yield " immediates=["
                for idx, item in enumerate(op_mapping.immediates):
                    if idx:
                        yield ", "
                    if not isinstance(item, type):
                        yield repr(item)
                    else:
                        yield item.__name__
                yield "],"
            if op_mapping.args:
                yield " args=["
                for idx, allowed_types_or_idx in enumerate(op_mapping.args):
                    if idx:
                        yield ", "
                    if isinstance(allowed_types_or_idx, int):
                        yield repr(allowed_types_or_idx)
                    else:  # noqa: PLR5501
                        if len(allowed_types_or_idx) == 1:
                            yield f"({pytype_repr(*allowed_types_or_idx)},)"
                        else:
                            yield "("
                            for idx2, allowed_type in enumerate(allowed_types_or_idx):
                                if idx2:
                                    yield ","
                                yield pytype_repr(allowed_type)
                            yield ")"
                yield "],"

            yield "),"
        yield "]"
        yield "),"

def build_awst_data(
    lang_spec: LanguageSpec,
    enums: list[str],
    function_ops: list[FunctionDef],
    class_ops: list[ClassDef],
) -> Iterable[str]:
    yield "import typing"
    yield "from collections.abc import Mapping, Sequence"
    yield "from puyapy.awst_build import pytypes"
    yield (
        "from puyapy.awst_build.intrinsic_models"
        " import FunctionOpMapping, OpMappingWithOverloads, PropertyOpMapping"
    )
    yield "ENUM_CLASSES: typing.Final[Mapping[str, Mapping[str, str]]] = dict("
    for enum_name in enums:
        yield f"{get_python_enum_class(enum_name)}=dict("
        for enum_value in lang_spec.arg_enums[enum_name]:
            # enum names currently match enum immediate values
            yield f'{enum_value.name}="{enum_value.name}",'
        yield "),"
    yield ")"
    yield ""
    yield "FUNC_TO_AST_MAPPER: typing.Final[Mapping[str, OpMappingWithOverloads]] = dict("
    for function_op in function_ops:
        yield "".join(build_op_specification_body(function_op))
    yield ")"

    yield (
        "NAMESPACE_CLASSES: "
        "typing.Final[Mapping[str, Mapping[str, PropertyOpMapping | OpMappingWithOverloads]]]"
        " = dict("
    )
    for class_op in class_ops:
        yield f"{class_op.name}=dict("
        for method in class_op.methods:
            yield "".join(build_op_specification_body(method))
        yield "),"
    yield ")"

def output_stub(
    lang_spec: LanguageSpec,
    enums: list[str],
    function_ops: list[FunctionDef],
    class_ops: list[ClassDef],
) -> None:
    references = ", ".join(
        sorted(
            str(pt).removeprefix("algopy.")
            for pt, lit_t in PYTYPE_TO_LITERAL.items()
            if str(pt).startswith("algopy.")
        )
    )
    stub: list[str] = [
        "import typing",
        "",
        f"from algopy import {references}",
    ]

    for arg_enum in enums:
        stub.extend(build_enum(lang_spec, arg_enum))

    for function in function_ops:
        stub.extend(build_method_stub(function))

    for class_op in class_ops:
        stub.extend(build_stub_class(class_op))

    stub_out_path = VCS_ROOT / "stubs" / "algopy-stubs" / f"{STUB_NAMESPACE}.pyi"
    stub_out_path.write_text("\n".join(stub), encoding="utf-8")
    subprocess.run(["ruff", "format", str(stub_out_path)], check=True, cwd=VCS_ROOT)

def pytype_stub_repr(pytype: pytypes.PyType) -> str:
    return str(pytype).replace("algopy.", "")

def output_awst_data(
    lang_spec: LanguageSpec,
    enums: list[str],
    function_ops: list[FunctionDef],
    class_ops: list[ClassDef],
) -> None:
    awst_data = build_awst_data(lang_spec, enums, function_ops, class_ops)

    awst_data_path = VCS_ROOT / "src" / "puyapy" / "awst_build" / "intrinsic_data.py"
    awst_data_path.write_text("\n".join(awst_data), encoding="utf-8")
    subprocess.run(["ruff", "format", str(awst_data_path)], check=True, cwd=VCS_ROOT)
    subprocess.run(["ruff", "check", "--fix", str(awst_data_path)], check=False, cwd=VCS_ROOT)

def _get_algorand_doc(op: str) -> str:
    return f"[`{op}`]({ALGORAND_OP_URL}#{op})"

if __name__ == "__main__":
    main()


#!/usr/bin/env python3
import contextlib
import enum
import json
import logging
import typing
from pathlib import Path

import attrs
import cattrs

logger = logging.getLogger(__name__)

STACK_INPUT_NAMES = "ABCDE"
STACK_OUTPUT_NAMES_FEW = "XYZ"  # 3 or less var
STACK_OUTPUT_NAMES_MANY = "WXYZ"  # 4 var

VARIABLE_SIZE_OPCODES = {
    "intcblock",
    "bytecblock",
    "pushbytes",
    "pushbytess",
    "pushint",
    "pushints",
    "switch",
    "match",
}

class NamedType(typing.TypedDict):
    """
    {
      "Name": "uint64",
      "Abbreviation": "i",
      "Bound": [
        0,
        18446744073709551615
      ],
      "AVMType": "uint64"
    },
    """

    Name: str
    Abbreviation: str
    AVMType: str

class ImmediateNote(typing.TypedDict, total=False):
    """
    {
      "Comment": "transaction field index",
      "Encoding": "uint8",
      "Name": "F",
      "Reference": "txn"
    }
    """

    Comment: str
    Encoding: str
    Name: str
    Reference: str

class Operation(typing.TypedDict, total=False):
    """
    {
      "Opcode": 0,
      "Name": "err",
      "Size": 1,
      "Doc": "Fail immediately.",
      "IntroducedVersion": 1,
      "Groups": [
        "Flow Control"
      ]
    },
    {
      "Opcode": 1,
      "Name": "sha256",
      "Args": [
        "[]byte"
      ],
      "Returns": [
        "[32]byte"
      ],
      "Size": 1,
      "Doc": "SHA256 hash of value A, yields [32]byte",
      "IntroducedVersion": 1,
      "Groups": [
        "Arithmetic"
      ]
    }
    """

    Doc: str
    Opcode: int
    Size: int
    Name: str
    IntroducedVersion: int
    Groups: list[str]
    Args: list[str]
    Returns: list[str]
    DocExtra: str
    ArgEnum: list[str]
    ArgEnumTypes: list[str]
    ArgEnumBytes: list[int]
    ArgModes: list[int]
    ArgEnumVersion: list[int]
    ImmediateNote: list[ImmediateNote]

    # the following values are not in the original langspec.json
    # these values are manually patched in during transform
    ArgEnumIsInput: bool
    Halts: bool
    # these values are output by a modified opdoc.go from go-algorand repo
    Cost: str
    ArgEnumDoc: list[str]
    Modes: int

class AlgorandLanguageSpec(typing.TypedDict):
    NamedTypes: list[NamedType]
    Ops: list[Operation]

class StackType(enum.StrEnum):
    uint64 = enum.auto()
    bytes = "[]byte"
    bytes_8 = "[8]byte"
    bytes_32 = "[32]byte"
    bytes_33 = "[33]byte"
    bytes_64 = "[64]byte"
    bytes_80 = "[80]byte"
    bytes_1232 = "[1232]byte"
    bytes_1793 = "[1793]byte"
    bool = enum.auto()
    address = enum.auto()
    address_or_index = enum.auto()
    any = enum.auto()
    bigint = enum.auto()
    box_name = "boxName"
    asset = enum.auto()
    application = enum.auto()
    state_key = "stateKey"

class RunMode(enum.StrEnum):
    app = enum.auto()
    sig = enum.auto()
    any = enum.auto()

@attrs.define
class StackValue:
    name: str
    """Name used to refer to this value in the Op.doc"""
    stack_type: StackType
    doc: str | None = None

@attrs.define
class ArgEnum:
    name: str
    doc: str | None
    stack_type: StackType | None
    mode: RunMode
    value: int
    min_avm_version: int

class ImmediateKind(enum.StrEnum):
    uint8 = enum.auto()
    int8 = enum.auto()
    label = enum.auto()
    varuint = enum.auto()
    bytes = enum.auto()

    # array types
    label_array = enum.auto()
    varuint_array = enum.auto()
    bytes_array = enum.auto()

    # not in original lang spec
    arg_enum = enum.auto()

@attrs.frozen(kw_only=True)
class Immediate:
    name: str
    """Name used to refer to this value in the Op.doc"""
    immediate_type: ImmediateKind
    arg_enum: str | None = None
    modifies_stack_input: int | None = None
    """Index of stack input type that this immediate modifies"""
    modifies_stack_output: int | None = None
    """Index of stack output type that this immediate modifies"""
    """field_group reference if immediate_type is field_group"""
    doc: str | None = None

@attrs.define
class Cost:
    value: int | None
    """Static cost of op, or None if cost is not static"""
    doc: str
    """Documentation describing how cost is calculated"""

@attrs.define
class Op:
    name: str
    """Name of op in TEAL"""
    code: int
    """Bytecode value"""
    size: int
    """Size in bytes of compiled op, 0 indicate size is variable"""
    doc: list[str]
    cost: Cost
    min_avm_version: int
    """AVM version op was introduced"""
    halts: bool
    mode: RunMode
    """True if this op halts the program"""
    groups: list[str] = attrs.field(factory=list)
    """Groups op belongs to"""
    stack_inputs: list[StackValue] = attrs.field(factory=list)
    """Inputs that come from the stack"""
    immediate_args: list[Immediate] = attrs.field(factory=list)
    """Arguments that are passed as immediate values in TEAL"""
    stack_outputs: list[StackValue] = attrs.field(factory=list)
    """Outputs left on the stack"""

@attrs.define
class LanguageSpec:
    ops: dict[str, Op] = attrs.field(factory=dict)
    arg_enums: dict[str, list[ArgEnum]] = attrs.field(factory=dict)

    @staticmethod
    def from_json(json: dict[str, typing.Any]) -> "LanguageSpec":
        return cattrs.structure(json, LanguageSpec)

    def to_json(self) -> dict[str, typing.Any]:
        return attrs.asdict(self)

def _patch_lang_spec(lang_spec: dict[str, typing.Any]) -> None:
    ops = {op["Name"]: op for op in lang_spec["Ops"]}

    # patch ops that use a stack type of any
    # for arguments that should be an Address or Address index
    for op_name in (
        "acct_params_get",
        "app_local_get",
        "app_local_put",
        "app_local_del",
        "app_local_get_ex",
        "app_opted_in",
        "asset_holding_get",
        "balance",
        "min_balance",
        "voter_params_get",
    ):
        _patch_arg_type(ops, op_name, 0, "any", "address_or_index")

    # patch ops that use a stack type of uint64
    # for arguments that should be an Application
    for op_name, arg_index in {
        "app_opted_in": 1,
        "app_global_get_ex": 0,
        "app_local_get_ex": 1,
        "app_params_get": 0,
    }.items():
        _patch_arg_type(ops, op_name, arg_index, "uint64", "application")

    # patch ops that use a stack type of uint64
    # for return types that should be a bool
    for op_name in [
        "!",
    ]:
        _patch_return_type(ops, op_name, 0, "uint64", "bool")

    # patch ops that use a stack type of uint64
    # for arguments that should be an Asset
    for op_name, arg_index in {
        "asset_holding_get": 1,
        "asset_params_get": 0,
    }.items():
        _patch_arg_type(ops, op_name, arg_index, "uint64", "asset")

    for op_name, arg_index in {
        "select": 2,
    }.items():
        _patch_arg_type(ops, op_name, arg_index, "uint64", "bool")

    # patch return bytes -> bigint
    for op_name in [
        "b+",
        "b*",
    ]:
        _patch_return_type(ops, op_name, 0, "[]byte", "bigint")

    # patch txn enum fields with asset and application types
    txn = ops["txn"]
    itxn_field = ops["itxn_field"]
    for op in (txn, itxn_field):
        for immediate in [
            "XferAsset",
            "ConfigAsset",
            "FreezeAsset",
        ]:
            _patch_arg_enum_type(op, immediate, "uint64", "asset")

        _patch_arg_enum_type(op, "ApplicationID", "uint64", "application")
    _patch_arg_enum_type(txn, "CreatedApplicationID", "uint64", "application")
    _patch_arg_enum_type(txn, "CreatedAssetID", "uint64", "asset")

    # patch txna enums
    txna = ops["txna"]
    _patch_arg_enum_type(txna, "Assets", "uint64", "asset")
    _patch_arg_enum_type(txna, "Applications", "uint64", "application")

    # patch global enums
    _patch_arg_enum_type(ops["global"], "CurrentApplicationID", "uint64", "application")

    # base64_decode has an ArgEnumTypes array when it probably shouldn't
    # as all stack outputs are bytes
    del ops["base64_decode"]["ArgEnumTypes"]

    # itxn_field reuses the same field group as txn, however it only uses a subset of fields
    # additionally ArgEnumTypes refers to the stack input types not the output types
    itxn_field = ops["itxn_field"]
    itxn_field["ImmediateNote"][0]["Reference"] = "itxn_field"
    itxn_field["ArgEnumIsInput"] = True
    # ops that never return encode this with a single return type of none
    # however currently this information is stripped when generating langspec.json
    ops["err"]["Halts"] = True
    ops["return"]["Halts"] = True

def _patch_arg_enum_type(
    op: dict[str, typing.Any], immediate: str, current_type: str, new_type: str
) -> None:
    arg_enum = op["ArgEnum"]
    assert immediate in arg_enum, f"Expected {immediate} arg enum for {op['Name']}"
    immediate_index = arg_enum.index(immediate)
    arg_enum_types = op["ArgEnumTypes"]
    assert (
        arg_enum_types[immediate_index] == current_type
    ), f"Expected {immediate} to be {current_type}"
    arg_enum_types[immediate_index] = new_type

def _patch_arg_type(
    ops: dict[str, typing.Any], op_name: str, arg_index: int, current_type: str, new_type: str
) -> None:
    op_args = ops[op_name]["Args"]
    assert (
        op_args[arg_index] == current_type
    ), f"Expected {op_name} arg {arg_index} to be {current_type}"
    op_args[arg_index] = new_type

def _patch_return_type(
    ops: dict[str, typing.Any], op_name: str, return_index: int, current_type: str, new_type: str
) -> None:
    returns = ops[op_name]["Returns"]
    assert (
        returns[return_index] == current_type
    ), f"Expected {op_name} return {return_index} to be {current_type}"
    returns[return_index] = new_type

def create_indexed_enum(op: Operation) -> list[ArgEnum]:
    enum_names = op["ArgEnum"]
    enum_types: list[str] | list[None] = op.get("ArgEnumTypes", [])
    enum_docs = op["ArgEnumDoc"]
    enum_bytes = op["ArgEnumBytes"]
    enum_modes = op["ArgModes"]
    enum_versions = op["ArgEnumVersion"]

    if not enum_types:
        enum_types = [None] * len(enum_names)

    result = list[ArgEnum]()
    for enum_name, enum_type, enum_doc, enum_mode, enum_byte, enum_version in zip(
        enum_names, enum_types, enum_docs, enum_modes, enum_bytes, enum_versions, strict=True
    ):
        stack_type = None if enum_type is None else StackType(enum_type)
        enum_value = ArgEnum(
            name=enum_name,
            doc=enum_doc if enum_doc else None,
            stack_type=stack_type,
            mode=_map_enum_mode(op["Modes"], enum_mode),
            value=enum_byte,
            min_avm_version=enum_version,
        )
        result.append(enum_value)
    return result

def _map_enum_mode(op_mode: int, arg_mode: int = 0) -> RunMode:
    mode = arg_mode or op_mode
    match mode:
        case 1:
            return RunMode.sig
        case 2:
            return RunMode.app
        case 3:
            return RunMode.any
        case _:
            raise ValueError("Unexpected run mode")

def transform_encoding(value: str) -> ImmediateKind:
    match value:
        case "uint8":
            result = ImmediateKind.uint8
        case "int8":
            result = ImmediateKind.int8
        case "int16 (big-endian)":
            result = ImmediateKind.label
        case "varuint":
            result = ImmediateKind.varuint
        case "varuint length, bytes":
            result = ImmediateKind.bytes
        case "varuint count, [varuint ...]":
            result = ImmediateKind.varuint_array
        case "varuint count, [varuint length, bytes ...]":
            result = ImmediateKind.bytes_array
        case "varuint count, [int16 (big-endian) ...]":
            result = ImmediateKind.label_array
        case _:
            raise ValueError(f"Unknown Encoding: {value}")
    return result

def transform_stack_args(op: Operation) -> list[StackValue]:
    result = list[StackValue]()
    args = op.get("Args", [])
    assert len(args) <= len(STACK_INPUT_NAMES), f"More args than expected for {op['Name']}"
    for index, arg_type in enumerate(op.get("Args", [])):
        name = STACK_INPUT_NAMES[index]
        stack_type = StackType(arg_type)
        result.append(StackValue(name=name, stack_type=stack_type))
    return result

def transform_immediates(
    arg_enums: dict[str, list[ArgEnum]],
    algorand_ops: dict[str, Operation],
    op: Operation,
) -> list[Immediate]:
    op_name = op["Name"]
    result = list[Immediate]()
    for immediate in op.get("ImmediateNote", []):
        arg_enum_reference = immediate.get("Reference")
        if arg_enum_reference is not None:
            arg_enum = op.get("ArgEnum")
            if arg_enum_reference not in arg_enums:
                try:
                    enum_op = algorand_ops[arg_enum_reference]
                except KeyError:
                    enum_op = op
                assert arg_enum, f"Expected enum for {op_name}"
                arg_enums[arg_enum_reference] = create_indexed_enum(enum_op)

            if arg_enum is not None:
                assert len(arg_enum) == len(
                    arg_enums[arg_enum_reference]
                ), f"Arg Enum lengths don't match for {op_name}"

        modifies_stack_input: int | None = None
        modifies_stack_output: int | None = None

        if arg_enum_reference and any(a.stack_type for a in arg_enums[arg_enum_reference]):
            assert all(a.stack_type for a in arg_enums[arg_enum_reference])
            if op.get("ArgEnumIsInput"):
                modifies_stack_input = 0
            else:
                modifies_stack_output = 0

        result.append(
            Immediate(
                name=immediate["Name"],
                immediate_type=(
                    transform_encoding(immediate["Encoding"])
                    if arg_enum_reference is None
                    else ImmediateKind.arg_enum
                ),
                modifies_stack_input=modifies_stack_input,
                modifies_stack_output=modifies_stack_output,
                arg_enum=arg_enum_reference,
                doc=immediate["Comment"],
            )
        )
    return result

def transform_returns(op: Operation) -> list[StackValue]:
    try:
        returns = op["Returns"]
    except KeyError:
        return []
    num_returns = len(returns)
    if num_returns <= len(STACK_OUTPUT_NAMES_FEW):
        return_argument_names = STACK_OUTPUT_NAMES_FEW
    elif num_returns <= len(STACK_OUTPUT_NAMES_MANY):
        return_argument_names = STACK_OUTPUT_NAMES_MANY
    else:
        raise AssertionError(f"More returns than expected for {op['Name']}")
    return [
        StackValue(
            name=name,
            stack_type=StackType(return_type),
        )
        for name, return_type in zip(return_argument_names, returns, strict=False)
    ]

def transform_doc(op: Operation) -> list[str]:
    doc = op["Doc"].splitlines()

    doc_extra = op.get("DocExtra")
    if doc_extra:
        doc.extend(doc_extra.splitlines())
    return doc

def get_immediate_encoded_size(immediate: Immediate) -> int:
    match immediate.immediate_type:
        case ImmediateKind.uint8 | ImmediateKind.int8 | ImmediateKind.arg_enum:
            return 1
        case ImmediateKind.label:
            return 2
        case ImmediateKind():
            return 0
        case _:
            raise ValueError(f"Cannot determine size of {immediate.immediate_type}")

def transform_cost(op: Operation) -> Cost:
    algorand_cost = op["Cost"]
    cost = Cost(value=None, doc=algorand_cost)
    with contextlib.suppress(ValueError):
        cost.value = int(algorand_cost)

    return cost

def transform_spec(lang_spec: AlgorandLanguageSpec) -> LanguageSpec:
    result = LanguageSpec()
    arg_enums = result.arg_enums
    algorand_ops = {o["Name"]: o for o in sorted(lang_spec["Ops"], key=lambda x: x["Name"])}
    for op_name, algorand_op in algorand_ops.items():
        op = Op(
            name=op_name,
            code=algorand_op["Opcode"],
            size=algorand_op["Size"],
            doc=transform_doc(algorand_op),
            cost=transform_cost(algorand_op),
            min_avm_version=algorand_op["IntroducedVersion"],
            groups=algorand_op["Groups"],
            immediate_args=transform_immediates(arg_enums, algorand_ops, algorand_op),
            stack_inputs=transform_stack_args(algorand_op),
            stack_outputs=transform_returns(algorand_op),
            halts=algorand_op.get("Halts", False),
            mode=_map_enum_mode(algorand_op["Modes"]),
        )
        validate_op(result, op)
        result.ops[op.name] = op
    return result

def validate_op(lang_spec: LanguageSpec, op: Op) -> None:
    # validate op size
    instruction_size = 0 if op.name in VARIABLE_SIZE_OPCODES else 1
    expected_size = (
        sum([get_immediate_encoded_size(a) for a in op.immediate_args]) + instruction_size
    )
    assert op.size == expected_size, f"Unexpected size for specified immediate args for {op.name}"

    # validate immediate modifiers
    for immediate in op.immediate_args:
        if immediate.immediate_type == ImmediateKind.arg_enum:
            assert immediate.arg_enum in lang_spec.arg_enums
            if immediate.modifies_stack_input is not None:
                assert immediate.modifies_stack_input < len(op.stack_inputs), (
                    f"Immediate for {op.name} references stack input "
                    f"that does not exist {immediate.modifies_stack_input}"
                )
            if immediate.modifies_stack_output is not None:
                assert immediate.modifies_stack_output < len(op.stack_outputs), (
                    f"Immediate for {op.name} references stack output "
                    f"that does not exist {immediate.modifies_stack_output}"
                )
        else:
            assert not immediate.arg_enum
            assert not immediate.modifies_stack_input
            assert not immediate.modifies_stack_output

def main() -> None:
    vcs_root = Path(__file__).parent.parent
    spec_path = vcs_root / "langspec.json"

    output_path = vcs_root / "langspec.puya.json"
    logger.info(f"Transforming {spec_path} to {output_path}")

    lang_spec_json = json.loads(spec_path.read_text(encoding="utf-8"))
    _patch_lang_spec(lang_spec_json)
    lang_spec = typing.cast(AlgorandLanguageSpec, lang_spec_json)

    puya_spec = transform_spec(lang_spec)
    puya_json = json.dumps(puya_spec.to_json(), indent=4)
    output_path.write_text(puya_json, encoding="utf-8")

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(message)s")
    main()


#!/usr/bin/env python3
import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

MYPY_REPO = "https://github.com/python/mypy.git"
VCS_ROOT = Path(__file__).parent.parent
TYPESHED_README = """
This is PuyaPy's custom typeshed, which is a curated subset of the official MyPy typeshed.
It only includes the required stubs used by PuyaPy as this speeds up MyPy's parsing speed
significantly.

However this means certain python modules such as `enum` or `dataclasses` cannot be used in
PuyaPy stubs unless this typeshed is updated.

The contents of the typeshed are populated by the `scripts/vendor_mypy.py` script, which is used
to vendor new versions of MyPy or to update the stubs included in this typeshed. So to add new
stubs, update that script and rerun.
""".strip()

def clone_branch(version: str) -> str:
    git_clone = f"git clone --depth=1 --branch={version} --single-branch {MYPY_REPO} ."
    print(f"Executing: {git_clone}")
    subprocess.run(git_clone.split(), check=True)

    git_hash = subprocess.run("git rev-parse HEAD".split(), capture_output=True, check=True).stdout
    assert git_hash is not None
    subprocess.run("rm -rf .git".split(), check=True)
    return git_hash.decode("utf8").strip()

def vendor_mypy(version: str) -> None:
    puya_src_dir = VCS_ROOT / "src" / "puyapy"
    vendor_dir = puya_src_dir / "_vendor"
    mypy_vendor = vendor_dir / "mypy"
    print(f"Vendoring mypy into: {mypy_vendor}")

    print("Removing existing mypy files...")
    shutil.rmtree(mypy_vendor, ignore_errors=True)

    print(f"Cloning mypy {version}...")
    with tempfile.TemporaryDirectory() as tmp_dir:
        os.chdir(tmp_dir)
        git_hash = clone_branch(version)
        print(f"Checked out mypy {version} @ {git_hash}")

        print(f"Copying mypy into {mypy_vendor}...")
        shutil.copytree(Path(tmp_dir) / "mypy", mypy_vendor)
        (mypy_vendor / ".version").write_text(f"{version}: {git_hash}")
    print("Updating custom typeshed")
    update_puya_typeshed(mypy_vendor / "typeshed", puya_src_dir / "_typeshed")

def update_puya_typeshed(mypy_typeshed: Path, puya_typeshed: Path) -> None:
    shutil.rmtree(puya_typeshed, ignore_errors=True)

    stubs = Path("stubs")
    stdlib = Path("stdlib")
    relative_to_copy = [
        # hard coded in mpyy/modulefinder.py, minimum requirements for mypy
        stubs / "mypy-extensions" / "mypy_extensions.pyi",
        stdlib / "VERSIONS",
        # hard coded in mpyy/build.py, minimum requirements for mypy
        stdlib / "builtins.pyi",
        stdlib / "typing.pyi",
        stdlib / "types.pyi",
        stdlib / "typing_extensions.pyi",
        stdlib / "_typeshed" / "__init__.pyi",
        stdlib / "_collections_abc.pyi",
        stdlib / "collections" / "abc.pyi",
        stdlib / "sys" / "__init__.pyi",
        stdlib / "abc.pyi",
        # needed for puyapy
        # stdlib / "enum.pyi"
    ]

    (puya_typeshed / stdlib).mkdir(exist_ok=True, parents=True)
    (puya_typeshed / stubs).mkdir(exist_ok=True, parents=True)
    for relative in relative_to_copy:
        copy_src = mypy_typeshed / relative
        copy_dst = puya_typeshed / relative
        if copy_src.is_dir():
            shutil.copytree(copy_src, copy_dst)
        else:
            copy_dst.parent.mkdir(exist_ok=True, parents=True)
            shutil.copy(copy_src, copy_dst)
    (puya_typeshed / stdlib / "collections" / "__init__.pyi").touch()
    (puya_typeshed / "README.md").write_text(TYPESHED_README)

if __name__ == "__main__":
    if len(sys.argv) > 1:
        vendor_mypy(version=sys.argv[1])
    else:
        print("Usage: python vendor_mypy.py <version>")
        print("e.g. python vendor_mypy.py v1.5.0")


from algopy import (
    Bytes,
    UInt64,
    subroutine,
    urange,
)
from algopy.op import (
    btoi,
    bzero,
    extract,
    extract_uint16,
    getbit,
    itob,
    replace,
    select_uint64,
    setbit_bytes,
    substring,
)

UINT16_SIZE = 2
UINT64_SIZE = 8
UINT16_OFFSET = UINT64_SIZE - UINT16_SIZE

@subroutine
def dynamic_array_pop_bit(array: Bytes) -> tuple[Bytes, Bytes]:
    """
    Pop the last item from an arc4 dynamic array of arc4 encoded boolean items

    array: The bytes for the source array

    returns: tuple of (The popped item, The updated bytes for the source array)
    """
    array_length = extract_uint16(array, 0)
    length_minus_1 = array_length - 1
    result = replace(array, 0, extract(itob(length_minus_1), UINT16_OFFSET, 0))
    popped_location = length_minus_1 + UINT16_SIZE * 8
    popped = setbit_bytes(b"\x00", 0, getbit(result, popped_location))
    result = setbit_bytes(result, popped_location, 0)
    result = substring(result, 0, UINT16_SIZE + ((length_minus_1 + 7) // 8))
    return popped, result

@subroutine
def dynamic_array_pop_fixed_size(array: Bytes, fixed_byte_size: UInt64) -> tuple[Bytes, Bytes]:
    """
    Pop the last item from an arc4 dynamic array of fixed sized items

    array: The bytes for the source array

    returns: tuple of (The popped item, The updated bytes for the source array)
    """
    array_length = extract_uint16(array, 0)
    length_minus_1 = array_length - 1
    result = replace(array, 0, extract(itob(length_minus_1), UINT16_OFFSET, 0))
    item_location = result.length - fixed_byte_size
    popped = extract(result, item_location, fixed_byte_size)
    result = substring(result, 0, item_location)
    return popped, result

@subroutine
def dynamic_array_pop_byte_length_head(array: Bytes) -> tuple[Bytes, Bytes]:
    """
    Pop the last item from an arc4 dynamic array of items that are prefixed with their length in
    bytes, e.g. arc4.String, arc4.DynamicBytes

    source: The bytes for the source array

    returns: tuple of (The popped item, The updated bytes for the source array)
    """
    array_length = extract_uint16(array, 0)
    length_minus_1 = array_length - 1
    popped_header_offset = length_minus_1 * UINT16_SIZE
    head_and_tail = extract(array, UINT16_SIZE, 0)
    popped_offset = extract_uint16(head_and_tail, popped_header_offset)

    popped = substring(head_and_tail, popped_offset, head_and_tail.length)
    head_and_tail = substring(head_and_tail, 0, popped_header_offset) + substring(
        head_and_tail, popped_header_offset + 2, popped_offset
    )

    updated = extract(
        itob(length_minus_1), UINT16_OFFSET, UINT16_SIZE
    ) + recalculate_head_for_elements_with_byte_length_head(
        array_head_and_tail=head_and_tail, length=length_minus_1, start_at_index=UInt64(0)
    )

    return popped, updated

@subroutine
def dynamic_array_pop_dynamic_element(array: Bytes) -> tuple[Bytes, Bytes]:
    """
    Pop the last item from an arc4 dynamic array of dynamically sized items

    array: The bytes for the source array

    returns: tuple of (The popped item, The updated bytes for the source array)
    """
    array_length = extract_uint16(array, 0)
    length_minus_1 = array_length - 1
    popped_header_offset = length_minus_1 * UINT16_SIZE
    head_and_tail = extract(array, UINT16_SIZE, 0)
    popped_offset = extract_uint16(head_and_tail, popped_header_offset)

    popped = substring(head_and_tail, popped_offset, head_and_tail.length)

    new_head = Bytes()
    for head_offset in urange(0, length_minus_1 * UINT16_SIZE, UINT16_SIZE):
        item_offset = extract_uint16(head_and_tail, head_offset)
        item_offset -= UINT16_SIZE
        new_head += extract(itob(item_offset), UINT16_OFFSET, UINT16_SIZE)

    updated = (
        extract(itob(length_minus_1), UINT16_OFFSET, UINT16_SIZE)
        + new_head
        + substring(head_and_tail, popped_header_offset + UINT16_SIZE, popped_offset)
    )

    return popped, updated

@subroutine
def dynamic_array_concat_bits(
    *, array: Bytes, new_items_bytes: Bytes, new_items_count: UInt64, is_packed: bool
) -> Bytes:
    """
    Concat data to an arc4 dynamic array of arc4 encoded boolean values

    array: The bytes for the source array
    new_items_bytes: Either the data portion of an arc4 packed array of booleans
                        or
                     a sparse array of concatenated arc4 booleans
    new_items_count: The count of new items being added
    is_packed: True if new_items_bytes represents a packed array, else False

    returns: The updated bytes for the source array
    """
    array_length = extract_uint16(array, 0)
    new_array_length = array_length + new_items_count
    new_array_length_b = extract(itob(new_array_length), UINT16_OFFSET, 0)
    result = replace(array, 0, new_array_length_b)
    current_bytes = (array_length + 7) // 8
    required_bytes = (new_array_length + 7) // 8
    if current_bytes < required_bytes:
        result += bzero(required_bytes - current_bytes)

    write_offset = array_length + 8 * UINT16_SIZE
    for i in urange(0, new_items_count, UInt64(1) if is_packed else UInt64(8)):
        result = setbit_bytes(result, write_offset, getbit(new_items_bytes, i))
        write_offset += 1

    return result

@subroutine
def dynamic_array_concat_byte_length_head(
    array: Bytes, new_items_bytes: Bytes, new_items_count: UInt64
) -> Bytes:
    """
    Replace a single item in an arc4 dynamic array of items that are prefixed with
    their byte length

    array: The bytes of the source array
    new_items_bytes: The bytes for all new items, concatenated
    new_items_counts: The count of new items being added

    returns: The updated bytes for the source array
    """
    array_length = extract_uint16(array, 0)
    new_length = array_length + new_items_count
    header_end = array_length * UINT16_SIZE + 2

    return extract(
        itob(new_length), UINT16_OFFSET, UINT16_SIZE
    ) + recalculate_head_for_elements_with_byte_length_head(
        array_head_and_tail=(
            substring(array, 2, header_end)
            + bzero(new_items_count * UINT16_SIZE)
            + substring(array, header_end, array.length)
            + new_items_bytes
        ),
        length=new_length,
        start_at_index=UInt64(0),
    )

@subroutine
def dynamic_array_concat_dynamic_element(
    *,
    array_items_count: UInt64,
    array_head_and_tail: Bytes,
    new_items_count: UInt64,
    new_head_and_tail: Bytes,
) -> Bytes:
    new_head = Bytes()
    item_offset_adjustment = new_items_count * UINT16_SIZE
    for head_offset in urange(0, array_items_count * UINT16_SIZE, UINT16_SIZE):
        item_offset = extract_uint16(array_head_and_tail, head_offset)
        new_head += extract(itob(item_offset_adjustment + item_offset), UINT16_OFFSET, UINT16_SIZE)

    item_offset_adjustment = array_head_and_tail.length
    for head_offset in urange(0, new_items_count * UINT16_SIZE, UINT16_SIZE):
        item_offset = extract_uint16(new_head_and_tail, head_offset)
        new_head += extract(itob(item_offset_adjustment + item_offset), UINT16_OFFSET, UINT16_SIZE)
    return (
        extract(itob(array_items_count + new_items_count), UINT16_OFFSET, UINT16_SIZE)
        + new_head
        + substring(
            array_head_and_tail, array_items_count * UINT16_SIZE, array_head_and_tail.length
        )
        + substring(new_head_and_tail, new_items_count * UINT16_SIZE, new_head_and_tail.length)
    )

@subroutine
def dynamic_array_replace_byte_length_head(array: Bytes, new_item: Bytes, index: UInt64) -> Bytes:
    """
    Replace a single item in an arc4 dynamic array of items that are prefixed with
    their byte length

    array: The bytes of the source array
    new_item: The bytes of the new item to be inserted
    index: The index to insert the new item at
    array_length: The length of the array

    returns: The updated bytes for the source array
    """
    size_b = substring(array, 0, UINT16_SIZE)
    array_length = btoi(size_b)
    return size_b + static_array_replace_byte_length_head(
        array_head_and_tail=extract(array, UINT16_SIZE, 0),
        new_item=new_item,
        index=index,
        array_length=array_length,
    )

@subroutine
def dynamic_array_replace_dynamic_element(source: Bytes, new_item: Bytes, index: UInt64) -> Bytes:
    size_b = substring(source, 0, UINT16_SIZE)
    array_length = btoi(size_b)
    return size_b + static_array_replace_dynamic_element(
        array_head_and_tail=extract(source, UINT16_SIZE, 0),
        new_item=new_item,
        index=index,
        array_length=array_length,
    )

@subroutine
def static_array_replace_dynamic_element(
    *, array_head_and_tail: Bytes, new_item: Bytes, index: UInt64, array_length: UInt64
) -> Bytes:
    original_offset = extract_uint16(array_head_and_tail, index * 2)
    next_item_offset = extract_uint16(array_head_and_tail, (index + 1) * 2)
    end_of_tail = array_head_and_tail.length
    is_before_end = array_length - index - 1
    end_offset = select_uint64(end_of_tail, next_item_offset, is_before_end)

    original_item_length = end_offset - original_offset
    new_item_length = new_item.length
    new_head_and_tail = (
        substring(array_head_and_tail, 0, original_offset)
        + new_item
        + substring(array_head_and_tail, end_offset, end_of_tail)
    )
    for head_offset in urange((index + 1) * 2, array_length * 2, 2):
        tail_offset = extract_uint16(new_head_and_tail, head_offset)
        tail_offset += new_item_length
        tail_offset -= original_item_length
        tail_offset_bytes = extract(itob(tail_offset), UINT16_OFFSET, UINT16_SIZE)
        new_head_and_tail = replace(new_head_and_tail, head_offset, tail_offset_bytes)
    return new_head_and_tail

@subroutine
def static_array_replace_byte_length_head(
    array_head_and_tail: Bytes, new_item: Bytes, index: UInt64, array_length: UInt64
) -> Bytes:
    """
    Replace a single item in an arc4 dynamic array of items that are prefixed with
    their byte length

    array_head_and_tail: The head and tail bytes of the source array
    new_item: The bytes of the new item to be inserted
    index: The index to insert the new item at
    array_length: The length of the array

    returns: The updated bytes for the source array
    """
    assert index < array_length, "Index out of bounds"
    offset_for_index = extract_uint16(array_head_and_tail, index * UINT16_SIZE)
    old_item_length = extract_uint16(array_head_and_tail, offset_for_index)
    old_item_end = offset_for_index + old_item_length + UINT16_SIZE
    return recalculate_head_for_elements_with_byte_length_head(
        array_head_and_tail=substring(array_head_and_tail, 0, offset_for_index)
        + new_item
        + substring(array_head_and_tail, old_item_end, array_head_and_tail.length),
        length=array_length,
        start_at_index=index,
    )

@subroutine
def recalculate_head_for_elements_with_byte_length_head(
    array_head_and_tail: Bytes, length: UInt64, start_at_index: UInt64
) -> Bytes:
    """
    Recalculates the offset values of an arc4 static array, where each item's head consists of
    its length in bytes as uint16

    array_data: The static array data
    length: The length of the static array
    start_at_index: Optionally start at a non-zero index for performance optimisation. The offset
                    at this index is assumed to be correct if start_at_index is not 0

    returns: The updated bytes for the source array
    """
    tail_offset = select_uint64(
        length * UINT16_SIZE,
        extract_uint16(array_head_and_tail, start_at_index * UINT16_SIZE),
        start_at_index,  # use length * UINT16_SIZE if 0 otherwise inspect head
    )

    for head_offset in urange(start_at_index * UINT16_SIZE, length * UINT16_SIZE, UINT16_SIZE):
        tail_offset_bytes = extract(itob(tail_offset), UINT16_OFFSET, UINT16_SIZE)
        array_head_and_tail = replace(array_head_and_tail, head_offset, tail_offset_bytes)
        tail_offset += extract_uint16(array_head_and_tail, tail_offset) + UINT16_SIZE
        head_offset += UINT16_SIZE
    return array_head_and_tail


from algopy import Bytes, UInt64, op, subroutine

@subroutine
def is_substring(item: Bytes, sequence: Bytes) -> bool:
    """
    Search for a shorter string in a larger one.
    """

    start = UInt64(0)
    while start + item.length <= sequence.length:
        if item == op.substring(sequence, start, start + item.length):
            return True
        start += 1
    return False


from algopy import (
    Bytes,
    OnCompleteAction,
    OpUpFeeSource,
    TransactionType,
    UInt64,
    op,
    subroutine,
)

@subroutine
def ensure_budget(required_budget: UInt64, fee_source: OpUpFeeSource) -> None:
    # A budget buffer is necessary to deal with an edge case of ensure_budget():
    #   if the current budget is equal to or only slightly higher than the
    #   required budget then it's possible for ensure_budget() to return with a
    #   current budget less than the required budget. The buffer prevents this
    #   from being the case.
    required_budget_with_buffer = required_budget + 10
    while required_budget_with_buffer > op.Global.opcode_budget():
        op.ITxnCreate.begin()
        op.ITxnCreate.set_type_enum(TransactionType.ApplicationCall)
        op.ITxnCreate.set_on_completion(OnCompleteAction.DeleteApplication)
        op.ITxnCreate.set_approval_program(Bytes.from_hex("068101"))
        op.ITxnCreate.set_clear_state_program(Bytes.from_hex("068101"))
        match fee_source:
            case OpUpFeeSource.GroupCredit:
                op.ITxnCreate.set_fee(0)
            case OpUpFeeSource.AppAccount:
                op.ITxnCreate.set_fee(op.Global.min_txn_fee)
            # case OpUpFeeSource.Any:
            # any means no fee set
        op.ITxnCreate.submit()


import argparse
from importlib.metadata import version
from pathlib import Path

import attrs

from puya.log import LogFormat, LogLevel, configure_logging
from puya.main import main

@attrs.define(kw_only=True)
class _PuyaCLIArgs:
    options: Path | None = None
    awst: Path | None = None
    source_annotations: Path | None = None
    log_level: LogLevel = LogLevel.info
    log_format: LogFormat = LogFormat.default

def cli() -> None:
    parser = argparse.ArgumentParser(
        prog="puya", formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    # TODO: use version of puya instead once package is split
    parser.add_argument("--version", action="version", version=f"%(prog)s {version('puyapy')}")
    parser.add_argument(
        "--log-level", type=LogLevel.from_string, choices=list(LogLevel), default=LogLevel.info
    )
    parser.add_argument(
        "--log-format",
        type=LogFormat.from_string,
        choices=list(LogFormat),
        default=LogFormat.default,
    )
    parser.add_argument("--options", type=Path, required=True)
    parser.add_argument("--awst", type=Path, required=True)
    parser.add_argument("--source-annotations", type=Path)
    parsed_args = _PuyaCLIArgs()
    parser.parse_args(namespace=parsed_args)
    configure_logging(min_log_level=parsed_args.log_level, log_format=parsed_args.log_format)

    assert parsed_args.options
    options_json = parsed_args.options.read_text("utf8")
    assert parsed_args.awst
    awst_json = parsed_args.awst.read_text("utf8")
    source_annotations_json = None
    if parsed_args.source_annotations:
        source_annotations_json = parsed_args.source_annotations.read_text("utf8")
    main(
        options_json=options_json,
        awst_json=awst_json,
        source_annotations_json=source_annotations_json,
    )

if __name__ == "__main__":
    cli()


ZERO_ADDRESS = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAY5HFKQ"
ENCODED_ADDRESS_LENGTH = 58
PUBLIC_KEY_HASH_LENGTH = 32
ADDRESS_CHECKSUM_LENGTH = 4
MAX_BIGUINT_BITS = 512
MAX_UINT64 = 2**64 - 1
MAX_BIGUINT_BYTES = MAX_BIGUINT_BITS // 8
MAX_BYTES_LENGTH = 4096
MAX_SCRATCH_SLOT_NUMBER = 255
MAX_GLOBAL_STATE_KEYS = 64
MAX_LOCAL_STATE_KEYS = 16
MAX_STATE_KEY_LENGTH = 64
MIN_BOX_KEY_LENGTH = 1
MAX_BOX_KEY_LENGTH = 64
MAX_TRANSACTION_GROUP_SIZE = 16
MAX_APP_PAGE_SIZE = 2048
HASH_PREFIX_PROGRAM = b"Program"
"""Represents the prefix added to a program before hashing e.g. for a LogicSigs address"""
# Which language versions does this version of puya support targeting
# This will typically just be the current mainnet version and potentially the vNext if it doesn't
# contain breaking changes
SUPPORTED_AVM_VERSIONS = [10, 11, 12]
# Which language version is currently deployed to mainnet
MAINNET_AVM_VERSION = 10


import base64
import json
import typing
from collections.abc import Collection, Mapping, Sequence

from puya import (
    artifact_metadata as md,
    log,
)
from puya.avm import OnCompletionAction
from puya.awst.nodes import (
    AppStorageKind,
    ARC4CreateOption,
)
from puya.parse import SourceLocation

OCA_ARC32_MAPPING = {
    OnCompletionAction.NoOp: "no_op",
    OnCompletionAction.OptIn: "opt_in",
    OnCompletionAction.CloseOut: "close_out",
    OnCompletionAction.ClearState: "clear_state",
    OnCompletionAction.UpdateApplication: "update_application",
    OnCompletionAction.DeleteApplication: "delete_application",
}

JSONValue: typing.TypeAlias = "str | int | float | bool | None | Sequence[JSONValue] | JSONDict"
JSONDict: typing.TypeAlias = Mapping[str, "JSONValue"]

logger = log.get_logger(__name__)

def _encode_source(teal_text: str) -> str:
    return base64.b64encode(teal_text.encode()).decode("utf-8")

def _encode_schema_declaration(state: md.ContractState) -> JSONDict:
    return {
        "type": state.storage_type.name,
        "key": state.key_or_prefix.decode("utf-8"),  # TODO: support not utf8 keys?
        "descr": state.description,
    }

def _encode_state_declaration(state: md.StateTotals) -> JSONDict:
    return {
        "global": {
            "num_byte_slices": state.global_bytes,
            "num_uints": state.global_uints,
        },
        "local": {
            "num_byte_slices": state.local_bytes,
            "num_uints": state.local_uints,
        },
    }

def _encode_schema(state: Collection[md.ContractState]) -> JSONDict:
    return {
        "declared": {
            s.name: _encode_schema_declaration(s) for s in sorted(state, key=lambda s: s.name)
        },
        "reserved": {},  # TODO?
    }

def _encode_call_config(method: md.ARC4Method) -> JSONDict:
    match method.create:
        case ARC4CreateOption.require:
            call_config = "CREATE"
        case ARC4CreateOption.allow:
            call_config = "ALL"
        case ARC4CreateOption.disallow:
            call_config = "CALL"
        case never:
            typing.assert_never(never)
    return {OCA_ARC32_MAPPING[oca]: call_config for oca in method.allowed_completion_types}

def _encode_bare_method_configs(methods: Sequence[md.ARC4BareMethod]) -> JSONDict:
    result: dict[str, JSONValue] = {}
    for method in methods:
        result.update(**_encode_call_config(method))
    return result

def _get_signature(method: md.ARC4ABIMethod) -> str:
    return f"{method.name}({','.join(m.type_ for m in method.args)}){method.returns.type_}"

def _encode_default_arg(arg: md.ARC4MethodArg, loc: SourceLocation | None) -> JSONDict | None:
    match arg.client_default:
        case None:
            return None
        case md.MethodArgDefaultConstant(data=constant_data, type_=constant_arc4_type):
            if constant_arc4_type == "string":
                string = constant_data[2:].decode("utf8")
                return {"source": "constant", "data": string}
            elif constant_arc4_type.startswith("uint"):
                number = int.from_bytes(constant_data, signed=False)
                return {"source": "constant", "data": number}
            else:
                logger.warning(
                    f"parameter {arg.name!r} has unsupported default constant type for ARC-32",
                    location=loc,
                )
                return None
        case md.MethodArgDefaultFromMethod(
            name=method_name, return_type=return_type, readonly=readonly
        ):
            return {
                "source": "abi-method",
                "data": {
                    "name": method_name,
                    "args": [],
                    "readonly": readonly,  # ARC-22
                    "returns": {"type": return_type},
                },
            }
        case md.MethodArgDefaultFromState(kind=kind, key=key):
            match kind:
                case AppStorageKind.app_global:
                    source_name = "global-state"
                case AppStorageKind.account_local:
                    source_name = "local-state"
                case AppStorageKind.box:
                    logger.error(
                        "default argument from box storage are not supported by ARC-32",
                        location=loc,
                    )
                case unexpected:
                    typing.assert_never(unexpected)
            return {
                "source": source_name,
                # TODO: handle non utf-8 bytes
                "data": key.decode("utf-8"),
            }

def _encode_arc32_method_hint(metadata: md.ContractMetaData, method: md.ARC4ABIMethod) -> JSONDict:
    structs = {a.name: metadata.structs[a.struct] for a in method.args if a.struct}
    if method.returns.struct:
        structs["output"] = metadata.structs[method.returns.struct]
    default_arguments = {
        arg.name: default
        for arg in method.args
        if (default := _encode_default_arg(arg, method.config_location)) is not None
    }
    return {
        # deprecated by ARC-22
        "read_only": True if method.readonly else None,
        "default_arguments": default_arguments or None,
        "call_config": _encode_call_config(method),
        "structs": _encode_arc32_method_structs(structs),
    }

def _encode_arc32_method_structs(structs: Mapping[str, md.ARC4Struct]) -> JSONDict | None:
    if len(structs):
        return {
            struct_purpose: {
                "name": struct_def.name,
                "elements": [[f.name, f.type] for f in struct_def.fields],
            }
            for struct_purpose, struct_def in structs.items()
        }
    return None

def _encode_arc32_hints(
    metadata: md.ContractMetaData, methods: list[md.ARC4ABIMethod]
) -> JSONDict:
    return {
        _get_signature(method): _encode_arc32_method_hint(metadata, method) for method in methods
    }

def _encode_abi_method(method: md.ARC4ABIMethod) -> JSONDict:
    return {
        "name": method.name,
        "args": [
            {
                "type": arg.type_,
                "name": arg.name,
                "desc": arg.desc,
            }
            for arg in method.args
        ],
        "readonly": method.readonly,  # ARC-22
        "returns": {
            "type": method.returns.type_,
            "desc": method.returns.desc,
        },
        "desc": method.desc,
    }

def _encode_arc4_contract(
    name: str, desc: str | None, methods: Sequence[md.ARC4ABIMethod]
) -> JSONDict:
    return {
        "name": name,
        "desc": desc,
        "methods": [_encode_abi_method(m) for m in methods],
        "networks": {},
    }

def _filter_none(value: JSONDict) -> JSONValue:
    if isinstance(value, dict):
        return {k: _filter_none(v) for k, v in value.items() if v is not None}
    if isinstance(value, list):
        return list(map(_filter_none, value))
    return value

def create_arc32_json(
    approval_program: str, clear_program: str, metadata: md.ContractMetaData
) -> str:
    bare_methods = [m for m in metadata.arc4_methods if isinstance(m, md.ARC4BareMethod)]
    abi_methods = [m for m in metadata.arc4_methods if isinstance(m, md.ARC4ABIMethod)]
    app_spec = {
        "hints": _encode_arc32_hints(metadata, abi_methods),
        "source": {
            "approval": _encode_source(approval_program),
            "clear": _encode_source(clear_program),
        },
        "state": _encode_state_declaration(metadata.state_totals),
        "schema": {
            "global": _encode_schema(metadata.global_state.values()),
            "local": _encode_schema(metadata.local_state.values()),
        },
        "contract": _encode_arc4_contract(metadata.name, metadata.description, abi_methods),
        "bare_call_config": _encode_bare_method_configs(bare_methods),
    }
    return json.dumps(_filter_none(app_spec), indent=4)