goget/internal/mirror/mirror.go

//go:build linux || freebsd
// +build linux freebsd

package mirror

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"codeberg.org/petrbalvin/goget/internal/core"
	format "codeberg.org/petrbalvin/goget/internal/format"
	"codeberg.org/petrbalvin/goget/internal/linkrewrite"
	"codeberg.org/petrbalvin/goget/internal/recursive"
	"codeberg.org/petrbalvin/goget/internal/transport"
	"golang.org/x/net/html"
)

// MirrorConfig holds configuration for website mirroring.
type MirrorConfig struct {
	BaseURL         *url.URL
	OutputDir       string
	MaxDepth        int // 0 = infinite
	FollowExternal  bool
	ExcludePatterns []string
	IncludePatterns []string
	Verbose         bool
	Parallel        int
	Delay           time.Duration
	UserAgent       string
	ConvertLinks    bool // Convert links for local viewing
	DownloadAssets  bool // Download CSS, JS, images
	PruneEmpty      bool // Remove empty directories after download
	RestrictFiles   bool // Only download files from original host
	Timeout         time.Duration
	RateLimiter     *transport.TokenBucket
	RespectRobots   bool // Respect robots.txt
	DryRun          bool // List URLs without writing files
}

// MirrorStats contains statistics about a mirroring operation.
type MirrorStats struct {
	TotalURLs      int64
	DownloadedURLs int64
	FailedURLs     int64
	SkippedURLs    int64
	ConvertedLinks int64
	TotalBytes     int64
	StartTime      time.Time
	EndTime        time.Time
}

// robotsPolicy stores robots.txt rules
type robotsPolicy struct {
	allowedPaths []string
	disallowed   []string
	crawlDelay   time.Duration
	userAgent    string
}

// RobotsChecker checks robots.txt rules
type RobotsChecker struct {
	policies  map[string]*robotsPolicy // host -> policy
	mu        sync.RWMutex
	userAgent string
}

// NewRobotsChecker creates a new robots.txt checker
func NewRobotsChecker(userAgent string) *RobotsChecker {
	return &RobotsChecker{
		policies:  make(map[string]*robotsPolicy),
		userAgent: userAgent,
	}
}

// CanFetch checks if URL can be fetched according to robots.txt
func (rc *RobotsChecker) CanFetch(u *url.URL) bool {
	rc.mu.RLock()
	policy, exists := rc.policies[u.Host]
	rc.mu.RUnlock()

	if !exists {
		return true // No policy means allowed
	}

	path := u.Path
	if path == "" {
		path = "/"
	}

	// Check disallowed first (most specific match wins)
	for _, disallowed := range policy.disallowed {
		if strings.HasPrefix(path, disallowed) {
			return false
		}
	}

	return true
}

// GetCrawlDelay returns the crawl delay for a host
func (rc *RobotsChecker) GetCrawlDelay(u *url.URL) time.Duration {
	rc.mu.RLock()
	policy, exists := rc.policies[u.Host]
	rc.mu.RUnlock()

	if !exists || policy.crawlDelay == 0 {
		return 0
	}
	return policy.crawlDelay
}

// FetchPolicy fetches and parses robots.txt for a host
func (rc *RobotsChecker) FetchPolicy(ctx context.Context, u *url.URL, client *http.Client) error {
	rc.mu.Lock()
	defer rc.mu.Unlock()

	// Already fetched
	if _, exists := rc.policies[u.Host]; exists {
		return nil
	}

	robotsURL := &url.URL{
		Scheme: u.Scheme,
		Host:   u.Host,
		Path:   "/robots.txt",
	}

	req, err := http.NewRequestWithContext(ctx, "GET", robotsURL.String(), nil)
	if err != nil {
		return err
	}
	req.Header.Set("User-Agent", rc.userAgent)

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

	if resp.StatusCode != http.StatusOK {
		// No robots.txt means everything is allowed
		rc.policies[u.Host] = &robotsPolicy{}
		return nil
	}

	policy := parseRobots(resp.Body, rc.userAgent)
	rc.policies[u.Host] = policy
	return nil
}

// parseRobots parses robots.txt content
func parseRobots(r io.Reader, userAgent string) *robotsPolicy {
	policy := &robotsPolicy{
		userAgent: userAgent,
	}

	scanner := bufio.NewScanner(r)
	var currentAgents []string
	var matchingAgents bool

	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}

		idx := strings.Index(line, ":")
		if idx == -1 {
			continue
		}

		directive := strings.ToLower(strings.TrimSpace(line[:idx]))
		value := strings.TrimSpace(line[idx+1:])

		switch directive {
		case "user-agent":
			if len(currentAgents) > 0 && !matchingAgents {
				currentAgents = nil
			}
			agent := strings.ToLower(value)
			currentAgents = append(currentAgents, agent)
			matchingAgents = agent == userAgent || agent == "*"
		case "disallow":
			if matchingAgents && value != "" {
				policy.disallowed = append(policy.disallowed, value)
			}
		case "allow":
			if matchingAgents && value != "" {
				policy.allowedPaths = append(policy.allowedPaths, value)
			}
		case "crawl-delay":
			if matchingAgents {
				if delay, err := strconv.ParseFloat(value, 64); err == nil {
					policy.crawlDelay = time.Duration(delay * float64(time.Second))
				}
			}
		}
	}

	if err := scanner.Err(); err != nil {
		// Scanner error, return what we parsed so far
	}

	return policy
}

// Mirror manages the website mirroring process.
type Mirror struct {
	config        *MirrorConfig
	proto         core.Protocol
	httpClient    *http.Client
	filter        *recursive.URLFilter
	rewriter      *linkrewrite.Rewriter
	robotsChecker *RobotsChecker
	visited       map[string]bool
	visitedMu     sync.Mutex
	queue         chan *mirrorTask
	stats         *MirrorStats
	wg            sync.WaitGroup
	ctx           context.Context
	cancel        context.CancelFunc
}

// mirrorTask represents a task for the mirror
type mirrorTask struct {
	url   *url.URL
	depth int
}

// NewMirror creates new mirror instance
func NewMirror(cfg *MirrorConfig, proto core.Protocol) *Mirror {
	ctx, cancel := context.WithCancel(context.Background())

	filter := recursive.NewURLFilter(cfg.BaseURL, recursive.URLFilterConfig{
		MaxDepth:        cfg.MaxDepth,
		FollowExternal:  cfg.FollowExternal,
		ExcludePatterns: cfg.ExcludePatterns,
		IncludePatterns: cfg.IncludePatterns,
	})

	rewriter := linkrewrite.New(cfg.BaseURL, cfg.OutputDir)

	robotsChecker := NewRobotsChecker(cfg.UserAgent)

	// Create HTTP client for robots.txt fetching
	httpClient := &http.Client{
		Timeout: 10 * time.Second,
		Transport: &http.Transport{
			DialContext: (&net.Dialer{
				Timeout: 5 * time.Second,
			}).DialContext,
			TLSHandshakeTimeout: 5 * time.Second,
		},
	}

	mirror := &Mirror{
		config:        cfg,
		proto:         proto,
		httpClient:    httpClient,
		filter:        filter,
		rewriter:      rewriter,
		robotsChecker: robotsChecker,
		visited:       make(map[string]bool),
		queue:         make(chan *mirrorTask, 500), // Larger buffer
		stats:         &MirrorStats{StartTime: time.Now()},
		ctx:           ctx,
		cancel:        cancel,
	}

	return mirror
}

// Download starts the mirroring process and returns the final statistics.
func (m *Mirror) Download(ctx context.Context) (*MirrorStats, error) {
	// Create output directory
	if err := os.MkdirAll(m.config.OutputDir, 0755); err != nil {
		return nil, fmt.Errorf("failed to create output directory: %w", err)
	}

	// Fetch robots.txt if enabled
	if m.config.RespectRobots {
		if err := m.robotsChecker.FetchPolicy(ctx, m.config.BaseURL, m.httpClient); err != nil {
			if m.config.Verbose {
				fmt.Fprintf(os.Stderr, "[mirror] Warning: failed to fetch robots.txt: %v\n", err)
			}
		}
	}

	// Add initial URL to queue
	m.visitedMu.Lock()
	m.visited[m.config.BaseURL.String()] = true
	m.visitedMu.Unlock()

	m.wg.Add(1)
	select {
	case m.queue <- &mirrorTask{url: m.config.BaseURL, depth: 0}:
	case <-ctx.Done():
		m.wg.Done()
		return m.stats, ctx.Err()
	}

	// Start worker goroutines
	workers := m.config.Parallel
	if workers <= 0 {
		workers = 8 // More workers for mirror mode
	}

	for i := 0; i < workers; i++ {
		go m.worker()
	}

	// Close queue on context cancellation to unblock workers
	go func() {
		<-m.ctx.Done()
		close(m.queue)
	}()

	// Wait for all tasks to complete
	m.wg.Wait()
	close(m.queue)

	m.stats.EndTime = time.Now()

	// Prune empty directories if requested
	if m.config.PruneEmpty {
		m.pruneEmptyDirectories()
	}

	return m.stats, nil
}

// worker processes mirror tasks from the queue.
func (m *Mirror) worker() {
	for {
		select {
		case task, ok := <-m.queue:
			if !ok {
				return
			}
			if m.ctx.Err() != nil {
				m.wg.Done()
				continue
			}
			m.processTask(task)
		case <-m.ctx.Done():
			return
		}
	}
}

// processTask handles a single mirror task (download, link extraction, etc.).
func (m *Mirror) processTask(task *mirrorTask) {
	defer m.wg.Done()

	// Check if cancelled
	if m.ctx.Err() != nil {
		return
	}

	// Check robots.txt
	if m.config.RespectRobots && !m.robotsChecker.CanFetch(task.url) {
		atomic.AddInt64(&m.stats.SkippedURLs, 1)
		if m.config.Verbose {
			fmt.Fprintf(os.Stderr, "[skip] %s (robots.txt)\n", task.url.String())
		}
		return
	}

	// Apply rate limiting
	if m.config.RateLimiter != nil {
		m.config.RateLimiter.Allow(1)
	}

	// Apply per-host crawl delay
	if delay := m.robotsChecker.GetCrawlDelay(task.url); delay > 0 {
		select {
		case <-time.After(delay):
		case <-m.ctx.Done():
			return
		}
	}

	outputPath := recursive.GetLocalPath(m.config.BaseURL, task.url, m.config.OutputDir)

	if m.config.DryRun {
		if m.config.Verbose {
			fmt.Fprintf(os.Stderr, "[dry-run] %s (depth %d)\n", task.url.String(), task.depth)
		}
		atomic.AddInt64(&m.stats.SkippedURLs, 1)
		return
	}

	// Download the URL
	if m.config.Verbose {
		fmt.Fprintf(os.Stderr, "[download] %s (depth %d)\n", task.url.String(), task.depth)
	}

	// Create parent directory
	if err := os.MkdirAll(filepath.Dir(outputPath), 0755); err != nil {
		atomic.AddInt64(&m.stats.FailedURLs, 1)
		if m.config.Verbose {
			fmt.Fprintf(os.Stderr, "[error] failed to create directory: %v\n", err)
		}
		return
	}

	// Create download request
	req := &core.DownloadRequest{
		URL:            task.url,
		Output:         outputPath,
		Resume:         true,
		Timeout:        m.config.Timeout,
		Verbose:        false,
		AutoDecompress: false, // Don't decompress HTML files
		Headers: map[string]string{
			"User-Agent": m.config.UserAgent,
		},
		Ctx: m.ctx,
	}

	// Execute download
	result, err := m.proto.Download(m.ctx, req)
	if err != nil {
		atomic.AddInt64(&m.stats.FailedURLs, 1)
		if m.config.Verbose {
			fmt.Fprintf(os.Stderr, "[error] %s: %v\n", task.url.String(), err)
		}
		return
	}

	atomic.AddInt64(&m.stats.DownloadedURLs, 1)
	atomic.AddInt64(&m.stats.TotalBytes, result.BytesDownloaded)

	if m.config.Verbose {
		fmt.Fprintf(os.Stderr, "[done] %s -> %s (%s)\n", task.url.String(), outputPath, format.Bytes(result.BytesDownloaded))
	}

	// Register downloaded URL for link rewriting
	m.rewriter.Register(task.url, outputPath)

	// Process content based on type
	contentType := m.detectContentType(outputPath, task.url.Path)

	switch contentType {
	case "html":
		// Extract links and queue them
		m.extractAndQueue(task.url, outputPath, task.depth+1)

		// Convert links if enabled
		if m.config.ConvertLinks {
			m.convertLinksInFile(outputPath)
		}
	case "css":
		// Extract URLs from CSS and queue them
		if m.config.DownloadAssets {
			m.extractCSSURLs(task.url, outputPath, task.depth+1)
		}
	}
}

// detectContentType determines the content type based on file extension and content sniffing.
func (m *Mirror) detectContentType(outputPath, urlPath string) string {
	ext := strings.ToLower(filepath.Ext(outputPath))

	// Check extension first
	switch ext {
	case ".html", ".htm", ".php", ".asp", ".aspx", ".xhtml":
		return "html"
	case ".css":
		return "css"
	case ".js", ".mjs":
		return "js"
	case ".png", ".jpg", ".jpeg", ".gif", ".svg", ".ico", ".webp", ".bmp", ".avif":
		return "image"
	case ".woff", ".woff2", ".ttf", ".eot", ".otf":
		return "font"
	case ".xml", ".rss", ".atom":
		return "xml"
	case ".json":
		return "json"
	}

	// Check if it's a directory index
	if ext == "" && (urlPath == "" || urlPath == "/" || strings.HasSuffix(urlPath, "/")) {
		return "html"
	}

	// Try content detection
	data, err := os.ReadFile(outputPath)
	if err != nil || len(data) == 0 {
		return "unknown"
	}

	snippet := data
	if len(snippet) > 512 {
		snippet = snippet[:512]
	}
	content := strings.ToLower(string(snippet))

	// HTML detection
	if strings.Contains(content, "<!doctype html") ||
		strings.Contains(content, "<html") ||
		strings.Contains(content, "<head") ||
		strings.Contains(content, "<body") {
		return "html"
	}

	// CSS detection
	if strings.Contains(content, "{") && strings.Contains(content, "}") &&
		(strings.Contains(content, ":") || strings.Contains(content, "@")) {
		return "css"
	}

	// XML detection
	if strings.HasPrefix(content, "<?xml") || strings.HasPrefix(content, "<rss") {
		return "xml"
	}

	return "unknown"
}

// extractAndQueue extracts links from HTML and adds new URLs to the processing queue.
func (m *Mirror) extractAndQueue(baseURL *url.URL, outputPath string, depth int) {
	data, err := os.ReadFile(outputPath)
	if err != nil {
		if m.config.Verbose {
			fmt.Fprintf(os.Stderr, "[error] failed to read %s for link extraction: %v\n", outputPath, err)
		}
		return
	}

	links, err := recursive.ExtractLinks(baseURL, data)
	if err != nil {
		if m.config.Verbose {
			fmt.Fprintf(os.Stderr, "[error] failed to parse links from %s: %v\n", outputPath, err)
		}
		return
	}

	added := 0
	for _, link := range links {
		// Check filter first (quick rejection)
		if !m.filter.ShouldDownload(link, depth) {
			continue
		}

		// Atomic check-and-add for visited
		m.visitedMu.Lock()
		if m.visited[link.String()] {
			m.visitedMu.Unlock()
			continue
		}
		m.visited[link.String()] = true
		m.visitedMu.Unlock()

		// Add to queue with context-aware send
		m.wg.Add(1)
		select {
		case m.queue <- &mirrorTask{url: link, depth: depth}:
			added++
		case <-m.ctx.Done():
			m.wg.Done()
			return
		}
	}

	if m.config.Verbose && added > 0 {
		fmt.Fprintf(os.Stderr, "[links] queued %d new URLs from %s\n", added, outputPath)
	}
}

// extractCSSURLs extracts URLs from CSS files and queues them for download.
func (m *Mirror) extractCSSURLs(baseURL *url.URL, cssPath string, depth int) {
	data, err := os.ReadFile(cssPath)
	if err != nil {
		if m.config.Verbose {
			fmt.Fprintf(os.Stderr, "[error] failed to read CSS %s: %v\n", cssPath, err)
		}
		return
	}

	urls := recursive.ExtractCSSURLs(baseURL, data)

	// Queue found URLs
	for _, resolved := range urls {
		// Check filter
		if !m.filter.ShouldDownload(resolved, depth) {
			continue
		}

		// Atomic check-and-add
		m.visitedMu.Lock()
		if m.visited[resolved.String()] {
			m.visitedMu.Unlock()
			continue
		}
		m.visited[resolved.String()] = true
		m.visitedMu.Unlock()

		m.wg.Add(1)
		select {
		case m.queue <- &mirrorTask{url: resolved, depth: depth}:
		case <-m.ctx.Done():
			m.wg.Done()
			return
		}
	}
}

// convertLinksInFile rewrites URLs in an HTML file for local viewing.
func (m *Mirror) convertLinksInFile(filePath string) {
	data, err := os.ReadFile(filePath)
	if err != nil {
		if m.config.Verbose {
			fmt.Fprintf(os.Stderr, "[error] failed to read %s for link conversion: %v\n", filePath, err)
		}
		return
	}

	doc, err := html.Parse(bytes.NewReader(data))
	if err != nil {
		if m.config.Verbose {
			fmt.Fprintf(os.Stderr, "[error] failed to parse HTML %s: %v\n", filePath, err)
		}
		return
	}

	converted := m.rewriter.RewriteLinks(doc, filePath)

	if converted > 0 {
		// Write modified HTML
		var buf bytes.Buffer
		if err := html.Render(&buf, doc); err != nil {
			if m.config.Verbose {
				fmt.Fprintf(os.Stderr, "[error] failed to render HTML %s: %v\n", filePath, err)
			}
			return
		}

		if err := os.WriteFile(filePath, buf.Bytes(), 0644); err != nil {
			if m.config.Verbose {
				fmt.Fprintf(os.Stderr, "[error] failed to write modified HTML %s: %v\n", filePath, err)
			}
			return
		}

		atomic.AddInt64(&m.stats.ConvertedLinks, int64(converted))

		if m.config.Verbose {
			fmt.Fprintf(os.Stderr, "[convert] converted %d links in %s\n", converted, filePath)
		}
	}
}

// pruneEmptyDirectories removes empty directories after mirroring is complete.
func (m *Mirror) pruneEmptyDirectories() {
	if m.config.Verbose {
		fmt.Fprintf(os.Stderr, "[prune] removing empty directories...\n")
	}

	// Collect all directories
	dirs := make([]string, 0)
	filepath.Walk(m.config.OutputDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}
		if info.IsDir() && path != m.config.OutputDir {
			dirs = append(dirs, path)
		}
		return nil
	})

	// Sort by depth (deepest first) for single-pass removal
	sort.Slice(dirs, func(i, j int) bool {
		return strings.Count(dirs[i], string(filepath.Separator)) > strings.Count(dirs[j], string(filepath.Separator))
	})

	// Remove empty directories bottom-up
	removed := 0
	for _, dir := range dirs {
		entries, err := os.ReadDir(dir)
		if err == nil && len(entries) == 0 {
			if err := os.Remove(dir); err == nil {
				removed++
				if m.config.Verbose {
					fmt.Fprintf(os.Stderr, "[prune] removed empty directory: %s\n", dir)
				}
			}
		}
	}

	if m.config.Verbose && removed > 0 {
		fmt.Fprintf(os.Stderr, "[prune] removed %d empty directories\n", removed)
	}
}

// Cancel stops the mirroring process.
func (m *Mirror) Cancel() {
	m.cancel()
}

// GetStats returns the current mirroring statistics.
func (m *Mirror) GetStats() *MirrorStats {
	return m.stats
}

// GetStatsSummary returns a human-readable summary of mirroring statistics.
func (s *MirrorStats) GetStatsSummary() string {
	duration := s.EndTime.Sub(s.StartTime)
	if duration == 0 {
		duration = time.Since(s.StartTime)
	}

	speed := float64(0)
	if duration.Seconds() > 0 {
		speed = float64(atomic.LoadInt64(&s.TotalBytes)) / duration.Seconds()
	}

	return fmt.Sprintf("Downloaded: %d files, %s total, %d links converted, %v elapsed, %.1f KB/s avg",
		atomic.LoadInt64(&s.DownloadedURLs),
		format.Bytes(atomic.LoadInt64(&s.TotalBytes)),
		atomic.LoadInt64(&s.ConvertedLinks),
		duration.Round(time.Second),
		speed/1024)
}