goget/internal/recursive/crawler.go

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

package recursive

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"math/rand/v2"
	"mime"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

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

// CrawlerConfig configures the crawler
type CrawlerConfig struct {
	MaxDepth        int
	FollowExternal  bool
	ExcludePatterns []string
	IncludePatterns []string
	AcceptPatterns  []string
	NoParent        bool
	PageRequisites  bool
	ConvertLinks    bool
	OutputDir       string
	Verbose         bool
	Parallel        int
	Delay           time.Duration
	RandomWait      bool
	UserAgent       string
	DryRun          bool
}

// CrawlerStats represents crawler statistics
type CrawlerStats struct {
	TotalURLs      int
	DownloadedURLs int
	FailedURLs     int
	SkippedURLs    int
	TotalBytes     int64
	StartTime      time.Time
	EndTime        time.Time
	mu             sync.Mutex
}

// Crawler performs recursive downloading
type Crawler struct {
	config    *CrawlerConfig
	proto     core.Protocol
	filter    *URLFilter
	visited   map[string]bool
	visitedMu sync.Mutex
	queue     chan *crawlTask
	queueOnce sync.Once
	stats     *CrawlerStats
	rewriter  *linkrewrite.Rewriter
	wg        sync.WaitGroup
	ctx       context.Context
	cancel    context.CancelFunc
}

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

// NewCrawler creates new crawler
func NewCrawler(cfg *CrawlerConfig, proto core.Protocol) *Crawler {
	ctx, cancel := context.WithCancel(context.Background())

	filter := NewURLFilter(nil, URLFilterConfig{
		MaxDepth:        cfg.MaxDepth,
		FollowExternal:  cfg.FollowExternal,
		ExcludePatterns: cfg.ExcludePatterns,
		IncludePatterns: cfg.IncludePatterns,
		AcceptPatterns:  cfg.AcceptPatterns,
		NoParent:        cfg.NoParent,
	})

	crawler := &Crawler{
		config:  cfg,
		proto:   proto,
		filter:  filter,
		visited: make(map[string]bool),
		queue:   make(chan *crawlTask, 100),
		stats:   &CrawlerStats{StartTime: time.Now()},
		ctx:     ctx,
		cancel:  cancel,
	}
	if cfg.ConvertLinks {
		crawler.rewriter = linkrewrite.New(nil, cfg.OutputDir)
	}

	return crawler
}

// Download recursively downloads the URL and all links
func (c *Crawler) Download(ctx context.Context, baseURL *url.URL) (*CrawlerStats, error) {
	c.filter.baseURL = baseURL
	if c.rewriter != nil {
		c.rewriter = linkrewrite.New(baseURL, c.config.OutputDir)
	}

	// Create output directory
	if err := os.MkdirAll(c.config.OutputDir, 0755); err != nil {
		return nil, fmt.Errorf("failed to create output directory: %w", err)
	}

	// Add to WaitGroup BEFORE adding to queue
	c.wg.Add(1)
	c.queue <- &crawlTask{url: baseURL, depth: 0}

	// Start worker goroutines
	workers := c.config.Parallel
	if workers <= 0 {
		workers = 4
	}

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

	// Close queue on context cancellation to unblock workers
	closeQueue := func() {
		c.queueOnce.Do(func() {
			close(c.queue)
		})
	}

	go func() {
		<-c.ctx.Done()
		closeQueue()
	}()

	// Wait for all tasks to complete
	c.wg.Wait()
	closeQueue()

	c.stats.EndTime = time.Now()
	return c.stats, nil
}

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

			// Delay between requests (with optional randomization)
			if c.config.Delay > 0 {
				d := c.config.Delay
				if c.config.RandomWait {
					d = randomizeDelay(d)
				}
				select {
				case <-time.After(d):
				case <-c.ctx.Done():
					return
				}
			}
		case <-c.ctx.Done():
			return
		}
	}
}

// processTask processes a single crawl task
func (c *Crawler) processTask(task *crawlTask) {
	defer c.wg.Done()

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

	// Check if already visited
	c.visitedMu.Lock()
	if c.visited[task.url.String()] {
		c.visitedMu.Unlock()
		c.stats.mu.Lock()
		c.stats.SkippedURLs++
		c.stats.mu.Unlock()
		return
	}
	c.visited[task.url.String()] = true
	c.visitedMu.Unlock()

	// Check filter
	if !c.filter.ShouldDownload(task.url, task.depth) {
		c.stats.mu.Lock()
		c.stats.SkippedURLs++
		c.stats.mu.Unlock()
		if c.config.Verbose {
			fmt.Fprintf(os.Stderr, "[skip] %s (filtered)\n", task.url.String())
		}
		return
	}

	outputPath := GetLocalPath(c.filter.baseURL, task.url, c.config.OutputDir)

	if c.config.DryRun {
		if c.config.Verbose {
			fmt.Fprintf(os.Stderr, "[dry-run] %s (depth %d)\n", task.url.String(), task.depth)
		}
		c.stats.mu.Lock()
		c.stats.SkippedURLs++
		c.stats.mu.Unlock()
		return
	}

	// Download the URL
	if c.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 {
		c.stats.mu.Lock()
		c.stats.FailedURLs++
		c.stats.mu.Unlock()
		if c.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:        30 * time.Minute,
		Verbose:        false,
		AutoDecompress: true,
		Headers: map[string]string{
			"User-Agent": c.config.UserAgent,
		},
		Ctx: c.ctx,
	}

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

	c.stats.mu.Lock()
	c.stats.DownloadedURLs++
	c.stats.TotalBytes += result.BytesDownloaded
	c.stats.mu.Unlock()

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

	if c.config.ConvertLinks && c.rewriter != nil {
		c.rewriter.Register(task.url, outputPath)
		c.rewriteLinks(outputPath, task.url)
	}

	// Check if content is HTML by reading file content
	if c.shouldParseForLinks(outputPath, task.url.Path) {
		if c.config.Verbose {
			fmt.Fprintf(os.Stderr, "[parse] %s looks like HTML, extracting links\n", outputPath)
		}
		c.extractAndQueue(task.url, outputPath, task.depth+1)
	}

	// Delay between requests
	if c.config.Delay > 0 {
		select {
		case <-time.After(c.config.Delay):
		case <-c.ctx.Done():
		}
	}
}

// Better HTML detection - check content, not just extension
func (c *Crawler) shouldParseForLinks(outputPath, urlPath string) bool {
	// First check extension (quick path)
	ext := strings.ToLower(filepath.Ext(outputPath))
	if ext == ".html" || ext == ".htm" || ext == ".php" || ext == ".asp" || ext == ".aspx" {
		return true
	}
	if ext == "" && (urlPath == "" || urlPath == "/" || strings.HasSuffix(urlPath, "/")) {
		return true
	}

	// If extension doesn't match, read only the first 512 bytes to check for HTML signature.
	f, err := os.Open(outputPath)
	if err != nil {
		return false
	}
	defer f.Close()

	snippet := make([]byte, 512)
	n, err := io.ReadFull(f, snippet)
	if err != nil && err != io.ErrUnexpectedEOF && err != io.EOF {
		return false
	}
	snippet = snippet[:n]
	if len(snippet) == 0 {
		return false
	}
	content := strings.ToLower(string(snippet))

	// Look for common HTML indicators
	if strings.Contains(content, "<!doctype html") ||
		strings.Contains(content, "<html") ||
		strings.Contains(content, "<head") ||
		strings.Contains(content, "<body") {
		return true
	}

	// Check Content-Type via mime sniffing (basic)
	if contentType := httpDetectContentType(snippet); contentType != "" {
		mediaType, _, err := mime.ParseMediaType(contentType)
		if err == nil && mediaType == "text/html" {
			return true
		}
	}

	return false
}

// httpDetectContentType basic MIME type detection from content
func httpDetectContentType(data []byte) string {
	if len(data) < 2 {
		return ""
	}
	if bytes.HasPrefix(data, []byte("<!")) || bytes.HasPrefix(data, []byte("<html")) || bytes.HasPrefix(data, []byte("<HTML")) {
		return "text/html"
	}
	return ""
}

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

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

	if c.config.Verbose {
		fmt.Fprintf(os.Stderr, "[links] found %d links in %s\n", len(links), outputPath)
	}

	// Add valid links to queue
	added := 0
	for _, link := range links {
		if c.filter.ShouldDownload(link, depth) {
			c.visitedMu.Lock()
			if !c.visited[link.String()] {
				c.visitedMu.Unlock()
				// Add to WaitGroup before adding to queue
				c.wg.Add(1)
				select {
				case c.queue <- &crawlTask{url: link, depth: depth}:
					added++
				case <-c.ctx.Done():
					c.wg.Done()
					return
				}
			} else {
				c.visitedMu.Unlock()
			}
		}
	}
	if c.config.Verbose && added > 0 {
		fmt.Fprintf(os.Stderr, "[queue] added %d new URLs to queue\n", added)
	}
}

// GetStats returns current statistics
func (c *Crawler) GetStats() *CrawlerStats {
	return c.stats
}

// Cancel cancels the crawler
func (c *Crawler) Cancel() {
	c.cancel()
}

// GetStatsSummary returns a summary of statistics
func (s *CrawlerStats) 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(s.TotalBytes) / duration.Seconds()
	}

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

// rewriteLinks rewrites absolute URLs to local paths in HTML files.
func (c *Crawler) rewriteLinks(outputPath string, sourceURL *url.URL) {
	if c.rewriter == nil {
		return
	}

	data, err := os.ReadFile(outputPath)
	if err != nil {
		return
	}

	doc, err := html.Parse(bytes.NewReader(data))
	if err != nil {
		return
	}

	if c.rewriter.RewriteLinks(doc, outputPath) == 0 {
		return
	}

	var buf bytes.Buffer
	if err := html.Render(&buf, doc); err != nil {
		return
	}
	_ = os.WriteFile(outputPath, buf.Bytes(), 0644)
	_ = sourceURL
}

// randomizeDelay returns a randomized delay between 0.5x and 1.5x.
func randomizeDelay(d time.Duration) time.Duration {
	factor := 0.5 + rand.Float64()
	return time.Duration(float64(d) * factor)
}