feat: refactor upload command and implement progress tracking

cmd/root.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/krau/SaveAny-Bot/cmd/upload"
 	"github.com/krau/SaveAny-Bot/config"
 	"github.com/spf13/cobra"
 )
@@ -16,6 +17,7 @@ var rootCmd = &cobra.Command{
 
 func init() {
 	config.RegisterFlags(rootCmd)
+	upload.Register(rootCmd)
 }
 
 func Execute(ctx context.Context) {

cmd/upload/cmd.go

@@ -1,8 +1,9 @@
-package cmd
+package upload
 
 import (
 	"context"
 	"fmt"
+	"io"
 	"os"
 	"path"
 	"path/filepath"
@@ -10,6 +11,7 @@ import (
 	"github.com/charmbracelet/log"
 	"github.com/krau/SaveAny-Bot/client/bot"
 	"github.com/krau/SaveAny-Bot/common/cache"
+	"github.com/krau/SaveAny-Bot/common/utils/ioutil"
 	"github.com/krau/SaveAny-Bot/common/utils/tgutil"
 	"github.com/krau/SaveAny-Bot/config"
 	"github.com/krau/SaveAny-Bot/database"
@@ -25,13 +27,14 @@ var uploadCmd = &cobra.Command{
 	RunE:  Upload,
 }
 
-func init() {
+func Register(root *cobra.Command) {
 	uploadCmd.Flags().StringP("file", "f", "", "file path to upload")
 	uploadCmd.MarkFlagRequired("file")
 	uploadCmd.Flags().StringP("storage", "s", "", "storage name to upload to")
 	uploadCmd.MarkFlagRequired("storage")
 	uploadCmd.Flags().StringP("dir", "d", "", "storage dir to upload to, default is the base_path of the storage")
-	rootCmd.AddCommand(uploadCmd)
+	uploadCmd.Flags().Bool("no-progress", false, "disable progress bar")
+	root.AddCommand(uploadCmd)
 }
 
 func Upload(cmd *cobra.Command, args []string) error {
@@ -47,6 +50,10 @@ func Upload(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
+	noProgress, err := cmd.Flags().GetBool("no-progress")
+	if err != nil {
+		return err
+	}
 
 	ctx := cmd.Context()
 	log := log.FromContext(ctx)
@@ -86,10 +93,36 @@ func Upload(cmd *cobra.Command, args []string) error {
 	ctx = context.WithValue(ctx, ctxkey.ContentLength, fileSize)
 	ctx = tgutil.ExtWithContext(ctx, bot.ExtContext())
 
+	// Create progress reader and UI
+	var reader io.Reader
+	var progressUI *UploadProgress
 	log.Info("Uploading file...", "file", fp, "to", storname, "as", uploadPath)
-	if err := stor.Save(ctx, file, uploadPath); err != nil {
+
+	if !noProgress && fileSize > 0 {
+		progressUI = NewUploadProgress(ctx, fileName, fileSize)
+		progressUI.Start()
+
+		reader = ioutil.NewProgressReader(file, fileSize, func(read int64, total int64) {
+			if total > 0 {
+				progressUI.UpdateProgress(float64(read) / float64(total))
+			}
+		})
+	} else {
+		reader = file
+	}
+
+	if err := stor.Save(ctx, reader, uploadPath); err != nil {
+		if progressUI != nil {
+			progressUI.SetError(err)
+			progressUI.Wait()
+		}
 		log.Fatal("Failed to upload file", "error", err)
 	}
+
+	if progressUI != nil {
+		progressUI.Done()
+		progressUI.Wait()
+	}
 	log.Info("File uploaded successfully")
 	return nil
 }

cmd/upload/progress_stub.go

@@ -0,0 +1,35 @@
+//go:build no_bubbletea
+
+package upload
+
+import "context"
+
+type uploadModel struct {
+}
+
+// UploadProgress manages the progress UI for uploads
+type UploadProgress struct {
+}
+
+// NewUploadProgress creates a new upload progress tracker
+func NewUploadProgress(ctx context.Context, fileName string, fileSize int64) *UploadProgress {
+	return &UploadProgress{}
+}
+
+// Start starts the progress UI in a goroutine and returns immediately
+func (up *UploadProgress) Start() {}
+
+// UpdateProgress updates the progress bar with a new percentage (0.0 - 1.0)
+func (up *UploadProgress) UpdateProgress(percent float64) {}
+
+// SetError sets an error and quits the progress UI
+func (up *UploadProgress) SetError(err error) {}
+
+// Done signals that the upload is complete
+func (up *UploadProgress) Done() {}
+
+// Wait waits for the progress UI to finish
+func (up *UploadProgress) Wait() {}
+
+// Quit quits the progress UI
+func (up *UploadProgress) Quit() {}

cmd/upload/progress_tea.go

@@ -0,0 +1,178 @@
+//go:build !no_bubbletea
+
+package upload
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/charmbracelet/bubbles/progress"
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/lipgloss"
+	"github.com/dustin/go-humanize"
+)
+
+var (
+	helpStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("#626262"))
+)
+
+// progressMsg is sent to update the progress bar
+type progressMsg float64
+
+// progressErrMsg is sent when an error occurs
+type progressErrMsg struct{ err error }
+
+// progressDoneMsg is sent when the upload is complete
+type progressDoneMsg struct{}
+
+// uploadModel is the bubbletea model for the upload progress UI
+type uploadModel struct {
+	progress  progress.Model
+	fileName  string
+	fileSize  int64
+	bytesRead int64
+	err       error
+	done      bool
+	quitting  bool
+	width     int
+}
+
+func newUploadModel(fileName string, fileSize int64) uploadModel {
+	p := progress.New(
+		progress.WithDefaultGradient(),
+		progress.WithWidth(50),
+	)
+	return uploadModel{
+		progress: p,
+		fileName: fileName,
+		fileSize: fileSize,
+	}
+}
+
+func (m uploadModel) Init() tea.Cmd {
+	return nil
+}
+
+func (m uploadModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	switch msg := msg.(type) {
+	case tea.WindowSizeMsg:
+		m.width = msg.Width
+		m.progress.Width = min(msg.Width-10, 80)
+		return m, nil
+
+	case progressMsg:
+		var cmds []tea.Cmd
+		percent := float64(msg)
+		m.bytesRead = int64(percent * float64(m.fileSize))
+
+		cmds = append(cmds, m.progress.SetPercent(percent))
+		return m, tea.Batch(cmds...)
+
+	case progressErrMsg:
+		m.err = msg.err
+		return m, tea.Quit
+
+	case progressDoneMsg:
+		m.done = true
+		m.progress.SetPercent(1.0)
+		return m, tea.Quit
+
+	case progress.FrameMsg:
+		// Don't process frame messages if we're done or quitting
+		if m.done || m.quitting {
+			return m, nil
+		}
+		progressModel, cmd := m.progress.Update(msg)
+		m.progress = progressModel.(progress.Model)
+		return m, cmd
+	}
+
+	return m, nil
+}
+
+func (m uploadModel) View() string {
+	if m.err != nil {
+		return fmt.Sprintf("\n  ❌ Error: %s\n\n", m.err.Error())
+	}
+
+	var sb strings.Builder
+	sb.WriteString("\n")
+
+	// File info
+	sb.WriteString(fmt.Sprintf("  📁 %s\n", m.fileName))
+	sb.WriteString(fmt.Sprintf("  📊 %s / %s\n\n",
+		humanize.Bytes(uint64(m.bytesRead)),
+		humanize.Bytes(uint64(m.fileSize)),
+	))
+
+	// Progress bar
+	sb.WriteString("  ")
+	sb.WriteString(m.progress.View())
+	sb.WriteString("\n\n")
+
+	if m.done {
+		sb.WriteString("  √ Upload complete!\n\n")
+	} else {
+		sb.WriteString(helpStyle.Render("  Press Ctrl+C to cancel"))
+		sb.WriteString("\n\n")
+	}
+
+	return sb.String()
+}
+
+// UploadProgress manages the progress UI for uploads
+type UploadProgress struct {
+	program *tea.Program
+	ctx     context.Context
+	cancel  context.CancelFunc
+}
+
+// NewUploadProgress creates a new upload progress tracker
+func NewUploadProgress(ctx context.Context, fileName string, fileSize int64) *UploadProgress {
+	model := newUploadModel(fileName, fileSize)
+	ctx, cancel := context.WithCancel(ctx)
+	p := tea.NewProgram(
+		model,
+		tea.WithoutSignalHandler(),
+		tea.WithContext(ctx),
+		tea.WithInput(nil), // Disable keyboard input, rely on context cancellation
+	)
+	return &UploadProgress{
+		program: p,
+		ctx:     ctx,
+		cancel:  cancel,
+	}
+}
+
+// Start starts the progress UI in a goroutine and returns immediately
+func (up *UploadProgress) Start() {
+	go func() {
+		up.program.Run()
+	}()
+}
+
+// UpdateProgress updates the progress bar with a new percentage (0.0 - 1.0)
+func (up *UploadProgress) UpdateProgress(percent float64) {
+	up.program.Send(progressMsg(percent))
+}
+
+// SetError sets an error and quits the progress UI
+func (up *UploadProgress) SetError(err error) {
+	up.program.Send(progressErrMsg{err: err})
+}
+
+// Done signals that the upload is complete
+func (up *UploadProgress) Done() {
+	up.program.Send(progressDoneMsg{})
+}
+
+// Wait waits for the progress UI to finish
+func (up *UploadProgress) Wait() {
+	up.program.Wait()
+}
+
+// Quit quits the progress UI
+func (up *UploadProgress) Quit() {
+	up.program.Quit()
+}