SaveAny-Bot/storage/rclone/rclone.go

package rclone

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os/exec"
	"path"
	"strings"
	"time"

	"github.com/charmbracelet/log"
	"github.com/krau/SaveAny-Bot/common/utils/fsutil"
	config "github.com/krau/SaveAny-Bot/config/storage"
	"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
	storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
	"github.com/krau/SaveAny-Bot/pkg/storagetypes"
)

type Rclone struct {
	config config.RcloneStorageConfig
	logger *log.Logger
}

func (r *Rclone) Init(ctx context.Context, cfg config.StorageConfig) error {
	rcloneConfig, ok := cfg.(*config.RcloneStorageConfig)
	if !ok {
		return fmt.Errorf("failed to cast rclone config")
	}
	if err := rcloneConfig.Validate(); err != nil {
		return err
	}
	r.config = *rcloneConfig
	r.logger = log.FromContext(ctx).WithPrefix(fmt.Sprintf("rclone[%s]", r.config.Name))

	// 检查 rclone 是否安装
	if _, err := exec.LookPath("rclone"); err != nil {
		return ErrRcloneNotFound
	}

	args := r.buildBaseArgs()
	args = append(args, "listremotes")
	cmd := exec.CommandContext(ctx, "rclone", args...)
	output, err := cmd.Output()
	if err != nil {
		r.logger.Errorf("Failed to list remotes: %v", err)
		return fmt.Errorf("failed to verify rclone: %w", err)
	}

	remoteName := strings.TrimSuffix(r.config.Remote, ":")

	found := false
	scanner := bufio.NewScanner(bytes.NewReader(output))
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		line = strings.TrimSuffix(line, ":")
		if line == remoteName {
			found = true
			break
		}
	}

	if !found {
		r.logger.Errorf("Remote %s not found in rclone config", r.config.Remote)
		return ErrRemoteNotFound
	}

	r.logger.Infof("Initialized rclone storage with remote: %s", r.config.Remote)
	return nil
}

func (r *Rclone) Type() storenum.StorageType {
	return storenum.Rclone
}

func (r *Rclone) Name() string {
	return r.config.Name
}

func (r *Rclone) buildBaseArgs() []string {
	var args []string
	if r.config.ConfigPath != "" {
		args = append(args, "--config", r.config.ConfigPath)
	}
	args = append(args, r.config.Flags...)
	return args
}

func (r *Rclone) getRemotePath(storagePath string) string {
	remote := r.config.Remote
	if !strings.HasSuffix(remote, ":") {
		remote += ":"
	}
	basePath := strings.TrimPrefix(r.config.BasePath, "/")
	fullPath := path.Join(basePath, storagePath)
	return remote + fullPath
}

func (r *Rclone) Save(ctx context.Context, reader io.Reader, storagePath string) error {
	r.logger.Infof("Saving file to %s", storagePath)

	candidate := storagePath
	if overwrite, _ := ctx.Value(ctxkey.OverwriteExisting).(bool); !overwrite {
		candidate = fsutil.UniquePath("", storagePath, func(c string) bool {
			return r.Exists(ctx, c)
		}, 100)
	}

	remotePath := r.getRemotePath(candidate)
	r.logger.Debugf("Remote path: %s", remotePath)

	// Use rclone rcat to read from stdin and upload
	args := r.buildBaseArgs()
	args = append(args, "rcat", remotePath)

	cmd := exec.CommandContext(ctx, "rclone", args...)
	cmd.Stdin = reader

	var stderr bytes.Buffer
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		r.logger.Errorf("Failed to save file: %v, stderr: %s", err, stderr.String())
		return fmt.Errorf("%w: %s", ErrFailedToSaveFile, stderr.String())
	}

	r.logger.Infof("Successfully saved file to %s", candidate)
	return nil
}

func (r *Rclone) Exists(ctx context.Context, storagePath string) bool {
	remotePath := r.getRemotePath(storagePath)

	args := r.buildBaseArgs()
	args = append(args, "lsf", remotePath)

	cmd := exec.CommandContext(ctx, "rclone", args...)
	err := cmd.Run()
	return err == nil
}

// lsjsonItem represents a single entry in the output of `rclone lsjson`
type lsjsonItem struct {
	Path     string `json:"Path"`
	Name     string `json:"Name"`
	Size     int64  `json:"Size"`
	MimeType string `json:"MimeType"`
	ModTime  string `json:"ModTime"`
	IsDir    bool   `json:"IsDir"`
}

// ListFiles implements storage.StorageListable
func (r *Rclone) ListFiles(ctx context.Context, dirPath string) ([]storagetypes.FileInfo, error) {
	r.logger.Infof("Listing files in %s", dirPath)

	remotePath := r.getRemotePath(dirPath)

	args := r.buildBaseArgs()
	args = append(args, "lsjson", remotePath)

	cmd := exec.CommandContext(ctx, "rclone", args...)
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		r.logger.Errorf("Failed to list files: %v, stderr: %s", err, stderr.String())
		return nil, fmt.Errorf("%w: %s", ErrFailedToListFiles, stderr.String())
	}

	var items []lsjsonItem
	if err := json.Unmarshal(stdout.Bytes(), &items); err != nil {
		r.logger.Errorf("Failed to parse lsjson output: %v", err)
		return nil, fmt.Errorf("failed to parse lsjson output: %w", err)
	}

	files := make([]storagetypes.FileInfo, 0, len(items))
	for _, item := range items {
		var modTime time.Time
		if item.ModTime != "" {
			parsedTime, err := time.Parse(time.RFC3339Nano, item.ModTime)
			if err != nil {
				r.logger.Warnf("Failed to parse mod time %q for %s: %v", item.ModTime, item.Name, err)
			} else {
				modTime = parsedTime
			}
		}

		files = append(files, storagetypes.FileInfo{
			Name:    item.Name,
			Path:    path.Join(dirPath, item.Name),
			Size:    item.Size,
			IsDir:   item.IsDir,
			ModTime: modTime,
		})
	}

	r.logger.Debugf("Found %d files/directories in %s", len(files), dirPath)
	return files, nil
}

// OpenFile implements storage.StorageReadable
func (r *Rclone) OpenFile(ctx context.Context, filePath string) (io.ReadCloser, int64, error) {
	r.logger.Infof("Opening file %s", filePath)

	remotePath := r.getRemotePath(filePath)

	size, err := r.getFileSize(ctx, remotePath)
	if err != nil {
		r.logger.Errorf("Failed to get file size: %v", err)
		return nil, 0, fmt.Errorf("%w: %v", ErrFailedToOpenFile, err)
	}

	args := r.buildBaseArgs()
	args = append(args, "cat", remotePath)

	cmd := exec.CommandContext(ctx, "rclone", args...)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, 0, fmt.Errorf("failed to create stdout pipe: %w", err)
	}

	if err := cmd.Start(); err != nil {
		return nil, 0, fmt.Errorf("failed to start rclone cat: %w", err)
	}

	reader := &rcloneCatReader{
		reader: stdout,
		cmd:    cmd,
		logger: r.logger,
	}

	r.logger.Debugf("Opened file %s (size: %d bytes)", filePath, size)
	return reader, size, nil
}

func (r *Rclone) getFileSize(ctx context.Context, remotePath string) (int64, error) {
	args := r.buildBaseArgs()
	args = append(args, "lsjson", remotePath)

	cmd := exec.CommandContext(ctx, "rclone", args...)
	var stdout bytes.Buffer
	cmd.Stdout = &stdout

	if err := cmd.Run(); err != nil {
		return 0, err
	}

	var items []lsjsonItem
	if err := json.Unmarshal(stdout.Bytes(), &items); err != nil {
		return 0, err
	}

	if len(items) > 0 {
		return items[0].Size, nil
	}
	return 0, nil
}

type rcloneCatReader struct {
	reader io.ReadCloser
	cmd    *exec.Cmd
	logger *log.Logger
}

func (r *rcloneCatReader) Read(p []byte) (n int, err error) {
	return r.reader.Read(p)
}

func (r *rcloneCatReader) Close() error {
	if err := r.reader.Close(); err != nil {
		r.logger.Warnf("Failed to close reader: %v", err)
	}
	// Kill the cat process so Wait cannot block on a hung pipe.
	if r.cmd.Process != nil {
		if err := r.cmd.Process.Kill(); err != nil {
			r.logger.Warnf("Failed to kill rclone cat process: %v", err)
		}
	}
	if err := r.cmd.Wait(); err != nil {
		r.logger.Warnf("rclone cat process exited with error: %v", err)
	}
	return nil
}