LZBlockRelay

read_enabled: public(bool)
read_channel: public(uint32)
mainnet_eid: public(uint32)
mainnet_block_view: public(address)

@external
def set_read_config(
    _is_enabled: bool, _read_channel: uint32, _mainnet_eid: uint32, _mainnet_view: address
):
    """
    @notice Configure read functionality
    @param _is_enabled Whether this contract can initiate reads
    @param _read_channel LZ read channel ID
    @param _mainnet_eid Mainnet endpoint ID
    @param _mainnet_view MainnetBlockView contract address
    """
    ownable._check_owner()

    assert _read_channel > OApp.READ_CHANNEL_THRESHOLD, "Invalid read channel"

    assert (_is_enabled and _mainnet_eid != 0 and _mainnet_view != empty(address)) or (
        not _is_enabled and _mainnet_eid == 0 and _mainnet_view == empty(address)
    ), "Invalid read config"

    # Clean up old peer if switching channels while read is enabled
    # This prevents leaving stale peer mappings when changing read channels
    if self.read_enabled and self.read_channel != _read_channel:
        OApp._setPeer(self.read_channel, convert(empty(address), bytes32))

    self.read_enabled = _is_enabled
    self.read_channel = _read_channel
    self.mainnet_eid = _mainnet_eid
    self.mainnet_block_view = _mainnet_view

    peer: bytes32 = convert(self, bytes32) if _is_enabled else convert(empty(address), bytes32)
    OApp._setPeer(_read_channel, peer)

interface IBlockOracle:
    def commit_block(block_number: uint256, block_hash: bytes32) -> bool: nonpayable
    def last_confirmed_block_number() -> uint256: view
    def get_block_hash(block_number: uint256) -> bytes32: view

from snekmate.auth import ownable

@external
def set_block_oracle(_oracle: address):
    """
    @notice Set the block oracle address
    @param _oracle Block oracle address
    """
    ownable._check_owner()

    self.block_oracle = IBlockOracle(_oracle)

from snekmate.auth import ownable

@external
def withdraw_eth(_amount: uint256):
    """
    @notice Withdraw ETH from contract
    @dev ETH can be accumulated from LZ refunds
    @param _amount Amount to withdraw
    """
    ownable._check_owner()

    assert self.balance >= _amount, "Insufficient balance"
    send(msg.sender, _amount)

from ..modules.oapp_vyper.src import OApp  # main module

exports: (
    OApp.endpoint,
    OApp.peers,
    OApp.setPeer,
    OApp.setDelegate,
    OApp.isComposeMsgSender,
    OApp.allowInitializePath,
    OApp.nextNonce,
)

# LayerZero EndpointV2 interface
interface ILayerZeroEndpointV2:
    def quote(_params: MessagingParams, _sender: address) -> MessagingFee: view
    def send(_params: MessagingParams, _refundAddress: address) -> MessagingReceipt: payable
    def setDelegate(_delegate: address): nonpayable
    def eid() -> uint32: view
    def lzToken() -> address: view

# The LayerZero endpoint associated with the given OApp
endpoint: public(immutable(ILayerZeroEndpointV2))

from ..modules.oapp_vyper.src import OApp  # main module

exports: (
    OApp.endpoint,
    OApp.peers,
    OApp.setPeer,
    OApp.setDelegate,
    OApp.isComposeMsgSender,
    OApp.allowInitializePath,
    OApp.nextNonce,
)

# Mapping to store peers associated with corresponding endpoints
peers: public(HashMap[uint32, bytes32])

from ..modules.oapp_vyper.src import OApp  # main module

exports: (
    OApp.endpoint,
    OApp.peers,
    OApp.setPeer,
    OApp.setDelegate,
    OApp.isComposeMsgSender,
    OApp.allowInitializePath,
    OApp.nextNonce,
)

event PeerSet:
    eid: uint32
    peer: bytes32

# Mapping to store peers associated with corresponding endpoints
peers: public(HashMap[uint32, bytes32])

@external
def setPeer(_eid: uint32, _peer: bytes32):
    """
    @notice Sets the peer address (OApp instance) for a corresponding endpoint.
    @param _eid The endpoint ID.
    @param _peer The address of the peer to be associated with the corresponding endpoint.
    @dev Only the owner/admin of the OApp can call this function.
    @dev Indicates that the peer is trusted to send LayerZero messages to this OApp.
    @dev Set this to bytes32(0) to remove the peer address.
    @dev Peer is a bytes32 to accommodate non-evm chains.
    """
    ownable._check_owner()

    self._setPeer(_eid, _peer)

@internal
def _setPeer(_eid: uint32, _peer: bytes32):
    """
    @notice Internal function to set peer address
    @param _eid The endpoint ID.
    @param _peer The address of the peer to be associated with the corresponding endpoint.
    """
    self.peers[_eid] = _peer

    log PeerSet(eid=_eid, peer=_peer)

from ..modules.oapp_vyper.src import OApp  # main module

exports: (
    OApp.endpoint,
    OApp.peers,
    OApp.setPeer,
    OApp.setDelegate,
    OApp.isComposeMsgSender,
    OApp.allowInitializePath,
    OApp.nextNonce,
)

@external
def set_peers(_eids: DynArray[uint32, MAX_N_BROADCAST], _peers: DynArray[address, MAX_N_BROADCAST]):
    """
    @notice Set peers for a corresponding endpoints. Batched version of OApp.setPeer that accept address (EVM only).
    @param _eids The endpoint IDs.
    @param _peers Addresses of the peers to be associated with the corresponding endpoints.
    """
    ownable._check_owner()

    assert len(_eids) == len(_peers), "Invalid peer arrays"
    for i: uint256 in range(0, len(_eids), bound=MAX_N_BROADCAST):
        OApp._setPeer(_eids[i], convert(_peers[i], bytes32))

event PeerSet:
    eid: uint32
    peer: bytes32

# Mapping to store peers associated with corresponding endpoints
peers: public(HashMap[uint32, bytes32])

@internal
def _setPeer(_eid: uint32, _peer: bytes32):
    """
    @notice Internal function to set peer address
    @param _eid The endpoint ID.
    @param _peer The address of the peer to be associated with the corresponding endpoint.
    """
    self.peers[_eid] = _peer

    log PeerSet(eid=_eid, peer=_peer)

from ..modules.oapp_vyper.src import OApp  # main module

exports: (
    OApp.endpoint,
    OApp.peers,
    OApp.setPeer,
    OApp.setDelegate,
    OApp.isComposeMsgSender,
    OApp.allowInitializePath,
    OApp.nextNonce,
)

# LayerZero EndpointV2 interface
interface ILayerZeroEndpointV2:
    def quote(_params: MessagingParams, _sender: address) -> MessagingFee: view
    def send(_params: MessagingParams, _refundAddress: address) -> MessagingReceipt: payable
    def setDelegate(_delegate: address): nonpayable
    def eid() -> uint32: view
    def lzToken() -> address: view

@external
def setDelegate(_delegate: address):
    """
    @notice Sets the delegate address for the OApp.
    @param _delegate The address of the delegate to be set.
    @dev Only the owner/admin of the OApp can call this function.
    @dev Provides the ability for a delegate to set configs, on behalf of the OApp,
    directly on the Endpoint contract.
    """
    ownable._check_owner()

    extcall endpoint.setDelegate(_delegate)

from ..modules.oapp_vyper.src import OApp  # main module

exports: (
    OApp.endpoint,
    OApp.peers,
    OApp.setPeer,
    OApp.setDelegate,
    OApp.isComposeMsgSender,
    OApp.allowInitializePath,
    OApp.nextNonce,
)

struct Origin:
    srcEid: uint32
    sender: bytes32
    nonce: uint64

@external
@view
def isComposeMsgSender(
    _origin: Origin, _message: Bytes[MAX_MESSAGE_SIZE], _sender: address
) -> bool:
    """
    @notice Indicates whether an address is an approved composeMsg sender to the Endpoint.
    @param _origin The origin information containing the source endpoint and sender address.
    @param _message The lzReceive payload.
    @param _sender The sender address.
    @return isSender Is a valid sender.
    """
    return _sender == self

from ..modules.oapp_vyper.src import OApp  # main module

exports: (
    OApp.endpoint,
    OApp.peers,
    OApp.setPeer,
    OApp.setDelegate,
    OApp.isComposeMsgSender,
    OApp.allowInitializePath,
    OApp.nextNonce,
)

struct Origin:
    srcEid: uint32
    sender: bytes32
    nonce: uint64

@external
@view
def allowInitializePath(_origin: Origin) -> bool:
    """
    @notice Checks if the path initialization is allowed based on the provided origin.
    @param _origin The origin information containing the source endpoint and sender address.
    @return Whether the path has been initialized.
    @dev This indicates to the endpoint that the OApp has enabled msgs for this particular path to be received.
    @dev This defaults to assuming if a peer has been set, its initialized.
    """
    return self.peers[_origin.srcEid] == _origin.sender

from ..modules.oapp_vyper.src import OApp  # main module

exports: (
    OApp.endpoint,
    OApp.peers,
    OApp.setPeer,
    OApp.setDelegate,
    OApp.isComposeMsgSender,
    OApp.allowInitializePath,
    OApp.nextNonce,
)

@external
@pure
def nextNonce(_srcEid: uint32, _sender: bytes32) -> uint64:
    """
    @notice Retrieves the next nonce for a given source endpoint and sender address.
    @dev Vyper-specific: If your app relies on ordered execution, you must change this function.
    @param _srcEid The source endpoint ID.
    @param _sender The sender address.
    @return nonce The next nonce.
    @dev The path nonce starts from 1. If 0 is returned it means that there is NO nonce ordered enforcement.
    @dev Is required by the off-chain executor to determine the OApp expects msg execution is ordered.
    @dev This is also enforced by the OApp.
    @dev By default this is NOT enabled. ie. nextNonce is hardcoded to return 0.
    """
    return 0

@external
@payable
def request_block_hash(
    _target_eids: DynArray[uint32, MAX_N_BROADCAST],
    _target_fees: DynArray[uint256, MAX_N_BROADCAST],
    _lz_receive_gas_limit: uint128,
    _read_gas_limit: uint128,
    _block_number: uint256 = 0,
):
    """
    @notice Request block hash from mainnet and broadcast to specified targets
    @param _target_eids List of chain IDs to broadcast to
    @param _target_fees List of fees per chain (must match _target_eids length)
    @param _lz_receive_gas_limit Gas limit for lzReceive (same for all targets)
    @param _read_gas_limit Gas limit for read operation
    @param _block_number Optional block number (0 means latest)
    @dev User must ensure msg.value is sufficient:
        - must cover read fee (quote_read_fee)
        - must cover broadcast fees (quote_broadcast_fees)
    """

    assert self.read_enabled, "Read not enabled"
    assert len(_target_eids) == len(_target_fees), "Length mismatch"

    self._request_block_hash(
        _block_number,
        _target_eids,
        _target_fees,
        _lz_receive_gas_limit,
        _read_gas_limit,
    )

@internal
@payable
def _request_block_hash(
    _block_number: uint256,
    _target_eids: DynArray[uint32, MAX_N_BROADCAST],
    _target_fees: DynArray[uint256, MAX_N_BROADCAST],
    _lz_receive_gas_limit: uint128,
    _read_gas_limit: uint128,
):
    """
    @notice Internal function to request block hash from mainnet and broadcast to specified targets
    @param _block_number Block number to request
    @param _target_eids Target EIDs to broadcast to
    @param _target_fees Target fees to pay per broadcast
    @param _lz_receive_gas_limit Gas limit for lzReceive
    @param _read_gas_limit Gas limit for read operation
    """

    # Store target EIDs and fees for lzReceive
    cached_targets: DynArray[BroadcastTarget, MAX_N_BROADCAST] = []
    sum_target_fees: uint256 = 0
    for i: uint256 in range(0, len(_target_eids), bound=MAX_N_BROADCAST):
        cached_targets.append(BroadcastTarget(eid=_target_eids[i], fee=_target_fees[i]))
        sum_target_fees += _target_fees[i]

    assert sum_target_fees <= msg.value, "Insufficient value" # dev: check is here because we sum here

    message: Bytes[OApp.MAX_MESSAGE_SIZE] = self._prepare_read_request(_block_number)

    # Create options using OptionsBuilder module
    options: Bytes[OptionsBuilder.MAX_OPTIONS_TOTAL_SIZE] = OptionsBuilder.newOptions()
    options = OptionsBuilder.addExecutorLzReadOption(
        options, _read_gas_limit, READ_RETURN_SIZE, convert(sum_target_fees, uint128)
    )

    # Send message
    fees: OApp.MessagingFee = OApp.MessagingFee(nativeFee=msg.value, lzTokenFee=0)
    # Fees = read fee + broadcast fees (value of read return message)
    receipt: OApp.MessagingReceipt = OApp._lzSend(
        self.read_channel, message, options, fees, msg.sender # dev: refund excess fee to sender
    )

    # Store targets for lzReceive using receipt.guid as key
    self.broadcast_data[receipt.guid] = BroadcastData(
        targets=cached_targets,
        gas_limit=_lz_receive_gas_limit,
    )

@external
@payable
def broadcast_latest_block(
    _target_eids: DynArray[uint32, MAX_N_BROADCAST],
    _target_fees: DynArray[uint256, MAX_N_BROADCAST],
    _lz_receive_gas_limit: uint128,
):
    """
    @notice Broadcast latest confirmed block hash to specified chains
    @param _target_eids List of chain IDs to broadcast to
    @param _target_fees List of fees per chain (must match _target_eids length)
    @param _lz_receive_gas_limit Gas limit for lzReceive (same for all targets)
    @dev Only broadcast what was received via lzRead to prevent potentially malicious hashes from other sources
    """

    assert self.read_enabled, "Can only broadcast from read-enabled chains"
    assert self.block_oracle != empty(IBlockOracle), "Oracle not configured"
    assert len(_target_eids) == len(_target_fees), "Length mismatch"

    # Get latest block from oracle
    block_number: uint256 = staticcall self.block_oracle.last_confirmed_block_number()
    block_hash: bytes32 = staticcall self.block_oracle.get_block_hash(block_number)
    assert block_hash != empty(bytes32), "No confirmed blocks"

    # Only broadcast if this block was received via lzRead
    assert self.received_blocks[block_number] == block_hash, "Unknown source"

    # Prepare broadcast targets
    broadcast_targets: DynArray[BroadcastTarget, MAX_N_BROADCAST] = []
    for i: uint256 in range(0, len(_target_eids), bound=MAX_N_BROADCAST):
        broadcast_targets.append(BroadcastTarget(eid=_target_eids[i], fee=_target_fees[i]))

    self._broadcast_block(
        block_number,
        block_hash,
        BroadcastData(targets=broadcast_targets, gas_limit=_lz_receive_gas_limit),
        msg.sender,
    )

@payable
@external
def lzReceive(
    _origin: OApp.Origin,
    _guid: bytes32,
    _message: Bytes[OApp.MAX_MESSAGE_SIZE],
    _executor: address,
    _extraData: Bytes[OApp.MAX_EXTRA_DATA_SIZE],
):
    """
    @notice Handle messages: read responses, and regular messages
    @dev Two types of messages:
        1. Read responses (from read channel)
        2. Regular messages (block hash broadcasts from other chains)
    @param _origin Origin information containing srcEid, sender, and nonce
    @param _guid Global unique identifier for the message
    @param _message The encoded message payload containing block number and hash
    @param _executor Address of the executor for the message
    @param _extraData Additional data passed by the executor
    """
    # Verify message source
    OApp._lzReceive(_origin, _guid, _message, _executor, _extraData)

    if _origin.srcEid == self.read_channel:
        # Only handle read response if read is enabled
        assert self.read_enabled, "Read not enabled"
        # Decode block hash and number from response
        block_number: uint256 = 0
        block_hash: bytes32 = empty(bytes32)
        block_number, block_hash = abi_decode(_message, (uint256, bytes32))
        if block_hash == empty(bytes32):
            return  # Invalid response

        # Store received block hash
        self.received_blocks[block_number] = block_hash

        # Commit block hash to oracle
        self._commit_block(block_number, block_hash)

        broadcast_data: BroadcastData = self.broadcast_data[_guid]

        if len(broadcast_data.targets) > 0:
            # Verify that attached value covers requested broadcast fees
            total_fee: uint256 = 0
            for target: BroadcastTarget in broadcast_data.targets:
                total_fee += target.fee
            assert msg.value >= total_fee, "Insufficient msg.value"

            # Perform broadcast
            self._broadcast_block(
                block_number,
                block_hash,
                broadcast_data,
                self, # dev: refund excess fee to self
            )
    else:
        # Regular message - decode and commit block hash
        block_number: uint256 = 0
        block_hash: bytes32 = empty(bytes32)
        block_number, block_hash = abi_decode(_message, (uint256, bytes32))
        self._commit_block(block_number, block_hash)

@external
@view
def quote_read_fee(
    _read_gas_limit: uint128,
    _value: uint128,
) -> uint256:
    """
    @notice Quote fee for reading block hash from mainnet
    @param _read_gas_limit Gas to be provided in return message
    @param _value Value to be provided in return message
    @return Fee in native tokens required for the read operation
    """
    assert self.read_enabled, "Read not enabled - call set_read_config"

    message: Bytes[OApp.MAX_MESSAGE_SIZE] = self._prepare_read_request(0) # dev: 0 for latest block

    # Create options using OptionsBuilder module
    options: Bytes[OptionsBuilder.MAX_OPTIONS_TOTAL_SIZE] = OptionsBuilder.newOptions()
    options = OptionsBuilder.addExecutorLzReadOption(
        options, _read_gas_limit, READ_RETURN_SIZE, _value
    )

    return OApp._quote(
        self.read_channel,
        message,
        options,
        False,
    ).nativeFee

@internal
@pure
def newOptions() -> Bytes[MAX_OPTIONS_TOTAL_SIZE]:
    """
    @notice Creates a new options container with type 3.
    @return options The newly created options container.
    """
    options: Bytes[MAX_OPTIONS_TOTAL_SIZE] = concat(convert(TYPE_3, bytes2), b"")

    return options

@internal
@pure
def addExecutorOption(
    _options: Bytes[MAX_OPTIONS_TOTAL_SIZE],
    _optionType: uint8,
    _option: Bytes[MAX_OPTION_SINGLE_SIZE],
) -> Bytes[MAX_OPTIONS_TOTAL_SIZE]:
    """
    @dev Adds an executor option to the existing options.
    @param _options The existing options container.
    @param _optionType The type of the executor option.
    @param _option The encoded data for the executor option.
    @return options The updated options container.
    """
    assert convert(slice(_options, 0, 2), uint16) == TYPE_3, "OApp: invalid option type"
    # Account for header bytes: 1 worker + 2 size + 1 type = 4 bytes
    assert (len(_options) + len(_option) + 4 <= MAX_OPTIONS_TOTAL_SIZE), "OApp: options size exceeded"

    return concat(
        convert(_options, Bytes[MAX_OPTIONS_TOTAL_SIZE - MAX_OPTION_SINGLE_SIZE - 4]), # downcast Bytes size, -4 for header
        convert(EXECUTOR_WORKER_ID, bytes1),
        convert(convert(len(_option) + 1, uint16), bytes2),  # +1 for optionType
        convert(_optionType, bytes1),
        _option,
    )

@external
@view
def quote_broadcast_fees(
    _target_eids: DynArray[uint32, MAX_N_BROADCAST],
    _lz_receive_gas_limit: uint128,
) -> DynArray[uint256, MAX_N_BROADCAST]:
    """
    @notice Quote fees for broadcasting block hash to specified targets
    @param _target_eids List of chain IDs to broadcast to
    @param _lz_receive_gas_limit Gas limit for lzReceive
    @return Array of fees per target chain (0 if target not configured)
    """
    # Prepare dummy broadcast message (uint256 number, bytes32 hash)
    message: Bytes[OApp.MAX_MESSAGE_SIZE] = abi_encode(empty(uint256), empty(bytes32))

    # Prepare array of fees per chain
    fees: DynArray[uint256, MAX_N_BROADCAST] = []

    # Prepare options (same for all targets)
    options: Bytes[OptionsBuilder.MAX_OPTIONS_TOTAL_SIZE] = OptionsBuilder.newOptions()
    options = OptionsBuilder.addExecutorLzReceiveOption(options, _lz_receive_gas_limit, 0)

    # Cycle through targets
    for eid: uint32 in _target_eids:
        target: bytes32 = OApp.peers[eid]  # Use peers directly
        if target == empty(bytes32):
            fees.append(0)
            continue

        # Get fee for target EID and append to array
        fees.append(OApp._quote(eid, message, options, False).nativeFee)

    return fees

read_enabled: public(bool)

read_channel: public(uint32)

mainnet_eid: public(uint32)

mainnet_block_view: public(address)

block_oracle: public(IBlockOracle)