nats-upload/main.go

package main

import (
	"context"
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/nats-io/nats.go"
	"github.com/nats-io/nats.go/jetstream"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"golang.org/x/mod/semver"
)

type Config struct {
	NatsURL         string `mapstructure:"nats"`
	BucketName      string `mapstructure:"bucket"`
	Directory       string `mapstructure:"dir"`
	Prefix          string `mapstructure:"prefix"`
	BinaryName      string `mapstructure:"binary"`
	NotifyTopic     string `mapstructure:"notify"`
	SkipNotify      bool   `mapstructure:"skip-notify"`
	Cleanup         int    `mapstructure:"cleanup"`
	CleanupAll      bool   `mapstructure:"cleanup-all"`
	GiteaURL        string `mapstructure:"gitea-url"`
	GiteaToken      string `mapstructure:"gitea-token"`
	CleanupFeatures bool   `mapstructure:"cleanup-features"`
	Repository      string `mapstructure:"repository"`
}

var rootCmd = &cobra.Command{
	Use:   "nats-upload",
	Short: "Upload binaries to NATS object store and cleanup old versions",
	RunE: func(cmd *cobra.Command, args []string) error {
		var cfg Config
		if err := viper.Unmarshal(&cfg); err != nil {
			return fmt.Errorf("failed to unmarshal config: %w", err)
		}
		if cfg.Directory == "" && cfg.Cleanup == 0 && !cfg.CleanupFeatures {
			return errors.New("directory path is required, or cleanup/cleanup-features must be enabled")
		}
		return runUploadAndCleanup(cmd.Context(), &cfg)
	},
}

var cleanCmd = &cobra.Command{
	Use:   "clean",
	Short: "Cleanup old versions in NATS object store",
	RunE: func(cmd *cobra.Command, args []string) error {
		var cfg Config
		if err := viper.Unmarshal(&cfg); err != nil {
			return fmt.Errorf("failed to unmarshal config: %w", err)
		}
		if cfg.Cleanup == 0 {
			return errors.New("cleanup count must be greater than 0")
		}
		return runCleanupOnly(cmd.Context(), &cfg)
	},
}

var cleanFeaturesCmd = &cobra.Command{
	Use:   "clean-features",
	Short: "Delete feature-branch binaries whose branches no longer exist in Gitea",
	RunE: func(cmd *cobra.Command, args []string) error {
		var cfg Config
		if err := viper.Unmarshal(&cfg); err != nil {
			return fmt.Errorf("failed to unmarshal config: %w", err)
		}
		return runCleanupFeatures(cmd.Context(), &cfg)
	},
}

func init() {
	cobra.OnInitialize(initConfig)

	rootCmd.PersistentFlags().String("nats", "nats://localhost:4222", "NATS server URL")
	rootCmd.PersistentFlags().String("bucket", "binaries", "Object store bucket name")
	rootCmd.PersistentFlags().String("binary", "", "Binary name (defaults to first binary found)")
	rootCmd.PersistentFlags().Int("cleanup", 2, "Keep only N most recent versions (0 disables cleanup)")
	rootCmd.PersistentFlags().Bool("cleanup-all", false, "Cleanup all binaries, not just current one")
	rootCmd.PersistentFlags().Bool("clean-all", false, "Alias for --cleanup-all")
	rootCmd.PersistentFlags().String("gitea-url", "", "Gitea server URL (for cleaning feature branches)")
	rootCmd.PersistentFlags().String("gitea-token", "", "Gitea API token (for cleaning feature branches)")
	rootCmd.PersistentFlags().Bool("cleanup-features", false, "Cleanup stale feature-branch binaries")
	rootCmd.PersistentFlags().String("repository", "", "Repository path (owner/repo) for Gitea branch API")

	rootCmd.Flags().String("dir", "upload", "Directory containing binaries to upload")
	rootCmd.Flags().String("prefix", "", "Prefix to strip from paths (like 'upload/')")
	rootCmd.Flags().String("notify", "binaries.update", "NATS topic to publish update notification")
	rootCmd.Flags().Bool("skip-notify", false, "Skip publishing update notification")

	rootPersistentFlags := rootCmd.PersistentFlags()
	for _, name := range []string{"nats", "bucket", "binary", "cleanup", "cleanup-all", "gitea-url", "gitea-token", "cleanup-features", "repository"} {
		if err := viper.BindPFlag(name, rootPersistentFlags.Lookup(name)); err != nil {
			log.Fatalf("error binding %s flag: %v", name, err)
		}
	}

	rootFlags := rootCmd.Flags()
	for _, name := range []string{"dir", "prefix", "notify", "skip-notify"} {
		if err := viper.BindPFlag(name, rootFlags.Lookup(name)); err != nil {
			log.Fatalf("error binding %s flag: %v", name, err)
		}
	}

	rootCmd.AddCommand(cleanCmd, cleanFeaturesCmd)
}

func initConfig() {
	viper.SetEnvPrefix("INPUT")
	viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_", ".", "_"))
	viper.AutomaticEnv()

	viper.RegisterAlias("nats_url", "nats")
	viper.RegisterAlias("source", "dir")
	viper.RegisterAlias("strip_prefix", "prefix")
	viper.RegisterAlias("notify_topic", "notify")
	viper.RegisterAlias("clean_all", "cleanup-all")
	viper.RegisterAlias("gitea_url", "gitea-url")
	viper.RegisterAlias("gitea_token", "gitea-token")
	viper.RegisterAlias("cleanup_features", "cleanup-features")
	viper.RegisterAlias("repository", "repository")

	_ = viper.BindEnv("nats", "INPUT_NATS", "INPUT_NATS_URL")
	_ = viper.BindEnv("dir", "INPUT_DIR", "INPUT_SOURCE")
	_ = viper.BindEnv("prefix", "INPUT_PREFIX", "INPUT_STRIP_PREFIX")
	_ = viper.BindEnv("notify", "INPUT_NOTIFY", "INPUT_NOTIFY_TOPIC")
	_ = viper.BindEnv("cleanup-all", "INPUT_CLEANUP_ALL", "INPUT_CLEAN_ALL")
	_ = viper.BindEnv("gitea-url", "INPUT_GITEA_URL", "INPUT_GITEA_URL")
	_ = viper.BindEnv("gitea-token", "INPUT_GITEA_TOKEN", "INPUT_GITEA_TOKEN")
	_ = viper.BindEnv("cleanup-features", "INPUT_CLEANUP_FEATURES", "INPUT_CLEANUP_FEATURES")
	_ = viper.BindEnv("repository", "INPUT_REPOSITORY", "INPUT_REPOSITORY", "GITHUB_REPOSITORY")
}

type NATSClient struct {
	Conn  *nats.Conn
	JS    jetstream.JetStream
	Store jetstream.ObjectStore
}

func getNATSConnection(ctx context.Context, cfg *Config) (*NATSClient, error) {
	nc, err := nats.Connect(cfg.NatsURL)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to NATS: %w", err)
	}

	js, err := jetstream.New(nc)
	if err != nil {
		nc.Close()
		return nil, fmt.Errorf("failed to create JetStream context: %w", err)
	}

	store, err := js.ObjectStore(ctx, cfg.BucketName)
	if err != nil {
		store, err = js.CreateObjectStore(ctx, jetstream.ObjectStoreConfig{
			Bucket:      cfg.BucketName,
			Description: "Binary storage for self-update",
		})
		if err != nil {
			nc.Close()
			return nil, fmt.Errorf("failed to get/create object store: %w", err)
		}
		log.Printf("Created object store: %s", cfg.BucketName)
	}

	return &NATSClient{
		Conn:  nc,
		JS:    js,
		Store: store,
	}, nil
}

func runUploadAndCleanup(ctx context.Context, cfg *Config) error {
	client, err := getNATSConnection(ctx, cfg)
	if err != nil {
		return err
	}
	defer client.Conn.Close()

	if cfg.Directory != "" {
		err := filepath.Walk(cfg.Directory, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}

			if info.IsDir() {
				return nil
			}

			data, err := os.ReadFile(path)
			if err != nil {
				return fmt.Errorf("failed to read %s: %w", path, err)
			}

			relPath, err := filepath.Rel(cfg.Directory, path)
			if err != nil {
				return fmt.Errorf("failed to get relative path: %w", err)
			}

			objectKey := relPath
			if cfg.Prefix != "" {
				objectKey = strings.TrimPrefix(relPath, cfg.Prefix)
			}

			objectKey = filepath.ToSlash(objectKey)

			if cfg.BinaryName == "" {
				parts := strings.Split(objectKey, "/")
				if len(parts) >= 2 {
					cfg.BinaryName = parts[0]
				}
			}

			log.Printf("Uploading %s as %s (%d bytes)", path, objectKey, len(data))

			_, err = client.Store.PutBytes(ctx, objectKey, data)
			if err != nil {
				return fmt.Errorf("failed to upload %s: %w", path, err)
			}

			log.Printf("✓ Uploaded %s", objectKey)
			return nil
		})

		if err != nil {
			return fmt.Errorf("failed to upload files: %w", err)
		}

		log.Printf("Successfully uploaded all files from %s to NATS object store '%s'", cfg.Directory, cfg.BucketName)
	}

	if cfg.Cleanup > 0 {
		log.Printf("Cleaning up old versions, keeping %d most recent", cfg.Cleanup)
		err := cleanupOldVersions(ctx, client.Store, cfg.BinaryName, cfg.Cleanup, cfg.CleanupAll)
		if err != nil {
			return fmt.Errorf("failed to cleanup old versions: %w", err)
		}
	}

	if cfg.CleanupFeatures {
		log.Printf("Cleaning up stale feature branch binaries...")
		if err := runCleanupFeatures(ctx, cfg); err != nil {
			return fmt.Errorf("failed to cleanup feature branches: %w", err)
		}
	}

	if !cfg.SkipNotify && cfg.NotifyTopic != "" {
		log.Printf("Publishing update notification to topic: %s", cfg.NotifyTopic)

		message := fmt.Sprintf("binaries updated in %s", cfg.BucketName)
		err := client.Conn.Publish(cfg.NotifyTopic, []byte(message))
		if err != nil {
			return fmt.Errorf("failed to publish notification: %w", err)
		}

		err = client.Conn.Flush()
		if err != nil {
			return fmt.Errorf("failed to flush notification: %w", err)
		}

		log.Printf("✓ Published update notification")
	}
	return nil
}

func runCleanupOnly(ctx context.Context, cfg *Config) error {
	client, err := getNATSConnection(ctx, cfg)
	if err != nil {
		return err
	}
	defer client.Conn.Close()

	log.Printf("Cleaning up old versions, keeping %d most recent", cfg.Cleanup)
	err = cleanupOldVersions(ctx, client.Store, cfg.BinaryName, cfg.Cleanup, cfg.CleanupAll)
	if err != nil {
		return fmt.Errorf("failed to cleanup old versions: %w", err)
	}
	return nil
}

func runCleanupFeatures(ctx context.Context, cfg *Config) error {
	client, err := getNATSConnection(ctx, cfg)
	if err != nil {
		return err
	}
	defer client.Conn.Close()

	log.Printf("Cleaning up stale feature branch binaries from bucket %s", cfg.BucketName)

	objects, err := client.Store.List(ctx)
	if err != nil {
		return fmt.Errorf("failed to list objects: %w", err)
	}

	type featureTag struct {
		binary    string
		objectKey string
		version   string
	}
	var featureObjects []featureTag

	for _, obj := range objects {
		parts := strings.Split(obj.Name, "/")
		if len(parts) < 3 {
			continue
		}

		version := parts[len(parts)-1]
		if !strings.HasPrefix(version, "feature-") {
			continue
		}

		featureObjects = append(featureObjects, featureTag{
			binary:    parts[0],
			objectKey: obj.Name,
			version:   version,
		})
	}

	if len(featureObjects) == 0 {
		log.Printf("No feature branch objects found")
		return nil
	}

	log.Printf("Found %d feature branch objects, checking active feature branches via git...", len(featureObjects))

	// Use git ls-remote to list remote branches (no token needed)
	cmd := exec.CommandContext(ctx, "git", "ls-remote", "--heads", "origin")
	output, err := cmd.Output()
	if err != nil {
		return fmt.Errorf("failed to list remote branches via git: %w", err)
	}

	activeBranches := make(map[string]bool)
	for _, line := range strings.Split(string(output), "\n") {
		// Format: "<sha>\trefs/heads/<branch>"
		parts := strings.Split(line, "\t")
		if len(parts) < 2 {
			continue
		}
		ref := parts[len(parts)-1]
		branchName := strings.TrimPrefix(ref, "refs/heads/")
		if strings.HasPrefix(branchName, "feature/") {
			activeBranches[branchName] = true
		}
	}

	log.Printf("Active feature branches: %d", len(activeBranches))
	for branch := range activeBranches {
		log.Printf("  - %s", branch)
	}

	var deleted int
	for _, fo := range featureObjects {
		branchPart := strings.TrimPrefix(fo.version, "feature-")
		branchName := "feature/" + strings.ReplaceAll(branchPart, "--", "/")

		if !activeBranches[branchName] {
			log.Printf("Deleting stale: %s (branch %s no longer exists)", fo.objectKey, branchName)
			err := client.Store.Delete(ctx, fo.objectKey)
			if err != nil {
				log.Printf("Failed to delete %s: %v", fo.objectKey, err)
			} else {
				deleted++
			}
		}
	}

	log.Printf("Cleanup complete: deleted %d stale feature branch objects", deleted)
	return nil
}
func cleanupOldVersions(ctx context.Context, store jetstream.ObjectStore, currentBinary string, keepCount int, cleanAll bool) error {
	objects, err := store.List(ctx)
	if err != nil {
		return fmt.Errorf("failed to list objects: %w", err)
	}

	versionsByPath := make(map[string][]*jetstream.ObjectInfo)

	for _, obj := range objects {
		parts := strings.Split(obj.Name, "/")
		if len(parts) < 3 {
			continue
		}

		binName := parts[0]
		arch := parts[1]
		pathKey := binName + "/" + arch

		if !cleanAll && currentBinary != "" && binName != currentBinary {
			continue
		}

		versionsByPath[pathKey] = append(versionsByPath[pathKey], obj)
	}

	for pathKey, versions := range versionsByPath {
		if len(versions) <= keepCount {
			log.Printf("Path %s has %d versions, keeping all", pathKey, len(versions))
			continue
		}

		sort.Slice(versions, func(i, j int) bool {
			versionI := filepath.Base(versions[i].Name)
			versionJ := filepath.Base(versions[j].Name)

			if !strings.HasPrefix(versionI, "v") {
				versionI = "v" + versionI
			}
			if !strings.HasPrefix(versionJ, "v") {
				versionJ = "v" + versionJ
			}

			return semver.Compare(versionI, versionJ) > 0
		})

		toDelete := versions[keepCount:]
		log.Printf("Path %s has %d versions, deleting %d old versions", pathKey, len(versions), len(toDelete))

		for _, obj := range toDelete {
			log.Printf("Deleting old version: %s", obj.Name)
			err := store.Delete(ctx, obj.Name)
			if err != nil && !errors.Is(err, jetstream.ErrObjectNotFound) {
				return fmt.Errorf("failed to delete %s: %w", obj.Name, err)
			}
			log.Printf("✓ Deleted %s", obj.Name)
		}
	}

	return nil
}
func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
	defer cancel()

	if err := rootCmd.ExecuteContext(ctx); err != nil {
		_, _ = fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}