bot.py

# -*- coding: utf-8 -*-
"""
Core bot class for Uomi testnet automation.
Loads wallets, executes tasks in sequence, and tracks results.
"""

from __future__ import annotations

import json
import logging
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

_log = logging.getLogger(__name__)

BASE_DIR = Path(__file__).resolve().parent
HISTORY_FILE = BASE_DIR / "history.json"


class UomiBot:
    """Testnet automation bot that runs configured tasks across wallets."""

    def __init__(self, config: Dict[str, Any]):
        self._config = config
        self._wallets = config.get("wallets", [])
        self._rpc_url = config.get("rpc_url", "")
        self._chain_id = config.get("chain_id", 0)
        self._gas_limit = config.get("gas_limit", 300000)
        self._gas_price_gwei = config.get("gas_price_gwei", 1)
        self._interval = config.get("task_interval_sec", 60)
        self._proxy = config.get("proxy") or None
        self._history = self._load_history()
        self._running = False

    def _load_history(self) -> Dict[str, Any]:
        if HISTORY_FILE.exists():
            try:
                with open(HISTORY_FILE, "r", encoding="utf-8") as fh:
                    return json.load(fh)
            except (json.JSONDecodeError, IOError):
                pass
        return {}

    def _save_history(self) -> None:
        try:
            with open(HISTORY_FILE, "w", encoding="utf-8") as fh:
                json.dump(self._history, fh, indent=2, ensure_ascii=False)
        except IOError as exc:
            _log.error("Failed to save history: %s", exc)

    def _get_web3(self):
        try:
            from web3 import Web3
        except ImportError:
            raise ImportError("web3 is required. Install with: pip install web3")

        if self._proxy:
            from web3 import HTTPProvider
            provider = HTTPProvider(
                self._rpc_url,
                request_kwargs={"proxies": {"http": self._proxy, "https": self._proxy}},
            )
        else:
            provider = Web3.HTTPProvider(self._rpc_url)

        w3 = Web3(provider)
        if not w3.is_connected():
            raise ConnectionError(f"Cannot connect to RPC: {self._rpc_url}")
        return w3

    def _get_wallet_history(self, address: str) -> Dict[str, Any]:
        if address not in self._history:
            self._history[address] = {
                "faucet_claims": 0,
                "staked_amount": "0",
                "governance_votes": 0,
                "last_activity": "",
            }
        return self._history[address]

    def _resolve_task(self, task_name: str):
        from tasks import get_task_executor
        return get_task_executor(task_name)

    def run(self, task_names: List[str]) -> None:
        self._running = True
        _log.info("Bot started with %d wallet(s) and %d task(s)", len(self._wallets), len(task_names))

        try:
            w3 = self._get_web3()
        except Exception as exc:
            _log.error("Failed to connect: %s", exc)
            print(f"  [!] Connection failed: {exc}")
            return

        cycle = 0
        while self._running:
            cycle += 1
            print(f"\n--- Cycle {cycle} ---")

            for wallet in self._wallets:
                address = wallet.get("address", "")
                pk = wallet.get("private_key", "")
                label = wallet.get("label", address[:8])

                if not pk:
                    print(f"  [{label}] Skipping: no private key")
                    continue

                print(f"  [{label}] Processing wallet {address[:8]}...{address[-4:]}")
                wallet_hist = self._get_wallet_history(address)

                for task_name in task_names:
                    try:
                        executor = self._resolve_task(task_name)
                        if executor is None:
                            print(f"    [{task_name}] Unknown task, skipping")
                            continue

                        print(f"    [{task_name}] Executing...")
                        result = executor(
                            w3=w3,
                            address=address,
                            private_key=pk,
                            chain_id=self._chain_id,
                            gas_limit=self._gas_limit,
                            gas_price_gwei=self._gas_price_gwei,
                            config=self._config,
                        )

                        if result.get("success"):
                            print(f"    [{task_name}] Success: {result.get('message', 'OK')}")
                            self._update_history(wallet_hist, task_name, result)
                        else:
                            print(f"    [{task_name}] Failed: {result.get('message', 'unknown error')}")

                    except Exception as exc:
                        print(f"    [{task_name}] Error: {exc}")
                        _log.exception("Task %s failed for %s", task_name, address[:8])

                wallet_hist["last_activity"] = time.strftime("%Y-%m-%d %H:%M:%S")

            self._save_history()
            print(f"\n  Waiting {self._interval}s before next cycle (Ctrl+C to stop)...")

            try:
                time.sleep(self._interval)
            except KeyboardInterrupt:
                self._running = False
                break

        print("\nBot stopped.")

    def _update_history(
        self, wallet_hist: Dict[str, Any], task_name: str, result: Dict[str, Any]
    ) -> None:
        if task_name == "faucet":
            wallet_hist["faucet_claims"] = wallet_hist.get("faucet_claims", 0) + 1
        elif task_name == "staking":
            amount = result.get("staked", "0")
            wallet_hist["staked_amount"] = amount
        elif task_name == "governance":
            wallet_hist["governance_votes"] = wallet_hist.get("governance_votes", 0) + 1

    def stop(self) -> None:
        self._running = False