universe/scripts/mm-bot/chains.go

package main

import (
	"context"
	"fmt"
	"log"
	"math/big"
	"sync"

	"github.com/luxfi/geth/ethclient"
)

// ChainClient holds a connection to a single EVM chain.
type ChainClient struct {
	Config ChainConfig
	Client *ethclient.Client
	AMM    *AMMClient
}

// ChainManager manages connections to multiple EVM chains.
type ChainManager struct {
	mu      sync.RWMutex
	clients map[string]*ChainClient // keyed by chain name (lowercase)
}

// NewChainManager creates a ChainManager and connects to all configured chains.
func NewChainManager(chains []ChainConfig) (*ChainManager, error) {
	cm := &ChainManager{
		clients: make(map[string]*ChainClient),
	}

	for _, cc := range chains {
		client, err := ethclient.Dial(cc.RPC)
		if err != nil {
			log.Printf("[%s] WARNING: failed to connect to %s: %v", cc.Name, cc.RPC, err)
			continue
		}

		// Verify chain ID
		ctx := context.Background()
		chainID, err := client.ChainID(ctx)
		if err != nil {
			log.Printf("[%s] WARNING: failed to get chain ID: %v", cc.Name, err)
			client.Close()
			continue
		}

		if chainID.Cmp(cc.ChainID) != 0 {
			log.Printf("[%s] WARNING: chain ID mismatch: expected %s, got %s",
				cc.Name, cc.ChainID.String(), chainID.String())
			client.Close()
			continue
		}

		amm := NewAMMClient(client, cc)
		cm.clients[cc.Name] = &ChainClient{
			Config: cc,
			Client: client,
			AMM:    amm,
		}

		log.Printf("[%s] Connected (chainID=%s, rpc=%s)", cc.Name, chainID.String(), cc.RPC)
	}

	if len(cm.clients) == 0 {
		return nil, fmt.Errorf("failed to connect to any chain")
	}

	return cm, nil
}

// Get returns the client for a named chain.
func (cm *ChainManager) Get(name string) *ChainClient {
	cm.mu.RLock()
	defer cm.mu.RUnlock()
	return cm.clients[name]
}

// All returns all connected chain clients.
func (cm *ChainManager) All() []*ChainClient {
	cm.mu.RLock()
	defer cm.mu.RUnlock()

	result := make([]*ChainClient, 0, len(cm.clients))
	for _, cc := range cm.clients {
		result = append(result, cc)
	}
	return result
}

// Close closes all chain connections.
func (cm *ChainManager) Close() {
	cm.mu.Lock()
	defer cm.mu.Unlock()

	for name, cc := range cm.clients {
		cc.Client.Close()
		log.Printf("[%s] Disconnected", name)
	}
}

// GetNativeBalance returns the native LUX balance on a chain.
func (cc *ChainClient) GetNativeBalance(ctx context.Context, wallet *Wallet) (*big.Int, error) {
	return cc.Client.BalanceAt(ctx, wallet.Address, nil)
}

// PrintBalances logs the native and token balances for a wallet on this chain.
func (cc *ChainClient) PrintBalances(ctx context.Context, wallet *Wallet) {
	native, err := cc.GetNativeBalance(ctx, wallet)
	if err != nil {
		log.Printf("[%s] %s native balance error: %v", cc.Config.Name, wallet.ShortAddr(), err)
		return
	}

	wlux, _ := cc.AMM.GetTokenBalance(ctx, cc.Config.WLUX, wallet.Address)
	leth, _ := cc.AMM.GetTokenBalance(ctx, cc.Config.LETH, wallet.Address)
	lbtc, _ := cc.AMM.GetTokenBalance(ctx, cc.Config.LBTC, wallet.Address)
	lusdc, _ := cc.AMM.GetTokenBalance(ctx, cc.Config.LUSDC, wallet.Address)

	log.Printf("[%s] %s balances: native=%s WLUX=%s LETH=%s LBTC=%s LUSDC=%s",
		cc.Config.Name,
		wallet.ShortAddr(),
		formatWei(native),
		formatWei(wlux),
		formatWei(leth),
		formatWei(lbtc),
		formatWei(lusdc),
	)
}

// formatWei formats a wei value as a human-readable string with 4 decimal places.
func formatWei(wei *big.Int) string {
	if wei == nil {
		return "0"
	}

	// Convert to float for display
	f := new(big.Float).SetInt(wei)
	divisor := new(big.Float).SetInt(new(big.Int).Exp(big.NewInt(10), big.NewInt(18), nil))
	f.Quo(f, divisor)

	result, _ := f.Float64()
	return fmt.Sprintf("%.4f", result)
}