mev-beta/pkg/arbitrage/calculator.go

package arbitrage

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

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

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

// CalculatorConfig contains configuration for profitability calculations
type CalculatorConfig struct {
	MinProfitWei     *big.Int // Minimum net profit in wei
	MinROI           float64  // Minimum ROI percentage (e.g., 0.05 = 5%)
	MaxPriceImpact   float64  // Maximum acceptable price impact (e.g., 0.10 = 10%)
	MaxGasPriceGwei  uint64   // Maximum gas price in gwei
	SlippageTolerance float64  // Slippage tolerance (e.g., 0.005 = 0.5%)
}

// DefaultCalculatorConfig returns default configuration
func DefaultCalculatorConfig() *CalculatorConfig {
	minProfit := new(big.Int).Mul(big.NewInt(5), new(big.Int).Exp(big.NewInt(10), big.NewInt(16), nil)) // 0.05 ETH

	return &CalculatorConfig{
		MinProfitWei:     minProfit,
		MinROI:           0.05,  // 5%
		MaxPriceImpact:   0.10,  // 10%
		MaxGasPriceGwei:  100,   // 100 gwei
		SlippageTolerance: 0.005, // 0.5%
	}
}

// Calculator calculates profitability of arbitrage opportunities
type Calculator struct {
	config      *CalculatorConfig
	logger      *slog.Logger
	gasEstimator *GasEstimator
}

// NewCalculator creates a new calculator
func NewCalculator(config *CalculatorConfig, gasEstimator *GasEstimator, logger *slog.Logger) *Calculator {
	if config == nil {
		config = DefaultCalculatorConfig()
	}

	return &Calculator{
		config:      config,
		gasEstimator: gasEstimator,
		logger:      logger.With("component", "calculator"),
	}
}

// CalculateProfitability calculates the profitability of a path
func (c *Calculator) CalculateProfitability(ctx context.Context, path *Path, inputAmount *big.Int, gasPrice *big.Int) (*Opportunity, error) {
	if len(path.Pools) == 0 {
		return nil, fmt.Errorf("path has no pools")
	}

	if inputAmount == nil || inputAmount.Sign() <= 0 {
		return nil, fmt.Errorf("invalid input amount")
	}

	startTime := time.Now()

	// Simulate the swap through each pool in the path
	currentAmount := new(big.Int).Set(inputAmount)
	pathSteps := make([]*PathStep, 0, len(path.Pools))

	totalPriceImpact := 0.0

	for i, pool := range path.Pools {
		tokenIn := path.Tokens[i]
		tokenOut := path.Tokens[i+1]

		// Calculate swap output
		amountOut, priceImpact, err := c.calculateSwapOutput(pool, tokenIn, tokenOut, currentAmount)
		if err != nil {
			c.logger.Warn("failed to calculate swap output",
				"pool", pool.Address.Hex(),
				"error", err,
			)
			return nil, fmt.Errorf("failed to calculate swap at pool %s: %w", pool.Address.Hex(), err)
		}

		// Create path step
		step := &PathStep{
			PoolAddress: pool.Address,
			Protocol:    pool.Protocol,
			TokenIn:     tokenIn,
			TokenOut:    tokenOut,
			AmountIn:    currentAmount,
			AmountOut:   amountOut,
			Fee:         pool.Fee,
		}

		// Calculate fee amount
		step.FeeAmount = c.calculateFeeAmount(currentAmount, pool.Fee, pool.Protocol)

		// Store V3-specific state if applicable
		if pool.Protocol == types.ProtocolUniswapV3 && pool.SqrtPriceX96 != nil {
			step.SqrtPriceX96Before = new(big.Int).Set(pool.SqrtPriceX96)

			// Calculate new price after swap
			zeroForOne := tokenIn == pool.Token0
			newPrice, err := c.calculateNewPriceV3(pool, currentAmount, zeroForOne)
			if err == nil {
				step.SqrtPriceX96After = newPrice
			}
		}

		pathSteps = append(pathSteps, step)
		totalPriceImpact += priceImpact

		// Update current amount for next hop
		currentAmount = amountOut
	}

	// Calculate profits
	outputAmount := currentAmount
	grossProfit := new(big.Int).Sub(outputAmount, inputAmount)

	// Estimate gas cost
	gasCost, err := c.gasEstimator.EstimateGasCost(ctx, path, gasPrice)
	if err != nil {
		c.logger.Warn("failed to estimate gas cost", "error", err)
		gasCost = big.NewInt(0)
	}

	// Calculate net profit
	netProfit := new(big.Int).Sub(grossProfit, gasCost)

	// Calculate ROI
	roi := 0.0
	if inputAmount.Sign() > 0 {
		inputFloat, _ := new(big.Float).SetInt(inputAmount).Float64()
		profitFloat, _ := new(big.Float).SetInt(netProfit).Float64()
		roi = profitFloat / inputFloat
	}

	// Average price impact across all hops
	avgPriceImpact := totalPriceImpact / float64(len(pathSteps))

	// Create opportunity
	opportunity := &Opportunity{
		ID:           fmt.Sprintf("%s-%d", path.Pools[0].Address.Hex(), time.Now().UnixNano()),
		Type:         path.Type,
		DetectedAt:   startTime,
		BlockNumber:  path.Pools[0].BlockNumber,
		Path:         pathSteps,
		InputToken:   path.Tokens[0],
		OutputToken:  path.Tokens[len(path.Tokens)-1],
		InputAmount:  inputAmount,
		OutputAmount: outputAmount,
		GrossProfit:  grossProfit,
		GasCost:      gasCost,
		NetProfit:    netProfit,
		ROI:          roi,
		PriceImpact:  avgPriceImpact,
		Priority:     c.calculatePriority(netProfit, roi),
		ExecuteAfter: time.Now(),
		ExpiresAt:    time.Now().Add(30 * time.Second), // 30 second expiration
		Executable:   c.isExecutable(netProfit, roi, avgPriceImpact),
	}

	c.logger.Debug("calculated profitability",
		"opportunityID", opportunity.ID,
		"inputAmount", inputAmount.String(),
		"outputAmount", outputAmount.String(),
		"grossProfit", grossProfit.String(),
		"netProfit", netProfit.String(),
		"roi", fmt.Sprintf("%.2f%%", roi*100),
		"priceImpact", fmt.Sprintf("%.2f%%", avgPriceImpact*100),
		"gasPrice", gasCost.String(),
		"executable", opportunity.Executable,
		"duration", time.Since(startTime),
	)

	return opportunity, nil
}

// calculateSwapOutput calculates the output amount for a swap
func (c *Calculator) calculateSwapOutput(pool *types.PoolInfo, tokenIn, tokenOut common.Address, amountIn *big.Int) (*big.Int, float64, error) {
	switch pool.Protocol {
	case types.ProtocolUniswapV2, types.ProtocolSushiSwap:
		return c.calculateSwapOutputV2(pool, tokenIn, tokenOut, amountIn)
	case types.ProtocolUniswapV3:
		return c.calculateSwapOutputV3(pool, tokenIn, tokenOut, amountIn)
	case types.ProtocolCurve:
		return c.calculateSwapOutputCurve(pool, tokenIn, tokenOut, amountIn)
	default:
		return nil, 0, fmt.Errorf("unsupported protocol: %s", pool.Protocol)
	}
}

// calculateSwapOutputV2 calculates output for UniswapV2-style pools
func (c *Calculator) calculateSwapOutputV2(pool *types.PoolInfo, tokenIn, tokenOut common.Address, amountIn *big.Int) (*big.Int, float64, error) {
	if pool.Reserve0 == nil || pool.Reserve1 == nil {
		return nil, 0, fmt.Errorf("pool has nil reserves")
	}

	// Determine direction
	var reserveIn, reserveOut *big.Int
	if tokenIn == pool.Token0 {
		reserveIn = pool.Reserve0
		reserveOut = pool.Reserve1
	} else if tokenIn == pool.Token1 {
		reserveIn = pool.Reserve1
		reserveOut = pool.Reserve0
	} else {
		return nil, 0, fmt.Errorf("token not in pool")
	}

	// Apply fee (0.3% = 9970/10000)
	fee := pool.Fee
	if fee == 0 {
		fee = 30 // Default 0.3%
	}

	// amountInWithFee = amountIn * (10000 - fee) / 10000
	amountInWithFee := new(big.Int).Mul(amountIn, big.NewInt(int64(10000-fee)))
	amountInWithFee.Div(amountInWithFee, big.NewInt(10000))

	// amountOut = (reserveOut * amountInWithFee) / (reserveIn + amountInWithFee)
	numerator := new(big.Int).Mul(reserveOut, amountInWithFee)
	denominator := new(big.Int).Add(reserveIn, amountInWithFee)
	amountOut := new(big.Int).Div(numerator, denominator)

	// Calculate price impact
	priceImpact := c.calculatePriceImpactV2(reserveIn, reserveOut, amountIn, amountOut)

	return amountOut, priceImpact, nil
}

// calculateSwapOutputV3 calculates output for UniswapV3 pools
func (c *Calculator) calculateSwapOutputV3(pool *types.PoolInfo, tokenIn, tokenOut common.Address, amountIn *big.Int) (*big.Int, float64, error) {
	if pool.SqrtPriceX96 == nil || pool.Liquidity == nil {
		return nil, 0, fmt.Errorf("pool missing V3 state")
	}

	zeroForOne := tokenIn == pool.Token0

	// Use V3 math utilities
	amountOut, priceAfter, err := parsers.CalculateSwapAmounts(
		pool.SqrtPriceX96,
		pool.Liquidity,
		amountIn,
		zeroForOne,
		pool.Fee,
	)
	if err != nil {
		return nil, 0, fmt.Errorf("V3 swap calculation failed: %w", err)
	}

	// Calculate price impact
	priceImpact := c.calculatePriceImpactV3(pool.SqrtPriceX96, priceAfter)

	return amountOut, priceImpact, nil
}

// calculateSwapOutputCurve calculates output for Curve pools
func (c *Calculator) calculateSwapOutputCurve(pool *types.PoolInfo, tokenIn, tokenOut common.Address, amountIn *big.Int) (*big.Int, float64, error) {
	// Simplified Curve calculation
	// In production, this should use the actual Curve StableSwap formula

	if pool.Reserve0 == nil || pool.Reserve1 == nil {
		return nil, 0, fmt.Errorf("pool has nil reserves")
	}

	// Determine direction (validate token is in pool)
	if tokenIn != pool.Token0 && tokenIn != pool.Token1 {
		return nil, 0, fmt.Errorf("token not in pool")
	}

	// Simplified: assume 1:1 swap with low slippage for stablecoins
	// TODO: Implement proper Curve StableSwap math using reserves and amp coefficient
	// This is a rough approximation - actual Curve math is more complex
	fee := pool.Fee
	if fee == 0 {
		fee = 4 // Default 0.04% for Curve
	}

	// Scale amounts to same decimals
	amountInScaled := amountIn
	if tokenIn == pool.Token0 {
		amountInScaled = types.ScaleToDecimals(amountIn, pool.Token0Decimals, 18)
	} else {
		amountInScaled = types.ScaleToDecimals(amountIn, pool.Token1Decimals, 18)
	}

	// Apply fee
	amountOutScaled := new(big.Int).Mul(amountInScaled, big.NewInt(int64(10000-fee)))
	amountOutScaled.Div(amountOutScaled, big.NewInt(10000))

	// Scale back to output token decimals
	var amountOut *big.Int
	if tokenOut == pool.Token0 {
		amountOut = types.ScaleToDecimals(amountOutScaled, 18, pool.Token0Decimals)
	} else {
		amountOut = types.ScaleToDecimals(amountOutScaled, 18, pool.Token1Decimals)
	}

	// Curve has very low price impact for stablecoins
	priceImpact := 0.001 // 0.1%

	return amountOut, priceImpact, nil
}

// calculateNewPriceV3 calculates the new sqrtPriceX96 after a swap
func (c *Calculator) calculateNewPriceV3(pool *types.PoolInfo, amountIn *big.Int, zeroForOne bool) (*big.Int, error) {
	_, priceAfter, err := parsers.CalculateSwapAmounts(
		pool.SqrtPriceX96,
		pool.Liquidity,
		amountIn,
		zeroForOne,
		pool.Fee,
	)
	return priceAfter, err
}

// calculatePriceImpactV2 calculates price impact for V2 swaps
func (c *Calculator) calculatePriceImpactV2(reserveIn, reserveOut, amountIn, amountOut *big.Int) float64 {
	// Price before swap
	priceBefore := new(big.Float).Quo(
		new(big.Float).SetInt(reserveOut),
		new(big.Float).SetInt(reserveIn),
	)

	// Price after swap
	newReserveIn := new(big.Int).Add(reserveIn, amountIn)
	newReserveOut := new(big.Int).Sub(reserveOut, amountOut)

	if newReserveIn.Sign() == 0 {
		return 1.0 // 100% impact
	}

	priceAfter := new(big.Float).Quo(
		new(big.Float).SetInt(newReserveOut),
		new(big.Float).SetInt(newReserveIn),
	)

	// Impact = |priceAfter - priceBefore| / priceBefore
	diff := new(big.Float).Sub(priceAfter, priceBefore)
	diff.Abs(diff)
	impact := new(big.Float).Quo(diff, priceBefore)

	impactFloat, _ := impact.Float64()
	return impactFloat
}

// calculatePriceImpactV3 calculates price impact for V3 swaps
func (c *Calculator) calculatePriceImpactV3(priceBefore, priceAfter *big.Int) float64 {
	if priceBefore.Sign() == 0 {
		return 1.0
	}

	priceBeforeFloat := new(big.Float).SetInt(priceBefore)
	priceAfterFloat := new(big.Float).SetInt(priceAfter)

	diff := new(big.Float).Sub(priceAfterFloat, priceBeforeFloat)
	diff.Abs(diff)
	impact := new(big.Float).Quo(diff, priceBeforeFloat)

	impactFloat, _ := impact.Float64()
	return impactFloat
}

// calculateFeeAmount calculates the fee paid in a swap
func (c *Calculator) calculateFeeAmount(amountIn *big.Int, feeBasisPoints uint32, protocol types.ProtocolType) *big.Int {
	if feeBasisPoints == 0 {
		return big.NewInt(0)
	}

	// Fee amount = amountIn * feeBasisPoints / 10000
	feeAmount := new(big.Int).Mul(amountIn, big.NewInt(int64(feeBasisPoints)))
	feeAmount.Div(feeAmount, big.NewInt(10000))

	return feeAmount
}

// calculatePriority calculates priority score for an opportunity
func (c *Calculator) calculatePriority(netProfit *big.Int, roi float64) int {
	// Priority based on both absolute profit and ROI
	// Higher profit and ROI = higher priority

	profitScore := 0
	if netProfit.Sign() > 0 {
		// Convert to ETH for scoring
		profitEth := new(big.Float).Quo(
			new(big.Float).SetInt(netProfit),
			new(big.Float).SetInt64(1e18),
		)
		profitEthFloat, _ := profitEth.Float64()
		profitScore = int(profitEthFloat * 100) // Scale to integer
	}

	roiScore := int(roi * 1000) // Scale to integer

	priority := profitScore + roiScore
	return priority
}

// isExecutable checks if an opportunity meets execution criteria
func (c *Calculator) isExecutable(netProfit *big.Int, roi, priceImpact float64) bool {
	// Check minimum profit
	if netProfit.Cmp(c.config.MinProfitWei) < 0 {
		return false
	}

	// Check minimum ROI
	if roi < c.config.MinROI {
		return false
	}

	// Check maximum price impact
	if priceImpact > c.config.MaxPriceImpact {
		return false
	}

	return true
}

// OptimizeInputAmount finds the optimal input amount for maximum profit
func (c *Calculator) OptimizeInputAmount(ctx context.Context, path *Path, gasPrice *big.Int, maxInput *big.Int) (*Opportunity, error) {
	c.logger.Debug("optimizing input amount",
		"path", fmt.Sprintf("%d pools", len(path.Pools)),
		"maxInput", maxInput.String(),
	)

	// Binary search for optimal input
	low := new(big.Int).Div(maxInput, big.NewInt(100))  // Start at 1% of max
	high := new(big.Int).Set(maxInput)
	bestOpp := (*Opportunity)(nil)

	iterations := 0
	maxIterations := 20

	for low.Cmp(high) < 0 && iterations < maxIterations {
		iterations++

		// Try mid point
		mid := new(big.Int).Add(low, high)
		mid.Div(mid, big.NewInt(2))

		opp, err := c.CalculateProfitability(ctx, path, mid, gasPrice)
		if err != nil {
			c.logger.Warn("optimization iteration failed", "error", err)
			break
		}

		if bestOpp == nil || opp.NetProfit.Cmp(bestOpp.NetProfit) > 0 {
			bestOpp = opp
		}

		// If profit is increasing, try larger amount
		// If profit is decreasing, try smaller amount
		if opp.NetProfit.Sign() > 0 && opp.PriceImpact < c.config.MaxPriceImpact {
			low = new(big.Int).Add(mid, big.NewInt(1))
		} else {
			high = new(big.Int).Sub(mid, big.NewInt(1))
		}
	}

	if bestOpp == nil {
		return nil, fmt.Errorf("failed to find profitable input amount")
	}

	c.logger.Info("optimized input amount",
		"iterations", iterations,
		"optimalInput", bestOpp.InputAmount.String(),
		"netProfit", bestOpp.NetProfit.String(),
		"roi", fmt.Sprintf("%.2f%%", bestOpp.ROI*100),
	)

	return bestOpp, nil
}