FragmentAPI/pyfragment/domains/tonapi/account.py

from __future__ import annotations

import base64
import logging
from typing import TYPE_CHECKING, Any

from ton_core import NetworkGlobalID
from tonutils.clients import TonapiClient
from tonutils.contracts.jetton import get_wallet_address_get_method, get_wallet_data_get_method
from tonutils.exceptions import ProviderResponseError

from pyfragment.core.constants import MIN_TON_BALANCE, MIN_USDT_BALANCE, USDT_TON_MASTER_ADDRESS, WALLET_CLASSES
from pyfragment.exceptions import WalletError
from pyfragment.models.wallet import WalletInfo

if TYPE_CHECKING:
    from pyfragment.client import FragmentClient


logger = logging.getLogger(__name__)


async def get_usdt_balance(ton: Any, wallet_address: str) -> float:
    """Return the USDT balance for a Fragment-linked TON wallet."""
    try:
        jetton_wallet_address = await get_wallet_address_get_method(
            client=ton,
            address=USDT_TON_MASTER_ADDRESS,
            owner_address=wallet_address,
        )
        wallet_data = await get_wallet_data_get_method(client=ton, address=jetton_wallet_address)
        raw_balance = int(wallet_data[0]) if wallet_data else 0
        return float(raw_balance) / 1_000_000.0
    except ProviderResponseError as exc:
        if exc.code == 404:
            logger.debug("No USDT jetton wallet found for '%s'; treating balance as 0", wallet_address)
            return 0.0
        logger.error("Failed to load USDT balance for wallet '%s': %s", wallet_address, exc, exc_info=True)
        raise WalletError(WalletError.USDT_BALANCE_CHECK_FAILED.format(exc=exc)) from exc
    except Exception as exc:
        logger.exception("Failed to load USDT balance for wallet '%s' due to an unexpected error", wallet_address)
        raise WalletError(WalletError.USDT_BALANCE_CHECK_FAILED.format(exc=exc)) from exc


async def check_ton_payment_balance(
    balance_ton: float,
    amount_ton: float,
    required_payment_amount: float | None,
) -> None:
    """Validate that the TON wallet can cover a TON-denominated payment."""
    tx_price_ton = amount_ton
    if required_payment_amount is not None and required_payment_amount > 0:
        tx_price_ton = max(tx_price_ton, required_payment_amount)

    required_ton = max(tx_price_ton, MIN_TON_BALANCE)
    if balance_ton < required_ton:
        logger.error(
            "Failed TON balance check: balance=%s TON, required=%s TON",
            round(balance_ton, 6),
            round(required_ton, 6),
        )
        raise WalletError(WalletError.LOW_TON_BALANCE.format(balance=balance_ton, required=required_ton))


async def check_usdt_payment_balance(
    balance_ton: float,
    required_payment_amount: float | None,
    ton: Any,
    wallet_address: str,
) -> None:
    """Validate that the wallet can cover a USDT-denominated payment."""
    if balance_ton < MIN_TON_BALANCE:
        logger.error(
            "Failed TON gas reserve check for USDT payment: balance=%s TON, required=%s TON",
            round(balance_ton, 6),
            MIN_TON_BALANCE,
        )
        raise WalletError(WalletError.LOW_TON_BALANCE.format(balance=balance_ton, required=MIN_TON_BALANCE))

    usdt_balance = await get_usdt_balance(ton, wallet_address)
    required_usdt = required_payment_amount if required_payment_amount is not None else MIN_USDT_BALANCE
    if usdt_balance < required_usdt:
        logger.error(
            "Failed USDT balance check for wallet '%s': balance=%s USDT, required=%s USDT",
            wallet_address,
            round(usdt_balance, 6),
            round(required_usdt, 6),
        )
        raise WalletError(WalletError.LOW_USDT_BALANCE.format(balance=usdt_balance, required=required_usdt))


async def get_account_info(client: FragmentClient) -> dict[str, Any]:
    """Build the wallet payload Fragment needs to prepare a transaction."""
    async with TonapiClient(network=NetworkGlobalID.MAINNET, api_key=client.api_key) as ton:
        try:
            wallet_cls = WALLET_CLASSES[client.wallet_version]
            wallet, pub_key, _, _ = wallet_cls.from_mnemonic(client=ton, mnemonic=client.seed)
            boc = wallet.state_init.serialize().to_boc()
            return {
                "address": wallet.address.to_str(False, False),
                "publicKey": pub_key.as_hex,
                "chain": "-239",
                "walletStateInit": base64.b64encode(boc).decode(),
            }
        except Exception as exc:
            logger.exception("Failed to build Fragment account info from the configured wallet")
            raise WalletError(WalletError.ACCOUNT_INFO_FAILED.format(exc=exc)) from exc


async def get_wallet_info(client: FragmentClient) -> WalletInfo:
    """Fetch the wallet address, chain state, and TON/USDT balances."""
    async with TonapiClient(network=NetworkGlobalID.MAINNET, api_key=client.api_key) as ton:
        try:
            wallet_cls = WALLET_CLASSES[client.wallet_version]
            wallet, _, _, _ = wallet_cls.from_mnemonic(client=ton, mnemonic=client.seed)
            await wallet.refresh()
            wallet_address = wallet.address.to_str(False, False)
            usdt_balance = await get_usdt_balance(ton, wallet_address)
            return WalletInfo(
                address=wallet.address.to_str(is_user_friendly=True, is_bounceable=False),
                state=wallet.state.value,
                ton_balance=round(wallet.balance / 1_000_000_000, 4),
                usdt_balance=round(usdt_balance, 4),
            )
        except Exception as exc:
            logger.exception("Failed to fetch wallet info from Tonapi")
            raise WalletError(WalletError.WALLET_INFO_FAILED.format(exc=exc)) from exc