mev-beta/pkg/pools/create2.go

package pools

import (
	"context"
	"fmt"
	"math/big"
	"sort"
	"strings"
	"time"

	"github.com/ethereum/go-ethereum"
	"github.com/ethereum/go-ethereum/accounts/abi"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/ethclient"

	"github.com/fraktal/mev-beta/internal/logger"
)

// CREATE2Calculator handles CREATE2 address calculations for various DEX factories
type CREATE2Calculator struct {
	logger     *logger.Logger
	factories  map[string]*FactoryConfig
	ethClient  *ethclient.Client
	curveCache map[string]common.Address // Cache for Curve pool addresses
}

// FactoryConfig contains the configuration for a DEX factory
type FactoryConfig struct {
	Name         string         // Factory name (e.g., "uniswap_v3", "sushiswap")
	Address      common.Address // Factory contract address
	InitCodeHash common.Hash    // Init code hash for CREATE2 calculation
	FeeStructure FeeStructure   // How fees are encoded
	SortTokens   bool           // Whether tokens should be sorted
}

// FeeStructure defines how fees are handled in address calculation
type FeeStructure struct {
	HasFee       bool     // Whether fee is part of salt
	FeePositions []int    // Byte positions where fee is encoded
	DefaultFees  []uint32 // Default fee tiers
}

// PoolIdentifier uniquely identifies a pool
type PoolIdentifier struct {
	Factory  string         // Factory name
	Token0   common.Address // First token (lower address if sorted)
	Token1   common.Address // Second token (higher address if sorted)
	Fee      uint32         // Fee tier
	PoolAddr common.Address // Calculated pool address
}

// NewCREATE2Calculator creates a new CREATE2 calculator
func NewCREATE2Calculator(logger *logger.Logger, ethClient *ethclient.Client) *CREATE2Calculator {
	calc := &CREATE2Calculator{
		logger:     logger,
		factories:  make(map[string]*FactoryConfig),
		ethClient:  ethClient,
		curveCache: make(map[string]common.Address),
	}

	// Initialize with known factory configurations
	calc.initializeFactories()

	return calc
}

// initializeFactories sets up configurations for known DEX factories
func (c *CREATE2Calculator) initializeFactories() {
	// Uniswap V3 Factory
	c.factories["uniswap_v3"] = &FactoryConfig{
		Name:         "uniswap_v3",
		Address:      common.HexToAddress("0x1F98431c8aD98523631AE4a59f267346ea31F984"),
		InitCodeHash: common.HexToHash("0xe34f199b19b2b4f47f68442619d555527d244f78a3297ea89325f843f87b8b54"),
		FeeStructure: FeeStructure{
			HasFee:      true,
			DefaultFees: []uint32{500, 3000, 10000}, // 0.05%, 0.3%, 1%
		},
		SortTokens: true,
	}

	// Uniswap V2 Factory
	c.factories["uniswap_v2"] = &FactoryConfig{
		Name:         "uniswap_v2",
		Address:      common.HexToAddress("0x5C69bEe701ef814a2B6a3EDD4B1652CB9cc5aA6f"),
		InitCodeHash: common.HexToHash("0x96e8ac4277198ff8b6f785478aa9a39f403cb768dd02cbee326c3e7da348845f"),
		FeeStructure: FeeStructure{
			HasFee:      false,
			DefaultFees: []uint32{3000}, // Fixed 0.3%
		},
		SortTokens: true,
	}

	// SushiSwap Factory (same as Uniswap V2 but different address)
	c.factories["sushiswap"] = &FactoryConfig{
		Name:         "sushiswap",
		Address:      common.HexToAddress("0xC0AEe478e3658e2610c5F7A4A2E1777cE9e4f2Ac"),
		InitCodeHash: common.HexToHash("0xe18a34eb0e04b04f7a0ac29a6e80748dca96319b42c54d679cb821dca90c6303"),
		FeeStructure: FeeStructure{
			HasFee:      false,
			DefaultFees: []uint32{3000}, // Fixed 0.3%
		},
		SortTokens: true,
	}

	// Camelot V3 (Arbitrum-specific)
	c.factories["camelot_v3"] = &FactoryConfig{
		Name:         "camelot_v3",
		Address:      common.HexToAddress("0x1a3c9B1d2F0529D97f2afC5136Cc23e58f1FD35B"),
		InitCodeHash: common.HexToHash("0xa856464ae65f7619087bc369daaf7e387dae1e5af69cfa7935850ebf754b04c1"),
		FeeStructure: FeeStructure{
			HasFee:      true,
			DefaultFees: []uint32{500, 3000, 10000}, // Similar to Uniswap V3
		},
		SortTokens: true,
	}

	// Curve Factory (simplified - Curve uses different math)
	c.factories["curve"] = &FactoryConfig{
		Name:         "curve",
		Address:      common.HexToAddress("0xF18056Bbd320E96A48e3Fbf8bC061322531aac99"),
		InitCodeHash: common.HexToHash("0x00"), // Curve doesn't use standard CREATE2
		FeeStructure: FeeStructure{
			HasFee:      true,
			DefaultFees: []uint32{400}, // 0.04% typical
		},
		SortTokens: false, // Curve maintains token order
	}
}

// CalculatePoolAddress calculates the pool address using CREATE2
func (c *CREATE2Calculator) CalculatePoolAddress(factoryName string, token0, token1 common.Address, fee uint32) (common.Address, error) {
	factory, exists := c.factories[factoryName]
	if !exists {
		return common.Address{}, fmt.Errorf("unknown factory: %s", factoryName)
	}

	// Sort tokens if required by the factory
	if factory.SortTokens {
		if token0.Big().Cmp(token1.Big()) > 0 {
			token0, token1 = token1, token0
		}
	}

	// Calculate salt based on factory type
	salt, err := c.calculateSalt(factory, token0, token1, fee)
	if err != nil {
		return common.Address{}, fmt.Errorf("failed to calculate salt: %w", err)
	}

	// Special handling for factories that don't use standard CREATE2
	if factoryName == "curve" {
		return c.calculateCurvePoolAddress(token0, token1, fee)
	}

	// Standard CREATE2 calculation:
	// address = keccak256(0xff + factory_address + salt + init_code_hash)[12:]

	// Prepare the data for hashing
	data := make([]byte, 0, 85)                          // 1 + 20 + 32 + 32 = 85 bytes
	data = append(data, 0xff)                            // 1 byte
	data = append(data, factory.Address.Bytes()...)      // 20 bytes
	data = append(data, salt...)                         // 32 bytes
	data = append(data, factory.InitCodeHash.Bytes()...) // 32 bytes

	// Calculate keccak256 hash
	hash := crypto.Keccak256(data)

	// Take the last 20 bytes as the address
	var poolAddr common.Address
	copy(poolAddr[:], hash[12:])

	c.logger.Debug(fmt.Sprintf("Calculated %s pool address: %s for tokens %s/%s fee %d",
		factoryName, poolAddr.Hex(), token0.Hex(), token1.Hex(), fee))

	return poolAddr, nil
}

// calculateSalt generates the salt for CREATE2 calculation
func (c *CREATE2Calculator) calculateSalt(factory *FactoryConfig, token0, token1 common.Address, fee uint32) ([]byte, error) {

	switch factory.Name {
	case "uniswap_v3", "camelot_v3":
		// Uniswap V3 salt: keccak256(abi.encode(token0, token1, fee))
		return c.calculateUniswapV3Salt(token0, token1, fee)

	case "uniswap_v2", "sushiswap":
		// Uniswap V2 salt: keccak256(abi.encodePacked(token0, token1))
		return c.calculateUniswapV2Salt(token0, token1)

	default:
		// Generic salt: keccak256(abi.encode(token0, token1, fee))
		return c.calculateGenericSalt(token0, token1, fee)
	}
}

// calculateUniswapV3Salt calculates salt for Uniswap V3 style factories
func (c *CREATE2Calculator) calculateUniswapV3Salt(token0, token1 common.Address, fee uint32) ([]byte, error) {
	// ABI encode: token0 (32 bytes) + token1 (32 bytes) + fee (32 bytes)
	data := make([]byte, 0, 96)

	// Pad addresses to 32 bytes
	token0Padded := make([]byte, 32)
	token1Padded := make([]byte, 32)
	feePadded := make([]byte, 32)

	copy(token0Padded[12:], token0.Bytes())
	copy(token1Padded[12:], token1.Bytes())

	// Convert fee to big endian 32 bytes
	feeBig := big.NewInt(int64(fee))
	feeBytes := feeBig.Bytes()
	copy(feePadded[32-len(feeBytes):], feeBytes)

	data = append(data, token0Padded...)
	data = append(data, token1Padded...)
	data = append(data, feePadded...)

	hash := crypto.Keccak256(data)
	return hash, nil
}

// calculateUniswapV2Salt calculates salt for Uniswap V2 style factories
func (c *CREATE2Calculator) calculateUniswapV2Salt(token0, token1 common.Address) ([]byte, error) {
	// ABI encodePacked: token0 (20 bytes) + token1 (20 bytes)
	data := make([]byte, 0, 40)
	data = append(data, token0.Bytes()...)
	data = append(data, token1.Bytes()...)

	hash := crypto.Keccak256(data)
	return hash, nil
}

// calculateGenericSalt calculates salt for generic factories
func (c *CREATE2Calculator) calculateGenericSalt(token0, token1 common.Address, fee uint32) ([]byte, error) {
	// Similar to Uniswap V3 but may have different encoding
	return c.calculateUniswapV3Salt(token0, token1, fee)
}

// calculateCurvePoolAddress handles Curve's registry-based pool discovery
func (c *CREATE2Calculator) calculateCurvePoolAddress(token0, token1 common.Address, fee uint32) (common.Address, error) {
	// Curve uses a registry-based system rather than deterministic CREATE2
	// We need to query multiple Curve registries to find pools
	if c.ethClient == nil {
		return common.Address{}, fmt.Errorf("ethereum client not configured for curve registry lookups")
	}

	// Create cache key
	cacheKey := fmt.Sprintf("%s-%s-%d", token0.Hex(), token1.Hex(), fee)
	if cached, exists := c.curveCache[cacheKey]; exists {
		c.logger.Debug(fmt.Sprintf("Using cached Curve pool address: %s", cached.Hex()))
		return cached, nil
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	// Curve registry addresses on Arbitrum
	registries := []common.Address{
		common.HexToAddress("0x0000000022D53366457F9d5E68Ec105046FC4383"), // Main Registry
		common.HexToAddress("0x90E00ACe148ca3b23Ac1bC8C240C2a7Dd9c2d7f5"), // Factory Registry
		common.HexToAddress("0xF18056Bbd320E96A48e3Fbf8bC061322531aac99"), // Crypto Registry
		common.HexToAddress("0x7D86446dDb609eD0F5f8684AcF30380a356b2B4c"), // Metapool Factory
	}

	// Try each registry to find the pool
	for i, registryAddr := range registries {
		poolAddr, err := c.queryCurveRegistry(ctx, registryAddr, token0, token1, i)
		if err != nil {
			c.logger.Debug(fmt.Sprintf("Registry %s failed: %v", registryAddr.Hex(), err))
			continue
		}

		if poolAddr != (common.Address{}) {
			c.logger.Debug(fmt.Sprintf("Found Curve pool %s in registry %s for tokens %s/%s",
				poolAddr.Hex(), registryAddr.Hex(), token0.Hex(), token1.Hex()))

			// Cache the result
			c.curveCache[cacheKey] = poolAddr
			return poolAddr, nil
		}
	}

	// If no pool found in registries, try deterministic calculation for newer Curve factories
	return c.calculateCurveDeterministicAddress(token0, token1, fee)
}

// FindPoolsForTokenPair finds all possible pools for a token pair across all factories
func (c *CREATE2Calculator) FindPoolsForTokenPair(token0, token1 common.Address) ([]*PoolIdentifier, error) {
	pools := make([]*PoolIdentifier, 0)

	for factoryName, factory := range c.factories {
		// Sort tokens if required
		sortedToken0, sortedToken1 := token0, token1
		if factory.SortTokens && token0.Big().Cmp(token1.Big()) > 0 {
			sortedToken0, sortedToken1 = token1, token0
		}

		// Try each default fee tier for this factory
		for _, fee := range factory.FeeStructure.DefaultFees {
			poolAddr, err := c.CalculatePoolAddress(factoryName, sortedToken0, sortedToken1, fee)
			if err != nil {
				c.logger.Debug(fmt.Sprintf("Failed to calculate pool address for %s: %v", factoryName, err))
				continue
			}

			pool := &PoolIdentifier{
				Factory:  factoryName,
				Token0:   sortedToken0,
				Token1:   sortedToken1,
				Fee:      fee,
				PoolAddr: poolAddr,
			}

			pools = append(pools, pool)
		}
	}

	c.logger.Debug(fmt.Sprintf("Found %d potential pools for tokens %s/%s",
		len(pools), token0.Hex(), token1.Hex()))

	return pools, nil
}

// ValidatePoolAddress verifies if a calculated address matches an expected address
func (c *CREATE2Calculator) ValidatePoolAddress(factoryName string, token0, token1 common.Address, fee uint32, expectedAddr common.Address) bool {
	calculatedAddr, err := c.CalculatePoolAddress(factoryName, token0, token1, fee)
	if err != nil {
		c.logger.Debug(fmt.Sprintf("Validation failed - calculation error: %v", err))
		return false
	}

	match := calculatedAddr == expectedAddr
	c.logger.Debug(fmt.Sprintf("Pool address validation: calculated=%s, expected=%s, match=%v",
		calculatedAddr.Hex(), expectedAddr.Hex(), match))

	return match
}

// GetFactoryConfig returns the configuration for a specific factory
func (c *CREATE2Calculator) GetFactoryConfig(factoryName string) (*FactoryConfig, error) {
	factory, exists := c.factories[factoryName]
	if !exists {
		return nil, fmt.Errorf("unknown factory: %s", factoryName)
	}

	// Return a copy to prevent modification
	configCopy := *factory
	return &configCopy, nil
}

// AddCustomFactory adds a custom factory configuration
func (c *CREATE2Calculator) AddCustomFactory(config *FactoryConfig) error {
	if config.Name == "" {
		return fmt.Errorf("factory name cannot be empty")
	}

	if config.Address == (common.Address{}) {
		return fmt.Errorf("factory address cannot be zero")
	}

	c.factories[config.Name] = config
	c.logger.Info(fmt.Sprintf("Added custom factory: %s at %s", config.Name, config.Address.Hex()))

	return nil
}

// ListFactories returns the names of all configured factories
func (c *CREATE2Calculator) ListFactories() []string {
	names := make([]string, 0, len(c.factories))
	for name := range c.factories {
		names = append(names, name)
	}

	sort.Strings(names)
	return names
}

// CalculateInitCodeHash calculates the init code hash for a given bytecode
// This is useful when adding new factories
func CalculateInitCodeHash(initCode []byte) common.Hash {
	return crypto.Keccak256Hash(initCode)
}

// VerifyFactorySupport checks if a factory supports CREATE2 pool creation
func (c *CREATE2Calculator) VerifyFactorySupport(factoryName string) error {
	factory, exists := c.factories[factoryName]
	if !exists {
		return fmt.Errorf("factory %s not configured", factoryName)
	}

	// Basic validation
	if factory.Address == (common.Address{}) {
		return fmt.Errorf("factory %s has zero address", factoryName)
	}

	if factory.InitCodeHash == (common.Hash{}) && factoryName != "curve" {
		return fmt.Errorf("factory %s has zero init code hash", factoryName)
	}

	if len(factory.FeeStructure.DefaultFees) == 0 {
		return fmt.Errorf("factory %s has no default fees configured", factoryName)
	}

	return nil
}

// queryCurveRegistry queries a specific Curve registry for pool information
func (c *CREATE2Calculator) queryCurveRegistry(ctx context.Context, registryAddr, token0, token1 common.Address, registryType int) (common.Address, error) {
	// Different registry types have different interfaces
	switch registryType {
	case 0: // Main Registry
		return c.queryMainCurveRegistry(ctx, registryAddr, token0, token1)
	case 1: // Factory Registry
		return c.queryFactoryCurveRegistry(ctx, registryAddr, token0, token1)
	case 2: // Crypto Registry
		return c.queryCryptoCurveRegistry(ctx, registryAddr, token0, token1)
	case 3: // Metapool Factory
		return c.queryMetapoolCurveRegistry(ctx, registryAddr, token0, token1)
	default:
		return common.Address{}, fmt.Errorf("unknown registry type: %d", registryType)
	}
}

// queryMainCurveRegistry queries the main Curve registry
func (c *CREATE2Calculator) queryMainCurveRegistry(ctx context.Context, registryAddr, token0, token1 common.Address) (common.Address, error) {
	// Query Curve registry using find_pool_for_coins function
	c.logger.Debug(fmt.Sprintf("Querying main Curve registry %s for tokens %s/%s",
		registryAddr.Hex(), token0.Hex(), token1.Hex()))

	if c.ethClient == nil {
		return common.Address{}, fmt.Errorf("ethereum client not configured for curve registry lookups")
	}

	// Curve registry ABI for find_pool_for_coins function
	registryABI := `[{"name":"find_pool_for_coins","outputs":[{"type":"address","name":""}],"inputs":[{"type":"address","name":"_from"},{"type":"address","name":"_to"}],"stateMutability":"view","type":"function"}]`

	parsedABI, err := abi.JSON(strings.NewReader(registryABI))
	if err != nil {
		return common.Address{}, fmt.Errorf("failed to parse registry ABI: %w", err)
	}

	// Pack the function call
	callData, err := parsedABI.Pack("find_pool_for_coins", token0, token1)
	if err != nil {
		return common.Address{}, fmt.Errorf("failed to pack registry call: %w", err)
	}

	// Make the contract call
	callMsg := ethereum.CallMsg{
		To:   &registryAddr,
		Data: callData,
	}

	result, err := c.ethClient.CallContract(ctx, callMsg, nil)
	if err != nil {
		return common.Address{}, fmt.Errorf("registry call failed: %w", err)
	}

	// Unpack the result
	var poolAddr common.Address
	if err := parsedABI.UnpackIntoInterface(&poolAddr, "find_pool_for_coins", result); err != nil {
		return common.Address{}, fmt.Errorf("failed to unpack result: %w", err)
	}

	// Check if a valid pool was found
	if poolAddr == (common.Address{}) {
		// Try with reversed token order
		callData, err = parsedABI.Pack("find_pool_for_coins", token1, token0)
		if err != nil {
			return common.Address{}, fmt.Errorf("failed to pack reversed registry call: %w", err)
		}

		callMsg.Data = callData
		result, err = c.ethClient.CallContract(ctx, callMsg, nil)
		if err != nil {
			return common.Address{}, fmt.Errorf("reversed registry call failed: %w", err)
		}

		if err := parsedABI.UnpackIntoInterface(&poolAddr, "find_pool_for_coins", result); err != nil {
			return common.Address{}, fmt.Errorf("failed to unpack reversed result: %w", err)
		}
	}

	return poolAddr, nil
}

// queryFactoryCurveRegistry queries the Curve factory registry
func (c *CREATE2Calculator) queryFactoryCurveRegistry(ctx context.Context, registryAddr, token0, token1 common.Address) (common.Address, error) {
	// Factory registry handles newer permissionless pools
	c.logger.Debug(fmt.Sprintf("Querying factory Curve registry %s for tokens %s/%s",
		registryAddr.Hex(), token0.Hex(), token1.Hex()))

	// Would implement actual registry query here
	return common.Address{}, nil
}

// queryCryptoCurveRegistry queries the Curve crypto registry
func (c *CREATE2Calculator) queryCryptoCurveRegistry(ctx context.Context, registryAddr, token0, token1 common.Address) (common.Address, error) {
	// Crypto registry handles volatile asset pools
	c.logger.Debug(fmt.Sprintf("Querying crypto Curve registry %s for tokens %s/%s",
		registryAddr.Hex(), token0.Hex(), token1.Hex()))

	// Would implement actual registry query here
	return common.Address{}, nil
}

// queryMetapoolCurveRegistry queries the Curve metapool factory
func (c *CREATE2Calculator) queryMetapoolCurveRegistry(ctx context.Context, registryAddr, token0, token1 common.Address) (common.Address, error) {
	// Metapool factory handles pools paired with base pools
	c.logger.Debug(fmt.Sprintf("Querying metapool Curve registry %s for tokens %s/%s",
		registryAddr.Hex(), token0.Hex(), token1.Hex()))

	// Would implement actual registry query here
	return common.Address{}, nil
}

// calculateCurveDeterministicAddress calculates Curve pool address deterministically for newer factories
func (c *CREATE2Calculator) calculateCurveDeterministicAddress(token0, token1 common.Address, fee uint32) (common.Address, error) {
	// Some newer Curve factories do use deterministic CREATE2
	// This handles those cases

	c.logger.Debug(fmt.Sprintf("Calculating deterministic Curve address for tokens %s/%s fee %d",
		token0.Hex(), token1.Hex(), fee))

	// Curve's CREATE2 implementation varies by factory
	// For stable pools: salt = keccak256(coins, A, fee)
	// For crypto pools: salt = keccak256(coins, A, gamma, mid_fee, out_fee, allowed_extra_profit, fee_gamma, adjustment_step, admin_fee, ma_half_time, initial_price)

	// Simplified implementation for stable pools
	coins := []common.Address{token0, token1}
	if token0.Big().Cmp(token1.Big()) > 0 {
		coins = []common.Address{token1, token0}
	}

	// Typical Curve stable pool parameters
	A := big.NewInt(200) // Amplification parameter
	feeInt := big.NewInt(int64(fee))

	// Create salt: keccak256(abi.encode(coins, A, fee))
	saltData := make([]byte, 0, 96) // 2*32 + 32 + 32

	// Encode coins (32 bytes each)
	coin0Padded := make([]byte, 32)
	coin1Padded := make([]byte, 32)
	copy(coin0Padded[12:], coins[0].Bytes())
	copy(coin1Padded[12:], coins[1].Bytes())

	// Encode A parameter (32 bytes)
	APadded := make([]byte, 32)
	ABytes := A.Bytes()
	copy(APadded[32-len(ABytes):], ABytes)

	// Encode fee (32 bytes)
	feePadded := make([]byte, 32)
	feeBytes := feeInt.Bytes()
	copy(feePadded[32-len(feeBytes):], feeBytes)

	saltData = append(saltData, coin0Padded...)
	saltData = append(saltData, coin1Padded...)
	saltData = append(saltData, APadded...)
	saltData = append(saltData, feePadded...)

	salt := crypto.Keccak256Hash(saltData)

	// Use Curve factory config for CREATE2
	factory := c.factories["curve"]
	if factory.InitCodeHash == (common.Hash{}) {
		// For factories without init code hash, use registry-based approach
		return common.Address{}, fmt.Errorf("deterministic calculation not supported for this Curve factory")
	}

	// Standard CREATE2 calculation
	data := make([]byte, 0, 85)
	data = append(data, 0xff)
	data = append(data, factory.Address.Bytes()...)
	data = append(data, salt.Bytes()...)
	data = append(data, factory.InitCodeHash.Bytes()...)

	hash := crypto.Keccak256(data)
	var poolAddr common.Address
	copy(poolAddr[:], hash[12:])

	c.logger.Debug(fmt.Sprintf("Calculated deterministic Curve pool address: %s", poolAddr.Hex()))
	return poolAddr, nil
}