path: root/internal/image/download.go

package image

import (
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

// DownloadOptions configures image download behavior
type DownloadOptions struct {
	OutputDir         string // Directory to save images
	OverwriteExisting bool   // Whether to overwrite existing files
	CreateDir         bool   // Create output directory if it doesn't exist
	FileNamePattern   string // Pattern for file naming (e.g., "{word}_{source}")
	MaxSizeBytes      int64  // Maximum file size to download (0 = no limit)
}

// DefaultDownloadOptions returns sensible defaults for image downloads
func DefaultDownloadOptions() *DownloadOptions {
	return &DownloadOptions{
		OutputDir:         "./images",
		OverwriteExisting: false,
		CreateDir:         true,
		FileNamePattern:   "{word}_{source}",
		MaxSizeBytes:      10 * 1024 * 1024, // 10MB
	}
}

// Downloader handles image downloads from search results.
// It requires ImageClient (ImageSearcher + AttributionProvider) because it
// both downloads images and embeds attribution text in saved files.
type Downloader struct {
	searcher ImageClient
	options  *DownloadOptions
}

// NewDownloader creates a new image downloader
func NewDownloader(searcher ImageClient, options *DownloadOptions) *Downloader {
	if options == nil {
		options = DefaultDownloadOptions()
	}
	return &Downloader{
		searcher: searcher,
		options:  options,
	}
}

// DownloadImage downloads a single image to the specified path.
func (d *Downloader) DownloadImage(ctx context.Context, result *SearchResult, outputPath string) (err error) {
	// Ensure directory exists
	dir := filepath.Dir(outputPath)
	if dir != "" && dir != "." {
		if err := os.MkdirAll(dir, 0755); err != nil {
			return fmt.Errorf("create output dir %q: %w", dir, err)
		}
	}

	// Check if file already exists
	if !d.options.OverwriteExisting {
		if _, err := os.Stat(outputPath); err == nil {
			return fmt.Errorf("output file exists: %s", outputPath)
		}
	}

	// Download the image
	reader, err := d.searcher.Download(ctx, result.URL)
	if err != nil {
		return fmt.Errorf("download %q: %w", result.URL, err)
	}
	defer func() {
		_ = reader.Close()
	}()

	// Create output file
	file, err := os.Create(outputPath)
	if err != nil {
		return fmt.Errorf("create output file %q: %w", outputPath, err)
	}
	defer func() {
		if closeErr := file.Close(); err == nil && closeErr != nil {
			err = fmt.Errorf("close output file %q: %w", outputPath, closeErr)
		}
	}()

	// Copy with size limit if specified
	var written int64
	if d.options.MaxSizeBytes > 0 {
		written, err = io.CopyN(file, reader, d.options.MaxSizeBytes)
		if err != nil && err != io.EOF {
			_ = os.Remove(outputPath) // Clean up on error
			return fmt.Errorf("write output file %q: %w", outputPath, err)
		}

		// Check if we hit the size limit
		if written == d.options.MaxSizeBytes {
			// Try to read one more byte to see if file is larger
			if _, err := reader.Read(make([]byte, 1)); err != io.EOF {
				_ = os.Remove(outputPath) // Clean up
				return fmt.Errorf("image exceeds max size %d bytes", d.options.MaxSizeBytes)
			}
		}

		// Ensure the file is fully written to disk before returning
		if err := file.Sync(); err != nil {
			return fmt.Errorf("sync output file %q: %w", outputPath, err)
		}
	} else {
		_, err = io.Copy(file, reader)
		if err != nil {
			_ = os.Remove(outputPath) // Clean up on error
			return fmt.Errorf("write output file %q: %w", outputPath, err)
		}
	}

	// Ensure the file is fully written to disk before returning
	if err := file.Sync(); err != nil {
		return fmt.Errorf("sync output file %q: %w", outputPath, err)
	}

	// Save attribution if required
	if attribution := d.searcher.GetAttribution(result); attribution != "" {
		attrPath := strings.TrimSuffix(outputPath, filepath.Ext(outputPath)) + "_attribution.txt"
		if err := os.WriteFile(attrPath, []byte(attribution), 0644); err != nil {
			// Non-fatal error - log but don't fail the download
			fmt.Fprintf(os.Stderr, "Warning: failed to save attribution: %v\n", err)
		}
	}

	return nil
}

// DownloadBestMatch downloads the best matching image for a query
func (d *Downloader) DownloadBestMatch(ctx context.Context, query string) (*SearchResult, string, error) {
	// Search for images
	opts := DefaultSearchOptions(query)
	opts.PerPage = 5 // Get top 5 results

	results, err := d.searcher.Search(ctx, opts)
	if err != nil {
		return nil, "", fmt.Errorf("search images: %w", err)
	}

	if len(results) == 0 {
		return nil, "", fmt.Errorf("no images found for %q", query)
	}

	// Try to download the first available image
	for i, result := range results {
		// Generate filename
		filename := d.generateFileName(query, &result, i)
		outputPath := filepath.Join(d.options.OutputDir, filename)

		// Try to download
		err := d.DownloadImage(ctx, &result, outputPath)
		if err == nil {
			return &result, outputPath, nil
		}

		// Log error and try next
		fmt.Fprintf(os.Stderr, "Warning: failed to download image %d: %v\n", i+1, err)
	}

	return nil, "", fmt.Errorf("no downloadable images found for %q", query)
}

// generateFileName creates a filename based on the pattern
func (d *Downloader) generateFileName(word string, result *SearchResult, index int) string {
	// Start with the pattern
	filename := d.options.FileNamePattern

	// Replace placeholders
	filename = strings.ReplaceAll(filename, "{word}", sanitizeFileName(word))
	filename = strings.ReplaceAll(filename, "{source}", result.Source)
	filename = strings.ReplaceAll(filename, "{id}", result.ID)
	filename = strings.ReplaceAll(filename, "{index}", fmt.Sprintf("%d", index))

	// Determine extension from URL
	ext := filepath.Ext(result.URL)
	if strings.HasPrefix(result.URL, "data:image/png") {
		ext = ".png"
	} else if ext == "" || len(ext) > 5 { // Probably not a real extension
		ext = ".jpg" // Default to jpg
	}

	// Add extension if not present
	if filepath.Ext(filename) == "" {
		filename += ext
	}

	return filename
}

// sanitizeFileName removes or replaces characters that are problematic in filenames
func sanitizeFileName(name string) string {
	// Replace common problematic characters
	replacer := strings.NewReplacer(
		"/", "_",
		"\\", "_",
		":", "_",
		"*", "_",
		"?", "_",
		"\"", "_",
		"<", "_",
		">", "_",
		"|", "_",
		" ", "_",
		".", "_",
	)

	sanitized := replacer.Replace(name)

	// Ensure the filename is not too long
	if len(sanitized) > 50 {
		sanitized = sanitized[:50]
	}

	return sanitized
}

// DownloadBestMatchWithOptions downloads the best matching image for given search options
func (d *Downloader) DownloadBestMatchWithOptions(ctx context.Context, opts *SearchOptions) (*SearchResult, string, error) {
	// Search for images
	searchOpts := *opts    // Copy to avoid modifying original
	searchOpts.PerPage = 5 // Get top 5 results

	results, err := d.searcher.Search(ctx, &searchOpts)
	if err != nil {
		return nil, "", fmt.Errorf("search images: %w", err)
	}

	if len(results) == 0 {
		return nil, "", fmt.Errorf("no images found for %q", opts.Query)
	}

	// Try to download the first available image
	for i, result := range results {
		// Generate filename
		filename := d.generateFileName(opts.Query, &result, i)
		outputPath := filepath.Join(d.options.OutputDir, filename)

		// Try to download
		err := d.DownloadImage(ctx, &result, outputPath)
		if err == nil {
			return &result, outputPath, nil
		}

		// Log error and try next
		fmt.Fprintf(os.Stderr, "Warning: failed to download image %d: %v\n", i+1, err)
	}

	return nil, "", fmt.Errorf("no downloadable images found for %q", opts.Query)
}