goget/internal/output/writer.go

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

package output

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"
	"syscall"
	"time"
)

// WriterConfig holds configuration for the Writer.
type WriterConfig struct {
	// Output is the destination file path (use "-" for stdout).
	Output string

	// BufferSize is the size of the write buffer.
	BufferSize int

	// Atomic enables atomic write via temp file + rename.
	Atomic bool

	// Resume enables support for resuming interrupted downloads.
	Resume bool

	// ProgressCallback is called on each write with the current progress.
	ProgressCallback func(current, total int64, speed float64)

	// Verbose enables verbose logging.
	Verbose bool

	// CreateDirs creates missing parent directories of Output via os.MkdirAll
	// before opening the file. Ignored when Output is "-" (stdout). The
	// equivalent of curl --create-dirs / wget --directory-prefix.
	CreateDirs bool
}

// DefaultWriterConfig returns a WriterConfig with sensible defaults.
func DefaultWriterConfig() *WriterConfig {
	return &WriterConfig{
		BufferSize: 32 * 1024, // 32KB
		Atomic:     true,
		Resume:     false,
		Verbose:    false,
	}
}

// Writer is a progress-aware file writer supporting atomic writes and resume.
type Writer struct {
	config    *WriterConfig
	file      *os.File
	tempFile  *os.File
	written   int64
	total     int64
	startTime time.Time
	mu        sync.Mutex
	callback  func(current, total int64, speed float64)
	ctx       context.Context
	cancel    context.CancelFunc
}

// NewWriter creates a new writer
func NewWriter(cfg *WriterConfig) (*Writer, error) {
	if cfg == nil {
		cfg = DefaultWriterConfig()
	}

	w := &Writer{
		config:    cfg,
		written:   0,
		startTime: time.Now(),
		callback:  cfg.ProgressCallback,
	}

	// Create context for cancellation.
	w.ctx, w.cancel = context.WithCancel(context.Background())

	// Determine output path
	outputPath := cfg.Output
	if outputPath == "" {
		outputPath = "-" // stdout
	}

	// Create parent directories on demand (curl --create-dirs / wget
	// --directory-prefix). Skip for stdout where there is no path.
	if cfg.CreateDirs && outputPath != "-" {
		dir := filepath.Dir(outputPath)
		if dir != "" && dir != "." {
			if err := os.MkdirAll(dir, 0755); err != nil {
				return nil, fmt.Errorf("failed to create output directory: %w", err)
			}
		}
	}

	// Atomic write - create temp file
	if cfg.Atomic && outputPath != "-" {
		dir := filepath.Dir(outputPath)
		name := filepath.Base(outputPath)
		tempName := fmt.Sprintf(".%s.tmp.%d", name, time.Now().UnixNano())
		tempPath := filepath.Join(dir, tempName)

		tempFile, err := os.OpenFile(tempPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
		if err != nil {
			return nil, fmt.Errorf("failed to create temp file: %w", err)
		}

		w.tempFile = tempFile
		w.file = tempFile
	} else if outputPath == "-" {
		w.file = os.Stdout
	} else {
		// Check if file exists for resume.
		if cfg.Resume {
			if info, err := os.Stat(outputPath); err == nil {
				w.written = info.Size()
				// Open in append mode
				file, err := os.OpenFile(outputPath, os.O_APPEND|os.O_WRONLY, 0644)
				if err != nil {
					return nil, fmt.Errorf("failed to open file for resume: %w", err)
				}
				w.file = file
			} else {
				// File doesn't exist, create new
				file, err := os.Create(outputPath)
				if err != nil {
					return nil, fmt.Errorf("failed to create file: %w", err)
				}
				w.file = file
			}
		} else {
			// Check if file exists
			if _, err := os.Stat(outputPath); err == nil {
				return nil, fmt.Errorf("file already exists: %s", outputPath)
			}

			file, err := os.Create(outputPath)
			if err != nil {
				return nil, fmt.Errorf("failed to create file: %w", err)
			}
			w.file = file
		}
	}

	return w, nil
}

// Write implements the io.Writer interface.
func (w *Writer) Write(p []byte) (int, error) {
	w.mu.Lock()
	defer w.mu.Unlock()

	n, err := w.file.Write(p)
	if err != nil {
		return n, err
	}

	w.written += int64(n)

	// Report progress
	if w.callback != nil {
		duration := time.Since(w.startTime)
		speed := float64(w.written) / duration.Seconds()
		w.callback(w.written, w.total, speed)
	}

	return n, nil
}

// SetTotal sets the expected total size for progress tracking.
func (w *Writer) SetTotal(total int64) {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.total = total
}

// SetProgressCallback sets the callback function for progress reporting.
func (w *Writer) SetProgressCallback(cb func(current, total int64, speed float64)) {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.callback = cb
}

// Written returns the number of bytes written so far.
func (w *Writer) Written() int64 {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.written
}

// Total returns the expected total size.
func (w *Writer) Total() int64 {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.total
}

// Close finalizes the write and performs the atomic rename if applicable.
func (w *Writer) Close() error {
	w.cancel()

	if w.file == nil {
		return nil
	}

	// Don't close stdout — it's shared across the process
	if w.file != os.Stdout {
		if err := w.file.Close(); err != nil {
			return err
		}
	}

	// Atomic rename if temp file was used
	if w.tempFile != nil && w.config.Output != "-" {
		tempPath := w.tempFile.Name()
		if err := os.Rename(tempPath, w.config.Output); err != nil {
			if errors.Is(err, syscall.EXDEV) {
				// Cross-filesystem rename is not supported by the OS.
				// Fall back to a copy + remove so the data still reach
				// the destination directory instead of being lost.
				if fallbackErr := copyFileAndRemove(tempPath, w.config.Output); fallbackErr != nil {
					return fmt.Errorf("atomic rename failed across filesystems and fallback copy failed: rename=%v copy=%w", err, fallbackErr)
				}
				return nil
			}
			// Any other error (permission denied, target is a directory,
			// target busy on Windows, etc.) — keep the temp file so the
			// user can recover the partial download manually.
			return fmt.Errorf("failed to rename temp file %q to %q: %w (temp file kept for recovery)", tempPath, w.config.Output, err)
		}
	}

	return nil
}

// copyFileAndRemove copies src to dst (truncating any existing file) and
// removes src on success. Used as a fallback for os.Rename when the temp
// file and the destination are on different filesystems.
func copyFileAndRemove(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

	out, err := os.OpenFile(dst, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}

	if _, err := io.Copy(out, in); err != nil {
		out.Close()
		return err
	}
	if err := out.Close(); err != nil {
		return err
	}

	return os.Remove(src)
}

// Cancel aborts the write and cleans up any temporary files.
func (w *Writer) Cancel() {
	w.cancel()

	if w.tempFile != nil {
		// Remove temp file on cancel
		tempPath := w.tempFile.Name()
		w.tempFile.Close()
		os.Remove(tempPath)
	}
}

// GetSpeed returns the average write speed in bytes per second.
func (w *Writer) GetSpeed() float64 {
	w.mu.Lock()
	defer w.mu.Unlock()
	duration := time.Since(w.startTime)
	if duration.Seconds() == 0 {
		return 0
	}
	return float64(w.written) / duration.Seconds()
}

// GetDuration returns the elapsed time since the writer was created.
func (w *Writer) GetDuration() time.Duration {
	return time.Since(w.startTime)
}

// GetProgress returns the progress as a percentage of the total size.
func (w *Writer) GetProgress() float64 {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.total == 0 {
		return 0
	}
	return float64(w.written) / float64(w.total) * 100
}