gohoarder/pkg/prewarming/worker.go

// Package prewarming runs background workers that pre-fetch hot packages
// to reduce first-request latency.
package prewarming

import (
	"context"
	"sync"
	"time"

	"github.com/lukaszraczylo/gohoarder/pkg/analytics"
	"github.com/lukaszraczylo/gohoarder/pkg/cache"
	"github.com/lukaszraczylo/gohoarder/pkg/network"
	"github.com/rs/zerolog/log"
)

// PackageInfo represents a package to pre-warm
type PackageInfo struct {
	Registry string
	Name     string
	Version  string
	Priority int
}

// Worker handles background pre-warming of popular packages
type Worker struct {
	cache         *cache.Manager
	analytics     *analytics.Engine
	client        *network.Client
	stopChan      chan struct{}
	wg            sync.WaitGroup
	interval      time.Duration
	maxConcurrent int
	enabled       bool
}

// Config holds pre-warming worker configuration
type Config struct {
	CacheManager  *cache.Manager
	Analytics     *analytics.Engine
	NetworkClient *network.Client
	Interval      time.Duration
	MaxConcurrent int
	TopPackages   int
	Enabled       bool
}

// NewWorker creates a new pre-warming worker
func NewWorker(cfg Config) *Worker {
	if cfg.Interval <= 0 {
		cfg.Interval = 1 * time.Hour
	}
	if cfg.MaxConcurrent <= 0 {
		cfg.MaxConcurrent = 5
	}
	if cfg.TopPackages <= 0 {
		cfg.TopPackages = 100
	}

	worker := &Worker{
		cache:         cfg.CacheManager,
		analytics:     cfg.Analytics,
		client:        cfg.NetworkClient,
		interval:      cfg.Interval,
		maxConcurrent: cfg.MaxConcurrent,
		enabled:       cfg.Enabled,
		stopChan:      make(chan struct{}),
	}

	if cfg.Enabled {
		log.Info().
			Dur("interval", cfg.Interval).
			Int("max_concurrent", cfg.MaxConcurrent).
			Msg("Pre-warming worker initialized")
	} else {
		log.Info().Msg("Pre-warming worker disabled")
	}

	return worker
}

// Start begins the pre-warming worker
func (w *Worker) Start(ctx context.Context) {
	if !w.enabled {
		log.Debug().Msg("Pre-warming worker is disabled, not starting")
		return
	}

	w.wg.Add(1)
	go w.run(ctx)
	log.Info().Msg("Pre-warming worker started")
}

// run is the main worker loop
func (w *Worker) run(ctx context.Context) {
	defer w.wg.Done()

	ticker := time.NewTicker(w.interval)
	defer ticker.Stop()

	// Run immediately on start
	w.prewarmPopularPackages(ctx)

	for {
		select {
		case <-ctx.Done():
			log.Info().Msg("Pre-warming worker stopping due to context cancellation")
			return
		case <-w.stopChan:
			log.Info().Msg("Pre-warming worker stopped")
			return
		case <-ticker.C:
			w.prewarmPopularPackages(ctx)
		}
	}
}

// prewarmPopularPackages fetches and caches popular packages
func (w *Worker) prewarmPopularPackages(ctx context.Context) {
	log.Info().Msg("Starting pre-warming cycle")

	// Get popular packages from analytics
	popularPackages := w.analytics.GetTopPackages(100)
	if len(popularPackages) == 0 {
		log.Debug().Msg("No popular packages found for pre-warming")
		return
	}

	// Get trending packages for additional candidates
	trendingPackages := w.analytics.GetTrendingPackages(50)

	// Combine and deduplicate
	packages := w.combinePackages(popularPackages, trendingPackages)

	log.Info().
		Int("packages", len(packages)).
		Msg("Identified packages for pre-warming")

	// Create work queue
	workChan := make(chan PackageInfo, len(packages))
	for _, pkg := range packages {
		workChan <- PackageInfo{
			Registry: pkg.Registry,
			Name:     pkg.Name,
			Version:  "latest", // Pre-warm latest version
			Priority: int(pkg.Downloads),
		}
	}
	close(workChan)

	// Start worker goroutines
	var wg sync.WaitGroup
	for i := 0; i < w.maxConcurrent; i++ {
		wg.Add(1)
		go func(workerID int) {
			defer wg.Done()
			w.processPackages(ctx, workerID, workChan)
		}(i)
	}

	wg.Wait()
	log.Info().Msg("Pre-warming cycle completed")
}

// processPackages processes packages from the work queue
func (w *Worker) processPackages(ctx context.Context, workerID int, workChan <-chan PackageInfo) {
	for pkg := range workChan {
		select {
		case <-ctx.Done():
			return
		default:
			w.prewarmPackage(ctx, pkg, workerID)
		}
	}
}

// prewarmPackage fetches and caches a single package
func (w *Worker) prewarmPackage(ctx context.Context, pkg PackageInfo, workerID int) {
	log.Debug().
		Int("worker", workerID).
		Str("registry", pkg.Registry).
		Str("package", pkg.Name).
		Str("version", pkg.Version).
		Msg("Pre-warming package")

	// Build URL based on registry
	url := w.buildPackageURL(pkg)
	if url == "" {
		log.Warn().
			Str("registry", pkg.Registry).
			Str("package", pkg.Name).
			Msg("Cannot build URL for registry")
		return
	}

	// Fetch package from upstream
	reqCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
	defer cancel()

	body, statusCode, err := w.client.Get(reqCtx, url, nil)
	if err != nil {
		log.Error().
			Err(err).
			Str("package", pkg.Name).
			Msg("Failed to fetch package for pre-warming")
		return
	}
	defer func() { _ = body.Close() }() // #nosec G104 -- Cleanup, error not critical

	if statusCode != 200 {
		log.Warn().
			Int("status", statusCode).
			Str("package", pkg.Name).
			Msg("Non-200 response for package")
		return
	}

	// Cache the package
	// In a real implementation, this would read the response body and store it
	log.Info().
		Str("package", pkg.Name).
		Str("version", pkg.Version).
		Msg("Successfully pre-warmed package")
}

// buildPackageURL builds the upstream URL for a package
func (w *Worker) buildPackageURL(pkg PackageInfo) string {
	// This is simplified - in reality, each registry has different URL patterns
	switch pkg.Registry {
	case "npm":
		return "https://registry.npmjs.org/" + pkg.Name
	case "pypi":
		return "https://pypi.org/simple/" + pkg.Name + "/"
	case "go":
		// Go modules use different URL patterns
		return "https://proxy.golang.org/" + pkg.Name + "/@latest"
	default:
		return ""
	}
}

// combinePackages merges popular and trending packages, removing duplicates
func (w *Worker) combinePackages(popular, trending []analytics.PopularPackage) []analytics.PopularPackage {
	seen := make(map[string]bool)
	result := make([]analytics.PopularPackage, 0, len(popular)+len(trending))

	for _, pkg := range popular {
		key := pkg.Registry + ":" + pkg.Name
		if !seen[key] {
			result = append(result, pkg)
			seen[key] = true
		}
	}

	for _, pkg := range trending {
		key := pkg.Registry + ":" + pkg.Name
		if !seen[key] {
			result = append(result, pkg)
			seen[key] = true
		}
	}

	return result
}

// Stop gracefully stops the pre-warming worker
func (w *Worker) Stop() {
	if !w.enabled {
		return
	}

	log.Info().Msg("Stopping pre-warming worker")
	close(w.stopChan)
	w.wg.Wait()
	log.Info().Msg("Pre-warming worker stopped")
}

// TriggerPrewarm manually triggers a pre-warming cycle
func (w *Worker) TriggerPrewarm(ctx context.Context) {
	if !w.enabled {
		log.Warn().Msg("Cannot trigger pre-warm: worker is disabled")
		return
	}

	log.Info().Msg("Manual pre-warming triggered")
	go w.prewarmPopularPackages(ctx)
}

// PrewarmPackage pre-warms a specific package
func (w *Worker) PrewarmPackage(ctx context.Context, registry, name, version string) error {
	if !w.enabled {
		log.Warn().Msg("Pre-warming worker is disabled")
		return nil
	}

	pkg := PackageInfo{
		Registry: registry,
		Name:     name,
		Version:  version,
		Priority: 100,
	}

	w.prewarmPackage(ctx, pkg, 0)
	return nil
}

// GetStatus returns the current status of the pre-warming worker
func (w *Worker) GetStatus() map[string]interface{} {
	return map[string]interface{}{
		"enabled":        w.enabled,
		"interval":       w.interval.String(),
		"max_concurrent": w.maxConcurrent,
	}
}