universe/scripts/mm-bot/amm.go

package main

import (
	"context"
	"crypto/ecdsa"
	"fmt"
	"math/big"
	"strings"
	"time"

	ethereum "github.com/luxfi/geth"
	"github.com/luxfi/geth/accounts/abi"
	"github.com/luxfi/geth/common"
	"github.com/luxfi/geth/core/types"
	gethcrypto "github.com/luxfi/geth/crypto"
	"github.com/luxfi/geth/ethclient"
)

// ABI JSON strings for AMM V2 Router and ERC20 interactions.
const routerABIJSON = `[
	{
		"name": "swapExactTokensForTokens",
		"type": "function",
		"inputs": [
			{"name": "amountIn", "type": "uint256"},
			{"name": "amountOutMin", "type": "uint256"},
			{"name": "path", "type": "address[]"},
			{"name": "to", "type": "address"},
			{"name": "deadline", "type": "uint256"}
		],
		"outputs": [{"type": "uint256[]"}],
		"stateMutability": "nonpayable"
	},
	{
		"name": "addLiquidity",
		"type": "function",
		"inputs": [
			{"name": "tokenA", "type": "address"},
			{"name": "tokenB", "type": "address"},
			{"name": "amountADesired", "type": "uint256"},
			{"name": "amountBDesired", "type": "uint256"},
			{"name": "amountAMin", "type": "uint256"},
			{"name": "amountBMin", "type": "uint256"},
			{"name": "to", "type": "address"},
			{"name": "deadline", "type": "uint256"}
		],
		"outputs": [
			{"type": "uint256"},
			{"type": "uint256"},
			{"type": "uint256"}
		],
		"stateMutability": "nonpayable"
	},
	{
		"name": "removeLiquidity",
		"type": "function",
		"inputs": [
			{"name": "tokenA", "type": "address"},
			{"name": "tokenB", "type": "address"},
			{"name": "liquidity", "type": "uint256"},
			{"name": "amountAMin", "type": "uint256"},
			{"name": "amountBMin", "type": "uint256"},
			{"name": "to", "type": "address"},
			{"name": "deadline", "type": "uint256"}
		],
		"outputs": [
			{"type": "uint256"},
			{"type": "uint256"}
		],
		"stateMutability": "nonpayable"
	},
	{
		"name": "getAmountsOut",
		"type": "function",
		"inputs": [
			{"name": "amountIn", "type": "uint256"},
			{"name": "path", "type": "address[]"}
		],
		"outputs": [{"type": "uint256[]"}],
		"stateMutability": "view"
	},
	{
		"name": "WETH",
		"type": "function",
		"inputs": [],
		"outputs": [{"type": "address"}],
		"stateMutability": "pure"
	}
]`

const erc20ABIJSON = `[
	{
		"name": "approve",
		"type": "function",
		"inputs": [
			{"name": "spender", "type": "address"},
			{"name": "amount", "type": "uint256"}
		],
		"outputs": [{"type": "bool"}],
		"stateMutability": "nonpayable"
	},
	{
		"name": "balanceOf",
		"type": "function",
		"inputs": [{"name": "account", "type": "address"}],
		"outputs": [{"type": "uint256"}],
		"stateMutability": "view"
	},
	{
		"name": "allowance",
		"type": "function",
		"inputs": [
			{"name": "owner", "type": "address"},
			{"name": "spender", "type": "address"}
		],
		"outputs": [{"type": "uint256"}],
		"stateMutability": "view"
	},
	{
		"name": "transfer",
		"type": "function",
		"inputs": [
			{"name": "to", "type": "address"},
			{"name": "amount", "type": "uint256"}
		],
		"outputs": [{"type": "bool"}],
		"stateMutability": "nonpayable"
	}
]`

const wluxABIJSON = `[
	{
		"name": "deposit",
		"type": "function",
		"inputs": [],
		"outputs": [],
		"stateMutability": "payable"
	},
	{
		"name": "withdraw",
		"type": "function",
		"inputs": [{"name": "wad", "type": "uint256"}],
		"outputs": [],
		"stateMutability": "nonpayable"
	}
]`

const factoryABIJSON = `[
	{
		"name": "getPair",
		"type": "function",
		"inputs": [
			{"name": "tokenA", "type": "address"},
			{"name": "tokenB", "type": "address"}
		],
		"outputs": [{"type": "address"}],
		"stateMutability": "view"
	},
	{
		"name": "allPairsLength",
		"type": "function",
		"inputs": [],
		"outputs": [{"type": "uint256"}],
		"stateMutability": "view"
	}
]`

const pairABIJSON = `[
	{
		"name": "getReserves",
		"type": "function",
		"inputs": [],
		"outputs": [
			{"type": "uint112"},
			{"type": "uint112"},
			{"type": "uint32"}
		],
		"stateMutability": "view"
	},
	{
		"name": "token0",
		"type": "function",
		"inputs": [],
		"outputs": [{"type": "address"}],
		"stateMutability": "view"
	},
	{
		"name": "token1",
		"type": "function",
		"inputs": [],
		"outputs": [{"type": "address"}],
		"stateMutability": "view"
	}
]`

var (
	routerABI  abi.ABI
	erc20ABI   abi.ABI
	wluxABI    abi.ABI
	factoryABI abi.ABI
	pairABI    abi.ABI
)

func init() {
	var err error
	routerABI, err = abi.JSON(strings.NewReader(routerABIJSON))
	if err != nil {
		panic("failed to parse router ABI: " + err.Error())
	}
	erc20ABI, err = abi.JSON(strings.NewReader(erc20ABIJSON))
	if err != nil {
		panic("failed to parse ERC20 ABI: " + err.Error())
	}
	wluxABI, err = abi.JSON(strings.NewReader(wluxABIJSON))
	if err != nil {
		panic("failed to parse WLUX ABI: " + err.Error())
	}
	factoryABI, err = abi.JSON(strings.NewReader(factoryABIJSON))
	if err != nil {
		panic("failed to parse factory ABI: " + err.Error())
	}
	pairABI, err = abi.JSON(strings.NewReader(pairABIJSON))
	if err != nil {
		panic("failed to parse pair ABI: " + err.Error())
	}
}

// TradePair represents a token pair for trading.
type TradePair struct {
	TokenA common.Address
	TokenB common.Address
	NameA  string
	NameB  string
}

// AMMClient wraps ethclient and provides AMM interaction methods.
type AMMClient struct {
	client   *ethclient.Client
	chainCfg ChainConfig
}

// NewAMMClient creates a new AMM client for the given chain.
func NewAMMClient(client *ethclient.Client, chainCfg ChainConfig) *AMMClient {
	return &AMMClient{
		client:   client,
		chainCfg: chainCfg,
	}
}

// addressFromKey derives the Ethereum address from a private key.
func addressFromKey(key *ecdsa.PrivateKey) common.Address {
	return gethcrypto.PubkeyToAddress(key.PublicKey)
}

// GetTokenBalance returns the ERC20 token balance of an address.
func (a *AMMClient) GetTokenBalance(ctx context.Context, token, account common.Address) (*big.Int, error) {
	data, err := erc20ABI.Pack("balanceOf", account)
	if err != nil {
		return nil, fmt.Errorf("pack balanceOf: %w", err)
	}

	result, err := a.client.CallContract(ctx, ethereum.CallMsg{
		To:   &token,
		Data: data,
	}, nil)
	if err != nil {
		return nil, fmt.Errorf("call balanceOf: %w", err)
	}

	out, err := erc20ABI.Unpack("balanceOf", result)
	if err != nil {
		return nil, fmt.Errorf("unpack balanceOf: %w", err)
	}

	return out[0].(*big.Int), nil
}

// GetAllowance returns the ERC20 allowance of owner for spender.
func (a *AMMClient) GetAllowance(ctx context.Context, token, owner, spender common.Address) (*big.Int, error) {
	data, err := erc20ABI.Pack("allowance", owner, spender)
	if err != nil {
		return nil, fmt.Errorf("pack allowance: %w", err)
	}

	result, err := a.client.CallContract(ctx, ethereum.CallMsg{
		To:   &token,
		Data: data,
	}, nil)
	if err != nil {
		return nil, fmt.Errorf("call allowance: %w", err)
	}

	out, err := erc20ABI.Unpack("allowance", result)
	if err != nil {
		return nil, fmt.Errorf("unpack allowance: %w", err)
	}

	return out[0].(*big.Int), nil
}

// ApproveToken sends an ERC20 approve transaction.
func (a *AMMClient) ApproveToken(ctx context.Context, key *ecdsa.PrivateKey, token, spender common.Address, amount *big.Int) (common.Hash, error) {
	data, err := erc20ABI.Pack("approve", spender, amount)
	if err != nil {
		return common.Hash{}, fmt.Errorf("pack approve: %w", err)
	}

	return a.sendTx(ctx, key, &token, big.NewInt(0), data, 100_000)
}

// GetAmountsOut queries the router for expected output amounts.
func (a *AMMClient) GetAmountsOut(ctx context.Context, amountIn *big.Int, path []common.Address) ([]*big.Int, error) {
	data, err := routerABI.Pack("getAmountsOut", amountIn, path)
	if err != nil {
		return nil, fmt.Errorf("pack getAmountsOut: %w", err)
	}

	router := a.chainCfg.AMMV2Router
	result, err := a.client.CallContract(ctx, ethereum.CallMsg{
		To:   &router,
		Data: data,
	}, nil)
	if err != nil {
		return nil, fmt.Errorf("call getAmountsOut: %w", err)
	}

	out, err := routerABI.Unpack("getAmountsOut", result)
	if err != nil {
		return nil, fmt.Errorf("unpack getAmountsOut: %w", err)
	}

	amounts := out[0].([]*big.Int)
	return amounts, nil
}

// SwapExactTokensForTokens executes a swap on the AMM V2 router.
func (a *AMMClient) SwapExactTokensForTokens(
	ctx context.Context,
	key *ecdsa.PrivateKey,
	amountIn *big.Int,
	amountOutMin *big.Int,
	path []common.Address,
	to common.Address,
) (common.Hash, error) {
	deadline := big.NewInt(time.Now().Unix() + 300) // 5 min deadline

	data, err := routerABI.Pack("swapExactTokensForTokens", amountIn, amountOutMin, path, to, deadline)
	if err != nil {
		return common.Hash{}, fmt.Errorf("pack swap: %w", err)
	}

	router := a.chainCfg.AMMV2Router
	return a.sendTx(ctx, key, &router, big.NewInt(0), data, 300_000)
}

// AddLiquidity adds liquidity to an AMM V2 pair.
func (a *AMMClient) AddLiquidity(
	ctx context.Context,
	key *ecdsa.PrivateKey,
	tokenA, tokenB common.Address,
	amountA, amountB *big.Int,
	to common.Address,
) (common.Hash, error) {
	deadline := big.NewInt(time.Now().Unix() + 300)

	// Use 99% of desired as minimum (1% slippage tolerance)
	minA := new(big.Int).Div(new(big.Int).Mul(amountA, big.NewInt(99)), big.NewInt(100))
	minB := new(big.Int).Div(new(big.Int).Mul(amountB, big.NewInt(99)), big.NewInt(100))

	data, err := routerABI.Pack("addLiquidity", tokenA, tokenB, amountA, amountB, minA, minB, to, deadline)
	if err != nil {
		return common.Hash{}, fmt.Errorf("pack addLiquidity: %w", err)
	}

	router := a.chainCfg.AMMV2Router
	return a.sendTx(ctx, key, &router, big.NewInt(0), data, 500_000)
}

// WrapLUX wraps native LUX into WLUX by calling deposit() with value.
func (a *AMMClient) WrapLUX(ctx context.Context, key *ecdsa.PrivateKey, amount *big.Int) (common.Hash, error) {
	data, err := wluxABI.Pack("deposit")
	if err != nil {
		return common.Hash{}, fmt.Errorf("pack deposit: %w", err)
	}

	wlux := a.chainCfg.WLUX
	return a.sendTx(ctx, key, &wlux, amount, data, 100_000)
}

// TransferNative sends a native LUX transfer (simple value transfer).
func (a *AMMClient) TransferNative(ctx context.Context, key *ecdsa.PrivateKey, to common.Address, amount *big.Int) (common.Hash, error) {
	return a.sendTx(ctx, key, &to, amount, nil, 21_000)
}

// TransferERC20 sends an ERC20 token transfer.
func (a *AMMClient) TransferERC20(ctx context.Context, key *ecdsa.PrivateKey, token, to common.Address, amount *big.Int) (common.Hash, error) {
	data, err := erc20ABI.Pack("transfer", to, amount)
	if err != nil {
		return common.Hash{}, fmt.Errorf("pack transfer: %w", err)
	}

	return a.sendTx(ctx, key, &token, big.NewInt(0), data, 100_000)
}

// GetPairReserves returns the reserves of a pair if it exists.
func (a *AMMClient) GetPairReserves(ctx context.Context, tokenA, tokenB common.Address) (*big.Int, *big.Int, error) {
	// First get the pair address
	pairAddr, err := a.getPairAddress(ctx, tokenA, tokenB)
	if err != nil {
		return nil, nil, err
	}
	if pairAddr == (common.Address{}) {
		return nil, nil, fmt.Errorf("pair does not exist")
	}

	data, err := pairABI.Pack("getReserves")
	if err != nil {
		return nil, nil, fmt.Errorf("pack getReserves: %w", err)
	}

	result, err := a.client.CallContract(ctx, ethereum.CallMsg{
		To:   &pairAddr,
		Data: data,
	}, nil)
	if err != nil {
		return nil, nil, fmt.Errorf("call getReserves: %w", err)
	}

	out, err := pairABI.Unpack("getReserves", result)
	if err != nil {
		return nil, nil, fmt.Errorf("unpack getReserves: %w", err)
	}

	reserve0 := out[0].(*big.Int)
	reserve1 := out[1].(*big.Int)
	return reserve0, reserve1, nil
}

// getPairAddress gets the pair contract address from the factory.
func (a *AMMClient) getPairAddress(ctx context.Context, tokenA, tokenB common.Address) (common.Address, error) {
	data, err := factoryABI.Pack("getPair", tokenA, tokenB)
	if err != nil {
		return common.Address{}, fmt.Errorf("pack getPair: %w", err)
	}

	factory := a.chainCfg.AMMV2Factory
	result, err := a.client.CallContract(ctx, ethereum.CallMsg{
		To:   &factory,
		Data: data,
	}, nil)
	if err != nil {
		return common.Address{}, fmt.Errorf("call getPair: %w", err)
	}

	out, err := factoryABI.Unpack("getPair", result)
	if err != nil {
		return common.Address{}, fmt.Errorf("unpack getPair: %w", err)
	}

	return out[0].(common.Address), nil
}

// sendTx signs and sends a transaction, then waits for the receipt.
func (a *AMMClient) sendTx(
	ctx context.Context,
	key *ecdsa.PrivateKey,
	to *common.Address,
	value *big.Int,
	data []byte,
	gasLimit uint64,
) (common.Hash, error) {
	from := addressFromKey(key)

	nonce, err := a.client.PendingNonceAt(ctx, from)
	if err != nil {
		return common.Hash{}, fmt.Errorf("get nonce: %w", err)
	}

	gasPrice, err := a.client.SuggestGasPrice(ctx)
	if err != nil {
		return common.Hash{}, fmt.Errorf("get gas price: %w", err)
	}

	tx := types.NewTx(&types.LegacyTx{
		Nonce:    nonce,
		To:       to,
		Value:    value,
		Gas:      gasLimit,
		GasPrice: gasPrice,
		Data:     data,
	})

	signer := types.NewEIP155Signer(a.chainCfg.ChainID)
	signedTx, err := types.SignTx(tx, signer, key)
	if err != nil {
		return common.Hash{}, fmt.Errorf("sign tx: %w", err)
	}

	if err := a.client.SendTransaction(ctx, signedTx); err != nil {
		return common.Hash{}, fmt.Errorf("send tx: %w", err)
	}

	// Wait for receipt with timeout
	receiptCtx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	hash := signedTx.Hash()
	for {
		receipt, err := a.client.TransactionReceipt(receiptCtx, hash)
		if err == nil {
			if receipt.Status == 0 {
				return hash, fmt.Errorf("tx reverted: %s", hash.Hex())
			}
			return hash, nil
		}

		select {
		case <-receiptCtx.Done():
			return hash, fmt.Errorf("receipt timeout for %s", hash.Hex())
		case <-time.After(2 * time.Second):
			// Poll again
		}
	}
}

// EnsureApproval checks and sets token approval if needed.
func (a *AMMClient) EnsureApproval(ctx context.Context, key *ecdsa.PrivateKey, token, spender common.Address, amount *big.Int) error {
	from := addressFromKey(key)

	allowance, err := a.GetAllowance(ctx, token, from, spender)
	if err != nil {
		return fmt.Errorf("check allowance: %w", err)
	}

	if allowance.Cmp(amount) >= 0 {
		return nil // Already approved enough
	}

	// Approve a large amount to avoid repeated approvals
	maxApproval := new(big.Int).Mul(amount, big.NewInt(1000))
	_, err = a.ApproveToken(ctx, key, token, spender, maxApproval)
	if err != nil {
		return fmt.Errorf("approve: %w", err)
	}

	return nil
}

// GetTradePairs returns all available trading pairs for a chain.
func GetTradePairs(cc ChainConfig) []TradePair {
	return []TradePair{
		{TokenA: cc.WLUX, TokenB: cc.LETH, NameA: "WLUX", NameB: "LETH"},
		{TokenA: cc.WLUX, TokenB: cc.LBTC, NameA: "WLUX", NameB: "LBTC"},
		{TokenA: cc.WLUX, TokenB: cc.LUSDC, NameA: "WLUX", NameB: "LUSDC"},
		{TokenA: cc.LETH, TokenB: cc.LUSDC, NameA: "LETH", NameB: "LUSDC"},
	}
}