mev-beta/pkg/transport/unified_provider_manager.go

package transport

import (
	"fmt"
	"io/ioutil"

	"github.com/ethereum/go-ethereum/ethclient"
	"gopkg.in/yaml.v3"
)

// UnifiedProviderManager manages all provider pools (read-only, execution, testing)
type UnifiedProviderManager struct {
	ReadOnlyPool    *ReadOnlyProviderPool
	ExecutionPool   *ExecutionProviderPool
	TestingPool     *TestingProviderPool
	config          ProvidersConfig
	providerConfigs map[string]ProviderConfig
}

// OperationMode defines the type of operation being performed
type OperationMode int

const (
	ModeReadOnly OperationMode = iota
	ModeExecution
	ModeTesting
)

// NewUnifiedProviderManager creates a new unified provider manager
func NewUnifiedProviderManager(configPath string) (*UnifiedProviderManager, error) {
	// Load configuration
	config, err := LoadProvidersConfig(configPath)
	if err != nil {
		return nil, fmt.Errorf("failed to load provider config: %w", err)
	}

	// Create provider configs map for easy lookup
	providerConfigs := make(map[string]ProviderConfig)
	for _, provider := range config.Providers {
		providerConfigs[provider.Name] = provider
	}

	manager := &UnifiedProviderManager{
		config:          config,
		providerConfigs: providerConfigs,
	}

	// Initialize provider pools
	if err := manager.initializePools(); err != nil {
		return nil, fmt.Errorf("failed to initialize provider pools: %w", err)
	}

	return manager, nil
}

// initializePools initializes all provider pools based on configuration
func (upm *UnifiedProviderManager) initializePools() error {
	var err error

	// Initialize read-only pool if configured
	if poolConfig, exists := upm.config.ProviderPools["read_only"]; exists {
		upm.ReadOnlyPool, err = NewReadOnlyProviderPool(poolConfig, upm.providerConfigs)
		if err != nil {
			return fmt.Errorf("failed to initialize read-only pool: %w", err)
		}
	}

	// Initialize execution pool if configured
	if poolConfig, exists := upm.config.ProviderPools["execution"]; exists {
		upm.ExecutionPool, err = NewExecutionProviderPool(poolConfig, upm.providerConfigs)
		if err != nil {
			return fmt.Errorf("failed to initialize execution pool: %w", err)
		}
	}

	// Initialize testing pool if configured
	if poolConfig, exists := upm.config.ProviderPools["testing"]; exists {
		upm.TestingPool, err = NewTestingProviderPool(poolConfig, upm.providerConfigs)
		if err != nil {
			return fmt.Errorf("failed to initialize testing pool: %w", err)
		}
	}

	return nil
}

// GetPoolForMode returns the appropriate provider pool for the given operation mode
func (upm *UnifiedProviderManager) GetPoolForMode(mode OperationMode) (ProviderPool, error) {
	switch mode {
	case ModeReadOnly:
		if upm.ReadOnlyPool == nil {
			return nil, fmt.Errorf("read-only pool not initialized")
		}
		return upm.ReadOnlyPool, nil
	case ModeExecution:
		if upm.ExecutionPool == nil {
			return nil, fmt.Errorf("execution pool not initialized")
		}
		return upm.ExecutionPool, nil
	case ModeTesting:
		if upm.TestingPool == nil {
			return nil, fmt.Errorf("testing pool not initialized")
		}
		return upm.TestingPool, nil
	default:
		return nil, fmt.Errorf("unknown operation mode: %d", mode)
	}
}

// GetReadOnlyHTTPClient returns an HTTP client optimized for read-only operations
func (upm *UnifiedProviderManager) GetReadOnlyHTTPClient() (*ethclient.Client, error) {
	if upm.ReadOnlyPool == nil {
		return nil, fmt.Errorf("read-only pool not initialized")
	}
	return upm.ReadOnlyPool.GetHTTPClient()
}

// GetReadOnlyWSClient returns a WebSocket client for real-time data
func (upm *UnifiedProviderManager) GetReadOnlyWSClient() (*ethclient.Client, error) {
	if upm.ReadOnlyPool == nil {
		return nil, fmt.Errorf("read-only pool not initialized")
	}
	return upm.ReadOnlyPool.GetWSClient()
}

// GetExecutionHTTPClient returns an HTTP client optimized for transaction execution
func (upm *UnifiedProviderManager) GetExecutionHTTPClient() (*ethclient.Client, error) {
	if upm.ExecutionPool == nil {
		return nil, fmt.Errorf("execution pool not initialized")
	}
	return upm.ExecutionPool.GetHTTPClient()
}

// GetTestingHTTPClient returns an HTTP client for testing (preferably Anvil)
func (upm *UnifiedProviderManager) GetTestingHTTPClient() (*ethclient.Client, error) {
	if upm.TestingPool == nil {
		return nil, fmt.Errorf("testing pool not initialized")
	}
	return upm.TestingPool.GetHTTPClient()
}

// GetAllStats returns statistics for all provider pools
func (upm *UnifiedProviderManager) GetAllStats() map[string]interface{} {
	stats := make(map[string]interface{})

	if upm.ReadOnlyPool != nil {
		stats["read_only"] = upm.ReadOnlyPool.GetStats()
	}

	if upm.ExecutionPool != nil {
		stats["execution"] = upm.ExecutionPool.GetStats()
	}

	if upm.TestingPool != nil {
		stats["testing"] = upm.TestingPool.GetStats()
	}

	// Add overall summary
	summary := map[string]interface{}{
		"total_pools":       len(stats),
		"pools_initialized": []string{},
	}

	for poolName := range stats {
		summary["pools_initialized"] = append(summary["pools_initialized"].([]string), poolName)
	}

	stats["summary"] = summary

	return stats
}

// CreateTestingSnapshot creates a snapshot in the testing environment
func (upm *UnifiedProviderManager) CreateTestingSnapshot() (string, error) {
	if upm.TestingPool == nil {
		return "", fmt.Errorf("testing pool not initialized")
	}
	return upm.TestingPool.CreateSnapshot()
}

// RevertTestingSnapshot reverts to a snapshot in the testing environment
func (upm *UnifiedProviderManager) RevertTestingSnapshot(snapshotID string) error {
	if upm.TestingPool == nil {
		return fmt.Errorf("testing pool not initialized")
	}
	return upm.TestingPool.RevertToSnapshot(snapshotID)
}

// Close shuts down all provider pools
func (upm *UnifiedProviderManager) Close() error {
	var errors []error

	if upm.ReadOnlyPool != nil {
		if err := upm.ReadOnlyPool.Close(); err != nil {
			errors = append(errors, fmt.Errorf("failed to close read-only pool: %w", err))
		}
	}

	if upm.ExecutionPool != nil {
		if err := upm.ExecutionPool.Close(); err != nil {
			errors = append(errors, fmt.Errorf("failed to close execution pool: %w", err))
		}
	}

	if upm.TestingPool != nil {
		if err := upm.TestingPool.Close(); err != nil {
			errors = append(errors, fmt.Errorf("failed to close testing pool: %w", err))
		}
	}

	if len(errors) > 0 {
		return fmt.Errorf("errors closing provider pools: %v", errors)
	}

	return nil
}

// LoadProvidersConfigFromFile loads configuration from a YAML file
func LoadProvidersConfigFromFile(path string) (ProvidersConfig, error) {
	var config ProvidersConfig

	// Read the YAML file
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return config, fmt.Errorf("failed to read config file %s: %w", path, err)
	}

	// Unmarshal the YAML data
	if err := yaml.Unmarshal(data, &config); err != nil {
		return config, fmt.Errorf("failed to parse YAML config: %w", err)
	}

	// Validate the configuration
	if err := validateProvidersConfig(&config); err != nil {
		return config, fmt.Errorf("invalid configuration: %w", err)
	}

	return config, nil
}

// validateProvidersConfig validates the provider configuration
func validateProvidersConfig(config *ProvidersConfig) error {
	if len(config.Providers) == 0 {
		return fmt.Errorf("no providers configured")
	}

	// Validate provider pools
	for poolName, poolConfig := range config.ProviderPools {
		if len(poolConfig.Providers) == 0 {
			return fmt.Errorf("provider pool '%s' has no providers", poolName)
		}

		// Check that all referenced providers exist
		providerNames := make(map[string]bool)
		for _, provider := range config.Providers {
			providerNames[provider.Name] = true
		}

		for _, providerName := range poolConfig.Providers {
			if !providerNames[providerName] {
				return fmt.Errorf("provider pool '%s' references unknown provider '%s'", poolName, providerName)
			}
		}
	}

	// Validate individual providers
	for i, provider := range config.Providers {
		if provider.Name == "" {
			return fmt.Errorf("provider %d has no name", i)
		}
		if provider.HTTPEndpoint == "" && provider.WSEndpoint == "" {
			return fmt.Errorf("provider %s has no endpoints", provider.Name)
		}
		if provider.RateLimit.RequestsPerSecond <= 0 {
			return fmt.Errorf("provider %s has invalid rate limit", provider.Name)
		}

		// Validate Anvil config if present
		if provider.Type == "anvil_fork" && provider.AnvilConfig == nil {
			return fmt.Errorf("provider %s is anvil_fork type but has no anvil_config", provider.Name)
		}
	}

	return nil
}

// GetProviderByName returns a specific provider configuration by name
func (upm *UnifiedProviderManager) GetProviderByName(name string) (ProviderConfig, bool) {
	config, exists := upm.providerConfigs[name]
	return config, exists
}

// GetProvidersByType returns all providers of a specific type
func (upm *UnifiedProviderManager) GetProvidersByType(providerType string) []ProviderConfig {
	var providers []ProviderConfig
	for _, provider := range upm.config.Providers {
		if provider.Type == providerType {
			providers = append(providers, provider)
		}
	}
	return providers
}

// GetProvidersByFeature returns all providers that support a specific feature
func (upm *UnifiedProviderManager) GetProvidersByFeature(feature string) []ProviderConfig {
	var providers []ProviderConfig
	for _, provider := range upm.config.Providers {
		for _, providerFeature := range provider.Features {
			if providerFeature == feature {
				providers = append(providers, provider)
				break
			}
		}
	}
	return providers
}