package main

import (
	"context"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"

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

func main() {
	if len(os.Args) < 2 {
		_, _ = fmt.Fprintf(os.Stderr, "Usage: %s <binary-name> [args...]\n", os.Args[0])
		os.Exit(1)
	}

	binaryName := os.Args[1]
	remainingArgs := os.Args[2:]

	natsURL := getEnv("NATS_URL", "nats://localhost:4222")
	bucketName := getEnv("NATS_BUCKET", "binaries")
	cacheDir := getEnv("NATS_CACHE_DIR", filepath.Join(os.TempDir(), "nats-exec-cache"))

	ctx := context.Background()

	nc, err := nats.Connect(natsURL)
	if err != nil {
		log.Fatalf("Failed to connect to NATS: %v", err)
	}
	defer nc.Close()

	js, err := jetstream.New(nc)
	if err != nil {
		log.Fatalf("Failed to create JetStream context: %v", err)
	}

	store, err := js.ObjectStore(ctx, bucketName)
	if err != nil {
		log.Fatalf("Failed to get object store %s: %v", bucketName, err)
	}

	// Format: binary/arch/version
	arch := runtime.GOARCH

	objects, err := store.List(ctx)
	if err != nil {
		log.Fatalf("Failed to list objects: %v", err)
	}

	var latestVersion string
	var latestKey string

	for _, info := range objects {
		// Key structure: <binaryName>/<arch>/<version>

		parts := strings.Split(info.Name, "/")
		if len(parts) != 3 {
			continue
		}
		if parts[0] != binaryName {
			continue
		}
		if parts[1] != arch {
			continue
		}

		version := parts[2]
		if !strings.HasPrefix(version, "v") {
			version = "v" + version
		}

		if semver.IsValid(version) {
			if latestVersion == "" || semver.Compare(version, latestVersion) > 0 {
				latestVersion = version
				latestKey = info.Name
			}
		}
	}

	if latestKey == "" {
		log.Fatalf("No version of %s found for arch %s", binaryName, arch)
	}

	localPath := filepath.Join(cacheDir, latestKey)

	if _, err := os.Stat(localPath); os.IsNotExist(err) {
		log.Printf("Downloading %s version %s...", binaryName, latestVersion)
		if err := downloadBinary(ctx, store, latestKey, localPath); err != nil {
			log.Fatalf("Failed to download binary: %v", err)
		}
	} else {
		log.Printf("Using cached %s version %s", binaryName, latestVersion)
	}

	if err := os.Chmod(localPath, 0755); err != nil {
		log.Fatalf("Failed to make binary executable: %v", err)
	}

	fullPath, err := filepath.Abs(localPath)
	if err != nil {
		log.Fatalf("Failed to get absolute path: %v", err)
	}

	env := os.Environ()
	args := append([]string{fullPath}, remainingArgs...)

	err = syscall.Exec(fullPath, args, env)
	if err != nil {
		log.Fatalf("Failed to exec %s: %v", fullPath, err)
	}
}

func downloadBinary(ctx context.Context, store jetstream.ObjectStore, key, localPath string) error {
	if err := os.MkdirAll(filepath.Dir(localPath), 0755); err != nil {
		return err
	}

	obj, err := store.Get(ctx, key)
	if err != nil {
		return err
	}

	f, err := os.OpenFile(localPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
	if err != nil {
		return err
	}
	defer f.Close()

	_, err = io.Copy(f, obj)
	return err
}

func getEnv(key, defaultValue string) string {
	if value, ok := os.LookupEnv(key); ok {
		return value
	}
	return defaultValue
}