go-jdenticon/cmd/jdenticon/batch.go

package main

import (
	"bufio"
	"context"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"

	"github.com/mattn/go-isatty"
	"github.com/schollz/progressbar/v3"
	"github.com/spf13/cobra"
	"gitea.dockr.co/kev/go-jdenticon/jdenticon"
)

const (
	// Maximum filename length to avoid filesystem issues
	maxFilenameLength = 200
)

var (
	// Keep a-z, A-Z, 0-9, underscore, hyphen, and period. Replace everything else.
	sanitizeRegex = regexp.MustCompile(`[^a-zA-Z0-9_.-]+`)
)

// batchJob represents a single identicon generation job
type batchJob struct {
	value      string
	outputPath string
	size       int
}

// batchStats tracks processing statistics atomically
type batchStats struct {
	processed int64
	failed    int64
}

// batchCmd represents the batch command
var batchCmd = &cobra.Command{
	Use:   "batch <input-file>",
	Short: "Generate multiple identicons from a list using concurrent processing",
	Long: `Generate multiple identicons from a list of values in a text file.

The input file should contain one value per line. Each value will be used to generate
an identicon saved to the output directory. The filename will be based on the input
value with special characters replaced.

Uses concurrent processing with a worker pool for optimal performance. The number
of concurrent workers can be controlled with the --concurrency flag.

Examples:
  jdenticon batch users.txt --output-dir ./avatars
  jdenticon batch emails.txt --output-dir ./avatars --format png --size 64
  jdenticon batch large-list.txt --output-dir ./avatars --concurrency 8`,
	Args: cobra.ExactArgs(1), // Validate argument count
	RunE: func(cmd *cobra.Command, args []string) error {
		return runConcurrentBatch(cmd, args)
	},
}

func init() {
	rootCmd.AddCommand(batchCmd)

	// Define local flags specific to 'batch'
	batchCmd.Flags().StringP("output-dir", "d", "", "Output directory for generated identicons (required)")
	_ = batchCmd.MarkFlagRequired("output-dir")

	// Concurrency control
	batchCmd.Flags().IntP("concurrency", "c", runtime.NumCPU(),
		fmt.Sprintf("Number of concurrent workers (default: %d)", runtime.NumCPU()))
}

// runConcurrentBatch executes the batch processing with concurrent workers
func runConcurrentBatch(cmd *cobra.Command, args []string) error {
	inputFile := args[0]
	outputDir, _ := cmd.Flags().GetString("output-dir")
	concurrency, _ := cmd.Flags().GetInt("concurrency")

	// Validate concurrency value
	if concurrency <= 0 {
		return fmt.Errorf("concurrency must be positive, got %d", concurrency)
	}

	// Get format from viper
	format, err := getFormatFromViper()
	if err != nil {
		return err
	}

	// Populate library config from root persistent flags
	config, size, err := populateConfigFromFlags()
	if err != nil {
		return err
	}

	// Create generator with custom config and larger cache for concurrent workload
	cacheSize := generatorCacheSize * concurrency // Scale cache with worker count
	generator, err := jdenticon.NewGeneratorWithConfig(config, cacheSize)
	if err != nil {
		return fmt.Errorf("failed to create generator: %w", err)
	}

	// Create output directory if it doesn't exist
	// #nosec G301 -- 0755 is standard for directories; CLI tool needs world-readable output
	if err = os.MkdirAll(outputDir, 0755); err != nil {
		return fmt.Errorf("failed to create output directory: %w", err)
	}

	// Set up graceful shutdown context
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	// Process the file and collect jobs
	jobs, total, err := prepareJobs(inputFile, outputDir, format, size)
	if err != nil {
		return err
	}

	if total == 0 {
		fmt.Fprintf(os.Stderr, "No valid entries found in input file\n")
		return nil
	}

	// Initialize progress reporting
	showProgress := isTTY(os.Stderr)
	var bar *progressbar.ProgressBar
	if showProgress {
		bar = createProgressBar(total)
	}

	// Execute concurrent processing
	stats, err := processConcurrentJobs(ctx, jobs, generator, format, concurrency, bar)
	if err != nil {
		return err
	}

	// Final status message
	processed := atomic.LoadInt64(&stats.processed)
	failed := atomic.LoadInt64(&stats.failed)

	if showProgress {
		fmt.Fprintf(os.Stderr, "\nBatch processing complete: %d succeeded, %d failed\n", processed, failed)
	} else {
		fmt.Fprintf(os.Stderr, "Batch processing complete: %d succeeded, %d failed\n", processed, failed)
	}

	if failed > 0 {
		return fmt.Errorf("some identicons failed to generate")
	}

	return nil
}

// prepareJobs reads the input file and creates a slice of jobs
func prepareJobs(inputFile, outputDir string, format FormatFlag, size int) ([]batchJob, int, error) {
	// #nosec G304 -- inputFile is provided via CLI flag, validated by cobra
	file, err := os.Open(inputFile)
	if err != nil {
		return nil, 0, fmt.Errorf("failed to open input file: %w", err)
	}
	defer file.Close()

	var jobs []batchJob
	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		value := strings.TrimSpace(scanner.Text())
		if value == "" {
			continue // Skip empty lines
		}

		// Generate filename from value (sanitize for filesystem)
		filename := sanitizeFilename(value)
		extension := ".svg"
		if format == FormatPNG {
			extension = ".png"
		}
		outputPath := filepath.Join(outputDir, filename+extension)

		jobs = append(jobs, batchJob{
			value:      value,
			outputPath: outputPath,
			size:       size,
		})
	}

	if err := scanner.Err(); err != nil {
		return nil, 0, fmt.Errorf("error reading input file: %w", err)
	}

	return jobs, len(jobs), nil
}

// createProgressBar creates a progress bar for the batch processing
func createProgressBar(total int) *progressbar.ProgressBar {
	return progressbar.NewOptions(total,
		progressbar.OptionSetWriter(os.Stderr),
		progressbar.OptionSetDescription("Processing identicons..."),
		progressbar.OptionEnableColorCodes(true),
		progressbar.OptionShowCount(),
		progressbar.OptionSetWidth(15),
		progressbar.OptionSetTheme(progressbar.Theme{
			Saucer:        "[green]=[reset]",
			SaucerHead:    "[green]>[reset]",
			SaucerPadding: " ",
			BarStart:      "[",
			BarEnd:        "]",
		}),
	)
}

// processConcurrentJobs executes jobs using a worker pool pattern
func processConcurrentJobs(ctx context.Context, jobs []batchJob, generator *jdenticon.Generator,
	format FormatFlag, concurrency int, bar *progressbar.ProgressBar) (*batchStats, error) {
	stats := &batchStats{}
	jobChan := make(chan batchJob, len(jobs))
	var wg sync.WaitGroup

	// Start worker goroutines
	for i := 0; i < concurrency; i++ {
		wg.Add(1)
		go func(workerID int) {
			defer wg.Done()
			batchWorker(ctx, workerID, jobChan, generator, format, stats, bar)
		}(i)
	}

	// Send jobs to workers
	go func() {
		defer close(jobChan)
		for _, job := range jobs {
			select {
			case jobChan <- job:
			case <-ctx.Done():
				return
			}
		}
	}()

	// Wait for all workers to complete
	wg.Wait()

	// Check if processing was interrupted
	if ctx.Err() != nil {
		return stats, fmt.Errorf("processing interrupted: %w", ctx.Err())
	}

	return stats, nil
}

// batchWorker processes jobs from the job channel
func batchWorker(ctx context.Context, workerID int, jobs <-chan batchJob, generator *jdenticon.Generator,
	format FormatFlag, stats *batchStats, bar *progressbar.ProgressBar) {
	for {
		select {
		case job, ok := <-jobs:
			if !ok {
				// Channel closed, no more jobs
				return
			}

			// Process the job with context for cancellation support
			if err := processJob(ctx, job, generator, format); err != nil {
				fmt.Fprintf(os.Stderr, "Worker %d failed to process %q: %v\n", workerID, job.value, err)
				atomic.AddInt64(&stats.failed, 1)
			} else {
				atomic.AddInt64(&stats.processed, 1)
			}

			// Update progress bar if available
			if bar != nil {
				_ = bar.Add(1)
			}

		case <-ctx.Done():
			// Shutdown signal received
			return
		}
	}
}

// processJob handles a single identicon generation job
func processJob(ctx context.Context, job batchJob, generator *jdenticon.Generator, format FormatFlag) error {
	// Generate identicon with context for cancellation support
	icon, err := generator.Generate(ctx, job.value, job.size)
	if err != nil {
		return fmt.Errorf("failed to generate identicon: %w", err)
	}

	// Generate output based on format
	result, err := renderIcon(icon, format)
	if err != nil {
		return fmt.Errorf("failed to render %s: %w", format, err)
	}

	// Write file
	// #nosec G306 -- 0644 is appropriate for generated image files (world-readable)
	if err := os.WriteFile(job.outputPath, result, 0644); err != nil {
		return fmt.Errorf("failed to write file: %w", err)
	}

	return nil
}

// sanitizeFilename converts a value to a safe filename using a whitelist approach
func sanitizeFilename(value string) string {
	// Replace @ separately for readability if desired
	filename := strings.ReplaceAll(value, "@", "_at_")

	// Replace all other invalid characters with a single underscore
	filename = sanitizeRegex.ReplaceAllString(filename, "_")

	// Limit length to avoid filesystem issues
	if len(filename) > maxFilenameLength {
		filename = filename[:maxFilenameLength]
	}

	return filename
}

// isTTY checks if the given file descriptor is a terminal
func isTTY(f *os.File) bool {
	return isatty.IsTerminal(f.Fd())
}