fastestmirror/internal/mirror/tester.go

package mirror

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"sort"
	"sync"
	"time"
)

// Mirror represents a package repository mirror
type Mirror struct {
	URL      string
	Country  string
	Score    float64
	Latency  time.Duration
	Speed    float64 // MB/s
	LastTest time.Time
	Error    error
}

// MirrorList manages a collection of mirrors
type MirrorList struct {
	Mirrors []Mirror
	mutex   sync.RWMutex
}

// TestResult contains the results of testing a mirror
type TestResult struct {
	Mirror  Mirror
	Success bool
	Error   error
}

// DefaultMirrors contains known mirrors for each distribution family
var DefaultMirrors = map[string][]string{
	"debian": {
		"http://deb.debian.org/debian/",
		"http://ftp.us.debian.org/debian/",
		"http://ftp.uk.debian.org/debian/",
		"http://ftp.de.debian.org/debian/",
		"http://mirror.kakao.com/debian/",
		"http://mirrors.kernel.org/debian/",
		"http://mirror.math.princeton.edu/pub/debian/",
		"http://debian.osuosl.org/debian/",
	},
	"arch": {
		"https://mirror.rackspace.com/archlinux/",
		"https://mirrors.kernel.org/archlinux/",
		"https://mirror.math.princeton.edu/pub/archlinux/",
		"https://mirrors.mit.edu/archlinux/",
		"https://mirror.cs.vt.edu/pub/ArchLinux/",
		"https://mirrors.liquidweb.com/archlinux/",
		"https://arch.mirror.constant.com/",
		"https://america.mirror.pkgbuild.com/",
	},
	"fedora": {
		"https://download.fedoraproject.org/pub/fedora/linux/",
		"https://mirrors.kernel.org/fedora/",
		"https://mirror.math.princeton.edu/pub/fedora/linux/",
		"https://mirrors.mit.edu/fedora/linux/",
		"https://fedora.mirror.constant.com/",
		"https://mirror.cs.vt.edu/pub/fedora/linux/",
		"https://mirrors.liquidweb.com/fedora/",
		"https://ftp.osuosl.org/pub/fedora/linux/",
	},
}

// NewMirrorList creates a new mirror list for the given distribution family
func NewMirrorList(family string) *MirrorList {
	ml := &MirrorList{
		Mirrors: make([]Mirror, 0),
	}
	return ml
}

// LoadMirrors fetches and loads mirrors for the distribution family
func (ml *MirrorList) LoadMirrors(ctx context.Context, family string) error {
	fetcher := NewMirrorFetcher()

	fmt.Printf("Fetching complete mirror list for %s...\n", family)

	mirrors, err := fetcher.FetchMirrors(ctx, family)
	if err != nil {
		fmt.Printf("Failed to fetch mirrors online, using fallback list: %v\n", err)
		mirrors = fetcher.getDefaultMirrors(family)
	}

	if len(mirrors) == 0 {
		return fmt.Errorf("no mirrors available for %s", family)
	}

	ml.mutex.Lock()
	defer ml.mutex.Unlock()

	ml.Mirrors = make([]Mirror, 0, len(mirrors))
	for _, mirrorURL := range mirrors {
		ml.Mirrors = append(ml.Mirrors, Mirror{
			URL: mirrorURL,
		})
	}

	fmt.Printf("Loaded %d mirrors for testing\n", len(ml.Mirrors))
	return nil
}

// TestMirrors tests all mirrors concurrently and updates their scores
func (ml *MirrorList) TestMirrors(ctx context.Context, timeout time.Duration) error {
	if len(ml.Mirrors) == 0 {
		return fmt.Errorf("no mirrors to test - this is awkward")
	}

	ml.mutex.Lock()
	defer ml.mutex.Unlock()

	results := make(chan TestResult, len(ml.Mirrors))
	var wg sync.WaitGroup

	fmt.Printf("Testing %d mirrors concurrently...\n", len(ml.Mirrors))

	// Test each mirror concurrently
	for i := range ml.Mirrors {
		wg.Add(1)
		go func(mirror *Mirror) {
			defer wg.Done()
			testResult := ml.testSingleMirror(ctx, mirror, timeout)
			results <- testResult
		}(&ml.Mirrors[i])
	}

	// Wait for all tests to complete
	go func() {
		wg.Wait()
		close(results)
	}()

	// Collect results and count successes
	successCount := 0
	failureCount := 0

	for result := range results {
		for i := range ml.Mirrors {
			if ml.Mirrors[i].URL == result.Mirror.URL {
				ml.Mirrors[i] = result.Mirror
				if result.Success && result.Mirror.Error == nil {
					successCount++
				} else {
					failureCount++
				}
				break
			}
		}
	}

	fmt.Printf("Testing complete: %d successful, %d failed\n", successCount, failureCount)

	// Sort by score (higher is better), but put failed mirrors at the end
	sort.Slice(ml.Mirrors, func(i, j int) bool {
		// If one has an error and the other doesn't, prioritize the working one
		if (ml.Mirrors[i].Error != nil) != (ml.Mirrors[j].Error != nil) {
			return ml.Mirrors[j].Error != nil // Working mirrors come first
		}
		// If both are working or both have errors, sort by score
		return ml.Mirrors[i].Score > ml.Mirrors[j].Score
	})

	return nil
}

// testSingleMirror tests a single mirror's speed and latency
func (ml *MirrorList) testSingleMirror(ctx context.Context, mirror *Mirror, timeout time.Duration) TestResult {
	testCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	// Test latency with HEAD request first
	latency, err := ml.testLatency(testCtx, mirror.URL)
	if err != nil {
		return TestResult{
			Mirror: Mirror{
				URL:      mirror.URL,
				Country:  mirror.Country,
				Error:    fmt.Errorf("latency test failed: %w", err),
				LastTest: time.Now(),
			},
			Success: false,
			Error:   err,
		}
	}

	// Validate that the mirror actually serves the required package databases
	err = ml.validateMirror(testCtx, mirror.URL)
	if err != nil {
		return TestResult{
			Mirror: Mirror{
				URL:      mirror.URL,
				Country:  mirror.Country,
				Latency:  latency,
				Error:    fmt.Errorf("validation failed: %w", err),
				LastTest: time.Now(),
			},
			Success: false,
			Error:   err,
		}
	}

	// Test download speed with a small file
	speed, err := ml.testSpeed(testCtx, mirror.URL)
	if err != nil {
		return TestResult{
			Mirror: Mirror{
				URL:      mirror.URL,
				Country:  mirror.Country,
				Latency:  latency,
				Error:    fmt.Errorf("speed test failed: %w", err),
				LastTest: time.Now(),
			},
			Success: false,
			Error:   err,
		}
	}

	// Calculate score based on speed and latency
	score := ml.calculateScore(speed, latency)

	return TestResult{
		Mirror: Mirror{
			URL:      mirror.URL,
			Country:  mirror.Country,
			Score:    score,
			Latency:  latency,
			Speed:    speed,
			LastTest: time.Now(),
			Error:    nil,
		},
		Success: true,
	}
}

// testLatency measures response time to a mirror
func (ml *MirrorList) testLatency(ctx context.Context, mirrorURL string) (time.Duration, error) {
	start := time.Now()

	req, err := http.NewRequestWithContext(ctx, "HEAD", mirrorURL, nil)
	if err != nil {
		return 0, err
	}

	client := &http.Client{
		Timeout: 2 * time.Second, // 2000ms timeout as requested
	}

	resp, err := client.Do(req)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 400 {
		return 0, fmt.Errorf("HTTP %d: mirror seems to be having a bad day", resp.StatusCode)
	}

	return time.Since(start), nil
}

// testSpeed measures download speed from a mirror
func (ml *MirrorList) testSpeed(ctx context.Context, mirrorURL string) (float64, error) {
	// Try to download a standard test file (like Release file for Debian)
	testPath := ml.getTestPath(mirrorURL)

	req, err := http.NewRequestWithContext(ctx, "GET", testPath, nil)
	if err != nil {
		return 0, err
	}

	client := &http.Client{
		Timeout: 15 * time.Second, // Reasonable timeout for downloading test files
	}

	start := time.Now()
	resp, err := client.Do(req)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 400 {
		return 0, fmt.Errorf("HTTP %d: couldn't download test file", resp.StatusCode)
	}

	// Read response and measure time
	bytesRead, err := io.Copy(io.Discard, resp.Body)
	if err != nil {
		return 0, err
	}

	duration := time.Since(start)
	if duration == 0 {
		return 0, fmt.Errorf("download completed too quickly to measure")
	}

	// Calculate speed in MB/s
	speed := float64(bytesRead) / duration.Seconds() / (1024 * 1024)
	return speed, nil
}

// validateMirror checks if the mirror actually serves the required package databases
func (ml *MirrorList) validateMirror(ctx context.Context, mirrorURL string) error {
	// For Arch mirrors, test multiple repositories to ensure they're properly formatted
	if contains(mirrorURL, "archlinux") {
		// Test all major repositories that pacman needs
		repositories := []string{"core", "extra", "multilib"}

		for _, repo := range repositories {
			testURL := mirrorURL + repo + "/os/x86_64/" + repo + ".db"
			err := ml.checkFileExists(ctx, testURL)
			if err != nil {
				return fmt.Errorf("repository %s not found: %w", repo, err)
			}
		}
		return nil
	}

	// For Debian mirrors, test Release file
	if contains(mirrorURL, "debian") {
		testURL := mirrorURL + "dists/stable/Release"
		return ml.checkFileExists(ctx, testURL)
	}

	// For Fedora, test repomd.xml
	if contains(mirrorURL, "fedora") {
		testURL := mirrorURL + "releases/39/Everything/x86_64/os/repodata/repomd.xml"
		return ml.checkFileExists(ctx, testURL)
	}

	// For unknown mirrors, just test basic connectivity
	return nil
}

// checkFileExists verifies that a specific file exists on the mirror
func (ml *MirrorList) checkFileExists(ctx context.Context, fileURL string) error {
	req, err := http.NewRequestWithContext(ctx, "HEAD", fileURL, nil)
	if err != nil {
		return err
	}

	client := &http.Client{
		Timeout: 5 * time.Second,
	}

	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode == 404 {
		return fmt.Errorf("required package database not found (404)")
	}

	if resp.StatusCode >= 400 {
		return fmt.Errorf("HTTP %d: mirror doesn't serve required files", resp.StatusCode)
	}

	return nil
}

// getTestPath returns an appropriate test file path for the mirror
func (ml *MirrorList) getTestPath(mirrorURL string) string {
	// For Debian-based mirrors, try Release file
	if contains(mirrorURL, "debian") {
		return mirrorURL + "dists/stable/Release"
	}

	// For Arch mirrors, try core.db
	if contains(mirrorURL, "archlinux") {
		return mirrorURL + "core/os/x86_64/core.db"
	}

	// For Fedora, try repomd.xml
	if contains(mirrorURL, "fedora") {
		return mirrorURL + "releases/39/Everything/x86_64/os/repodata/repomd.xml"
	}

	// Fallback: just try the base URL
	return mirrorURL
}

// calculateScore calculates a score based on speed and latency
func (ml *MirrorList) calculateScore(speed float64, latency time.Duration) float64 {
	// Higher speed is better, lower latency is better
	// Normalize and combine metrics
	speedScore := speed * 10                                 // Weight speed heavily
	latencyScore := 1000.0 / float64(latency.Milliseconds()) // Lower latency = higher score

	return speedScore + latencyScore
}

// GetBest returns the best performing mirror
func (ml *MirrorList) GetBest() *Mirror {
	ml.mutex.RLock()
	defer ml.mutex.RUnlock()

	if len(ml.Mirrors) == 0 {
		return nil
	}

	// Return the first mirror (should be highest scoring after sorting)
	for _, mirror := range ml.Mirrors {
		if mirror.Error == nil {
			return &mirror
		}
	}

	return nil
}

// GetTop returns the top N mirrors (only successful ones)
func (ml *MirrorList) GetTop(n int) []Mirror {
	ml.mutex.RLock()
	defer ml.mutex.RUnlock()

	result := make([]Mirror, 0, n)
	count := 0

	for _, mirror := range ml.Mirrors {
		if mirror.Error == nil && count < n {
			result = append(result, mirror)
			count++
		}
	}

	return result
}

// GetAll returns all mirrors with their test results
func (ml *MirrorList) GetAll() []Mirror {
	ml.mutex.RLock()
	defer ml.mutex.RUnlock()

	result := make([]Mirror, len(ml.Mirrors))
	copy(result, ml.Mirrors)
	return result
}

// contains is a helper function to check if a string contains a substring
func contains(s, substr string) bool {
	return len(s) >= len(substr) &&
		(s == substr ||
			s[:len(substr)] == substr ||
			s[len(s)-len(substr):] == substr ||
			findInString(s, substr))
}

func findInString(s, substr string) bool {
	for i := 0; i <= len(s)-len(substr); i++ {
		if s[i:i+len(substr)] == substr {
			return true
		}
	}
	return false
}