mev-beta/pkg/execution/flashloan.go

package execution

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

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto"

	"github.com/your-org/mev-bot/pkg/arbitrage"
)

// Aave V3 Pool address on Arbitrum
var AaveV3PoolAddress = common.HexToAddress("0x794a61358D6845594F94dc1DB02A252b5b4814aD")

// WETH address on Arbitrum
var WETHAddress = common.HexToAddress("0x82aF49447D8a07e3bd95BD0d56f35241523fBab1")

// FlashloanProvider represents different flashloan providers
type FlashloanProvider string

const (
	FlashloanProviderAaveV3     FlashloanProvider = "aave_v3"
	FlashloanProviderUniswapV3  FlashloanProvider = "uniswap_v3"
	FlashloanProviderUniswapV2  FlashloanProvider = "uniswap_v2"
)

// FlashloanConfig contains configuration for flashloans
type FlashloanConfig struct {
	// Provider preferences (ordered by preference)
	PreferredProviders []FlashloanProvider

	// Fee configuration
	AaveV3FeeBPS     uint16 // Aave V3 fee in basis points (default: 9 = 0.09%)
	UniswapV3FeeBPS  uint16 // Uniswap V3 flash fee (pool dependent)
	UniswapV2FeeBPS  uint16 // Uniswap V2 flash swap fee (30 bps)

	// Execution contract
	ExecutorContract common.Address // Custom contract that receives flashloan callback
}

// DefaultFlashloanConfig returns default configuration
func DefaultFlashloanConfig() *FlashloanConfig {
	return &FlashloanConfig{
		PreferredProviders: []FlashloanProvider{
			FlashloanProviderAaveV3,
			FlashloanProviderUniswapV3,
			FlashloanProviderUniswapV2,
		},
		AaveV3FeeBPS:    9,  // 0.09%
		UniswapV3FeeBPS: 0,  // No fee for flash swaps (pay in swap)
		UniswapV2FeeBPS: 30, // 0.3% (0.25% fee + 0.05% protocol)
	}
}

// FlashloanManager manages flashloan operations
type FlashloanManager struct {
	config *FlashloanConfig
	logger *slog.Logger

	// Provider-specific encoders
	aaveV3Encoder     *AaveV3FlashloanEncoder
	uniswapV3Encoder  *UniswapV3FlashloanEncoder
	uniswapV2Encoder  *UniswapV2FlashloanEncoder
}

// NewFlashloanManager creates a new flashloan manager
func NewFlashloanManager(config *FlashloanConfig, logger *slog.Logger) *FlashloanManager {
	if config == nil {
		config = DefaultFlashloanConfig()
	}

	return &FlashloanManager{
		config:           config,
		logger:           logger.With("component", "flashloan_manager"),
		aaveV3Encoder:    NewAaveV3FlashloanEncoder(),
		uniswapV3Encoder: NewUniswapV3FlashloanEncoder(),
		uniswapV2Encoder: NewUniswapV2FlashloanEncoder(),
	}
}

// FlashloanRequest represents a flashloan request
type FlashloanRequest struct {
	Token    common.Address
	Amount   *big.Int
	Provider FlashloanProvider
	Params   []byte // Additional parameters to pass to callback
}

// FlashloanTransaction represents an encoded flashloan transaction
type FlashloanTransaction struct {
	To       common.Address
	Data     []byte
	Value    *big.Int
	Provider FlashloanProvider
	Fee      *big.Int
}

// BuildFlashloanTransaction builds a flashloan transaction for an opportunity
func (fm *FlashloanManager) BuildFlashloanTransaction(
	ctx context.Context,
	opp *arbitrage.Opportunity,
	swapCalldata []byte,
) (*FlashloanTransaction, error) {
	fm.logger.Debug("building flashloan transaction",
		"opportunityID", opp.ID,
		"inputAmount", opp.InputAmount.String(),
	)

	// Determine best flashloan provider
	provider, err := fm.selectProvider(ctx, opp.InputToken, opp.InputAmount)
	if err != nil {
		return nil, fmt.Errorf("failed to select provider: %w", err)
	}

	fm.logger.Debug("selected flashloan provider", "provider", provider)

	// Build flashloan transaction
	var tx *FlashloanTransaction

	switch provider {
	case FlashloanProviderAaveV3:
		tx, err = fm.buildAaveV3Flashloan(opp, swapCalldata)

	case FlashloanProviderUniswapV3:
		tx, err = fm.buildUniswapV3Flashloan(opp, swapCalldata)

	case FlashloanProviderUniswapV2:
		tx, err = fm.buildUniswapV2Flashloan(opp, swapCalldata)

	default:
		return nil, fmt.Errorf("unsupported flashloan provider: %s", provider)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to build flashloan: %w", err)
	}

	fm.logger.Info("flashloan transaction built",
		"provider", provider,
		"amount", opp.InputAmount.String(),
		"fee", tx.Fee.String(),
	)

	return tx, nil
}

// buildAaveV3Flashloan builds an Aave V3 flashloan transaction
func (fm *FlashloanManager) buildAaveV3Flashloan(
	opp *arbitrage.Opportunity,
	swapCalldata []byte,
) (*FlashloanTransaction, error) {
	// Calculate fee
	fee := fm.calculateFee(opp.InputAmount, fm.config.AaveV3FeeBPS)

	// Encode flashloan call
	to, data, err := fm.aaveV3Encoder.EncodeFlashloan(
		[]common.Address{opp.InputToken},
		[]*big.Int{opp.InputAmount},
		fm.config.ExecutorContract,
		swapCalldata,
	)

	if err != nil {
		return nil, fmt.Errorf("failed to encode Aave V3 flashloan: %w", err)
	}

	return &FlashloanTransaction{
		To:       to,
		Data:     data,
		Value:    big.NewInt(0),
		Provider: FlashloanProviderAaveV3,
		Fee:      fee,
	}, nil
}

// buildUniswapV3Flashloan builds a Uniswap V3 flash swap transaction
func (fm *FlashloanManager) buildUniswapV3Flashloan(
	opp *arbitrage.Opportunity,
	swapCalldata []byte,
) (*FlashloanTransaction, error) {
	// Uniswap V3 flash swaps don't have a separate fee
	// The fee is paid as part of the swap
	fee := big.NewInt(0)

	// Get pool address for the flashloan token
	// In production, we'd query the pool with highest liquidity
	poolAddress := opp.Path[0].PoolAddress

	// Encode flash swap
	to, data, err := fm.uniswapV3Encoder.EncodeFlash(
		opp.InputToken,
		opp.InputAmount,
		poolAddress,
		fm.config.ExecutorContract,
		swapCalldata,
	)

	if err != nil {
		return nil, fmt.Errorf("failed to encode Uniswap V3 flash: %w", err)
	}

	return &FlashloanTransaction{
		To:       to,
		Data:     data,
		Value:    big.NewInt(0),
		Provider: FlashloanProviderUniswapV3,
		Fee:      fee,
	}, nil
}

// buildUniswapV2Flashloan builds a Uniswap V2 flash swap transaction
func (fm *FlashloanManager) buildUniswapV2Flashloan(
	opp *arbitrage.Opportunity,
	swapCalldata []byte,
) (*FlashloanTransaction, error) {
	// Calculate fee
	fee := fm.calculateFee(opp.InputAmount, fm.config.UniswapV2FeeBPS)

	// Get pool address
	poolAddress := opp.Path[0].PoolAddress

	// Encode flash swap
	to, data, err := fm.uniswapV2Encoder.EncodeFlash(
		opp.InputToken,
		opp.InputAmount,
		poolAddress,
		fm.config.ExecutorContract,
		swapCalldata,
	)

	if err != nil {
		return nil, fmt.Errorf("failed to encode Uniswap V2 flash: %w", err)
	}

	return &FlashloanTransaction{
		To:       to,
		Data:     data,
		Value:    big.NewInt(0),
		Provider: FlashloanProviderUniswapV2,
		Fee:      fee,
	}, nil
}

// selectProvider selects the best flashloan provider
func (fm *FlashloanManager) selectProvider(
	ctx context.Context,
	token common.Address,
	amount *big.Int,
) (FlashloanProvider, error) {
	// For now, use the first preferred provider
	// In production, we'd check availability and fees for each

	if len(fm.config.PreferredProviders) == 0 {
		return "", fmt.Errorf("no flashloan providers configured")
	}

	// Use first preferred provider
	return fm.config.PreferredProviders[0], nil
}

// calculateFee calculates the flashloan fee
func (fm *FlashloanManager) calculateFee(amount *big.Int, feeBPS uint16) *big.Int {
	// fee = amount * feeBPS / 10000
	fee := new(big.Int).Mul(amount, big.NewInt(int64(feeBPS)))
	fee.Div(fee, big.NewInt(10000))
	return fee
}

// CalculateTotalCost calculates the total cost including fee
func (fm *FlashloanManager) CalculateTotalCost(amount *big.Int, feeBPS uint16) *big.Int {
	fee := fm.calculateFee(amount, feeBPS)
	total := new(big.Int).Add(amount, fee)
	return total
}

// AaveV3FlashloanEncoder encodes Aave V3 flashloan calls
type AaveV3FlashloanEncoder struct {
	poolAddress common.Address
}

// NewAaveV3FlashloanEncoder creates a new Aave V3 flashloan encoder
func NewAaveV3FlashloanEncoder() *AaveV3FlashloanEncoder {
	return &AaveV3FlashloanEncoder{
		poolAddress: AaveV3PoolAddress,
	}
}

// EncodeFlashloan encodes an Aave V3 flashloan call
func (e *AaveV3FlashloanEncoder) EncodeFlashloan(
	assets []common.Address,
	amounts []*big.Int,
	receiverAddress common.Address,
	params []byte,
) (common.Address, []byte, error) {
	// flashLoan(address receivingAddress, address[] assets, uint256[] amounts, uint256[] modes, address onBehalfOf, bytes params, uint16 referralCode)
	methodID := crypto.Keccak256([]byte("flashLoan(address,address[],uint256[],uint256[],address,bytes,uint16)"))[:4]

	// For simplicity, this is a basic implementation
	// In production, we'd need to properly encode all dynamic arrays

	data := make([]byte, 0)
	data = append(data, methodID...)

	// receivingAddress
	data = append(data, padLeft(receiverAddress.Bytes(), 32)...)

	// Offset to assets array (7 * 32 bytes)
	data = append(data, padLeft(big.NewInt(224).Bytes(), 32)...)

	// Offset to amounts array (calculated based on assets length)
	assetsOffset := 224 + 32 + (32 * len(assets))
	data = append(data, padLeft(big.NewInt(int64(assetsOffset)).Bytes(), 32)...)

	// Offset to modes array
	modesOffset := assetsOffset + 32 + (32 * len(amounts))
	data = append(data, padLeft(big.NewInt(int64(modesOffset)).Bytes(), 32)...)

	// onBehalfOf (receiver address)
	data = append(data, padLeft(receiverAddress.Bytes(), 32)...)

	// Offset to params
	paramsOffset := modesOffset + 32 + (32 * len(assets))
	data = append(data, padLeft(big.NewInt(int64(paramsOffset)).Bytes(), 32)...)

	// referralCode (0)
	data = append(data, padLeft(big.NewInt(0).Bytes(), 32)...)

	// Assets array
	data = append(data, padLeft(big.NewInt(int64(len(assets))).Bytes(), 32)...)
	for _, asset := range assets {
		data = append(data, padLeft(asset.Bytes(), 32)...)
	}

	// Amounts array
	data = append(data, padLeft(big.NewInt(int64(len(amounts))).Bytes(), 32)...)
	for _, amount := range amounts {
		data = append(data, padLeft(amount.Bytes(), 32)...)
	}

	// Modes array (0 = no debt, we repay in same transaction)
	data = append(data, padLeft(big.NewInt(int64(len(assets))).Bytes(), 32)...)
	for range assets {
		data = append(data, padLeft(big.NewInt(0).Bytes(), 32)...)
	}

	// Params bytes
	data = append(data, padLeft(big.NewInt(int64(len(params))).Bytes(), 32)...)
	data = append(data, params...)

	// Pad params to 32-byte boundary
	remainder := len(params) % 32
	if remainder != 0 {
		padding := make([]byte, 32-remainder)
		data = append(data, padding...)
	}

	return e.poolAddress, data, nil
}

// UniswapV3FlashloanEncoder encodes Uniswap V3 flash calls
type UniswapV3FlashloanEncoder struct{}

// NewUniswapV3FlashloanEncoder creates a new Uniswap V3 flashloan encoder
func NewUniswapV3FlashloanEncoder() *UniswapV3FlashloanEncoder {
	return &UniswapV3FlashloanEncoder{}
}

// EncodeFlash encodes a Uniswap V3 flash call
func (e *UniswapV3FlashloanEncoder) EncodeFlash(
	token common.Address,
	amount *big.Int,
	poolAddress common.Address,
	recipient common.Address,
	data []byte,
) (common.Address, []byte, error) {
	// flash(address recipient, uint256 amount0, uint256 amount1, bytes data)
	methodID := crypto.Keccak256([]byte("flash(address,uint256,uint256,bytes)"))[:4]

	calldata := make([]byte, 0)
	calldata = append(calldata, methodID...)

	// recipient
	calldata = append(calldata, padLeft(recipient.Bytes(), 32)...)

	// amount0 or amount1 (depending on which token in the pool)
	// For simplicity, assume token0
	calldata = append(calldata, padLeft(amount.Bytes(), 32)...)
	calldata = append(calldata, padLeft(big.NewInt(0).Bytes(), 32)...)

	// Offset to data bytes
	calldata = append(calldata, padLeft(big.NewInt(128).Bytes(), 32)...)

	// Data length
	calldata = append(calldata, padLeft(big.NewInt(int64(len(data))).Bytes(), 32)...)

	// Data
	calldata = append(calldata, data...)

	// Padding
	remainder := len(data) % 32
	if remainder != 0 {
		padding := make([]byte, 32-remainder)
		calldata = append(calldata, padding...)
	}

	return poolAddress, calldata, nil
}

// UniswapV2FlashloanEncoder encodes Uniswap V2 flash swap calls
type UniswapV2FlashloanEncoder struct{}

// NewUniswapV2FlashloanEncoder creates a new Uniswap V2 flashloan encoder
func NewUniswapV2FlashloanEncoder() *UniswapV2FlashloanEncoder {
	return &UniswapV2FlashloanEncoder{}
}

// EncodeFlash encodes a Uniswap V2 flash swap call
func (e *UniswapV2FlashloanEncoder) EncodeFlash(
	token common.Address,
	amount *big.Int,
	poolAddress common.Address,
	recipient common.Address,
	data []byte,
) (common.Address, []byte, error) {
	// swap(uint amount0Out, uint amount1Out, address to, bytes data)
	methodID := crypto.Keccak256([]byte("swap(uint256,uint256,address,bytes)"))[:4]

	calldata := make([]byte, 0)
	calldata = append(calldata, methodID...)

	// amount0Out or amount1Out (depending on which token)
	// For simplicity, assume token0
	calldata = append(calldata, padLeft(amount.Bytes(), 32)...)
	calldata = append(calldata, padLeft(big.NewInt(0).Bytes(), 32)...)

	// to (recipient)
	calldata = append(calldata, padLeft(recipient.Bytes(), 32)...)

	// Offset to data bytes
	calldata = append(calldata, padLeft(big.NewInt(128).Bytes(), 32)...)

	// Data length
	calldata = append(calldata, padLeft(big.NewInt(int64(len(data))).Bytes(), 32)...)

	// Data
	calldata = append(calldata, data...)

	// Padding
	remainder := len(data) % 32
	if remainder != 0 {
		padding := make([]byte, 32-remainder)
		calldata = append(calldata, padding...)
	}

	return poolAddress, calldata, nil
}