mev-beta/orig/pkg/dex/balancer.go

package dex

import (
	"context"
	"fmt"
	"math/big"
	"strings"

	"github.com/ethereum/go-ethereum"
	"github.com/ethereum/go-ethereum/accounts/abi"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/ethclient"
)

// BalancerDecoder implements DEXDecoder for Balancer
type BalancerDecoder struct {
	*BaseDecoder
	vaultABI abi.ABI
	poolABI  abi.ABI
}

// Balancer Vault ABI (minimal)
const balancerVaultABI = `[
	{
		"name": "swap",
		"type": "function",
		"inputs": [
			{
				"name": "singleSwap",
				"type": "tuple",
				"components": [
					{"name": "poolId", "type": "bytes32"},
					{"name": "kind", "type": "uint8"},
					{"name": "assetIn", "type": "address"},
					{"name": "assetOut", "type": "address"},
					{"name": "amount", "type": "uint256"},
					{"name": "userData", "type": "bytes"}
				]
			},
			{
				"name": "funds",
				"type": "tuple",
				"components": [
					{"name": "sender", "type": "address"},
					{"name": "fromInternalBalance", "type": "bool"},
					{"name": "recipient", "type": "address"},
					{"name": "toInternalBalance", "type": "bool"}
				]
			},
			{"name": "limit", "type": "uint256"},
			{"name": "deadline", "type": "uint256"}
		],
		"outputs": [{"name": "amountCalculated", "type": "uint256"}]
	},
	{
		"name": "getPoolTokens",
		"type": "function",
		"inputs": [{"name": "poolId", "type": "bytes32"}],
		"outputs": [
			{"name": "tokens", "type": "address[]"},
			{"name": "balances", "type": "uint256[]"},
			{"name": "lastChangeBlock", "type": "uint256"}
		],
		"stateMutability": "view"
	}
]`

// Balancer Pool ABI (minimal)
const balancerPoolABI = `[
	{
		"name": "getPoolId",
		"type": "function",
		"inputs": [],
		"outputs": [{"name": "", "type": "bytes32"}],
		"stateMutability": "view"
	},
	{
		"name": "getNormalizedWeights",
		"type": "function",
		"inputs": [],
		"outputs": [{"name": "", "type": "uint256[]"}],
		"stateMutability": "view"
	},
	{
		"name": "getSwapFeePercentage",
		"type": "function",
		"inputs": [],
		"outputs": [{"name": "", "type": "uint256"}],
		"stateMutability": "view"
	}
]`

// Balancer Vault address on Arbitrum
var BalancerVaultAddress = common.HexToAddress("0xBA12222222228d8Ba445958a75a0704d566BF2C8")

// NewBalancerDecoder creates a new Balancer decoder
func NewBalancerDecoder(client *ethclient.Client) *BalancerDecoder {
	vaultABI, _ := abi.JSON(strings.NewReader(balancerVaultABI))
	poolABI, _ := abi.JSON(strings.NewReader(balancerPoolABI))

	return &BalancerDecoder{
		BaseDecoder: NewBaseDecoder(ProtocolBalancer, client),
		vaultABI:    vaultABI,
		poolABI:     poolABI,
	}
}

// DecodeSwap decodes a Balancer swap transaction
func (d *BalancerDecoder) DecodeSwap(tx *types.Transaction) (*SwapInfo, error) {
	data := tx.Data()
	if len(data) < 4 {
		return nil, fmt.Errorf("transaction data too short")
	}

	method, err := d.vaultABI.MethodById(data[:4])
	if err != nil {
		return nil, fmt.Errorf("failed to get method: %w", err)
	}

	if method.Name != "swap" {
		return nil, fmt.Errorf("unsupported method: %s", method.Name)
	}

	params := make(map[string]interface{})
	if err := method.Inputs.UnpackIntoMap(params, data[4:]); err != nil {
		return nil, fmt.Errorf("failed to unpack params: %w", err)
	}

	// Extract singleSwap struct
	singleSwap := params["singleSwap"].(struct {
		PoolId   [32]byte
		Kind     uint8
		AssetIn  common.Address
		AssetOut common.Address
		Amount   *big.Int
		UserData []byte
	})

	funds := params["funds"].(struct {
		Sender              common.Address
		FromInternalBalance bool
		Recipient           common.Address
		ToInternalBalance   bool
	})

	return &SwapInfo{
		Protocol:  ProtocolBalancer,
		TokenIn:   singleSwap.AssetIn,
		TokenOut:  singleSwap.AssetOut,
		AmountIn:  singleSwap.Amount,
		AmountOut: params["limit"].(*big.Int),
		Recipient: funds.Recipient,
		Deadline:  params["deadline"].(*big.Int),
		Fee:       big.NewInt(25), // 0.25% typical
	}, nil
}

// GetPoolReserves fetches current pool reserves for Balancer
func (d *BalancerDecoder) GetPoolReserves(ctx context.Context, client *ethclient.Client, poolAddress common.Address) (*PoolReserves, error) {
	// Get pool ID
	poolIdData, err := client.CallContract(ctx, ethereum.CallMsg{
		To:   &poolAddress,
		Data: d.poolABI.Methods["getPoolId"].ID,
	}, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to get pool ID: %w", err)
	}
	poolId := [32]byte{}
	copy(poolId[:], poolIdData)

	// Get pool tokens and balances from Vault
	getPoolTokensCalldata, err := d.vaultABI.Pack("getPoolTokens", poolId)
	if err != nil {
		return nil, fmt.Errorf("failed to pack getPoolTokens: %w", err)
	}

	tokensData, err := client.CallContract(ctx, ethereum.CallMsg{
		To:   &BalancerVaultAddress,
		Data: getPoolTokensCalldata,
	}, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to get pool tokens: %w", err)
	}

	var result struct {
		Tokens          []common.Address
		Balances        []*big.Int
		LastChangeBlock *big.Int
	}
	if err := d.vaultABI.UnpackIntoInterface(&result, "getPoolTokens", tokensData); err != nil {
		return nil, fmt.Errorf("failed to unpack pool tokens: %w", err)
	}

	if len(result.Tokens) < 2 {
		return nil, fmt.Errorf("pool has less than 2 tokens")
	}

	// Get weights
	weightsData, err := client.CallContract(ctx, ethereum.CallMsg{
		To:   &poolAddress,
		Data: d.poolABI.Methods["getNormalizedWeights"].ID,
	}, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to get weights: %w", err)
	}

	var weights []*big.Int
	if err := d.poolABI.UnpackIntoInterface(&weights, "getNormalizedWeights", weightsData); err != nil {
		return nil, fmt.Errorf("failed to unpack weights: %w", err)
	}

	// Get swap fee
	feeData, err := client.CallContract(ctx, ethereum.CallMsg{
		To:   &poolAddress,
		Data: d.poolABI.Methods["getSwapFeePercentage"].ID,
	}, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to get swap fee: %w", err)
	}
	fee := new(big.Int).SetBytes(feeData)

	return &PoolReserves{
		Token0:      result.Tokens[0],
		Token1:      result.Tokens[1],
		Reserve0:    result.Balances[0],
		Reserve1:    result.Balances[1],
		Protocol:    ProtocolBalancer,
		PoolAddress: poolAddress,
		Fee:         fee,
		Weights:     weights,
	}, nil
}

// CalculateOutput calculates expected output for Balancer weighted pools
func (d *BalancerDecoder) CalculateOutput(amountIn *big.Int, reserves *PoolReserves, tokenIn common.Address) (*big.Int, error) {
	if amountIn == nil || amountIn.Sign() <= 0 {
		return nil, fmt.Errorf("invalid amountIn")
	}

	if reserves.Weights == nil || len(reserves.Weights) < 2 {
		return nil, fmt.Errorf("missing pool weights")
	}

	var balanceIn, balanceOut, weightIn, weightOut *big.Int
	if tokenIn == reserves.Token0 {
		balanceIn = reserves.Reserve0
		balanceOut = reserves.Reserve1
		weightIn = reserves.Weights[0]
		weightOut = reserves.Weights[1]
	} else if tokenIn == reserves.Token1 {
		balanceIn = reserves.Reserve1
		balanceOut = reserves.Reserve0
		weightIn = reserves.Weights[1]
		weightOut = reserves.Weights[0]
	} else {
		return nil, fmt.Errorf("tokenIn not in pool")
	}

	if balanceIn.Sign() == 0 || balanceOut.Sign() == 0 {
		return nil, fmt.Errorf("insufficient liquidity")
	}

	// Balancer weighted pool formula:
	// amountOut = balanceOut * (1 - (balanceIn / (balanceIn + amountIn))^(weightIn/weightOut))
	// Simplified approximation for demonstration

	// Apply fee
	fee := reserves.Fee
	if fee == nil {
		fee = big.NewInt(25) // 0.25% = 25 basis points
	}

	amountInAfterFee := new(big.Int).Mul(amountIn, new(big.Int).Sub(big.NewInt(10000), fee))
	amountInAfterFee.Div(amountInAfterFee, big.NewInt(10000))

	// Simplified calculation: use ratio of weights
	// amountOut ≈ amountIn * (balanceOut/balanceIn) * (weightOut/weightIn)
	amountOut := new(big.Int).Mul(amountInAfterFee, balanceOut)
	amountOut.Div(amountOut, balanceIn)

	// Adjust by weight ratio (simplified)
	amountOut.Mul(amountOut, weightOut)
	amountOut.Div(amountOut, weightIn)

	// For production: Implement full weighted pool math with exponentiation
	// amountOut = balanceOut * (1 - (balanceIn / (balanceIn + amountInAfterFee))^(weightIn/weightOut))

	return amountOut, nil
}

// CalculatePriceImpact calculates price impact for Balancer
func (d *BalancerDecoder) CalculatePriceImpact(amountIn *big.Int, reserves *PoolReserves, tokenIn common.Address) (float64, error) {
	if amountIn == nil || amountIn.Sign() <= 0 {
		return 0, nil
	}

	var balanceIn *big.Int
	if tokenIn == reserves.Token0 {
		balanceIn = reserves.Reserve0
	} else {
		balanceIn = reserves.Reserve1
	}

	if balanceIn.Sign() == 0 {
		return 1.0, nil
	}

	// Price impact for weighted pools is lower than constant product
	amountInFloat := new(big.Float).SetInt(amountIn)
	balanceFloat := new(big.Float).SetInt(balanceIn)

	ratio := new(big.Float).Quo(amountInFloat, balanceFloat)

	// Weighted pools have better capital efficiency
	impact := new(big.Float).Mul(ratio, big.NewFloat(0.8))

	impactValue, _ := impact.Float64()
	return impactValue, nil
}

// GetQuote gets a price quote for Balancer
func (d *BalancerDecoder) GetQuote(ctx context.Context, client *ethclient.Client, tokenIn, tokenOut common.Address, amountIn *big.Int) (*PriceQuote, error) {
	// TODO: Implement pool lookup via Balancer subgraph or on-chain registry
	return nil, fmt.Errorf("GetQuote not yet implemented for Balancer")
}

// IsValidPool checks if a pool is a valid Balancer pool
func (d *BalancerDecoder) IsValidPool(ctx context.Context, client *ethclient.Client, poolAddress common.Address) (bool, error) {
	// Try to call getPoolId() - if it succeeds, it's a Balancer pool
	_, err := client.CallContract(ctx, ethereum.CallMsg{
		To:   &poolAddress,
		Data: d.poolABI.Methods["getPoolId"].ID,
	}, nil)

	return err == nil, nil
}