mev-beta/pkg/arbitrage/path_finder.go

package arbitrage

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

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

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

// PathFinderConfig contains configuration for path finding
type PathFinderConfig struct {
	MaxHops         int      // Maximum number of hops (2-4)
	MinLiquidity    *big.Int // Minimum liquidity per pool
	AllowedProtocols []types.ProtocolType
	MaxPathsPerPair int // Maximum paths to return per token pair
}

// DefaultPathFinderConfig returns default configuration
func DefaultPathFinderConfig() *PathFinderConfig {
	return &PathFinderConfig{
		MaxHops:      4,
		MinLiquidity: new(big.Int).Mul(big.NewInt(10000), new(big.Int).Exp(big.NewInt(10), big.NewInt(18), nil)), // 10,000 tokens
		AllowedProtocols: []types.ProtocolType{
			types.ProtocolUniswapV2,
			types.ProtocolUniswapV3,
			types.ProtocolSushiSwap,
			types.ProtocolCurve,
		},
		MaxPathsPerPair: 10,
	}
}

// PathFinder finds arbitrage paths between tokens
type PathFinder struct {
	cache  cache.PoolCache
	config *PathFinderConfig
	logger *slog.Logger
}

// NewPathFinder creates a new path finder
func NewPathFinder(cache cache.PoolCache, config *PathFinderConfig, logger *slog.Logger) *PathFinder {
	if config == nil {
		config = DefaultPathFinderConfig()
	}

	return &PathFinder{
		cache:  cache,
		config: config,
		logger: logger.With("component", "path_finder"),
	}
}

// Path represents a route through multiple pools
type Path struct {
	Tokens []common.Address
	Pools  []*types.PoolInfo
	Type   OpportunityType
}

// FindTwoPoolPaths finds simple two-pool arbitrage paths (A→B→A)
func (pf *PathFinder) FindTwoPoolPaths(ctx context.Context, tokenA, tokenB common.Address) ([]*Path, error) {
	pf.logger.Debug("finding two-pool paths",
		"tokenA", tokenA.Hex(),
		"tokenB", tokenB.Hex(),
	)

	// Get all pools containing tokenA and tokenB
	poolsAB, err := pf.cache.GetByTokenPair(ctx, tokenA, tokenB)
	if err != nil {
		return nil, fmt.Errorf("failed to get pools: %w", err)
	}

	// Filter by liquidity and protocols
	validPools := pf.filterPools(poolsAB)
	if len(validPools) < 2 {
		return nil, fmt.Errorf("insufficient pools for two-pool arbitrage: need at least 2, found %d", len(validPools))
	}

	paths := make([]*Path, 0)

	// Generate all pairs of pools
	for i := 0; i < len(validPools); i++ {
		for j := i + 1; j < len(validPools); j++ {
			pool1 := validPools[i]
			pool2 := validPools[j]

			// Two-pool arbitrage: buy on pool1, sell on pool2
			path := &Path{
				Tokens: []common.Address{tokenA, tokenB, tokenA},
				Pools:  []*types.PoolInfo{pool1, pool2},
				Type:   OpportunityTypeTwoPool,
			}
			paths = append(paths, path)

			// Also try reverse: buy on pool2, sell on pool1
			reversePath := &Path{
				Tokens: []common.Address{tokenA, tokenB, tokenA},
				Pools:  []*types.PoolInfo{pool2, pool1},
				Type:   OpportunityTypeTwoPool,
			}
			paths = append(paths, reversePath)
		}
	}

	pf.logger.Debug("found two-pool paths",
		"count", len(paths),
	)

	if len(paths) > pf.config.MaxPathsPerPair {
		paths = paths[:pf.config.MaxPathsPerPair]
	}

	return paths, nil
}

// FindTriangularPaths finds triangular arbitrage paths (A→B→C→A)
func (pf *PathFinder) FindTriangularPaths(ctx context.Context, tokenA common.Address) ([]*Path, error) {
	pf.logger.Debug("finding triangular paths",
		"tokenA", tokenA.Hex(),
	)

	// Get all pools containing tokenA
	poolsWithA, err := pf.cache.GetPoolsByToken(ctx, tokenA)
	if err != nil {
		return nil, fmt.Errorf("failed to get pools with tokenA: %w", err)
	}

	poolsWithA = pf.filterPools(poolsWithA)
	if len(poolsWithA) < 2 {
		return nil, fmt.Errorf("insufficient pools for triangular arbitrage")
	}

	paths := make([]*Path, 0)
	visited := make(map[string]bool)

	// For each pair of pools containing tokenA
	for i := 0; i < len(poolsWithA) && len(paths) < pf.config.MaxPathsPerPair; i++ {
		for j := i + 1; j < len(poolsWithA) && len(paths) < pf.config.MaxPathsPerPair; j++ {
			pool1 := poolsWithA[i]
			pool2 := poolsWithA[j]

			// Get the other tokens in each pool
			tokenB := pf.getOtherToken(pool1, tokenA)
			tokenC := pf.getOtherToken(pool2, tokenA)

			if tokenB == tokenC {
				continue // This would be a two-pool path
			}

			// Check if there's a pool connecting tokenB and tokenC
			poolsBC, err := pf.cache.GetByTokenPair(ctx, tokenB, tokenC)
			if err != nil {
				continue
			}

			poolsBC = pf.filterPools(poolsBC)
			if len(poolsBC) == 0 {
				continue
			}

			// For each connecting pool, create a triangular path
			for _, poolBC := range poolsBC {
				// Create path signature to avoid duplicates
				pathSig := fmt.Sprintf("%s-%s-%s", pool1.Address.Hex(), poolBC.Address.Hex(), pool2.Address.Hex())
				if visited[pathSig] {
					continue
				}
				visited[pathSig] = true

				path := &Path{
					Tokens: []common.Address{tokenA, tokenB, tokenC, tokenA},
					Pools:  []*types.PoolInfo{pool1, poolBC, pool2},
					Type:   OpportunityTypeTriangular,
				}
				paths = append(paths, path)

				if len(paths) >= pf.config.MaxPathsPerPair {
					break
				}
			}
		}
	}

	pf.logger.Debug("found triangular paths",
		"count", len(paths),
	)

	return paths, nil
}

// FindMultiHopPaths finds multi-hop arbitrage paths (up to MaxHops)
func (pf *PathFinder) FindMultiHopPaths(ctx context.Context, startToken, endToken common.Address, maxHops int) ([]*Path, error) {
	if maxHops < 2 || maxHops > pf.config.MaxHops {
		return nil, fmt.Errorf("invalid maxHops: must be between 2 and %d", pf.config.MaxHops)
	}

	pf.logger.Debug("finding multi-hop paths",
		"startToken", startToken.Hex(),
		"endToken", endToken.Hex(),
		"maxHops", maxHops,
	)

	paths := make([]*Path, 0)
	visited := make(map[string]bool)

	// BFS to find paths
	type searchNode struct {
		currentToken common.Address
		pools        []*types.PoolInfo
		tokens       []common.Address
		visited      map[common.Address]bool
	}

	queue := make([]*searchNode, 0)

	// Initialize with pools containing startToken
	startPools, err := pf.cache.GetPoolsByToken(ctx, startToken)
	if err != nil {
		return nil, fmt.Errorf("failed to get start pools: %w", err)
	}
	startPools = pf.filterPools(startPools)

	for _, pool := range startPools {
		nextToken := pf.getOtherToken(pool, startToken)
		if nextToken == (common.Address{}) {
			continue
		}

		visitedTokens := make(map[common.Address]bool)
		visitedTokens[startToken] = true

		queue = append(queue, &searchNode{
			currentToken: nextToken,
			pools:        []*types.PoolInfo{pool},
			tokens:       []common.Address{startToken, nextToken},
			visited:      visitedTokens,
		})
	}

	// BFS search
	for len(queue) > 0 && len(paths) < pf.config.MaxPathsPerPair {
		node := queue[0]
		queue = queue[1:]

		// Check if we've reached the end token
		if node.currentToken == endToken {
			// Found a path!
			pathSig := pf.getPathSignature(node.pools)
			if !visited[pathSig] {
				visited[pathSig] = true

				path := &Path{
					Tokens: node.tokens,
					Pools:  node.pools,
					Type:   OpportunityTypeMultiHop,
				}
				paths = append(paths, path)
			}
			continue
		}

		// Don't exceed max hops
		if len(node.pools) >= maxHops {
			continue
		}

		// Get pools containing current token
		nextPools, err := pf.cache.GetPoolsByToken(ctx, node.currentToken)
		if err != nil {
			continue
		}
		nextPools = pf.filterPools(nextPools)

		// Explore each next pool
		for _, pool := range nextPools {
			nextToken := pf.getOtherToken(pool, node.currentToken)
			if nextToken == (common.Address{}) {
				continue
			}

			// Don't revisit tokens (except endToken)
			if node.visited[nextToken] && nextToken != endToken {
				continue
			}

			// Create new search node
			newVisited := make(map[common.Address]bool)
			for k, v := range node.visited {
				newVisited[k] = v
			}
			newVisited[node.currentToken] = true

			newPools := make([]*types.PoolInfo, len(node.pools))
			copy(newPools, node.pools)
			newPools = append(newPools, pool)

			newTokens := make([]common.Address, len(node.tokens))
			copy(newTokens, node.tokens)
			newTokens = append(newTokens, nextToken)

			queue = append(queue, &searchNode{
				currentToken: nextToken,
				pools:        newPools,
				tokens:       newTokens,
				visited:      newVisited,
			})
		}
	}

	pf.logger.Debug("found multi-hop paths",
		"count", len(paths),
	)

	return paths, nil
}

// FindAllArbitragePaths finds all types of arbitrage paths for a token
func (pf *PathFinder) FindAllArbitragePaths(ctx context.Context, token common.Address) ([]*Path, error) {
	pf.logger.Debug("finding all arbitrage paths",
		"token", token.Hex(),
	)

	allPaths := make([]*Path, 0)

	// Find triangular paths
	triangular, err := pf.FindTriangularPaths(ctx, token)
	if err != nil {
		pf.logger.Warn("failed to find triangular paths", "error", err)
	} else {
		allPaths = append(allPaths, triangular...)
	}

	// Find two-pool paths with common pairs
	commonTokens := pf.getCommonTokens(ctx, token)
	for _, otherToken := range commonTokens {
		twoPools, err := pf.FindTwoPoolPaths(ctx, token, otherToken)
		if err != nil {
			continue
		}
		allPaths = append(allPaths, twoPools...)
	}

	pf.logger.Info("found all arbitrage paths",
		"token", token.Hex(),
		"totalPaths", len(allPaths),
	)

	return allPaths, nil
}

// filterPools filters pools by liquidity and protocol
func (pf *PathFinder) filterPools(pools []*types.PoolInfo) []*types.PoolInfo {
	filtered := make([]*types.PoolInfo, 0, len(pools))

	for _, pool := range pools {
		// Check if protocol is allowed
		allowed := false
		for _, proto := range pf.config.AllowedProtocols {
			if pool.Protocol == proto {
				allowed = true
				break
			}
		}
		if !allowed {
			continue
		}

		// Check minimum liquidity
		if pf.config.MinLiquidity != nil && pool.Liquidity != nil {
			if pool.Liquidity.Cmp(pf.config.MinLiquidity) < 0 {
				continue
			}
		}

		// Check if pool is active
		if !pool.IsActive {
			continue
		}

		filtered = append(filtered, pool)
	}

	return filtered
}

// getOtherToken returns the other token in a pool
func (pf *PathFinder) getOtherToken(pool *types.PoolInfo, token common.Address) common.Address {
	if pool.Token0 == token {
		return pool.Token1
	}
	if pool.Token1 == token {
		return pool.Token0
	}
	return common.Address{}
}

// getPathSignature creates a unique signature for a path
func (pf *PathFinder) getPathSignature(pools []*types.PoolInfo) string {
	sig := ""
	for i, pool := range pools {
		if i > 0 {
			sig += "-"
		}
		sig += pool.Address.Hex()
	}
	return sig
}

// getCommonTokens returns commonly traded tokens for finding two-pool paths
func (pf *PathFinder) getCommonTokens(ctx context.Context, baseToken common.Address) []common.Address {
	// In a real implementation, this would return the most liquid tokens
	// For now, return a hardcoded list of common Arbitrum tokens

	// WETH
	weth := common.HexToAddress("0x82aF49447D8a07e3bd95BD0d56f35241523fBab1")
	// USDC
	usdc := common.HexToAddress("0xFF970A61A04b1cA14834A43f5dE4533eBDDB5CC8")
	// USDT
	usdt := common.HexToAddress("0xFd086bC7CD5C481DCC9C85ebE478A1C0b69FCbb9")
	// DAI
	dai := common.HexToAddress("0xDA10009cBd5D07dd0CeCc66161FC93D7c9000da1")
	// ARB
	arb := common.HexToAddress("0x912CE59144191C1204E64559FE8253a0e49E6548")

	commonTokens := []common.Address{weth, usdc, usdt, dai, arb}

	// Filter out the base token itself
	filtered := make([]common.Address, 0)
	for _, token := range commonTokens {
		if token != baseToken {
			filtered = append(filtered, token)
		}
	}

	return filtered
}