mev-beta/pkg/dex/registry.go

package dex

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

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

// Registry manages all supported DEX protocols
type Registry struct {
	dexes  map[DEXProtocol]*DEXInfo
	mu     sync.RWMutex
	client *ethclient.Client
}

// NewRegistry creates a new DEX registry
func NewRegistry(client *ethclient.Client) *Registry {
	return &Registry{
		dexes:  make(map[DEXProtocol]*DEXInfo),
		client: client,
	}
}

// Register adds a DEX to the registry
func (r *Registry) Register(info *DEXInfo) error {
	if info == nil {
		return fmt.Errorf("DEX info cannot be nil")
	}
	if info.Decoder == nil {
		return fmt.Errorf("DEX decoder cannot be nil for %s", info.Name)
	}

	r.mu.Lock()
	defer r.mu.Unlock()

	r.dexes[info.Protocol] = info
	return nil
}

// Get retrieves a DEX by protocol
func (r *Registry) Get(protocol DEXProtocol) (*DEXInfo, error) {
	r.mu.RLock()
	defer r.mu.RUnlock()

	dex, exists := r.dexes[protocol]
	if !exists {
		return nil, fmt.Errorf("DEX protocol %s not registered", protocol)
	}
	if !dex.Active {
		return nil, fmt.Errorf("DEX protocol %s is not active", protocol)
	}
	return dex, nil
}

// GetAll returns all registered DEXes
func (r *Registry) GetAll() []*DEXInfo {
	r.mu.RLock()
	defer r.mu.RUnlock()

	dexes := make([]*DEXInfo, 0, len(r.dexes))
	for _, dex := range r.dexes {
		if dex.Active {
			dexes = append(dexes, dex)
		}
	}
	return dexes
}

// GetActiveDEXCount returns the number of active DEXes
func (r *Registry) GetActiveDEXCount() int {
	r.mu.RLock()
	defer r.mu.RUnlock()

	count := 0
	for _, dex := range r.dexes {
		if dex.Active {
			count++
		}
	}
	return count
}

// Deactivate deactivates a DEX
func (r *Registry) Deactivate(protocol DEXProtocol) error {
	r.mu.Lock()
	defer r.mu.Unlock()

	dex, exists := r.dexes[protocol]
	if !exists {
		return fmt.Errorf("DEX protocol %s not registered", protocol)
	}
	dex.Active = false
	return nil
}

// Activate activates a DEX
func (r *Registry) Activate(protocol DEXProtocol) error {
	r.mu.Lock()
	defer r.mu.Unlock()

	dex, exists := r.dexes[protocol]
	if !exists {
		return fmt.Errorf("DEX protocol %s not registered", protocol)
	}
	dex.Active = true
	return nil
}

// GetBestQuote finds the best price quote across all DEXes
func (r *Registry) GetBestQuote(ctx context.Context, tokenIn, tokenOut common.Address, amountIn *big.Int) (*PriceQuote, error) {
	dexes := r.GetAll()
	if len(dexes) == 0 {
		return nil, fmt.Errorf("no active DEXes registered")
	}

	type result struct {
		quote *PriceQuote
		err   error
	}

	results := make(chan result, len(dexes))

	// Query all DEXes in parallel
	for _, dex := range dexes {
		go func(d *DEXInfo) {
			quote, err := d.Decoder.GetQuote(ctx, r.client, tokenIn, tokenOut, amountIn)
			results <- result{quote: quote, err: err}
		}(dex)
	}

	// Collect results and find best quote
	var bestQuote *PriceQuote
	for i := 0; i < len(dexes); i++ {
		res := <-results
		if res.err != nil {
			continue // Skip failed quotes
		}
		if bestQuote == nil || res.quote.ExpectedOut.Cmp(bestQuote.ExpectedOut) > 0 {
			bestQuote = res.quote
		}
	}

	if bestQuote == nil {
		return nil, fmt.Errorf("no valid quotes found for %s -> %s", tokenIn.Hex(), tokenOut.Hex())
	}

	return bestQuote, nil
}

// FindArbitrageOpportunities finds arbitrage opportunities across DEXes
func (r *Registry) FindArbitrageOpportunities(ctx context.Context, tokenA, tokenB common.Address, amountIn *big.Int) ([]*ArbitragePath, error) {
	dexes := r.GetAll()
	if len(dexes) < 2 {
		return nil, fmt.Errorf("need at least 2 active DEXes for arbitrage, have %d", len(dexes))
	}

	opportunities := make([]*ArbitragePath, 0)

	// Simple 2-DEX arbitrage: Buy on DEX A, sell on DEX B
	for i, dexA := range dexes {
		for j, dexB := range dexes {
			if i >= j {
				continue // Avoid duplicate comparisons
			}

			// Get quote from DEX A (buy)
			quoteA, err := dexA.Decoder.GetQuote(ctx, r.client, tokenA, tokenB, amountIn)
			if err != nil {
				continue
			}

			// Get quote from DEX B (sell)
			quoteB, err := dexB.Decoder.GetQuote(ctx, r.client, tokenB, tokenA, quoteA.ExpectedOut)
			if err != nil {
				continue
			}

			// Calculate profit
			profit := new(big.Int).Sub(quoteB.ExpectedOut, amountIn)
			gasCost := new(big.Int).SetUint64((quoteA.GasEstimate + quoteB.GasEstimate) * 21000) // Rough estimate
			netProfit := new(big.Int).Sub(profit, gasCost)

			// Only consider profitable opportunities
			if netProfit.Sign() > 0 {
				profitETH := new(big.Float).Quo(
					new(big.Float).SetInt(netProfit),
					new(big.Float).SetInt(big.NewInt(1e18)),
				)
				profitFloat, _ := profitETH.Float64()

				roi := new(big.Float).Quo(
					new(big.Float).SetInt(netProfit),
					new(big.Float).SetInt(amountIn),
				)
				roiFloat, _ := roi.Float64()

				path := &ArbitragePath{
					Hops: []*PathHop{
						{
							DEX:         dexA.Protocol,
							PoolAddress: quoteA.PoolAddress,
							TokenIn:     tokenA,
							TokenOut:    tokenB,
							AmountIn:    amountIn,
							AmountOut:   quoteA.ExpectedOut,
							Fee:         quoteA.Fee,
						},
						{
							DEX:         dexB.Protocol,
							PoolAddress: quoteB.PoolAddress,
							TokenIn:     tokenB,
							TokenOut:    tokenA,
							AmountIn:    quoteA.ExpectedOut,
							AmountOut:   quoteB.ExpectedOut,
							Fee:         quoteB.Fee,
						},
					},
					TotalProfit: profit,
					ProfitETH:   profitFloat,
					ROI:         roiFloat,
					GasCost:     gasCost,
					NetProfit:   netProfit,
					Confidence:  0.8, // Base confidence for 2-hop arbitrage
				}

				opportunities = append(opportunities, path)
			}
		}
	}

	return opportunities, nil
}

// InitializeArbitrumDEXes initializes all Arbitrum DEXes
func (r *Registry) InitializeArbitrumDEXes() error {
	// UniswapV3
	uniV3 := &DEXInfo{
		Protocol:       ProtocolUniswapV3,
		Name:           "Uniswap V3",
		RouterAddress:  common.HexToAddress("0xE592427A0AEce92De3Edee1F18E0157C05861564"),
		FactoryAddress: common.HexToAddress("0x1F98431c8aD98523631AE4a59f267346ea31F984"),
		Fee:            big.NewInt(30), // 0.3% default
		PricingModel:   PricingConcentrated,
		Decoder:        NewUniswapV3Decoder(r.client),
		Active:         true,
	}
	if err := r.Register(uniV3); err != nil {
		return fmt.Errorf("failed to register UniswapV3: %w", err)
	}

	// SushiSwap
	sushi := &DEXInfo{
		Protocol:       ProtocolSushiSwap,
		Name:           "SushiSwap",
		RouterAddress:  common.HexToAddress("0x1b02dA8Cb0d097eB8D57A175b88c7D8b47997506"),
		FactoryAddress: common.HexToAddress("0xc35DADB65012eC5796536bD9864eD8773aBc74C4"),
		Fee:            big.NewInt(30), // 0.3%
		PricingModel:   PricingConstantProduct,
		Decoder:        NewSushiSwapDecoder(r.client),
		Active:         true,
	}
	if err := r.Register(sushi); err != nil {
		return fmt.Errorf("failed to register SushiSwap: %w", err)
	}

	// Curve - PRODUCTION READY
	curve := &DEXInfo{
		Protocol:       ProtocolCurve,
		Name:           "Curve",
		RouterAddress:  common.HexToAddress("0x0000000000000000000000000000000000000000"), // Curve uses individual pools
		FactoryAddress: common.HexToAddress("0xb17b674D9c5CB2e441F8e196a2f048A81355d031"), // Curve Factory on Arbitrum
		Fee:            big.NewInt(4), // 0.04% typical
		PricingModel:   PricingStableSwap,
		Decoder:        NewCurveDecoder(r.client),
		Active:         true, // ACTIVATED
	}
	if err := r.Register(curve); err != nil {
		return fmt.Errorf("failed to register Curve: %w", err)
	}

	// Balancer - PRODUCTION READY
	balancer := &DEXInfo{
		Protocol:       ProtocolBalancer,
		Name:           "Balancer",
		RouterAddress:  common.HexToAddress("0xBA12222222228d8Ba445958a75a0704d566BF2C8"), // Balancer Vault
		FactoryAddress: common.HexToAddress("0x0000000000000000000000000000000000000000"), // Uses Vault
		Fee:            big.NewInt(25), // 0.25% typical
		PricingModel:   PricingWeighted,
		Decoder:        NewBalancerDecoder(r.client),
		Active:         true, // ACTIVATED
	}
	if err := r.Register(balancer); err != nil {
		return fmt.Errorf("failed to register Balancer: %w", err)
	}

	return nil
}