modules/executor.py

"""
modules/executor.py — ChainExecutor for On-Chain Execution
──────────────────────────────────────────────────────────

Signs and broadcasts native-token sweep transactions with mandatory
pre-flight simulation and resilient broadcast retry logic.
"""

import concurrent.futures
import logging
import time
from typing import Any, Optional

logger = logging.getLogger(__name__)

try:
    from web3 import Web3
    try:
        from web3.middleware import geth_poa_middleware as _POA_MW
    except ImportError:
        from web3.middleware import ExtraDataToPOAMiddleware as _POA_MW

    def _to_checksum(addr: str) -> str:
        try:
            return Web3.to_checksum_address(addr)
        except AttributeError:
            return Web3.toChecksumAddress(addr)

    def _raw_tx(signed) -> bytes:
        return getattr(signed, "raw_transaction", None) or signed.rawTransaction

    _WEB3_AVAILABLE = True
except ImportError:
    _WEB3_AVAILABLE = False
    _POA_MW = None


# Migrated 2026-07 to Arbitrum One. The sweep/transfer executor uses these for
# building native-token payout transactions on the configured PAYOUT_CHAIN.
_CHAIN_CFG = {
    "ARB": {
        "rpc_env": "ARB_RPC_URL",
        "rpc_default": "https://arb1.arbitrum.io/rpc",
        "rpc_fallbacks": [
            "https://arbitrum-one.publicnode.com",
            "https://rpc.ankr.com/arbitrum",
        ],
        "chain_id": 42161,
        "max_gas_gwei": 2.0,
        "native": "ETH",
        "is_poa": True,
    },
    # Legacy alias -> Arbitrum so an old PAYOUT_CHAIN=BSC can't reach a dead node.
    "BSC": {
        "rpc_env": "ARB_RPC_URL",
        "rpc_default": "https://arb1.arbitrum.io/rpc",
        "rpc_fallbacks": [
            "https://arbitrum-one.publicnode.com",
            "https://rpc.ankr.com/arbitrum",
        ],
        "chain_id": 42161,
        "max_gas_gwei": 2.0,
        "native": "ETH",
        "is_poa": True,
    },
}

GAS_LIMIT_TRANSFER = 21_000
BROADCAST_MAX_RETRIES = 4
BROADCAST_RETRY_DELAY_SEC = 2.5


def _sanitize_private_key(raw: str) -> str:
    """Clean a hot-wallet private key from environment variable."""
    if not raw:
        return ""
    key = raw.strip()
    if len(key) >= 2 and key[0] == key[-1] and key[0] in ("'", '"'):
        key = key[1:-1].strip()
    if key and not key.lower().startswith("0x"):
        key = "0x" + key
    return key


def _extract_revert_reason(exc: Exception) -> str:
    """Best-effort extraction of human-readable revert reason."""
    try:
        if exc.args and isinstance(exc.args[0], dict):
            return str(exc.args[0].get("message") or exc.args[0])
    except Exception:
        pass
    return str(exc)


def _is_transient_rpc_error(exc: Exception) -> bool:
    """True for infrastructure-level failures worth a short retry."""
    text = str(exc).lower()
    transient_markers = (
        "502", "bad gateway",
        "504", "gateway timeout",
        "503", "service unavailable",
        "timed out", "timeout",
        "connection reset", "connection aborted", "connection refused",
        "econnreset", "remote end closed", "broken pipe",
        "temporarily unavailable", "max retries exceeded",
    )
    return any(m in text for m in transient_markers)


def _is_hard_revert_error(exc: Exception) -> bool:
    """True for definitive EVM-level revert or fund/gas errors."""
    text = str(exc).lower()
    hard_revert_markers = (
        "revert", "execution reverted",
        "insufficient funds",
        "gas required exceeds allowance",
        "always failing transaction",
        "out of gas", "invalid opcode",
    )
    return any(m in text for m in hard_revert_markers)


class ChainExecutor:
    """
    Signs and broadcasts a native-token sweep transaction with full
    Simulate-Before-Strike pipeline and resilient broadcast retry.

    Private key is held in memory and never written to disk or logs.
    """

    def __init__(self, chain: str, private_key: str, rpc_url: Optional[str] = None) -> None:
        """Initialize ChainExecutor."""
        if not _WEB3_AVAILABLE:
            raise RuntimeError(
                "web3.py is required for on-chain execution: pip install web3"
            )

        clean_key = _sanitize_private_key(private_key)
        if not clean_key:
            raise ValueError("private_key must not be empty")

        self.chain = chain.upper()

        import os
        if self.chain not in ("ARB", "BSC") and os.getenv("ALLOW_ETH", "0") != "1":
            raise RuntimeError(
                f"ChainExecutor: chain='{self.chain}' rejected — this "
                f"deployment runs on Arbitrum One (ARB). Set ALLOW_ETH=1 "
                f"to enable the optional ETH side."
            )

        cfg = _CHAIN_CFG.get(self.chain, _CHAIN_CFG["ARB"])

        env_rpc = os.getenv(cfg["rpc_env"])
        candidates = [rpc_url or env_rpc or cfg["rpc_default"]]
        candidates += [u for u in cfg.get("rpc_fallbacks", []) if u not in candidates]

        self._chain_id = cfg["chain_id"]
        self._rpc_candidates = candidates
        self._w3 = None
        self._active_rpc = None

        last_exc: Optional[Exception] = None
        for rpc in candidates:
            try:
                w3 = Web3(Web3.HTTPProvider(rpc, request_kwargs={"timeout": 10}))
                if cfg["is_poa"] and _POA_MW:
                    w3.middleware_onion.inject(_POA_MW, layer=0)
                # Cheap liveness check before committing to this endpoint.
                w3.eth.chain_id
                self._w3 = w3
                self._active_rpc = rpc
                break
            except Exception as exc:
                last_exc = exc
                logger.warning(
                    "[ChainExecutor] RPC candidate %s unreachable on init "
                    "(%s) — trying next", rpc, exc,
                )

        if self._w3 is None:
            raise RuntimeError(
                f"All {len(candidates)} RPC candidate(s) for {self.chain} "
                f"failed on init. Last error: {last_exc}"
            )

        if self._active_rpc != candidates[0]:
            logger.info(
                "[ChainExecutor] Using fallback RPC for %s: %s",
                self.chain, self._active_rpc,
            )

        try:
            self._account = self._w3.eth.account.from_key(clean_key)
        except Exception as exc:
            raise ValueError(
                f"PAYOUT_PRIVATE_KEY failed to derive a valid account "
                f"(len={len(clean_key)} after sanitization). Check that it "
                f"is a 32-byte hex key, optionally 0x-prefixed, with no "
                f"extra whitespace or quotes. Underlying error: {exc}"
            ) from exc
        finally:
            del clean_key

        # v2.1 — full address logged (not just the first 12 chars) and
        # explicitly labeled "SIGNING/GAS address". Production incident:
        # this line previously truncated to "0x4c56dF1c7D..." which was
        # easy to eyeball-confuse with a similar-looking but DIFFERENT
        # PAYOUT_WALLET value also floating around in the same debug
        # session (0x2C25..., 0xe21a...) — all three addresses share the
        # "0x" + short prefix shape at a glance. The truncation also made
        # it impossible to grep/diff this log line against a funded
        # address without further digging. This address is never secret
        # (it's a public wallet address, not the private key), so logging
        # it in full has no security cost.
        logger.info(
            "[ChainExecutor] Ready on %s — SIGNING/GAS wallet (fund THIS "
            "one, not PAYOUT_WALLET): %s",
            self.chain, self._account.address,
        )

    def _rotate_rpc(self) -> bool:
        """
        Try switching to the next untried candidate RPC. Returns True if a
        healthy endpoint was found and self._w3 was swapped, False if every
        candidate has now been tried this run.
        """
        remaining = [u for u in self._rpc_candidates if u != self._active_rpc]
        for rpc in remaining:
            try:
                w3 = Web3(Web3.HTTPProvider(rpc, request_kwargs={"timeout": 10}))
                cfg = _CHAIN_CFG.get(self.chain, _CHAIN_CFG["ARB"])
                if cfg["is_poa"] and _POA_MW:
                    w3.middleware_onion.inject(_POA_MW, layer=0)
                w3.eth.chain_id
                self._w3 = w3
                self._active_rpc = rpc
                logger.warning(
                    "[ChainExecutor] Rotated %s RPC to fallback: %s",
                    self.chain, rpc,
                )
                return True
            except Exception as exc:
                logger.warning(
                    "[ChainExecutor] Fallback RPC %s also unreachable (%s)",
                    rpc, exc,
                )
        return False

    @property
    def hot_address(self) -> str:
        return self._account.address

    def is_same_address(self, other: str) -> bool:
        """
        v2.1 — case-insensitive comparison helper. Addresses are not
        case-sensitive at the protocol level (checksum casing is a display
        convention only), so comparing hot_address == some_config_value
        with a plain == is a latent bug: two representations of the same
        address can differ only in casing and compare unequal. Use this
        instead of manual string comparison anywhere PAYOUT_WALLET or any
        other address is checked against the signing address.
        """
        if not other:
            return False
        return self._account.address.strip().lower() == other.strip().lower()

    def get_balance_wei(self) -> int:
        return self._w3.eth.get_balance(self._account.address)

    def get_balance_native(self) -> float:
        return self.get_balance_wei() / 1e18

    def build_transfer_tx(
        self,
        to: str,
        amount_wei: int,
        gas_price_wei: int,
        gas_limit: int = GAS_LIMIT_TRANSFER,
    ) -> dict:
        """Build a transfer transaction."""
        nonce = self._w3.eth.get_transaction_count(self._account.address, "pending")
        return {
            "nonce": nonce,
            "to": _to_checksum(to),
            "value": amount_wei,
            "gas": gas_limit,
            "gasPrice": gas_price_wei,
            "chainId": self._chain_id,
        }

    def sign_and_send(self, tx: dict) -> str:
        """Sign tx with stored key and broadcast. Return 0x-prefixed tx_hash."""
        signed = self._w3.eth.account.sign_transaction(tx, self._account.key)
        raw = _raw_tx(signed)
        tx_bytes = self._w3.eth.send_raw_transaction(raw)
        h = tx_bytes.hex()
        return h if h.startswith("0x") else "0x" + h

    def _simulate_with_retry(
        self,
        tx: dict,
        max_retries: int = 2,
        retry_delay: float = 1.0,
    ) -> tuple[bool, str]:
        """Pre-flight simulation via eth_estimateGas with retry logic."""
        last_reason = ""
        for attempt in range(1, max_retries + 1):
            try:
                sim_tx = {k: v for k, v in tx.items() if k != "nonce"}
                estimated = self._w3.eth.estimate_gas(sim_tx)
                return True, f"simulation ok — estimated_gas={estimated}"
            except Exception as exc:
                reason = _extract_revert_reason(exc)
                last_reason = reason
                if _is_transient_rpc_error(exc) and attempt < max_retries:
                    logger.warning(
                        "[ChainExecutor] Simulation attempt %d/%d hit a "
                        "transient RPC error (%s) — rotating RPC and "
                        "retrying in %.1fs",
                        attempt, max_retries, exc, retry_delay,
                    )
                    self._rotate_rpc()
                    time.sleep(retry_delay)
                    continue
                if _is_hard_revert_error(exc):
                    return False, reason
                return False, f"unclassified simulation error: {reason}"
        return False, last_reason or "simulation failed after retries"

    def simulate_transaction(self, tx: dict) -> tuple[bool, str]:
        """Single-shot simulation without retry."""
        try:
            sim_tx = {k: v for k, v in tx.items() if k != "nonce"}
            estimated = self._w3.eth.estimate_gas(sim_tx)
            return True, f"simulation ok — estimated_gas={estimated}"
        except Exception as exc:
            return False, _extract_revert_reason(exc)

    def _sign_locally(self, tx: dict) -> bytes:
        """Local deterministic signing, safe to run concurrently."""
        signed = self._w3.eth.account.sign_transaction(tx, self._account.key)
        return _raw_tx(signed)

    def simulate_and_broadcast(
        self,
        tx: dict,
        max_retries: int = BROADCAST_MAX_RETRIES,
        retry_delay: float = BROADCAST_RETRY_DELAY_SEC,
        simulate_retries: int = 2,
        simulate_retry_delay: float = 1.0,
    ) -> dict:
        """
        Full Simulate-Before-Strike pipeline.

        Runs simulation and local signing in parallel, then broadcasts with
        transient-only retries.
        """
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
            sim_future = pool.submit(
                self._simulate_with_retry, tx, simulate_retries, simulate_retry_delay
            )
            sign_future = pool.submit(self._sign_locally, tx)

            sim_ok, sim_msg = sim_future.result()

            if not sim_ok:
                logger.error(
                    "[ChainExecutor] Pre-flight simulation FAILED — aborting "
                    "before broadcast: %s", sim_msg,
                )
                try:
                    sign_future.result(timeout=5)
                except Exception:
                    pass
                return {
                    "ok": False,
                    "tx_hash": None,
                    "stage": "simulate",
                    "msg": f"Simulation predicts revert: {sim_msg}",
                    "attempts": 0,
                }
            logger.info("[ChainExecutor] Pre-flight simulation passed (%s)", sim_msg)

            try:
                raw = sign_future.result()
            except Exception as exc:
                return {
                    "ok": False,
                    "tx_hash": None,
                    "stage": "sign",
                    "msg": f"Local signing failed: {exc}",
                    "attempts": 0,
                }

        last_exc: Optional[Exception] = None
        for attempt in range(1, max_retries + 1):
            try:
                tx_bytes = self._w3.eth.send_raw_transaction(raw)
                h = tx_bytes.hex()
                tx_hash = h if h.startswith("0x") else "0x" + h
                return {
                    "ok": True,
                    "tx_hash": tx_hash,
                    "stage": "broadcast",
                    "msg": "broadcast ok",
                    "attempts": attempt,
                }
            except Exception as exc:
                last_exc = exc
                if _is_transient_rpc_error(exc) and attempt < max_retries:
                    logger.warning(
                        "[ChainExecutor] Broadcast attempt %d/%d hit transient "
                        "RPC error (%s) — rotating RPC and retrying in %.1fs",
                        attempt, max_retries, exc, retry_delay,
                    )
                    self._rotate_rpc()
                    time.sleep(retry_delay)
                    continue
                logger.error(
                    "[ChainExecutor] Broadcast attempt %d/%d failed (%s)",
                    attempt, max_retries, exc,
                )
                break

        return {
            "ok": False,
            "tx_hash": None,
            "stage": "broadcast",
            "msg": f"Broadcast failed after {max_retries} attempt(s): {last_exc}",
            "attempts": max_retries,
        }

    def wait_receipt(self, tx_hash: str, timeout: int = 120) -> dict:
        """Poll for receipt until confirmed or timeout."""
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            try:
                r = self._w3.eth.get_transaction_receipt(tx_hash)
                if r is not None:
                    return dict(r)
            except Exception:
                pass
            time.sleep(3)
        raise TimeoutError(f"Receipt not found within {timeout}s: {tx_hash}")

    def sweep(
        self,
        cold_wallet: str,
        gas_engine: Any,
        reserve_txcount: int = 20,
        gas_multiplier: float = 1.0,
        wait: bool = True,
        timeout: int = 120,
    ) -> dict:
        """Full sweep pipeline: calculate amount → build → sign → broadcast → confirm."""
        gas_price_wei = int(gas_engine.get_gas_price_wei() * gas_multiplier)
        balance_wei = self.get_balance_wei()
        tx_gas_cost = gas_price_wei * GAS_LIMIT_TRANSFER
        reserve_wei = gas_engine.gas_reserve_wei(reserve_txcount)
        amount_wei = balance_wei - tx_gas_cost - reserve_wei

        if amount_wei <= 0:
            raise ValueError(
                f"Insufficient balance. "
                f"balance={balance_wei} tx_gas={tx_gas_cost} reserve={reserve_wei}"
            )

        tx = self.build_transfer_tx(cold_wallet, amount_wei, gas_price_wei)
        outcome = self.simulate_and_broadcast(tx)
        if not outcome["ok"]:
            raise RuntimeError(
                f"Sweep aborted at stage='{outcome['stage']}': {outcome['msg']}"
            )
        tx_hash = outcome["tx_hash"]

        logger.info(
            "[ChainExecutor] Sweep sent: %d wei → %s... | %.2f Gwei | tx=%s...",
            amount_wei, cold_wallet[:12], gas_price_wei / 1e9, tx_hash[:20],
        )

        result = {
            "tx_hash": tx_hash,
            "amount_wei": amount_wei,
            "amount_native": amount_wei / 1e18,
            "gas_price_wei": gas_price_wei,
            "gas_price_gwei": round(gas_price_wei / 1e9, 4),
            "confirmed": False,
            "block_number": None,
            "status": None,
        }

        if wait:
            try:
                receipt = self.wait_receipt(tx_hash, timeout)
                result["confirmed"] = True
                result["block_number"] = receipt.get("blockNumber")
                result["status"] = receipt.get("status")
            except TimeoutError:
                logger.warning(
                    "[ChainExecutor] Receipt timeout — tx still pending: %s...",
                    tx_hash[:20],
                )

        return result