mev-beta/pkg/math/exchange_pricing.go

package math

import (
	"fmt"
	"math/big"
)

// ExchangeType represents different DEX protocols on Arbitrum
type ExchangeType string

const (
	ExchangeUniswapV3 ExchangeType = "uniswap_v3"
	ExchangeUniswapV2 ExchangeType = "uniswap_v2"
	ExchangeSushiSwap ExchangeType = "sushiswap"
	ExchangeCamelot   ExchangeType = "camelot"
	ExchangeBalancer  ExchangeType = "balancer"
	ExchangeTraderJoe ExchangeType = "traderjoe"
	ExchangeRamses    ExchangeType = "ramses"
	ExchangeCurve     ExchangeType = "curve"
	ExchangeKyber     ExchangeType = "kyber"
	ExchangeUniswapV4 ExchangeType = "uniswap_v4"
)

// ExchangePricer interface for exchange-specific price calculations
type ExchangePricer interface {
	GetSpotPrice(poolData *PoolData) (*UniversalDecimal, error)
	CalculateAmountOut(amountIn *UniversalDecimal, poolData *PoolData) (*UniversalDecimal, error)
	CalculateAmountIn(amountOut *UniversalDecimal, poolData *PoolData) (*UniversalDecimal, error)
	CalculatePriceImpact(amountIn *UniversalDecimal, poolData *PoolData) (*UniversalDecimal, error)
	GetMinimumLiquidity(poolData *PoolData) (*UniversalDecimal, error)
	ValidatePoolData(poolData *PoolData) error
}

// PoolData represents universal pool data structure
type PoolData struct {
	Address      string
	ExchangeType ExchangeType
	Token0       TokenInfo
	Token1       TokenInfo
	Reserve0     *UniversalDecimal
	Reserve1     *UniversalDecimal
	Fee          *UniversalDecimal // Fee as percentage (e.g., 0.003 for 0.3%)

	// Uniswap V3 specific
	SqrtPriceX96 *big.Int
	Tick         *big.Int
	Liquidity    *big.Int

	// Curve specific
	A *big.Int // Amplification coefficient

	// Balancer specific
	Weights     []*UniversalDecimal // Token weights
	SwapFeeRate *UniversalDecimal   // Swap fee rate
}

// TokenInfo represents token metadata
type TokenInfo struct {
	Address  string
	Symbol   string
	Decimals uint8
}

// ExchangePricingEngine manages all exchange-specific pricing logic
type ExchangePricingEngine struct {
	decimalConverter *DecimalConverter
	pricers          map[ExchangeType]ExchangePricer
}

// NewExchangePricingEngine creates a new pricing engine with all exchange support
func NewExchangePricingEngine() *ExchangePricingEngine {
	dc := NewDecimalConverter()

	engine := &ExchangePricingEngine{
		decimalConverter: dc,
		pricers:          make(map[ExchangeType]ExchangePricer),
	}

	// Register all exchange pricers
	engine.pricers[ExchangeUniswapV3] = NewUniswapV3Pricer(dc)
	engine.pricers[ExchangeUniswapV2] = NewUniswapV2Pricer(dc)
	engine.pricers[ExchangeSushiSwap] = NewSushiSwapPricer(dc)
	engine.pricers[ExchangeCamelot] = NewCamelotPricer(dc)
	engine.pricers[ExchangeBalancer] = NewBalancerPricer(dc)
	engine.pricers[ExchangeTraderJoe] = NewTraderJoePricer(dc)
	engine.pricers[ExchangeRamses] = NewRamsesPricer(dc)
	engine.pricers[ExchangeCurve] = NewCurvePricer(dc)

	return engine
}

// GetExchangePricer returns the appropriate pricer for an exchange
func (engine *ExchangePricingEngine) GetExchangePricer(exchangeType ExchangeType) (ExchangePricer, error) {
	pricer, exists := engine.pricers[exchangeType]
	if !exists {
		return nil, fmt.Errorf("unsupported exchange type: %s", exchangeType)
	}
	return pricer, nil
}

// CalculateSpotPrice gets spot price from any exchange
func (engine *ExchangePricingEngine) CalculateSpotPrice(poolData *PoolData) (*UniversalDecimal, error) {
	pricer, err := engine.GetExchangePricer(poolData.ExchangeType)
	if err != nil {
		return nil, err
	}

	return pricer.GetSpotPrice(poolData)
}

// UniswapV3Pricer implements Uniswap V3 concentrated liquidity pricing
type UniswapV3Pricer struct {
	dc *DecimalConverter
}

func NewUniswapV3Pricer(dc *DecimalConverter) *UniswapV3Pricer {
	return &UniswapV3Pricer{dc: dc}
}

func (p *UniswapV3Pricer) GetSpotPrice(poolData *PoolData) (*UniversalDecimal, error) {
	if poolData.SqrtPriceX96 == nil {
		return nil, fmt.Errorf("missing sqrtPriceX96 for Uniswap V3 pool")
	}

	// Use cached function for optimized calculation
	// Convert sqrtPriceX96 to actual price using cached constants
	// price = sqrtPriceX96^2 / 2^192
	price := SqrtPriceX96ToPriceCached(poolData.SqrtPriceX96)

	// Adjust for decimal differences between tokens
	if poolData.Token0.Decimals != poolData.Token1.Decimals {
		decimalDiff := int(poolData.Token1.Decimals) - int(poolData.Token0.Decimals)
		adjustment := new(big.Float).SetInt(new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(decimalDiff)), nil))
		price.Mul(price, adjustment)
	}

	// Convert back to big.Int with appropriate precision
	priceInt := new(big.Int)
	priceScaled := new(big.Float).Mul(price, new(big.Float).SetInt(p.dc.getScalingFactor(18)))
	priceScaled.Int(priceInt)

	return NewUniversalDecimal(priceInt, 18, fmt.Sprintf("%s/%s", poolData.Token1.Symbol, poolData.Token0.Symbol))
}

func (p *UniswapV3Pricer) CalculateAmountOut(amountIn *UniversalDecimal, poolData *PoolData) (*UniversalDecimal, error) {
	// Uniswap V3 concentrated liquidity calculation
	// This is a simplified version - production would need full tick math

	if poolData.Liquidity == nil || poolData.Liquidity.Sign() == 0 {
		return nil, fmt.Errorf("insufficient liquidity in Uniswap V3 pool")
	}

	// Apply fee
	feeAmount, err := p.dc.Multiply(amountIn, poolData.Fee, amountIn.Decimals, "FEE")
	if err != nil {
		return nil, fmt.Errorf("error calculating fee: %w", err)
	}

	amountInAfterFee, err := p.dc.Subtract(amountIn, feeAmount)
	if err != nil {
		return nil, fmt.Errorf("error subtracting fee: %w", err)
	}

	// Simplified constant product formula for demonstration
	// Real implementation would use tick mathematics
	numerator, err := p.dc.Multiply(amountInAfterFee, poolData.Reserve1, poolData.Reserve1.Decimals, "TEMP")
	if err != nil {
		return nil, err
	}

	denominator, err := p.dc.Add(poolData.Reserve0, amountInAfterFee)
	if err != nil {
		return nil, err
	}

	return p.dc.Divide(numerator, denominator, poolData.Token1.Decimals, poolData.Token1.Symbol)
}

func (p *UniswapV3Pricer) CalculateAmountIn(amountOut *UniversalDecimal, poolData *PoolData) (*UniversalDecimal, error) {
	// Reverse calculation for Uniswap V3
	if poolData.Reserve1.IsZero() || amountOut.Value.Cmp(poolData.Reserve1.Value) >= 0 {
		return nil, fmt.Errorf("insufficient liquidity for requested output amount")
	}

	// Simplified reverse calculation
	numerator, err := p.dc.Multiply(poolData.Reserve0, amountOut, poolData.Reserve0.Decimals, "TEMP")
	if err != nil {
		return nil, err
	}

	denominator, err := p.dc.Subtract(poolData.Reserve1, amountOut)
	if err != nil {
		return nil, err
	}

	amountInBeforeFee, err := p.dc.Divide(numerator, denominator, poolData.Token0.Decimals, "TEMP")
	if err != nil {
		return nil, err
	}

	// Add fee
	feeMultiplier, err := p.dc.FromString("1", 18, "FEE_MULT")
	if err != nil {
		return nil, err
	}

	oneMinusFee, err := p.dc.Subtract(feeMultiplier, poolData.Fee)
	if err != nil {
		return nil, err
	}

	return p.dc.Divide(amountInBeforeFee, oneMinusFee, poolData.Token0.Decimals, poolData.Token0.Symbol)
}

func (p *UniswapV3Pricer) CalculatePriceImpact(amountIn *UniversalDecimal, poolData *PoolData) (*UniversalDecimal, error) {
	// Calculate price before trade
	priceBefore, err := p.GetSpotPrice(poolData)
	if err != nil {
		return nil, fmt.Errorf("error getting spot price: %w", err)
	}

	// Calculate amount out
	amountOut, err := p.CalculateAmountOut(amountIn, poolData)
	if err != nil {
		return nil, fmt.Errorf("error calculating amount out: %w", err)
	}

	// Calculate effective price
	effectivePrice, err := p.dc.Divide(amountOut, amountIn, 18, "EFFECTIVE_PRICE")
	if err != nil {
		return nil, fmt.Errorf("error calculating effective price: %w", err)
	}

	// Calculate price impact as percentage
	priceDiff, err := p.dc.Subtract(priceBefore, effectivePrice)
	if err != nil {
		return nil, fmt.Errorf("error calculating price difference: %w", err)
	}

	return p.dc.CalculatePercentage(priceDiff, priceBefore)
}

func (p *UniswapV3Pricer) GetMinimumLiquidity(poolData *PoolData) (*UniversalDecimal, error) {
	if poolData.Liquidity == nil {
		return NewUniversalDecimal(big.NewInt(0), 18, "LIQUIDITY")
	}

	return NewUniversalDecimal(poolData.Liquidity, 18, "LIQUIDITY")
}

func (p *UniswapV3Pricer) ValidatePoolData(poolData *PoolData) error {
	if poolData.SqrtPriceX96 == nil {
		return fmt.Errorf("Uniswap V3 pool missing sqrtPriceX96")
	}
	if poolData.Liquidity == nil {
		return fmt.Errorf("Uniswap V3 pool missing liquidity")
	}
	if poolData.Fee == nil {
		return fmt.Errorf("Uniswap V3 pool missing fee")
	}
	return nil
}

// UniswapV2Pricer implements Uniswap V2 / SushiSwap constant product pricing
type UniswapV2Pricer struct {
	dc *DecimalConverter
}

func NewUniswapV2Pricer(dc *DecimalConverter) *UniswapV2Pricer {
	return &UniswapV2Pricer{dc: dc}
}

func (p *UniswapV2Pricer) GetSpotPrice(poolData *PoolData) (*UniversalDecimal, error) {
	if poolData.Reserve0.IsZero() {
		return nil, fmt.Errorf("zero reserve0 in constant product pool")
	}

	return p.dc.Divide(poolData.Reserve1, poolData.Reserve0, 18, fmt.Sprintf("%s/%s", poolData.Token1.Symbol, poolData.Token0.Symbol))
}

func (p *UniswapV2Pricer) CalculateAmountOut(amountIn *UniversalDecimal, poolData *PoolData) (*UniversalDecimal, error) {
	// Uniswap V2 constant product formula: x * y = k
	// amountOut = (amountIn * 997 * reserveOut) / (reserveIn * 1000 + amountIn * 997)

	// Apply fee (0.3% = 997/1000 remaining)
	feeNumerator, _ := p.dc.FromString("997", 0, "FEE_NUM")
	feeDenominator, _ := p.dc.FromString("1000", 0, "FEE_DEN")

	amountInWithFee, err := p.dc.Multiply(amountIn, feeNumerator, amountIn.Decimals, "TEMP")
	if err != nil {
		return nil, err
	}

	numerator, err := p.dc.Multiply(amountInWithFee, poolData.Reserve1, poolData.Reserve1.Decimals, "TEMP")
	if err != nil {
		return nil, err
	}

	reserveInScaled, err := p.dc.Multiply(poolData.Reserve0, feeDenominator, poolData.Reserve0.Decimals, "TEMP")
	if err != nil {
		return nil, err
	}

	denominator, err := p.dc.Add(reserveInScaled, amountInWithFee)
	if err != nil {
		return nil, err
	}

	return p.dc.Divide(numerator, denominator, poolData.Token1.Decimals, poolData.Token1.Symbol)
}

func (p *UniswapV2Pricer) CalculateAmountIn(amountOut *UniversalDecimal, poolData *PoolData) (*UniversalDecimal, error) {
	// Reverse calculation for constant product
	feeNumerator, _ := p.dc.FromString("1000", 0, "FEE_NUM")
	feeDenominator, _ := p.dc.FromString("997", 0, "FEE_DEN")

	numerator, err := p.dc.Multiply(poolData.Reserve0, amountOut, poolData.Reserve0.Decimals, "TEMP")
	if err != nil {
		return nil, err
	}

	numeratorWithFee, err := p.dc.Multiply(numerator, feeNumerator, numerator.Decimals, "TEMP")
	if err != nil {
		return nil, err
	}

	denominator, err := p.dc.Subtract(poolData.Reserve1, amountOut)
	if err != nil {
		return nil, err
	}

	denominatorWithFee, err := p.dc.Multiply(denominator, feeDenominator, denominator.Decimals, "TEMP")
	if err != nil {
		return nil, err
	}

	return p.dc.Divide(numeratorWithFee, denominatorWithFee, poolData.Token0.Decimals, poolData.Token0.Symbol)
}

func (p *UniswapV2Pricer) CalculatePriceImpact(amountIn *UniversalDecimal, poolData *PoolData) (*UniversalDecimal, error) {
	// Similar to Uniswap V3 implementation
	priceBefore, err := p.GetSpotPrice(poolData)
	if err != nil {
		return nil, err
	}

	amountOut, err := p.CalculateAmountOut(amountIn, poolData)
	if err != nil {
		return nil, err
	}

	effectivePrice, err := p.dc.Divide(amountOut, amountIn, 18, "EFFECTIVE_PRICE")
	if err != nil {
		return nil, err
	}

	priceDiff, err := p.dc.Subtract(priceBefore, effectivePrice)
	if err != nil {
		return nil, err
	}

	return p.dc.CalculatePercentage(priceDiff, priceBefore)
}

func (p *UniswapV2Pricer) GetMinimumLiquidity(poolData *PoolData) (*UniversalDecimal, error) {
	// Geometric mean of reserves
	product, err := p.dc.Multiply(poolData.Reserve0, poolData.Reserve1, 18, "LIQUIDITY")
	if err != nil {
		return nil, err
	}

	// Simplified square root - in production use precise sqrt algorithm
	sqrt := new(big.Int).Sqrt(product.Value)
	return NewUniversalDecimal(sqrt, 18, "LIQUIDITY")
}

func (p *UniswapV2Pricer) ValidatePoolData(poolData *PoolData) error {
	if poolData.Reserve0 == nil || poolData.Reserve1 == nil {
		return fmt.Errorf("missing reserves for constant product pool")
	}
	if poolData.Reserve0.IsZero() || poolData.Reserve1.IsZero() {
		return fmt.Errorf("zero reserves in constant product pool")
	}
	return nil
}

// SushiSwapPricer uses same logic as Uniswap V2
type SushiSwapPricer struct {
	*UniswapV2Pricer
}

func NewSushiSwapPricer(dc *DecimalConverter) *SushiSwapPricer {
	return &SushiSwapPricer{NewUniswapV2Pricer(dc)}
}

// CamelotPricer - Algebra-based DEX on Arbitrum
type CamelotPricer struct {
	*UniswapV3Pricer
}

func NewCamelotPricer(dc *DecimalConverter) *CamelotPricer {
	return &CamelotPricer{NewUniswapV3Pricer(dc)}
}

// BalancerPricer - Weighted pool implementation
type BalancerPricer struct {
	dc *DecimalConverter
}

func NewBalancerPricer(dc *DecimalConverter) *BalancerPricer {
	return &BalancerPricer{dc: dc}
}

func (p *BalancerPricer) GetSpotPrice(poolData *PoolData) (*UniversalDecimal, error) {
	if len(poolData.Weights) < 2 {
		return nil, fmt.Errorf("insufficient weights for Balancer pool")
	}

	// Balancer spot price = (reserveOut/weightOut) / (reserveIn/weightIn)
	reserveOutWeighted, err := p.dc.Divide(poolData.Reserve1, poolData.Weights[1], 18, "TEMP")
	if err != nil {
		return nil, err
	}

	reserveInWeighted, err := p.dc.Divide(poolData.Reserve0, poolData.Weights[0], 18, "TEMP")
	if err != nil {
		return nil, err
	}

	return p.dc.Divide(reserveOutWeighted, reserveInWeighted, 18, fmt.Sprintf("%s/%s", poolData.Token1.Symbol, poolData.Token0.Symbol))
}

func (p *BalancerPricer) CalculateAmountOut(amountIn *UniversalDecimal, poolData *PoolData) (*UniversalDecimal, error) {
	// Simplified Balancer calculation - production needs full weighted math
	return p.GetSpotPrice(poolData)
}

func (p *BalancerPricer) CalculateAmountIn(amountOut *UniversalDecimal, poolData *PoolData) (*UniversalDecimal, error) {
	spotPrice, err := p.GetSpotPrice(poolData)
	if err != nil {
		return nil, err
	}

	return p.dc.Divide(amountOut, spotPrice, poolData.Token0.Decimals, poolData.Token0.Symbol)
}

func (p *BalancerPricer) CalculatePriceImpact(amountIn *UniversalDecimal, poolData *PoolData) (*UniversalDecimal, error) {
	// Placeholder - would implement Balancer-specific price impact
	return NewUniversalDecimal(big.NewInt(0), 4, "PERCENT")
}

func (p *BalancerPricer) GetMinimumLiquidity(poolData *PoolData) (*UniversalDecimal, error) {
	return NewUniversalDecimal(big.NewInt(0), 18, "LIQUIDITY")
}

func (p *BalancerPricer) ValidatePoolData(poolData *PoolData) error {
	if len(poolData.Weights) < 2 {
		return fmt.Errorf("Balancer pool missing weights")
	}
	return nil
}

// Placeholder implementations for other exchanges
func NewTraderJoePricer(dc *DecimalConverter) *UniswapV2Pricer { return NewUniswapV2Pricer(dc) }
func NewRamsesPricer(dc *DecimalConverter) *UniswapV3Pricer    { return NewUniswapV3Pricer(dc) }

// CurvePricer - Stable swap implementation
type CurvePricer struct {
	dc *DecimalConverter
}

func NewCurvePricer(dc *DecimalConverter) *CurvePricer {
	return &CurvePricer{dc: dc}
}

func (p *CurvePricer) GetSpotPrice(poolData *PoolData) (*UniversalDecimal, error) {
	// Curve stable swap pricing - simplified version
	if poolData.A == nil {
		return nil, fmt.Errorf("missing amplification coefficient for Curve pool")
	}

	// For stable swaps, price should be close to 1:1
	return NewUniversalDecimal(p.dc.getScalingFactor(18), 18, fmt.Sprintf("%s/%s", poolData.Token1.Symbol, poolData.Token0.Symbol))
}

func (p *CurvePricer) CalculateAmountOut(amountIn *UniversalDecimal, poolData *PoolData) (*UniversalDecimal, error) {
	// Simplified stable swap calculation
	// Real implementation would use Newton's method for stable swap invariant
	feeAmount, err := p.dc.Multiply(amountIn, poolData.Fee, amountIn.Decimals, "FEE")
	if err != nil {
		return nil, err
	}

	return p.dc.Subtract(amountIn, feeAmount)
}

func (p *CurvePricer) CalculateAmountIn(amountOut *UniversalDecimal, poolData *PoolData) (*UniversalDecimal, error) {
	// Reverse stable swap calculation
	feeMultiplier, _ := p.dc.FromString("1", 18, "FEE_MULT")
	oneMinusFee, err := p.dc.Subtract(feeMultiplier, poolData.Fee)
	if err != nil {
		return nil, err
	}

	return p.dc.Divide(amountOut, oneMinusFee, poolData.Token0.Decimals, poolData.Token0.Symbol)
}

func (p *CurvePricer) CalculatePriceImpact(amountIn *UniversalDecimal, poolData *PoolData) (*UniversalDecimal, error) {
	// Curve pools have minimal price impact for stable pairs
	return NewUniversalDecimal(big.NewInt(1000), 4, "PERCENT") // 0.1%
}

func (p *CurvePricer) GetMinimumLiquidity(poolData *PoolData) (*UniversalDecimal, error) {
	return NewUniversalDecimal(big.NewInt(0), 18, "LIQUIDITY")
}

func (p *CurvePricer) ValidatePoolData(poolData *PoolData) error {
	if poolData.A == nil {
		return fmt.Errorf("Curve pool missing amplification coefficient")
	}
	return nil
}