modules/contract_manager.py

"""
modules/contract_manager.py — Garden Angel On-Chain Execution Layer v2.4
──────────────────────────────────────────────────────────────────────────────
v2.4 changelog (this revision):
  - NEW — "mdex" added to ROUTER_MAP. Confirmed via BscScan's own "Mdex:
    Router" label (user-supplied screenshot, 2026-07-10), matching MDEX's
    official URL (bsc.mdex.com) — 0x7DAe51BD3E3376B8c7c4900E9107f12Be3AF1bA8.
    Kept in sync with scanner.py v18.10's _BSC_ROUTER_ADDRESSES.

v2.3 changelog:
  - FIXED withdraw ABI/selector mismatch. The deployed FlashArbitrageV2 exposes
    withdraw(address token) — ONE argument — which sends the full token balance
    to the contract owner. This module previously declared and called a
    two-argument withdraw(address token, uint256 amount), so /withdraw encoded a
    selector that does not exist on-chain and always reverted, making stuck-token
    recovery impossible. The ABI entry and ContractManager.withdraw() now use the
    1-arg form. withdraw() still ACCEPTS an optional `amount` for backward
    compatibility with older callers, but ignores it (the contract always sweeps
    the full balance). VERIFY against your actual deployed contract on BscScan
    before relying on this — if your deployment really is 2-arg, revert this.

v2.2 changelog:
  - FIXED CHAIN MISMATCH: replaced Ethereum-mainnet AssetRegistry/ROUTER_MAP
    entries with verified BSC mainnet addresses (Binance-Peg tokens are 18
    decimals across the board; do not copy ETH-side 8/6-decimal values).

Required env vars:
    RPC_URL
    FLASH_ARBITRAGE_CONTRACT_ADDRESS
    PAYOUT_PRIVATE_KEY   (optional — read-only mode without it)
"""

from __future__ import annotations

import json
import logging
import os
from decimal import Decimal
from typing import Optional

from web3 import Web3
from web3.exceptions import ContractLogicError, TransactionNotFound  # noqa: F401  (re-exported for callers)

from modules.web3_client import Web3Client, get_web3_client

logger = logging.getLogger("garden_angel.contract_manager")


# ---------------------------------------------------------------------------
# FlashArbitrageV2 contract interface
# ---------------------------------------------------------------------------
# Matches the deployed contract exactly: startArbitrage takes a single
# `intermediateToken` address per leg, and withdraw takes a single `token`
# address (full-balance sweep to owner). There is no quoteRoundTrip view
# function on-chain.

FLASH_ARBITRAGE_V2_ABI = json.loads("""
[
  {"inputs": [], "name": "owner", "outputs": [{"internalType": "address", "name": "", "type": "address"}], "stateMutability": "view", "type": "function"},
  {"inputs": [], "name": "aavePool", "outputs": [{"internalType": "address", "name": "", "type": "address"}], "stateMutability": "view", "type": "function"},
  {"inputs": [
      {"internalType": "address", "name": "asset", "type": "address"},
      {"internalType": "uint256", "name": "amount", "type": "uint256"},
      {"internalType": "address", "name": "routerBuy", "type": "address"},
      {"internalType": "address", "name": "routerSell", "type": "address"},
      {"internalType": "address", "name": "intermediateToken", "type": "address"},
      {"internalType": "uint256", "name": "minIntermediateOut", "type": "uint256"},
      {"internalType": "uint256", "name": "minFinalOut", "type": "uint256"},
      {"internalType": "uint256", "name": "minProfit", "type": "uint256"}
    ], "name": "startArbitrage", "outputs": [], "stateMutability": "nonpayable", "type": "function"},
  {"inputs": [
      {"internalType": "address", "name": "token", "type": "address"}
    ], "name": "withdraw", "outputs": [], "stateMutability": "nonpayable", "type": "function"},
  {"inputs": [], "name": "withdrawETH", "outputs": [], "stateMutability": "nonpayable", "type": "function"}
]
""")

# Minimal ERC-20 ABI slice — just enough to read decimals() live as a
# fallback/verification path when an asset isn't in the static registry.
ERC20_DECIMALS_ABI = json.loads("""
[
  {"inputs": [], "name": "decimals", "outputs": [{"internalType": "uint8", "name": "", "type": "uint8"}], "stateMutability": "view", "type": "function"}
]
""")


# ---------------------------------------------------------------------------
# AssetRegistry — decimals lookup (checksummed BSC mainnet addresses)
# ---------------------------------------------------------------------------
# Binance-Peg tokens on BSC use 18 decimals across the board, even for assets
# whose Ethereum-native counterpart uses fewer (BTCB=18 vs mainnet WBTC=8).

class AssetRegistry:
    """Static decimals lookup keyed by checksummed asset address, with a live
    on-chain fallback (ERC-20 decimals()) for anything not listed here."""

    # Migrated 2026-07 to Arbitrum One. Unlike BSC's uniform 18-decimals,
    # Arbitrum uses the tokens' native decimals — USDC/USDT are 6, WBTC is 8.
    # Anything not listed falls back to an on-chain decimals() query.
    _STATIC: dict[str, int] = {
        Web3.to_checksum_address("0x82aF49447D8a07e3bd95BD0d56f35241523fBab1"): 18,  # WETH
        Web3.to_checksum_address("0x2f2a2543B76A4166549F7aaB2e75Bef0aefC5B0f"): 8,   # WBTC (8!)
        Web3.to_checksum_address("0x912CE59144191C1204E64559FE8253a0e49E6548"): 18,  # ARB
        Web3.to_checksum_address("0xfc5A1A6EB076a2C7aD06eD22C90d7E710E35ad0a"): 18,  # GMX
        Web3.to_checksum_address("0xf97f4df75117a78c1A5a0DBb814Af92458539FB4"): 18,  # LINK
        Web3.to_checksum_address("0x5979D7b546E38E414F7E9822514be443A4800529"): 18,  # wstETH
        Web3.to_checksum_address("0xaf88d065e77c8cC2239327C5EDb3A432268e5831"): 6,   # USDC native (6!)
        Web3.to_checksum_address("0xFF970A61A04b1cA14834A43f5dE4533eBDDB5CC8"): 6,   # USDC.e bridged (6!)
        Web3.to_checksum_address("0xFd086bC7CD5C481DCC9C85ebE478A1C0b69FCbb9"): 6,   # USDT (6!)
        Web3.to_checksum_address("0xDA10009cBd5D07dd0CeCc66161FC93D7c9000da1"): 18,  # DAI
    }

    def __init__(self, w3: Web3):
        self._w3 = w3
        self._cache: dict[str, int] = dict(self._STATIC)

    def decimals(self, asset_address: str) -> int:
        addr = Web3.to_checksum_address(asset_address)
        if addr in self._cache:
            return self._cache[addr]

        logger.warning(
            "AssetRegistry: %s not in static registry — querying decimals() on-chain.",
            addr,
        )
        token = self._w3.eth.contract(address=addr, abi=ERC20_DECIMALS_ABI)
        d = token.functions.decimals().call()
        self._cache[addr] = d
        return d

    def scale_amount(self, asset_address: str, human_amount: Decimal | float | int) -> int:
        """Convert a human-readable token amount into integer base units using
        that asset's actual decimals rather than an assumed 18."""
        d = self.decimals(asset_address)
        return int(Decimal(str(human_amount)) * (Decimal(10) ** d))


# ---------------------------------------------------------------------------
# ROUTER_MAP — DEX name -> checksummed Router02 address (BSC mainnet)
# ---------------------------------------------------------------------------

# ROUTER_MAP — DEX name -> checksummed Router address (Arbitrum One mainnet).
# Only Uniswap-V2-style routers (swapExactTokensForTokens) belong here; that
# is the ABI shape LocalExecutor builds its swaps against. Kept in sync with
# scanner.py's _BSC_ROUTER_ADDRESSES / _BSC_EXECUTABLE_DEXES.
ROUTER_MAP: dict[str, str] = {
    "sushiswap": Web3.to_checksum_address("0x1b02dA8Cb0d097eB8D57A175b88c7D8b47997506"),  # SushiSwap V2 Router
    "camelot":   Web3.to_checksum_address("0xc873fEcbd354f5A56E00E710B90EF4201db2448d"),  # Camelot V2 Router
    # Uniswap V3 / Camelot V3 use different interfaces (exactInputSingle /
    # Algebra) — deliberately NOT mapped here under the V2-shaped ABI. See
    # V3_ROUTER_MAP below and FlashArbitrageV3.sol for the V3-shaped path.
}

# ---------------------------------------------------------------------------
# Uniswap V3 (Arbitrum One) — quoting/execution against a SEPARATE deployed
# contract (FlashArbitrageV3.sol), not the V2-only FlashArbitrageV2. See that
# contract's own header for why this can't just be added to ROUTER_MAP above.
# ---------------------------------------------------------------------------
V3_ROUTER_MAP: dict[str, str] = {
    "uniswapv3": Web3.to_checksum_address("0x68b3465833fb72A70ecDF485E0e4C7bD8665Fc45"),  # SwapRouter02
}
UNISWAP_V3_QUOTER_V2 = Web3.to_checksum_address("0x61fFE014bA17989E743c5F6cB21bF9697530B21e")
UNISWAP_V3_FEE_TIERS: tuple[int, ...] = (100, 500, 3000, 10000)


def resolve_v3_router(dex_name: str) -> str:
    key = dex_name.strip().lower()
    if key not in V3_ROUTER_MAP:
        raise ValueError(
            f"Unknown V3 DEX '{dex_name}' — not in V3_ROUTER_MAP. "
            f"Known: {', '.join(sorted(V3_ROUTER_MAP))}."
        )
    return V3_ROUTER_MAP[key]


FLASH_ARBITRAGE_V3_ABI = json.loads("""
[
  {"inputs": [], "name": "owner", "outputs": [{"internalType": "address", "name": "", "type": "address"}], "stateMutability": "view", "type": "function"},
  {"inputs": [], "name": "aavePool", "outputs": [{"internalType": "address", "name": "", "type": "address"}], "stateMutability": "view", "type": "function"},
  {"inputs": [
      {"internalType": "address", "name": "asset", "type": "address"},
      {"internalType": "uint256", "name": "amount", "type": "uint256"},
      {"internalType": "address", "name": "routerBuy", "type": "address"},
      {"internalType": "bool", "name": "buyIsV3", "type": "bool"},
      {"internalType": "uint24", "name": "buyFeeTier", "type": "uint24"},
      {"internalType": "address", "name": "routerSell", "type": "address"},
      {"internalType": "bool", "name": "sellIsV3", "type": "bool"},
      {"internalType": "uint24", "name": "sellFeeTier", "type": "uint24"},
      {"internalType": "address", "name": "intermediateToken", "type": "address"},
      {"internalType": "uint256", "name": "minIntermediateOut", "type": "uint256"},
      {"internalType": "uint256", "name": "minFinalOut", "type": "uint256"},
      {"internalType": "uint256", "name": "minProfit", "type": "uint256"}
    ], "name": "startArbitrageMixed", "outputs": [], "stateMutability": "nonpayable", "type": "function"},
  {"inputs": [
      {"internalType": "address", "name": "token", "type": "address"}
    ], "name": "withdraw", "outputs": [], "stateMutability": "nonpayable", "type": "function"},
  {"inputs": [], "name": "withdrawETH", "outputs": [], "stateMutability": "nonpayable", "type": "function"}
]
""")


def resolve_router(dex_name: str) -> str:
    """Look up a router address by DEX name. Raises rather than guessing if the
    name isn't mapped — an unresolved router must never fall back to a default."""
    key = dex_name.strip().lower()
    if key not in ROUTER_MAP:
        raise ValueError(
            f"Unknown DEX '{dex_name}' — not in ROUTER_MAP. "
            f"Known: {', '.join(sorted(ROUTER_MAP))}. "
            "Add it explicitly rather than assuming a default router."
        )
    return ROUTER_MAP[key]


class ContractManager:
    """Contract-specific logic for FlashArbitrageV2. Connection and signing are
    delegated to a Web3Client instance."""

    def __init__(self, client: Optional[Web3Client] = None):
        try:
            contract_address = os.environ["FLASH_ARBITRAGE_CONTRACT_ADDRESS"]
        except KeyError as exc:
            raise RuntimeError(
                f"ContractManager requires {exc.args[0]} to be set — "
                "on-chain commands (/chainstatus, /owner, /execute, "
                "/withdraw, /withdraweth) will not work until it is."
            ) from exc

        self.client = client or get_web3_client()
        self.w3 = self.client.w3
        self.contract = self.w3.eth.contract(
            address=Web3.to_checksum_address(contract_address),
            abi=FLASH_ARBITRAGE_V2_ABI,
        )
        self.assets = AssetRegistry(self.w3)

    # --- read calls ---------------------------------------------------

    def get_owner(self) -> str:
        return self.contract.functions.owner().call()

    def get_aave_pool(self) -> str:
        return self.contract.functions.aavePool().call()

    def get_eth_balance(self) -> Decimal:
        wei = self.w3.eth.get_balance(self.contract.address)
        return Decimal(wei) / Decimal(10**18)

    # --- write calls ----------------------------------------------------
    # Gated on-chain by onlyOwner. The signing wallet used by Web3Client MUST be
    # the same address as the contract's owner or every call below reverts.

    def start_arbitrage(
        self,
        asset,
        amount,
        router_buy,
        router_sell,
        intermediate_token,
        min_intermediate_out,
        min_final_out,
        min_profit,
    ) -> str:
        """Matches the deployed 8-arg startArbitrage exactly."""
        fn = self.contract.functions.startArbitrage(
            Web3.to_checksum_address(asset),
            amount,
            Web3.to_checksum_address(router_buy),
            Web3.to_checksum_address(router_sell),
            Web3.to_checksum_address(intermediate_token),
            min_intermediate_out,
            min_final_out,
            min_profit,
        )
        return self.client.send(fn)

    def start_arbitrage_by_name(
        self,
        asset: str,
        human_amount: Decimal | float | int,
        dex_buy: str,
        dex_sell: str,
        intermediate_token: str,
        min_intermediate_out: int,
        min_final_out: int,
        min_profit: int,
    ) -> str:
        """Convenience wrapper: resolves DEX names via ROUTER_MAP and scales
        human_amount via AssetRegistry."""
        router_buy = resolve_router(dex_buy)
        router_sell = resolve_router(dex_sell)
        amount = self.assets.scale_amount(asset, human_amount)

        return self.start_arbitrage(
            asset, amount,
            router_buy, router_sell,
            intermediate_token,
            min_intermediate_out, min_final_out, min_profit,
        )

    def withdraw(self, token: str, amount: Optional[int] = None) -> str:
        """Withdraw the FULL balance of `token` to the contract owner.

        The deployed FlashArbitrageV2.withdraw takes ONLY (address token) and
        always sweeps the full balance to owner. `amount` is accepted for
        backward compatibility with older callers but is IGNORED on-chain — the
        previous 2-arg encoding reverted with a selector mismatch and could not
        recover funds. VERIFY against your deployed contract on BscScan.
        """
        if amount is not None:
            logger.info(
                "ContractManager.withdraw: `amount` is ignored — the deployed "
                "1-arg withdraw sweeps the full %s balance to owner.", token,
            )
        fn = self.contract.functions.withdraw(Web3.to_checksum_address(token))
        return self.client.send(fn)

    def withdraw_eth(self) -> str:
        fn = self.contract.functions.withdrawETH()
        return self.client.send(fn)


class ContractManagerV3:
    """Contract-specific logic for FlashArbitrageV3 (mixed V2/V3 legs) — a
    SEPARATE, optional deployment from ContractManager's FlashArbitrageV2.
    Constructing this does NOT require FLASH_ARBITRAGE_CONTRACT_ADDRESS; it
    requires FLASH_ARBITRAGE_V3_CONTRACT_ADDRESS instead, and is only ever
    constructed when that's actually set (see get_contract_manager_v3())."""

    def __init__(self, client: Optional[Web3Client] = None):
        try:
            contract_address = os.environ["FLASH_ARBITRAGE_V3_CONTRACT_ADDRESS"]
        except KeyError as exc:
            raise RuntimeError(
                f"ContractManagerV3 requires {exc.args[0]} to be set — deploy "
                "Contrect/FlashArbitrageV3.sol yourself first, then set this."
            ) from exc

        self.client = client or get_web3_client()
        self.w3 = self.client.w3
        self.contract = self.w3.eth.contract(
            address=Web3.to_checksum_address(contract_address),
            abi=FLASH_ARBITRAGE_V3_ABI,
        )
        self.assets = AssetRegistry(self.w3)

    def get_owner(self) -> str:
        return self.contract.functions.owner().call()

    def start_arbitrage_mixed(
        self,
        asset,
        amount,
        router_buy, buy_is_v3: bool, buy_fee_tier: int,
        router_sell, sell_is_v3: bool, sell_fee_tier: int,
        intermediate_token,
        min_intermediate_out,
        min_final_out,
        min_profit,
    ) -> str:
        fn = self.contract.functions.startArbitrageMixed(
            Web3.to_checksum_address(asset),
            amount,
            Web3.to_checksum_address(router_buy), bool(buy_is_v3), int(buy_fee_tier),
            Web3.to_checksum_address(router_sell), bool(sell_is_v3), int(sell_fee_tier),
            Web3.to_checksum_address(intermediate_token),
            min_intermediate_out,
            min_final_out,
            min_profit,
        )
        return self.client.send(fn)

    def withdraw(self, token: str) -> str:
        fn = self.contract.functions.withdraw(Web3.to_checksum_address(token))
        return self.client.send(fn)

    def withdraw_eth(self) -> str:
        fn = self.contract.functions.withdrawETH()
        return self.client.send(fn)


_contract_manager: Optional[ContractManager] = None
_contract_manager_v3: Optional[ContractManagerV3] = None
_contract_manager_v3_unavailable = False


def get_contract_manager() -> ContractManager:
    global _contract_manager
    if _contract_manager is None:
        _contract_manager = ContractManager()
    return _contract_manager


def get_contract_manager_v3() -> Optional[ContractManagerV3]:
    """Returns None (never raises) when FLASH_ARBITRAGE_V3_CONTRACT_ADDRESS
    isn't set — this is an opt-in feature, unlike get_contract_manager()'s
    required V2 contract. Cached the same way, including the "unavailable"
    outcome, so a missing env var doesn't retry construction every call."""
    global _contract_manager_v3, _contract_manager_v3_unavailable
    if _contract_manager_v3 is not None:
        return _contract_manager_v3
    if _contract_manager_v3_unavailable:
        return None
    try:
        _contract_manager_v3 = ContractManagerV3()
        return _contract_manager_v3
    except RuntimeError as exc:
        logger.info("[ContractManagerV3] not available: %s", exc)
        _contract_manager_v3_unavailable = True
        return None