CLIProxyAPI/internal/api/modules/amp/amp.go

// Package amp implements the Amp CLI routing module, providing OAuth-based
// integration with Amp CLI for ChatGPT and Anthropic subscriptions.
package amp

import (
	"fmt"
	"net/http/httputil"
	"strings"
	"sync"

	"github.com/gin-gonic/gin"
	"github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules"
	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
	sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
	log "github.com/sirupsen/logrus"
)

// Option configures the AmpModule.
type Option func(*AmpModule)

// AmpModule implements the RouteModuleV2 interface for Amp CLI integration.
// It provides:
//   - Reverse proxy to Amp control plane for OAuth/management
//   - Provider-specific route aliases (/api/provider/{provider}/...)
//   - Automatic gzip decompression for misconfigured upstreams
//   - Model mapping for routing unavailable models to alternatives
type AmpModule struct {
	secretSource    SecretSource
	proxy           *httputil.ReverseProxy
	accessManager   *sdkaccess.Manager
	authMiddleware_ gin.HandlerFunc
	modelMapper     *DefaultModelMapper
	enabled         bool
	registerOnce    sync.Once

	// configMu protects lastConfig for partial reload comparison
	configMu   sync.RWMutex
	lastConfig *config.AmpCode
}

// New creates a new Amp routing module with the given options.
// This is the preferred constructor using the Option pattern.
//
// Example:
//
//	ampModule := amp.New(
//	    amp.WithAccessManager(accessManager),
//	    amp.WithAuthMiddleware(authMiddleware),
//	    amp.WithSecretSource(customSecret),
//	)
func New(opts ...Option) *AmpModule {
	m := &AmpModule{
		secretSource: nil, // Will be created on demand if not provided
	}
	for _, opt := range opts {
		opt(m)
	}
	return m
}

// NewLegacy creates a new Amp routing module using the legacy constructor signature.
// This is provided for backwards compatibility.
//
// DEPRECATED: Use New with options instead.
func NewLegacy(accessManager *sdkaccess.Manager, authMiddleware gin.HandlerFunc) *AmpModule {
	return New(
		WithAccessManager(accessManager),
		WithAuthMiddleware(authMiddleware),
	)
}

// WithSecretSource sets a custom secret source for the module.
func WithSecretSource(source SecretSource) Option {
	return func(m *AmpModule) {
		m.secretSource = source
	}
}

// WithAccessManager sets the access manager for the module.
func WithAccessManager(am *sdkaccess.Manager) Option {
	return func(m *AmpModule) {
		m.accessManager = am
	}
}

// WithAuthMiddleware sets the authentication middleware for provider routes.
func WithAuthMiddleware(middleware gin.HandlerFunc) Option {
	return func(m *AmpModule) {
		m.authMiddleware_ = middleware
	}
}

// Name returns the module identifier
func (m *AmpModule) Name() string {
	return "amp-routing"
}

// Register sets up Amp routes if configured.
// This implements the RouteModuleV2 interface with Context.
// Routes are registered only once via sync.Once for idempotent behavior.
func (m *AmpModule) Register(ctx modules.Context) error {
	settings := ctx.Config.AmpCode
	upstreamURL := strings.TrimSpace(settings.UpstreamURL)

	// Determine auth middleware (from module or context)
	auth := m.getAuthMiddleware(ctx)

	// Use registerOnce to ensure routes are only registered once
	var regErr error
	m.registerOnce.Do(func() {
		// Initialize model mapper from config (for routing unavailable models to alternatives)
		m.modelMapper = NewModelMapper(settings.ModelMappings)

		// Store initial config for partial reload comparison
		settingsCopy := settings
		m.lastConfig = &settingsCopy

		// Always register provider aliases - these work without an upstream
		m.registerProviderAliases(ctx.Engine, ctx.BaseHandler, auth)

		// If no upstream URL, skip proxy routes but provider aliases are still available
		if upstreamURL == "" {
			log.Debug("amp upstream proxy disabled (no upstream URL configured)")
			log.Debug("amp provider alias routes registered")
			m.enabled = false
			return
		}

		// Create secret source with precedence: config > env > file
		// Cache secrets for 5 minutes to reduce file I/O
		if m.secretSource == nil {
			m.secretSource = NewMultiSourceSecret(settings.UpstreamAPIKey, 0 /* default 5min */)
		}

		// Create reverse proxy with gzip handling via ModifyResponse
		proxy, err := createReverseProxy(upstreamURL, m.secretSource)
		if err != nil {
			regErr = fmt.Errorf("failed to create amp proxy: %w", err)
			return
		}

		m.proxy = proxy
		m.enabled = true

		// Register management proxy routes (requires upstream)
		// Restrict to localhost by default for security (prevents drive-by browser attacks)
		handler := proxyHandler(proxy)
		m.registerManagementRoutes(ctx.Engine, ctx.BaseHandler, handler, settings.RestrictManagementToLocalhost)

		log.Infof("amp upstream proxy enabled for: %s", upstreamURL)
		log.Debug("amp provider alias routes registered")
	})

	return regErr
}

// getAuthMiddleware returns the authentication middleware, preferring the
// module's configured middleware, then the context middleware, then a fallback.
func (m *AmpModule) getAuthMiddleware(ctx modules.Context) gin.HandlerFunc {
	if m.authMiddleware_ != nil {
		return m.authMiddleware_
	}
	if ctx.AuthMiddleware != nil {
		return ctx.AuthMiddleware
	}
	// Fallback: no authentication (should not happen in production)
	log.Warn("amp module: no auth middleware provided, allowing all requests")
	return func(c *gin.Context) {
		c.Next()
	}
}

// OnConfigUpdated handles configuration updates with partial reload support.
// Only updates components that have actually changed to avoid unnecessary work.
// URL changes still require restart (logged as warning).
func (m *AmpModule) OnConfigUpdated(cfg *config.Config) error {
	newSettings := cfg.AmpCode

	// Get previous config for comparison
	m.configMu.RLock()
	oldSettings := m.lastConfig
	m.configMu.RUnlock()

	// Track what changed for logging
	var changes []string

	// Check model mappings change
	modelMappingsChanged := m.hasModelMappingsChanged(oldSettings, &newSettings)
	if modelMappingsChanged {
		if m.modelMapper != nil {
			m.modelMapper.UpdateMappings(newSettings.ModelMappings)
			changes = append(changes, "model-mappings")
			if m.enabled {
				log.Infof("amp config partial reload: model mappings updated (%d entries)", len(newSettings.ModelMappings))
			}
		} else if m.enabled {
			log.Warnf("amp model mapper not initialized, skipping model mapping update")
		}
	}

	if m.enabled {
		// Check upstream URL change (requires restart)
		newUpstreamURL := strings.TrimSpace(newSettings.UpstreamURL)
		oldUpstreamURL := ""
		if oldSettings != nil {
			oldUpstreamURL = strings.TrimSpace(oldSettings.UpstreamURL)
		}

		if newUpstreamURL == "" && oldUpstreamURL != "" {
			log.Warn("amp upstream URL removed from config, restart required to disable")
		} else if newUpstreamURL != oldUpstreamURL {
			changes = append(changes, "upstream-url(restart required)")
			log.Warnf("amp config: upstream-url changed (%s -> %s), restart required", oldUpstreamURL, newUpstreamURL)
		}

		// Check API key change
		apiKeyChanged := m.hasAPIKeyChanged(oldSettings, &newSettings)
		if apiKeyChanged {
			if m.secretSource != nil {
				if ms, ok := m.secretSource.(*MultiSourceSecret); ok {
					ms.UpdateExplicitKey(newSettings.UpstreamAPIKey)
					ms.InvalidateCache()
					changes = append(changes, "upstream-api-key")
					log.Debug("amp config partial reload: secret cache invalidated")
				}
			}
		}

		// Check restrict-management-to-localhost change (requires restart)
		if oldSettings != nil && oldSettings.RestrictManagementToLocalhost != newSettings.RestrictManagementToLocalhost {
			changes = append(changes, "restrict-management-to-localhost(restart required)")
			log.Warnf("amp config: restrict-management-to-localhost changed (%t -> %t), restart required",
				oldSettings.RestrictManagementToLocalhost, newSettings.RestrictManagementToLocalhost)
		}
	}

	// Store current config for next comparison
	m.configMu.Lock()
	settingsCopy := newSettings // copy struct
	m.lastConfig = &settingsCopy
	m.configMu.Unlock()

	// Log summary if any changes detected
	if len(changes) > 0 {
		log.Debugf("amp config partial reload completed: %v", changes)
	} else {
		log.Debug("amp config checked: no changes detected")
	}

	return nil
}

// hasModelMappingsChanged compares old and new model mappings.
func (m *AmpModule) hasModelMappingsChanged(old *config.AmpCode, new *config.AmpCode) bool {
	if old == nil {
		return len(new.ModelMappings) > 0
	}

	if len(old.ModelMappings) != len(new.ModelMappings) {
		return true
	}

	// Build map for efficient comparison
	oldMap := make(map[string]string, len(old.ModelMappings))
	for _, mapping := range old.ModelMappings {
		oldMap[strings.TrimSpace(mapping.From)] = strings.TrimSpace(mapping.To)
	}

	for _, mapping := range new.ModelMappings {
		from := strings.TrimSpace(mapping.From)
		to := strings.TrimSpace(mapping.To)
		if oldTo, exists := oldMap[from]; !exists || oldTo != to {
			return true
		}
	}

	return false
}

// hasAPIKeyChanged compares old and new API keys.
func (m *AmpModule) hasAPIKeyChanged(old *config.AmpCode, new *config.AmpCode) bool {
	oldKey := ""
	if old != nil {
		oldKey = strings.TrimSpace(old.UpstreamAPIKey)
	}
	newKey := strings.TrimSpace(new.UpstreamAPIKey)
	return oldKey != newKey
}

// GetModelMapper returns the model mapper instance (for testing/debugging).
func (m *AmpModule) GetModelMapper() *DefaultModelMapper {
	return m.modelMapper
}