Copy recursive

2 files changed, 309 insertions, 79 deletions

diff --git a/ccp.go b/ccp.go
index 2dc59c7..e2cdacf 100644
--- a/ccp.go
+++ b/ccp.go
@@ -1,18 +1,23 @@
 package main
 
 import (
-	"errors"
 	"flag"
 	"fmt"
-	"io"
 	"os"
+	"path/filepath"
+	"strings"
+	"sync/atomic"
 	"time"
 
 	"github.com/charmbracelet/bubbles/progress"
 	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/lipgloss"
+	"gitlab.com/rhogenson/ccp/internal/cp"
 	"gitlab.com/rhogenson/deque"
 )
 
+var f = flag.Bool("f", false, "if an existing destination file cannot be opened, remove it and try again")
+
 type measurement struct {
 	t time.Time
 	i int64
@@ -20,44 +25,98 @@ type measurement struct {
 
 type model struct {
 	progress progress.Model
+	msgs     chan tea.Msg
+
+	srcs []string
+	dst  string
 
 	max          int64
+	current      atomic.Int64
 	measurements deque.Deque[measurement]
-}
-
-func (m *model) Init() tea.Cmd {
-	return nil
+	copyingFiles map[string]bool
+	copyingFile  string
+	errs         []string
 }
 
 type (
-	progressMsg int
-	doneMsg     struct{}
+	tickMsg struct{}
+
+	maxMsg       int64
+	fileStartMsg string
+	fileDoneMsg  struct {
+		name string
+		err  error
+	}
+	doneMsg struct{}
 )
 
+func (m *model) listen() tea.Cmd {
+	return func() tea.Msg {
+		return <-m.msgs
+	}
+}
+
+func tick() tea.Cmd {
+	return tea.Tick(10*time.Millisecond, func(time.Time) tea.Msg { return tickMsg{} })
+}
+
+func (m *model) Init() tea.Cmd {
+	return tea.Batch(
+		func() tea.Msg {
+			cp.Copy(m.srcs, m.dst, *f, m)
+			return doneMsg{}
+		},
+		m.listen(),
+		tick())
+}
+
 func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	switch msg := msg.(type) {
-	case tea.WindowSizeMsg:
-		m.progress.Width = msg.Width - 4
-		return m, nil
-	case progressMsg:
-		n := int64(msg)
-		if m.measurements.Len() > 0 {
-			n += m.measurements.At(m.measurements.Len() - 1).i
+	case maxMsg:
+		m.max = int64(msg)
+		return m, m.listen()
+	case fileStartMsg:
+		name := string(msg)
+		m.copyingFiles[name] = true
+		if m.copyingFile == "" {
+			m.copyingFile = name
 		}
-		if m.measurements.Len() > 0 && m.measurements.At(m.measurements.Len()-1).i > n {
-			m.measurements.Reset()
+		return m, m.listen()
+	case fileDoneMsg:
+		delete(m.copyingFiles, msg.name)
+		if m.copyingFile == msg.name {
+			m.copyingFile = ""
+			for name := range m.copyingFiles {
+				m.copyingFile = name
+				break
+			}
 		}
-		now := time.Now()
-		for m.measurements.Len() > 2 && now.Sub(m.measurements.At(0).t) > 2*time.Minute {
-			m.measurements.PopFront()
+		if msg.err != nil {
+			m.errs = append(m.errs, msg.err.Error())
 		}
-		m.measurements.PushBack(measurement{now, n})
-		return m, m.progress.SetPercent(float64(n) / float64(m.max))
+		return m, m.listen()
 	case doneMsg:
 		return m, tea.Sequence(
-			m.progress.SetPercent(1),
+			m.progress.SetPercent(float64(m.current.Load())/float64(m.max)),
 			tea.Tick(500*time.Millisecond, func(time.Time) tea.Msg { return nil }),
 			tea.Quit)
+
+	case tickMsg:
+		now := time.Now()
+		for m.measurements.Len() > 1 && now.Sub(m.measurements.At(0).t) > 2*time.Minute {
+			m.measurements.PopFront()
+		}
+		n := m.current.Load()
+		m.measurements.PushBack(measurement{now, n})
+		cmds := []tea.Cmd{tick()}
+		if m.max > 0 {
+			cmds = append(cmds, m.progress.SetPercent(float64(n)/float64(m.max)))
+		}
+		return m, tea.Batch(cmds...)
+
+	case tea.WindowSizeMsg:
+		m.progress.Width = msg.Width - 4
+		return m, nil
 	// FrameMsg is sent when the progress bar wants to animate itself
 	case progress.FrameMsg:
 		progressModel, cmd := m.progress.Update(msg)
@@ -68,9 +127,15 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	}
 }
 
+var warningStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("3")).Render
+
 func (m *model) View() string {
+	copying := ""
+	if m.copyingFile != "" {
+		copying = "Copying " + m.copyingFile + "..."
+	}
 	etaStr := "calculating..."
-	if m.measurements.Len() > 1 {
+	if m.max > 0 && m.measurements.Len() > 1 {
 		first := m.measurements.At(0)
 		last := m.measurements.At(m.measurements.Len() - 1)
 		deltaT := last.t.Sub(first.t)
@@ -80,71 +145,64 @@ func (m *model) View() string {
 		}
 	}
 	return "\n" +
+		"  " + copying + "\n" +
 		"  " + m.progress.View() + "\n" +
-		"  " + "ETA: " + etaStr + "\n"
+		"  " + "ETA: " + etaStr + "\n\n" +
+		warningStyle(strings.Join(m.errs, "\n")) + "\n"
 }
 
-func copy(out, in *os.File, progress func(int64)) error {
-	for {
-		n, err := io.CopyN(out, in, 1024*1024)
-		if n > 0 {
-			progress(n)
-		}
-		if err != nil {
-			if err == io.EOF {
-				break
-			}
-			return err
-		}
-	}
-	if err := out.Close(); err != nil {
-		return err
-	}
-	return nil
+func (m *model) Max(n int64) {
+	m.msgs <- maxMsg(n)
 }
 
-func run() error {
-	args := flag.Args()
-	if len(args) != 2 {
-		return errors.New("usage error")
-	}
-	in, err := os.Open(args[0])
-	if err != nil {
-		return err
-	}
-	defer in.Close()
-	stat, err := in.Stat()
-	if err != nil {
-		return err
-	}
-	out, err := os.OpenFile(args[1], os.O_WRONLY|os.O_CREATE|os.O_TRUNC, stat.Mode().Perm())
-	if err != nil {
-		return err
-	}
-	defer out.Close()
-	p := tea.NewProgram(&model{
-		progress: progress.New(progress.WithDefaultGradient(), progress.WithoutPercentage()),
-		max:      stat.Size() + 1,
-	}, tea.WithInput(nil), tea.WithOutput(os.Stderr))
-	var copyErr error
-	go func() {
-		if copyErr = copy(out, in, func(n int64) { p.Send(progressMsg(n)) }); copyErr != nil {
-			p.Send(tea.QuitMsg{})
-			return
-		}
-		p.Send(doneMsg{})
-	}()
-	_, teaErr := p.Run()
-	if copyErr != nil {
-		return copyErr
-	}
-	return teaErr
+func (m *model) Progress(n int64) {
+	m.current.Add(n)
+}
+
+func (m *model) FileStart(name string) {
+	m.msgs <- fileStartMsg(name)
+}
+
+func (m *model) FileDone(name string, err error) {
+	m.msgs <- fileDoneMsg{name, err}
 }
 
 func main() {
+	flag.Usage = func() {
+		fmt.Fprintf(os.Stderr, `Usage: ccp [OPTION]... SOURCE DEST
+  or:  ccp [OPTION]... SOURCE... DIRECTORY
+
+Copy SOURCE to DEST, or multiple SOURCE(s) to DIRECTORY.
+
+`)
+		flag.PrintDefaults()
+	}
 	flag.Parse()
 
-	if err := run(); err != nil {
+	args := flag.Args()
+	if len(args) < 2 {
+		fmt.Fprintln(os.Stderr, "usage error")
+		os.Exit(2)
+	}
+	srcs, dst := args[:len(args)-1], args[len(args)-1]
+	if len(srcs) == 1 {
+		if stat, err := os.Stat(dst); err == nil && stat.IsDir() {
+			dst = filepath.Join(dst, filepath.Base(srcs[0]))
+		}
+	}
+	m := &model{
+		progress:     progress.New(progress.WithDefaultGradient(), progress.WithoutPercentage()),
+		msgs:         make(chan tea.Msg),
+		copyingFiles: make(map[string]bool),
+
+		srcs: srcs,
+		dst:  dst,
+	}
+	if _, err := tea.NewProgram(m, tea.WithInput(nil), tea.WithOutput(os.Stderr)).Run(); err != nil {
 		fmt.Fprintln(os.Stderr, err)
+		os.Exit(1)
+	}
+	if len(m.errs) > 0 {
+		os.Exit(1)
 	}
 }
diff --git a/internal/cp/cp.go b/internal/cp/cp.go
new file mode 100644
index 0000000..c19627f
--- /dev/null
+++ b/internal/cp/cp.go
@@ -0,0 +1,172 @@
+package cp
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"io/fs"
+	"os"
+	"path/filepath"
+)
+
+type Progress interface {
+	Max(int64)
+	Progress(int64)
+	FileStart(string)
+	FileDone(string, error)
+}
+
+func size(files []string) int64 {
+	var n int64 = 0
+	for _, f := range files {
+		filepath.WalkDir(f, func(path string, d fs.DirEntry, err error) error {
+			if err != nil {
+				return nil
+			}
+			switch d.Type() {
+			case 0: // regular file
+				stat, err := d.Info()
+				if err != nil {
+					return nil
+				}
+				n += 1 + stat.Size()
+			case os.ModeSymlink, os.ModeDir:
+				n++
+			}
+			return nil
+		})
+	}
+	return n
+}
+
+func fileExists(path string) bool {
+	_, err := os.Lstat(path)
+	return !errors.Is(err, os.ErrNotExist)
+}
+
+func copyRegularFile(src, dst string, force bool, progress Progress) (err error) {
+	progress.FileStart(src)
+
+	in, err := os.Open(src)
+	if err != nil {
+		return err
+	}
+	defer in.Close()
+	stat, err := in.Stat()
+	if err != nil {
+		return err
+	}
+	out, err := os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, stat.Mode().Perm())
+	if err != nil {
+		if !force || !fileExists(dst) {
+			return err
+		}
+		if err := os.RemoveAll(dst); err != nil {
+			return err
+		}
+		out, err = os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, stat.Mode().Perm())
+		if err != nil {
+			return err
+		}
+	}
+	defer out.Close()
+	for {
+		n, err := io.CopyN(out, in, 1024*1024)
+		if n > 0 {
+			progress.Progress(n)
+		}
+		if err != nil {
+			if err == io.EOF {
+				break
+			}
+			return err
+		}
+	}
+	if err := out.Close(); err != nil {
+		return err
+	}
+	progress.Progress(1)
+	return nil
+}
+
+func copySpecialFile(src string, d fs.DirEntry, dst string, force bool, progress Progress) error {
+	switch d.Type() {
+	case fs.ModeSymlink:
+		target, err := os.Readlink(src)
+		if err != nil {
+			return err
+		}
+		if err := os.Symlink(target, dst); err != nil {
+			if !force || !fileExists(dst) {
+				return err
+			}
+			if err := os.RemoveAll(dst); err != nil {
+				return err
+			}
+			if err := os.Symlink(target, dst); err != nil {
+				return err
+			}
+		}
+		progress.Progress(1)
+	case fs.ModeDir:
+		stat, err := d.Info()
+		if err != nil {
+			return err
+		}
+		if err := os.Mkdir(dst, stat.Mode().Perm()); err != nil {
+			if !force || !fileExists(dst) {
+				return err
+			}
+			if err := os.RemoveAll(dst); err != nil {
+				return err
+			}
+			if err := os.Mkdir(dst, stat.Mode().Perm()); err != nil {
+				return err
+			}
+		}
+		progress.Progress(1)
+	default:
+		return fmt.Errorf("%s: unknown file type %s", src, d.Type())
+	}
+	return nil
+}
+
+func Copy(srcs []string, dstRoot string, force bool, progress Progress) {
+	go func() {
+		progress.Max(size(srcs))
+	}()
+	const maxConcurrency = 10
+	// sem acts as a semaphore to limit the number of concurrent file copies
+	sem := make(chan struct{}, maxConcurrency)
+	for _, srcRoot := range srcs {
+		filepath.WalkDir(srcRoot, func(src string, d fs.DirEntry, err error) error {
+			if err != nil {
+				progress.FileDone(src, err)
+				return nil
+			}
+			relPath, err := filepath.Rel(srcRoot, src)
+			if err != nil {
+				progress.FileDone(src, err)
+				return nil
+			}
+			dst := filepath.Join(dstRoot, relPath)
+			if d.Type().IsRegular() {
+				sem <- struct{}{}
+				go func() {
+					defer func() { <-sem }()
+					err := copyRegularFile(src, dst, force, progress)
+					progress.FileDone(src, err)
+				}()
+			} else {
+				if err := copySpecialFile(src, d, dst, force, progress); err != nil {
+					progress.FileDone(src, err)
+					return nil
+				}
+			}
+			return nil
+		})
+	}
+	for range maxConcurrency {
+		sem <- struct{}{}
+	}
+}