import base64
import itertools
import typing
from collections import defaultdict
from collections.abc import Iterable, Mapping, Sequence
from importlib.metadata import version as metadata_version

from cattrs.preconf.json import make_converter
from packaging import version

from puya import (
    arc56_models as models,
    log,
)
from puya.artifact_metadata import (
    ARC4ABIMethod,
    ARC4BareMethod,
    ARC4Struct,
    ContractMetaData,
    ContractState,
    MethodArgDefault,
    MethodArgDefaultConstant,
    MethodArgDefaultFromMethod,
    MethodArgDefaultFromState,
)
from puya.awst.nodes import (
    AppStorageKind,
    ARC4CreateOption,
)
from puya.compilation_artifacts import CompiledProgram, DebugInfo
from puya.utils import unique

# TODO: use puya once the backend is shipped as separate package
_ALGOPY_VERSION = version.parse(metadata_version("puyapy"))
logger = log.get_logger(__name__)

def create_arc56_json(
    *,
    metadata: ContractMetaData,
    approval_program: CompiledProgram,
    clear_program: CompiledProgram,
    template_prefix: str,
) -> str:
    assert approval_program.debug_info is not None
    assert clear_program.debug_info is not None

    converter = make_converter(omit_if_default=True)
    bare_methods = [m for m in metadata.arc4_methods if isinstance(m, ARC4BareMethod)]
    abi_methods = [m for m in metadata.arc4_methods if isinstance(m, ARC4ABIMethod)]

    # use shorter name for structs unless there is a collision
    aliases = _StructAliases(metadata.structs.values())

    schema = metadata.state_totals
    app_spec = models.Contract(
        arcs=(22, 28),
        name=metadata.name,
        desc=metadata.description,
        networks={},
        structs={
            aliases.resolve(n): [
                models.StructField(
                    name=e.name,
                    type=aliases.resolve(e.struct) or e.type,
                )
                for e in s.fields
            ]
            for n, s in metadata.structs.items()
        },
        methods=[
            models.Method(
                name=m.name,
                desc=m.desc,
                args=[
                    models.MethodArg(
                        type=a.type_,
                        name=a.name,
                        desc=a.desc,
                        struct=aliases.resolve(a.struct),
                        defaultValue=_encode_default_arg(a.client_default),
                    )
                    for a in m.args
                ],
                returns=models.MethodReturns(
                    type=m.returns.type_,
                    desc=m.returns.desc,
                    struct=aliases.resolve(m.returns.struct),
                ),
                actions=_method_actions(m),
                readonly=m.readonly,
                events=[_struct_to_event(aliases, struct) for struct in m.events],
                # left for users to fill in for now
                recommendations=models.MethodRecommendations(
                    innerTransactionCount=None,
                    boxes=None,
                    accounts=None,
                    apps=None,
                    assets=None,
                ),
            )
            for m in abi_methods
        ],
        state=models.ContractState(
            schema={
                "global": {"ints": schema.global_uints, "bytes": schema.global_bytes},
                "local": {"ints": schema.local_uints, "bytes": schema.local_bytes},
            },
            keys={
                "global": _storage_keys(aliases, metadata.global_state),
                "local": _storage_keys(aliases, metadata.local_state),
                "box": _storage_keys(aliases, metadata.boxes),
            },
            maps={
                # note: at present there is no way of defining global/local maps
                "global": _storage_maps(aliases, metadata.global_state),
                "local": _storage_maps(aliases, metadata.local_state),
                "box": _storage_maps(aliases, metadata.boxes),
            },
        ),
        bareActions=_combine_actions(list(map(_method_actions, bare_methods))),
        sourceInfo={
            "approval": models.ProgramSourceInfo(
                sourceInfo=_get_source_info(approval_program.debug_info),
                pcOffsetMethod="cblocks" if approval_program.debug_info.op_pc_offset else "none",
            ),
            "clear": models.ProgramSourceInfo(
                sourceInfo=_get_source_info(clear_program.debug_info),
                pcOffsetMethod="cblocks" if clear_program.debug_info.op_pc_offset else "none",
            ),
        },
        source={
            "approval": _encode_str(approval_program.teal_src),
            "clear": _encode_str(clear_program.teal_src),
        },
        byteCode=(
            {
                "approval": _encode_bytes(approval_program.bytecode),
                "clear": _encode_bytes(clear_program.bytecode),
            }
            if approval_program.bytecode and clear_program.bytecode
            else None
        ),
        compilerInfo=(
            _compiler_info() if approval_program.bytecode and clear_program.bytecode else None
        ),
        events=[
            _struct_to_event(aliases, struct)
            for struct in unique(
                e for m in metadata.arc4_methods if isinstance(m, ARC4ABIMethod) for e in m.events
            )
        ],
        templateVariables={
            name.removeprefix(template_prefix): models.TemplateVariable(
                type=aliases.resolve(metadata.template_variable_types[name]),
                value=(
                    _encode_bytes(value.to_bytes(length=8) if isinstance(value, int) else value)
                    if value is not None
                    else None
                ),
            )
            for program in (approval_program, clear_program)
            for name, value in program.template_variables.items()
        },
        # TODO: provide a way for contracts to declare "public" scratch vars
        scratchVariables=None,
    )
    return converter.dumps(app_spec, indent=4)

def _get_source_info(debug_info: DebugInfo) -> Sequence[models.SourceInfo]:
    errors = defaultdict[str, list[int]](list)
    for pc, event in debug_info.pc_events.items():
        if error := event.get("error"):
            errors[error].append(pc)
    return [
        models.SourceInfo(
            pc=errors[error],
            errorMessage=error,
        )
        for error in sorted(errors)
    ]

class _StructAliases:
    def __init__(self, structs: Iterable[ARC4Struct]) -> None:
        alias_to_fullname = dict[str, str]()
        for struct in structs:
            alias = (
                struct.fullname
                if struct.name in alias_to_fullname or struct.name in models.AVMType
                else struct.name
            )
            alias_to_fullname[alias] = struct.fullname
        self.aliases = {v: k for k, v in alias_to_fullname.items()}

    @typing.overload
    def resolve(self, struct: str) -> str: ...

    @typing.overload
    def resolve(self, struct: None) -> None: ...

    def resolve(self, struct: str | None) -> str | None:
        if struct is None:
            return None
        return self.aliases.get(struct, struct)

def _struct_to_event(structs: _StructAliases, struct: ARC4Struct) -> models.Event:
    return models.Event(
        name=structs.resolve(struct.name),
        desc=struct.desc,
        args=[
            models.EventArg(
                name=f.name,
                type=f.type,
                struct=structs.resolve(f.struct),
            )
            for f in struct.fields
        ],
    )

def _storage_keys(
    structs: _StructAliases, state: Mapping[str, ContractState]
) -> models.StorageKeys:
    return {
        n: models.StorageKey(
            desc=m.description,
            keyType=structs.resolve(m.arc56_key_type),
            valueType=structs.resolve(m.arc56_value_type),
            key=_encode_bytes(m.key_or_prefix),
        )
        for n, m in state.items()
        if not m.is_map
    }

def _storage_maps(
    structs: _StructAliases, state: Mapping[str, ContractState]
) -> models.StorageMaps:
    return {
        n: models.StorageMap(
            desc=m.description,
            keyType=structs.resolve(m.arc56_key_type),
            valueType=structs.resolve(m.arc56_value_type),
            prefix=_encode_bytes(m.key_or_prefix),
        )
        for n, m in state.items()
        if m.is_map
    }

def _method_actions(method: ARC4BareMethod | ARC4ABIMethod) -> models.MethodActions:
    return models.MethodActions(
        create=[
            oca.name
            for oca in method.allowed_completion_types
            if method.create != ARC4CreateOption.disallow and allowed_create_oca(oca.name)
        ],
        call=[
            oca.name
            for oca in method.allowed_completion_types
            if method.create != ARC4CreateOption.require and allowed_call_oca(oca.name)
        ],
    )

def _encode_default_arg(default: MethodArgDefault | None) -> models.MethodArgDefaultValue | None:
    match default:
        case None:
            return None
        case MethodArgDefaultConstant(data=data, type_=type_string):
            return models.MethodArgDefaultValue(
                data=_encode_bytes(data),
                type=type_string,
                source=models.DefaultValueSource.literal,
            )
        case MethodArgDefaultFromState(key=key, kind=kind, key_type=key_type):
            match kind:
                case AppStorageKind.account_local:
                    source = models.DefaultValueSource.local_
                case AppStorageKind.app_global:
                    source = models.DefaultValueSource.global_
                case AppStorageKind.box:
                    source = models.DefaultValueSource.box
                case unexpected:
                    typing.assert_never(unexpected)
            return models.MethodArgDefaultValue(
                data=_encode_bytes(key), type=key_type, source=source
            )
        case MethodArgDefaultFromMethod(signature=signature):
            return models.MethodArgDefaultValue(
                data=signature,
                source=models.DefaultValueSource.method,
            )
        case unexpected:
            typing.assert_never(unexpected)

def _combine_actions(actions: Sequence[models.MethodActions]) -> models.MethodActions:
    return models.MethodActions(
        create=sorted(set(itertools.chain.from_iterable(a.create for a in actions))),
        call=sorted(set(itertools.chain.from_iterable(a.call for a in actions))),
    )

def allowed_create_oca(
    oca: str,
) -> typing.TypeGuard[typing.Literal["NoOp", "OptIn", "DeleteApplication"]]:
    return oca in ("NoOp", "OptIn", "DeleteApplication")

def allowed_call_oca(
    oca: str,
) -> typing.TypeGuard[
    typing.Literal["NoOp", "OptIn", "CloseOut", "UpdateApplication", "DeleteApplication"]
]:
    return oca in ("NoOp", "OptIn", "CloseOut", "UpdateApplication", "DeleteApplication")

def _encode_str(value: str) -> str:
    return _encode_bytes(value.encode("utf8"))

def _encode_bytes(value: bytes) -> str:
    return base64.b64encode(value).decode("utf-8")

def _compiler_info() -> models.CompilerInfo:
    return models.CompilerInfo(
        compiler="puya",
        compilerVersion=models.CompilerVersion(
            major=_ALGOPY_VERSION.major,
            minor=_ALGOPY_VERSION.minor,
            patch=_ALGOPY_VERSION.micro,
            commitHash=None,
        ),
    )


# ruff: noqa: N815
import enum
import typing
from collections.abc import Mapping, Sequence

import attrs

ABIType = str
"""An ABI-encoded type"""
StructName = str
"""The name of a defined struct"""
ProgramType = typing.Literal["approval", "clear"]

class AVMType(enum.StrEnum):
    """A native AVM type"""

    bytes = "AVMBytes"
    """Raw byteslice without the length prefixed that is specified in ARC-4"""
    string = "AVMString"
    """A utf-8 string without the length prefix that is specified in ARC-4"""
    uint64 = "AVMUint64"
    """A 64-bit unsigned integer"""

@attrs.frozen
class SourceInfo:
    pc: Sequence[int]
    """The program counter value(s). Could be offset if pcOffsetMethod is not 'none'"""
    errorMessage: str
    """A human-readable string that describes the error when the program fails at the given PC"""

@attrs.frozen
class ProgramSourceInfo:
    sourceInfo: Sequence[SourceInfo]
    """The source information for the program"""
    pcOffsetMethod: typing.Literal["none", "cblocks"]
    """
    How the program counter offset is calculated
    none: The pc values in sourceInfo are not offset
    cblocks: The pc values in sourceInfo are offset by the PC of the first op after the
             last cblock at the top of the program
    """

@attrs.frozen(kw_only=True)
class StorageKey:
    """Describes a single key in app storage"""

    desc: str | None = None
    """Description of what this storage key holds"""
    keyType: ABIType | AVMType | StructName
    """The type of the key"""
    valueType: ABIType | AVMType | StructName
    """The type of the value"""
    key: str
    """The base64-encoded key"""

@attrs.frozen(kw_only=True)
class StorageMap:
    """Describes a mapping of key-value pairs in storage"""

    desc: str | None = None
    """Description of what the key-value pairs in this mapping hold"""
    keyType: ABIType | AVMType | StructName
    """The type of the keys in the map"""
    valueType: ABIType | AVMType | StructName
    """The type of the values in the map"""
    prefix: str | None = None
    """The base64-encoded prefix of the map keys"""

@attrs.frozen
class StructField:
    """Information about a single field in a struct"""

    name: str
    """The name of the struct field"""
    type: ABIType | StructName | Sequence["StructField"]
    """The type of the struct field's value"""

@attrs.frozen
class EventArg:
    type: ABIType
    """
    The type of the argument.
    The `struct` field should also be checked to determine if this arg is a struct.
    """
    name: str | None = None
    """Optional, user-friendly name for the argument"""
    desc: str | None = None
    """Optional, user-friendly description for the argument"""
    struct: StructName | None = None
    """
    If the type is a struct, the name of the struct
    Note: this is a separate field to maintain backwards compatability with ARC-23
    """

@attrs.frozen(kw_only=True)
class Event:
    name: str
    """The name of the event"""
    desc: str | None = None
    """Optional, user-friendly description for the event"""
    args: Sequence[EventArg]
    """The arguments of the event, in order"""

class DefaultValueSource(enum.Enum):
    box = "box"
    """The data key signifies the box key to read the value from"""
    global_ = "global"
    """The data key signifies the global state key to read the value from"""
    local_ = "local"
    """The data key signifies the local state key to read the value from (for the sender)"""
    literal = "literal"
    """the value is a literal and should be passed directly as the argument"""
    method = "method"
    """
    The utf8 signature of the method in this contract to call to get the default value.
    If the method has arguments, they all must have default values.
    The method **MUST** be readonly so simulate can be used to get the default value.
    """

@attrs.frozen(kw_only=True)
class MethodArgDefaultValue:
    source: DefaultValueSource
    """Where the default value is coming from"""
    type: ABIType | AVMType | None = None
    """
    How the data is encoded.
    This is the encoding for the data provided here, not the arg type.
    Not relevant if source is method
    """
    data: str
    """Base64 encoded bytes, base64 ARC4 encoded uint64, or UTF-8 method selector"""

@attrs.frozen(kw_only=True)
class MethodArg:
    type: ABIType
    """
    The type of the argument.
    The `struct` field should also be checked to determine if this arg is a struct.
    """
    struct: StructName | None = None
    """
    If the type is a struct, the name of the struct.
    Note: this is a separate field to maintain backwards compatability with ARC-4
    """
    name: str | None = None
    """Optional, user-friendly name for the argument"""
    desc: str | None = None
    """Optional, user-friendly description for the argument"""
    defaultValue: MethodArgDefaultValue | None = None

@attrs.frozen(kw_only=True)
class MethodReturns:
    type: ABIType
    """
    The type of the return value, or "void" to indicate no return value.
    The `struct` field should also be checked to determine if this return value is a struct.
    """
    struct: StructName | None = None
    """
    If the type is a struct, the name of the struct
    """
    desc: str | None = None
    """Optional, user-friendly description for the return value"""

@attrs.frozen
class MethodActions:
    """An action is a combination of call/create and an OnComplete"""

    create: Sequence[typing.Literal["NoOp", "OptIn", "DeleteApplication"]]
    """OnCompletes this method allows when appID === 0"""
    call: Sequence[
        typing.Literal["NoOp", "OptIn", "CloseOut", "UpdateApplication", "DeleteApplication"]
    ]
    """OnCompletes this method allows when appID !== 0"""

@attrs.frozen(kw_only=True)
class MethodBoxRecommendation:
    app: int | None = None
    """The app ID for the box"""
    key: str
    """The base64 encoded box key"""
    readBytes: int
    """The number of bytes being read from the box"""
    writeBytes: int
    """The number of bytes being written to the box"""

@attrs.frozen(kw_only=True)
class MethodRecommendations:
    innerTransactionCount: int | None = None
    """The number of inner transactions the caller should cover the fees for"""
    boxes: MethodBoxRecommendation | None = None
    """Recommended box references to include"""
    accounts: Sequence[str] | None = None
    """Recommended foreign accounts"""
    apps: Sequence[int] | None = None
    """Recommended foreign apps"""
    assets: Sequence[int] | None = None
    """Recommended foreign assets"""

@attrs.frozen(kw_only=True)
class Method:
    """
    Describes a method in the contract.
    This interface is an extension of the interface described in ARC-4
    """

    name: str
    """The name of the method"""
    desc: str | None = None
    """Optional, user-friendly description for the method"""
    args: Sequence[MethodArg]
    """The arguments of the method, in order"""
    returns: MethodReturns
    """Information about the method's return value"""
    actions: MethodActions
    """Allowed actions for this method"""
    readonly: bool
    """If this method does not write anything to the ledger (ARC-22)"""
    events: Sequence[Event] = ()
    """ARC-28 events that MAY be emitted by this method"""
    recommendations: MethodRecommendations | None = None
    """Information that clients can use when calling the method"""

@attrs.frozen
class Network:
    appID: int
    """The app ID of the deployed contract in this network"""

class SchemaSizes(typing.TypedDict):
    ints: int
    bytes: int

ContractSchema = typing.TypedDict("ContractSchema", {"global": SchemaSizes, "local": SchemaSizes})
StorageMaps = Mapping[str, StorageMap]
StorageKeys = Mapping[str, StorageKey]
ContractStorage = typing.TypedDict(
    "ContractStorage", {"global": StorageMaps, "local": StorageMaps, "box": StorageMaps}
)
ContractKeys = typing.TypedDict(
    "ContractKeys", {"global": StorageKeys, "local": StorageKeys, "box": StorageKeys}
)

@attrs.frozen
class ContractState:
    schema: ContractSchema
    """
    Defines the values that should be used for GlobalNumUint, GlobalNumByteSlice, LocalNumUint,
    and LocalNumByteSlice when creating the application
    """
    keys: ContractKeys
    """Mapping of human-readable names to StorageKey objects"""
    maps: ContractStorage
    """Mapping of human-readable names to StorageMap objects"""

@attrs.frozen(kw_only=True)
class CompilerVersion:
    major: int
    minor: int
    patch: int
    commitHash: str | None = None

@attrs.frozen
class CompilerInfo:
    compiler: str
    """The name of the compiler"""
    compilerVersion: CompilerVersion

@attrs.frozen(kw_only=True)
class TemplateVariable:
    type: ABIType | AVMType | StructName
    """The type of the template variable"""
    value: str | None = None
    """If given, the the base64 encoded value used for the given app/program"""

@attrs.frozen
class ScratchVariable:
    slot: int
    type: ABIType | AVMType | StructName

@attrs.frozen(kw_only=True)
class Contract:
    """
    Describes the entire contract.
    This interface is an extension of the interface described in ARC-4
    """

    arcs: Sequence[int] = ()
    """
    The ARCs used and/or supported by this contract.
    All contracts implicitly support ARC-4 and ARC-56
    """
    name: str
    """A user-friendly name for the contract"""
    desc: str | None = None
    """Optional, user-friendly description for the interface"""
    networks: Mapping[str, Network] | None = None
    """
    Optional object listing the contract instances across different networks.
    The key is the base64 genesis hash of the network, and the value contains
    information about the deployed contract in the network indicated by the
    key. A key containing the human-readable name of the network MAY be
    included, but the corresponding genesis hash key MUST also be define
    """
    structs: Mapping[str, Sequence[StructField]]
    """
    Named structs use by the application.
    Each struct field appears in the same order as ABI encoding
    """
    methods: Sequence[Method]
    """All of the methods that the contract implements"""
    state: ContractState | None = None
    bareActions: MethodActions | None = None
    """Supported bare actions for the contract"""
    sourceInfo: Mapping[ProgramType, ProgramSourceInfo] | None = None
    """Information about the TEAL programs"""
    source: Mapping[ProgramType, str] | None = None
    """
    The pre-compiled TEAL that may contain template variables.
    MUST be omitted if included as part of ARC23
    """
    byteCode: Mapping[ProgramType, str] | None = None
    """
    The compiled bytecode for the application.
    MUST be omitted if included as part of ARC23
    """
    compilerInfo: CompilerInfo | None = None
    """
    Information used to get the given byteCode and/or PC values in sourceInfo.
    MUST be given if byteCode or PC values are present
    """
    events: Sequence[Event] | None = None
    """ARC-28 events that MAY be emitted by this contract"""
    templateVariables: Mapping[str, TemplateVariable] | None = None
    """
    A mapping of template variable names as they appear in the teal (not including TMPL_ prefix)
    to their respective types and values (if applicable)
    """
    scratchVariables: Mapping[str, ScratchVariable] | None = None
    """The scratch variables used during runtime"""


import re
import typing
from collections.abc import Sequence

import attrs
from immutabledict import immutabledict

from puya import avm
from puya.avm import AVMType
from puya.awst import nodes as awst_nodes
from puya.parse import SourceLocation
from puya.program_refs import ContractReference, LogicSigReference

@attrs.frozen
class ARC4StructField:
    name: str
    type: str
    struct: str | None

@attrs.frozen(kw_only=True)
class ARC4Struct:
    fullname: str
    desc: str | None = None
    fields: Sequence[ARC4StructField] = attrs.field(
        default=(), converter=tuple[ARC4StructField, ...]
    )

    @property
    def name(self) -> str:
        return re.split(r"\W", self.fullname)[-1]

@attrs.frozen(kw_only=True)
class MethodArgDefaultConstant:
    data: bytes
    type_: str

@attrs.frozen(kw_only=True)
class MethodArgDefaultFromState:
    kind: awst_nodes.AppStorageKind
    key: bytes
    key_type: str

@attrs.frozen(kw_only=True)
class MethodArgDefaultFromMethod:
    name: str
    return_type: str
    readonly: bool

    @property
    def signature(self) -> str:
        return f"{self.name}(){self.return_type}"

MethodArgDefault = (
    MethodArgDefaultConstant | MethodArgDefaultFromState | MethodArgDefaultFromMethod
)

@attrs.frozen
class ARC4MethodArg:
    name: str
    type_: str
    struct: str | None
    desc: str | None = attrs.field(hash=False)
    client_default: MethodArgDefault | None

@attrs.frozen
class ARC4Returns:
    type_: str
    struct: str | None
    desc: str | None = attrs.field(hash=False)

@attrs.frozen(kw_only=True)
class ARC4ABIMethod:
    id: str
    desc: str | None = attrs.field(hash=False)
    args: Sequence[ARC4MethodArg] = attrs.field(converter=tuple[ARC4MethodArg, ...])
    returns: ARC4Returns
    events: Sequence[ARC4Struct] = attrs.field(converter=tuple[ARC4Struct, ...])
    _config: awst_nodes.ARC4ABIMethodConfig

    @property
    def name(self) -> str:
        return self._config.name

    @property
    def allowed_completion_types(self) -> Sequence[avm.OnCompletionAction]:
        return self._config.allowed_completion_types

    @property
    def create(self) -> awst_nodes.ARC4CreateOption:
        return self._config.create

    @property
    def readonly(self) -> bool:
        return self._config.readonly

    @property
    def config_location(self) -> SourceLocation:
        return self._config.source_location

    @property
    def signature(self) -> str:
        return f"{self.name}({','.join(a.type_ for a in self.args)}){self.returns.type_}"

@attrs.frozen(kw_only=True)
class ARC4BareMethod:
    id: str
    desc: str | None = attrs.field(hash=False)
    _config: awst_nodes.ARC4BareMethodConfig

    @property
    def allowed_completion_types(self) -> Sequence[avm.OnCompletionAction]:
        return self._config.allowed_completion_types

    @property
    def create(self) -> awst_nodes.ARC4CreateOption:
        return self._config.create

    @property
    def config_location(self) -> SourceLocation:
        return self._config.source_location

ARC4Method = ARC4BareMethod | ARC4ABIMethod

@attrs.define(eq=False)
class ContractState:
    name: str
    source_location: SourceLocation
    key_or_prefix: bytes
    """Key value as bytes, or prefix if it is a map"""
    arc56_key_type: str
    arc56_value_type: str
    storage_type: typing.Literal[AVMType.uint64, AVMType.bytes]
    description: str | None
    is_map: bool
    """State describes a map"""

@attrs.frozen(kw_only=True)
class LogicSignatureMetaData:
    ref: LogicSigReference
    name: str
    description: str | None

@attrs.frozen
class StateTotals:
    global_uints: int
    local_uints: int
    global_bytes: int
    local_bytes: int

@attrs.frozen(kw_only=True)
class ContractMetaData:
    ref: ContractReference
    name: str
    description: str | None
    global_state: immutabledict[str, ContractState]
    local_state: immutabledict[str, ContractState]
    boxes: immutabledict[str, ContractState]
    state_totals: StateTotals
    arc4_methods: Sequence[ARC4Method]
    structs: immutabledict[str, ARC4Struct]
    template_variable_types: immutabledict[str, str]
    """Mapping of template variable names to their ARC-56 type"""

    @property
    def is_arc4(self) -> bool:
        return bool(self.arc4_methods)  # TODO: should this be an explicit flag instead?


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

import attrs

from puya import log
from puya.errors import CodeError
from puya.ir.models import CompiledContractReference, CompiledLogicSigReference, ModuleArtifact
from puya.ir.visitor import IRTraverser
from puya.parse import SourceLocation
from puya.program_refs import ContractReference, LogicSigReference

logger = log.get_logger(__name__)

@attrs.frozen(eq=False)
class Artifact:
    ir: ModuleArtifact
    depends_on: dict[ContractReference | LogicSigReference, SourceLocation | None] = attrs.field(
        factory=dict
    )

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

    @typing.final
    @property
    def source_location(self) -> SourceLocation:
        return self.ir.source_location

@attrs.define
class ArtifactCompilationSorter(IRTraverser):
    """
    Sorts IR artifacts so that programs that depend on the byte code of other programs
    are processed after their dependencies
    """

    artifacts: Mapping[ContractReference | LogicSigReference, Artifact]
    artifact: Artifact

    @classmethod
    def sort(
        cls,
        all_ir: Sequence[ModuleArtifact],
    ) -> Iterable[Artifact]:
        all_artifacts = {artifact.metadata.ref: Artifact(ir=artifact) for artifact in all_ir}
        artifacts = list(all_artifacts.values())
        for artifact in artifacts:
            reference_collector = cls(
                artifacts=all_artifacts,
                artifact=artifact,
            )
            for program in artifact.ir.all_programs():
                for subroutine in program.subroutines:
                    reference_collector.visit_all_blocks(subroutine.body)
        sorter = graphlib.TopologicalSorter(
            {artifact: [all_artifacts[n] for n in artifact.depends_on] for artifact in artifacts}
        )
        try:
            result = list(sorter.static_order())
        except graphlib.CycleError as ex:
            artifact_cycle: Sequence[Artifact] = ex.args[1]
            *_, before, last = artifact_cycle
            cycle_loc = last.depends_on[before.id]
            programs = " -> ".join(a.id for a in reversed(artifact_cycle))
            raise CodeError(f"cyclical program reference: {programs}", cycle_loc) from None
        return result

    def visit_compiled_contract_reference(self, const: CompiledContractReference) -> None:
        self._check_reference(const.artifact)
        # add unique references with source_location
        # will re-add a reference if current location is None
        if self.artifact.depends_on.get(const.artifact) is None:
            self.artifact.depends_on[const.artifact] = const.source_location

    def visit_compiled_logicsig_reference(self, const: CompiledLogicSigReference) -> None:
        self._check_reference(const.artifact)
        # add unique references with source_location
        # will re-add a reference if current location is None
        if self.artifact.depends_on.get(const.artifact) is None:
            self.artifact.depends_on[const.artifact] = const.source_location

    def _check_reference(self, ref: ContractReference | LogicSigReference) -> None:
        if ref not in self.artifacts:
            logger.critical(f"missing reference: {ref}")


import enum

@enum.unique
class AVMType(enum.Flag):
    bytes = enum.auto()
    uint64 = enum.auto()
    any = bytes | uint64

# values and names are matched to AVM definitions
class OnCompletionAction(enum.IntEnum):
    NoOp = 0
    OptIn = 1
    CloseOut = 2
    ClearState = 3
    UpdateApplication = 4
    DeleteApplication = 5

class TransactionType(enum.IntEnum):
    pay = 1
    keyreg = 2
    acfg = 3
    axfer = 4
    afrz = 5
    appl = 6


from immutabledict import immutabledict

from puya.awst import wtypes
from puya.errors import CodeError
from puya.parse import SourceLocation

def wtype_to_arc4(wtype: wtypes.WType, loc: SourceLocation | None = None) -> str:
    match wtype:
        case wtypes.ARC4Type(arc4_name=arc4_name):
            return arc4_name
        case (
            wtypes.void_wtype
            | wtypes.asset_wtype
            | wtypes.account_wtype
            | wtypes.application_wtype
        ):
            return wtype.name
        case wtypes.WGroupTransaction(transaction_type=transaction_type):
            return transaction_type.name if transaction_type else "txn"
    converted = maybe_avm_to_arc4_equivalent_type(wtype)
    if converted is None:
        raise CodeError(f"not an ARC4 type or native equivalent: {wtype}", loc)
    return wtype_to_arc4(converted, loc)

def maybe_avm_to_arc4_equivalent_type(wtype: wtypes.WType) -> wtypes.ARC4Type | None:
    match wtype:
        case wtypes.bool_wtype:
            return wtypes.arc4_bool_wtype
        case wtypes.uint64_wtype:
            return wtypes.ARC4UIntN(n=64, source_location=None)
        case wtypes.biguint_wtype:
            return wtypes.ARC4UIntN(n=512, source_location=None)
        case wtypes.bytes_wtype:
            return wtypes.ARC4DynamicArray(
                element_type=wtypes.arc4_byte_alias, native_type=wtype, source_location=None
            )
        case wtypes.string_wtype:
            return wtypes.arc4_string_alias
        case wtypes.WTuple(types=tuple_item_types) as wtuple:
            arc4_item_types = []
            for t in tuple_item_types:
                if isinstance(t, wtypes.ARC4Type):
                    arc4_item_types.append(t)
                else:
                    converted = maybe_avm_to_arc4_equivalent_type(t)
                    if converted is None:
                        return None
                    arc4_item_types.append(converted)
            if wtuple.fields:
                return wtypes.ARC4Struct(
                    name=wtuple.name,
                    desc=wtuple.desc,
                    frozen=True,
                    fields=immutabledict(zip(wtuple.fields, arc4_item_types, strict=True)),
                )
            else:
                return wtypes.ARC4Tuple(types=arc4_item_types, source_location=None)
        case _:
            return None


import typing

from puya.awst import nodes as awst_nodes
from puya.awst.function_traverser import FunctionTraverser
from puya.awst.visitors import ContractMemberVisitor, RootNodeVisitor

class AWSTTraverser(FunctionTraverser, RootNodeVisitor[None], ContractMemberVisitor[None]):
    @typing.override
    def visit_subroutine(self, statement: awst_nodes.Subroutine) -> None:
        statement.body.accept(self)

    @typing.override
    def visit_contract(self, statement: awst_nodes.Contract) -> None:
        for storage in statement.app_state:
            storage.accept(self)
        for method in statement.all_methods:
            method.accept(self)

    @typing.override
    def visit_logic_signature(self, statement: awst_nodes.LogicSignature) -> None:
        statement.program.accept(self)

    @typing.override
    def visit_app_storage_definition(self, defn: awst_nodes.AppStorageDefinition) -> None:
        defn.key.accept(self)

    @typing.override
    def visit_contract_method(self, statement: awst_nodes.ContractMethod) -> None:
        statement.body.accept(self)


import typing

import puya.awst.visitors
from puya.awst import nodes as awst_nodes

class FunctionTraverser(
    puya.awst.visitors.ExpressionVisitor[None],
    puya.awst.visitors.StatementVisitor[None],
):
    @typing.override
    def visit_assignment_statement(self, statement: awst_nodes.AssignmentStatement) -> None:
        statement.target.accept(self)
        statement.value.accept(self)

    @typing.override
    def visit_copy(self, expr: awst_nodes.Copy) -> None:
        expr.value.accept(self)

    @typing.override
    def visit_goto(self, statement: awst_nodes.Goto) -> None:
        pass

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

    @typing.override
    def visit_uint64_binary_operation(self, expr: awst_nodes.UInt64BinaryOperation) -> None:
        expr.left.accept(self)
        expr.right.accept(self)

    @typing.override
    def visit_biguint_binary_operation(self, expr: awst_nodes.BigUIntBinaryOperation) -> None:
        expr.left.accept(self)
        expr.right.accept(self)

    @typing.override
    def visit_reversed(self, expr: awst_nodes.Reversed) -> None:
        if isinstance(expr.expr, awst_nodes.Expression):
            expr.expr.accept(self)

    @typing.override
    def visit_integer_constant(self, expr: awst_nodes.IntegerConstant) -> None:
        pass

    @typing.override
    def visit_decimal_constant(self, expr: awst_nodes.DecimalConstant) -> None:
        pass

    @typing.override
    def visit_bool_constant(self, expr: awst_nodes.BoolConstant) -> None:
        pass

    @typing.override
    def visit_bytes_constant(self, expr: awst_nodes.BytesConstant) -> None:
        pass

    @typing.override
    def visit_string_constant(self, expr: awst_nodes.StringConstant) -> None:
        pass

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

    @typing.override
    def visit_compiled_contract(self, expr: awst_nodes.CompiledContract) -> None:
        for value in expr.template_variables.values():
            value.accept(self)

    @typing.override
    def visit_compiled_logicsig(self, expr: awst_nodes.CompiledLogicSig) -> None:
        for value in expr.template_variables.values():
            value.accept(self)

    @typing.override
    def visit_arc4_decode(self, expr: awst_nodes.ARC4Decode) -> None:
        expr.value.accept(self)

    @typing.override
    def visit_arc4_encode(self, expr: awst_nodes.ARC4Encode) -> None:
        expr.value.accept(self)

    @typing.override
    def visit_array_concat(self, expr: awst_nodes.ArrayConcat) -> None:
        expr.left.accept(self)
        expr.right.accept(self)

    @typing.override
    def visit_array_pop(self, expr: awst_nodes.ArrayPop) -> None:
        expr.base.accept(self)

    @typing.override
    def visit_array_extend(self, expr: awst_nodes.ArrayExtend) -> None:
        expr.base.accept(self)
        expr.other.accept(self)

    @typing.override
    def visit_method_constant(self, expr: awst_nodes.MethodConstant) -> None:
        pass

    @typing.override
    def visit_address_constant(self, expr: awst_nodes.AddressConstant) -> None:
        pass

    @typing.override
    def visit_numeric_comparison_expression(
        self, expr: awst_nodes.NumericComparisonExpression
    ) -> None:
        expr.lhs.accept(self)
        expr.rhs.accept(self)

    @typing.override
    def visit_var_expression(self, expr: awst_nodes.VarExpression) -> None:
        pass

    @typing.override
    def visit_assert_expression(self, expr: awst_nodes.AssertExpression) -> None:
        if expr.condition is not None:
            expr.condition.accept(self)

    @typing.override
    def visit_checked_maybe(self, expr: awst_nodes.CheckedMaybe) -> None:
        expr.expr.accept(self)

    @typing.override
    def visit_intrinsic_call(self, call: awst_nodes.IntrinsicCall) -> None:
        for arg in call.stack_args:
            arg.accept(self)

    @typing.override
    def visit_puya_lib_call(self, call: awst_nodes.PuyaLibCall) -> None:
        for arg in call.args:
            arg.value.accept(self)

    @typing.override
    def visit_group_transaction_reference(self, ref: awst_nodes.GroupTransactionReference) -> None:
        ref.index.accept(self)

    @typing.override
    def visit_create_inner_transaction(self, call: awst_nodes.CreateInnerTransaction) -> None:
        for expr in call.fields.values():
            expr.accept(self)

    @typing.override
    def visit_update_inner_transaction(self, call: awst_nodes.UpdateInnerTransaction) -> None:
        call.itxn.accept(self)
        for value in call.fields.values():
            value.accept(self)

    @typing.override
    def visit_submit_inner_transaction(self, call: awst_nodes.SubmitInnerTransaction) -> None:
        for expr in call.itxns:
            expr.accept(self)

    @typing.override
    def visit_inner_transaction_field(self, itxn_field: awst_nodes.InnerTransactionField) -> None:
        itxn_field.itxn.accept(self)
        if itxn_field.array_index:
            itxn_field.array_index.accept(self)

    @typing.override
    def visit_tuple_expression(self, expr: awst_nodes.TupleExpression) -> None:
        for item in expr.items:
            item.accept(self)

    @typing.override
    def visit_tuple_item_expression(self, expr: awst_nodes.TupleItemExpression) -> None:
        expr.base.accept(self)

    @typing.override
    def visit_field_expression(self, expr: awst_nodes.FieldExpression) -> None:
        expr.base.accept(self)

    @typing.override
    def visit_slice_expression(self, expr: awst_nodes.SliceExpression) -> None:
        expr.base.accept(self)
        if isinstance(expr.begin_index, awst_nodes.Expression):
            expr.begin_index.accept(self)
        if isinstance(expr.end_index, awst_nodes.Expression):
            expr.end_index.accept(self)

    @typing.override
    def visit_intersection_slice_expression(
        self, expr: awst_nodes.IntersectionSliceExpression
    ) -> None:
        expr.base.accept(self)
        if isinstance(expr.begin_index, awst_nodes.Expression):
            expr.begin_index.accept(self)
        if isinstance(expr.end_index, awst_nodes.Expression):
            expr.end_index.accept(self)

    @typing.override
    def visit_index_expression(self, expr: awst_nodes.IndexExpression) -> None:
        expr.base.accept(self)
        expr.index.accept(self)

    @typing.override
    def visit_conditional_expression(self, expr: awst_nodes.ConditionalExpression) -> None:
        expr.condition.accept(self)
        expr.true_expr.accept(self)
        expr.false_expr.accept(self)

    @typing.override
    def visit_single_evaluation(self, expr: awst_nodes.SingleEvaluation) -> None:
        expr.source.accept(self)

    @typing.override
    def visit_app_state_expression(self, expr: awst_nodes.AppStateExpression) -> None:
        expr.key.accept(self)

    @typing.override
    def visit_app_account_state_expression(
        self, expr: awst_nodes.AppAccountStateExpression
    ) -> None:
        expr.key.accept(self)
        expr.account.accept(self)

    @typing.override
    def visit_new_array(self, expr: awst_nodes.NewArray) -> None:
        for element in expr.values:
            element.accept(self)

    @typing.override
    def visit_new_struct(self, expr: awst_nodes.NewStruct) -> None:
        for element in expr.values.values():
            element.accept(self)

    @typing.override
    def visit_bytes_comparison_expression(
        self, expr: awst_nodes.BytesComparisonExpression
    ) -> None:
        expr.lhs.accept(self)
        expr.rhs.accept(self)

    @typing.override
    def visit_subroutine_call_expression(self, expr: awst_nodes.SubroutineCallExpression) -> None:
        for arg in expr.args:
            arg.value.accept(self)

    @typing.override
    def visit_bytes_binary_operation(self, expr: awst_nodes.BytesBinaryOperation) -> None:
        expr.left.accept(self)
        expr.right.accept(self)

    @typing.override
    def visit_boolean_binary_operation(self, expr: awst_nodes.BooleanBinaryOperation) -> None:
        expr.left.accept(self)
        expr.right.accept(self)

    @typing.override
    def visit_uint64_unary_operation(self, expr: awst_nodes.UInt64UnaryOperation) -> None:
        expr.expr.accept(self)

    @typing.override
    def visit_bytes_unary_operation(self, expr: awst_nodes.BytesUnaryOperation) -> None:
        expr.expr.accept(self)

    @typing.override
    def visit_not_expression(self, expr: awst_nodes.Not) -> None:
        expr.expr.accept(self)

    @typing.override
    def visit_block(self, statement: awst_nodes.Block) -> None:
        for stmt in statement.body:
            stmt.accept(self)

    @typing.override
    def visit_if_else(self, statement: awst_nodes.IfElse) -> None:
        statement.condition.accept(self)
        statement.if_branch.accept(self)
        if statement.else_branch:
            statement.else_branch.accept(self)

    @typing.override
    def visit_switch(self, statement: awst_nodes.Switch) -> None:
        statement.value.accept(self)
        for case, block in statement.cases.items():
            case.accept(self)
            block.accept(self)
        if statement.default_case:
            statement.default_case.accept(self)

    @typing.override
    def visit_while_loop(self, statement: awst_nodes.WhileLoop) -> None:
        statement.condition.accept(self)
        statement.loop_body.accept(self)

    @typing.override
    def visit_loop_exit(self, statement: awst_nodes.LoopExit) -> None:
        pass

    @typing.override
    def visit_return_statement(self, statement: awst_nodes.ReturnStatement) -> None:
        if statement.value is not None:
            statement.value.accept(self)

    @typing.override
    def visit_loop_continue(self, statement: awst_nodes.LoopContinue) -> None:
        pass

    @typing.override
    def visit_expression_statement(self, statement: awst_nodes.ExpressionStatement) -> None:
        statement.expr.accept(self)

    @typing.override
    def visit_template_var(self, statement: awst_nodes.TemplateVar) -> None:
        pass

    @typing.override
    def visit_uint64_augmented_assignment(
        self, statement: awst_nodes.UInt64AugmentedAssignment
    ) -> None:
        statement.target.accept(self)
        statement.value.accept(self)

    @typing.override
    def visit_biguint_augmented_assignment(
        self, statement: awst_nodes.BigUIntAugmentedAssignment
    ) -> None:
        statement.target.accept(self)
        statement.value.accept(self)

    @typing.override
    def visit_bytes_augmented_assignment(
        self, statement: awst_nodes.BytesAugmentedAssignment
    ) -> None:
        statement.target.accept(self)
        statement.value.accept(self)

    @typing.override
    def visit_for_in_loop(self, statement: awst_nodes.ForInLoop) -> None:
        statement.sequence.accept(self)
        statement.items.accept(self)
        statement.loop_body.accept(self)

    @typing.override
    def visit_reinterpret_cast(self, expr: awst_nodes.ReinterpretCast) -> None:
        expr.expr.accept(self)

    @typing.override
    def visit_enumeration(self, expr: awst_nodes.Enumeration) -> None:
        expr.expr.accept(self)

    @typing.override
    def visit_state_get_ex(self, expr: awst_nodes.StateGetEx) -> None:
        expr.field.accept(self)

    @typing.override
    def visit_state_delete(self, statement: awst_nodes.StateDelete) -> None:
        statement.field.accept(self)

    @typing.override
    def visit_state_get(self, expr: awst_nodes.StateGet) -> None:
        expr.field.accept(self)
        expr.default.accept(self)

    @typing.override
    def visit_state_exists(self, expr: awst_nodes.StateExists) -> None:
        expr.field.accept(self)

    @typing.override
    def visit_box_value_expression(self, expr: awst_nodes.BoxValueExpression) -> None:
        expr.key.accept(self)

    @typing.override
    def visit_biguint_postfix_unary_operation(
        self, expr: awst_nodes.BigUIntPostfixUnaryOperation
    ) -> None:
        expr.target.accept(self)

    @typing.override
    def visit_uint64_postfix_unary_operation(
        self, expr: awst_nodes.UInt64PostfixUnaryOperation
    ) -> None:
        expr.target.accept(self)

    @typing.override
    def visit_arc4_router(self, expr: awst_nodes.ARC4Router) -> None:
        pass

    @typing.override
    def visit_range(self, node: awst_nodes.Range) -> None:
        node.start.accept(self)
        node.stop.accept(self)
        node.step.accept(self)

    @typing.override
    def visit_emit(self, expr: awst_nodes.Emit) -> None:
        expr.value.accept(self)