mev-beta/pkg/pools/cache.go

package pools

import (
	"encoding/json"
	"fmt"
	"os"
	"sync"
	"time"

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

// PoolInfo represents information about a discovered pool
type PoolInfo struct {
	Address      common.Address `json:"address"`
	Protocol     string         `json:"protocol"`
	Version      string         `json:"version"`
	Type         string         `json:"type"` // "pool", "router", "vault"
	Token0       common.Address `json:"token0,omitempty"`
	Token1       common.Address `json:"token1,omitempty"`
	FirstSeen    time.Time      `json:"first_seen"`
	LastSeen     time.Time      `json:"last_seen"`
	SwapCount    uint64         `json:"swap_count"`
	IsVerified   bool           `json:"is_verified"`
}

// PoolCache manages discovered pools with thread-safe operations
type PoolCache struct {
	pools      map[common.Address]*PoolInfo
	mu         sync.RWMutex
	logger     log.Logger
	saveFile   string
	autoSave   bool
	saveEvery  int           // Save after N new pools
	newPools   int           // Counter for new pools since last save
	saveTicker *time.Ticker
}

// NewPoolCache creates a new pool cache
func NewPoolCache(saveFile string, autoSave bool, logger log.Logger) *PoolCache {
	cache := &PoolCache{
		pools:     make(map[common.Address]*PoolInfo),
		logger:    logger,
		saveFile:  saveFile,
		autoSave:  autoSave,
		saveEvery: 100, // Save every 100 new pools
		newPools:  0,
	}

	// Load existing pools from file
	if err := cache.Load(); err != nil {
		logger.Warn("failed to load pool cache", "error", err)
	}

	// Start periodic save if autosave enabled
	if autoSave {
		cache.saveTicker = time.NewTicker(5 * time.Minute)
		go cache.periodicSave()
	}

	return cache
}

// AddOrUpdate adds a new pool or updates an existing one
func (c *PoolCache) AddOrUpdate(pool *PoolInfo) bool {
	c.mu.Lock()
	defer c.mu.Unlock()

	existing, exists := c.pools[pool.Address]

	if exists {
		// Update existing pool
		existing.LastSeen = time.Now()
		existing.SwapCount++

		// Update tokens if they were unknown before
		if existing.Token0 == (common.Address{}) && pool.Token0 != (common.Address{}) {
			existing.Token0 = pool.Token0
		}
		if existing.Token1 == (common.Address{}) && pool.Token1 != (common.Address{}) {
			existing.Token1 = pool.Token1
		}

		return false // Not a new pool
	}

	// Add new pool
	pool.FirstSeen = time.Now()
	pool.LastSeen = time.Now()
	pool.SwapCount = 1
	c.pools[pool.Address] = pool

	c.newPools++

	// Auto-save if threshold reached
	if c.autoSave && c.newPools >= c.saveEvery {
		go c.Save() // Save in background
		c.newPools = 0
	}

	c.logger.Info("🆕 NEW POOL DISCOVERED",
		"address", pool.Address.Hex(),
		"protocol", pool.Protocol,
		"version", pool.Version,
		"type", pool.Type,
		"total_pools", len(c.pools),
	)

	return true // New pool
}

// Exists checks if a pool address is already in the cache
func (c *PoolCache) Exists(address common.Address) bool {
	c.mu.RLock()
	defer c.mu.RUnlock()

	_, exists := c.pools[address]
	return exists
}

// Get retrieves pool info by address
func (c *PoolCache) Get(address common.Address) (*PoolInfo, bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()

	pool, exists := c.pools[address]
	return pool, exists
}

// GetAll returns all pools (copy)
func (c *PoolCache) GetAll() []*PoolInfo {
	c.mu.RLock()
	defer c.mu.RUnlock()

	pools := make([]*PoolInfo, 0, len(c.pools))
	for _, pool := range c.pools {
		pools = append(pools, pool)
	}

	return pools
}

// GetByProtocol returns all pools for a specific protocol
func (c *PoolCache) GetByProtocol(protocol string) []*PoolInfo {
	c.mu.RLock()
	defer c.mu.RUnlock()

	pools := make([]*PoolInfo, 0)
	for _, pool := range c.pools {
		if pool.Protocol == protocol {
			pools = append(pools, pool)
		}
	}

	return pools
}

// Count returns the total number of pools
func (c *PoolCache) Count() int {
	c.mu.RLock()
	defer c.mu.RUnlock()

	return len(c.pools)
}

// Save writes the pool cache to disk
func (c *PoolCache) Save() error {
	c.mu.RLock()
	defer c.mu.RUnlock()

	if c.saveFile == "" {
		return fmt.Errorf("no save file configured")
	}

	// Convert to slice for JSON encoding
	pools := make([]*PoolInfo, 0, len(c.pools))
	for _, pool := range c.pools {
		pools = append(pools, pool)
	}

	data, err := json.MarshalIndent(pools, "", "  ")
	if err != nil {
		return fmt.Errorf("failed to marshal pools: %w", err)
	}

	if err := os.WriteFile(c.saveFile, data, 0644); err != nil {
		return fmt.Errorf("failed to write pool cache: %w", err)
	}

	c.logger.Info("💾 Pool cache saved", "file", c.saveFile, "pools", len(pools))
	return nil
}

// Load reads the pool cache from disk
func (c *PoolCache) Load() error {
	if c.saveFile == "" {
		return fmt.Errorf("no save file configured")
	}

	data, err := os.ReadFile(c.saveFile)
	if err != nil {
		if os.IsNotExist(err) {
			c.logger.Info("No existing pool cache found, starting fresh")
			return nil
		}
		return fmt.Errorf("failed to read pool cache: %w", err)
	}

	var pools []*PoolInfo
	if err := json.Unmarshal(data, &pools); err != nil {
		return fmt.Errorf("failed to unmarshal pools: %w", err)
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	c.pools = make(map[common.Address]*PoolInfo)
	for _, pool := range pools {
		c.pools[pool.Address] = pool
	}

	c.logger.Info("📖 Pool cache loaded", "file", c.saveFile, "pools", len(pools))
	return nil
}

// periodicSave saves the cache periodically
func (c *PoolCache) periodicSave() {
	for range c.saveTicker.C {
		if err := c.Save(); err != nil {
			c.logger.Error("periodic save failed", "error", err)
		}
	}
}

// Stop stops the periodic save and performs a final save
func (c *PoolCache) Stop() {
	if c.saveTicker != nil {
		c.saveTicker.Stop()
	}

	if c.autoSave {
		if err := c.Save(); err != nil {
			c.logger.Error("final save failed", "error", err)
		}
	}
}

// Stats returns cache statistics
func (c *PoolCache) Stats() map[string]interface{} {
	c.mu.RLock()
	defer c.mu.RUnlock()

	protocolCounts := make(map[string]int)
	totalSwaps := uint64(0)

	for _, pool := range c.pools {
		key := pool.Protocol
		if pool.Version != "" {
			key = pool.Protocol + "-" + pool.Version
		}
		protocolCounts[key]++
		totalSwaps += pool.SwapCount
	}

	return map[string]interface{}{
		"total_pools":     len(c.pools),
		"total_swaps":     totalSwaps,
		"protocol_counts": protocolCounts,
	}
}