cli/main.go

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"
)


type Config struct {
	ServerURL string `json:"server_url"`
	APIKey    string `json:"api_key"`
}

func main() {
	if len(os.Args) < 2 {
		printUsageAndExit()
	}

	command := os.Args[1]

	// Help check
	if command == "-h" || command == "--help" || command == "help" {
		printUsageAndExit()
	}

	switch command {
	case "save":
		// User can run either: hf-save <dirs>... OR hf-save save <dirs>...
		// We'll support both, but the primary expectation is: hf-save outputs logs
		if len(os.Args) < 3 {
			fmt.Println("Error: No directories specified to save.")
			os.Exit(1)
		}
		runSave(os.Args[2:])
	case "list":
		// user runs hf-list or hf-save list
		date := ""
		if len(os.Args) > 2 {
			date = os.Args[2]
		}
		runList(date)
	case "mount":
		// hf-mount outputs OR hf-mount 2026-06-15 16-35-20_outputs
		if len(os.Args) < 3 {
			fmt.Println("Error: Please specify what to restore.")
			os.Exit(1)
		}
		runMount(os.Args[2:])
	default:
		// If command is not recognized, treat all arguments as directories to save
		// This enables direct usage like: hf-save outputs logs
		runSave(os.Args[1:])
	}
}

func printUsageAndExit() {
	fmt.Println(`hf-save: Extreme fast backup/restore tool for ephemeral GPU instances

Usage:
  hf-save <dir1> [dir2]...   Save directories to the backup server
  hf-list                    List available backup dates
  hf-list <date>             List all backups on a specific date (e.g. hf-list 2026-06-15)
  hf-mount <dir_name>        Restore latest backup containing directory name
  hf-mount <date> <run>      Restore specific backup run (e.g. hf-mount 2026-06-15 16-35-20_outputs)

Aliases / Subcommands:
  hf-save save <dir>...
  hf-save list [date]
  hf-save mount <args>...
`)
	os.Exit(0)
}

func loadConfig() Config {
	home, err := os.UserHomeDir()
	if err != nil {
		fmt.Printf("Error determining user home directory: %v\n", err)
		os.Exit(1)
	}

	configPath := filepath.Join(home, ".config", "hf-save", "config.json")
	file, err := os.Open(configPath)
	if err != nil {
		fmt.Printf("Error: Configuration file not found at %s. Please run the init script first.\n", configPath)
		os.Exit(1)
	}
	defer file.Close()

	var cfg Config
	if err := json.NewDecoder(file).Decode(&cfg); err != nil {
		fmt.Printf("Error parsing configuration: %v\n", err)
		os.Exit(1)
	}

	cfg.ServerURL = strings.TrimSuffix(cfg.ServerURL, "/")
	return cfg
}

type uploadFile struct {
	relPath string
	absPath string
}

func ensureServerAwake(serverURL string) {
	client := &http.Client{Timeout: 5 * time.Second}
	fmt.Print("Checking connection to backup server... ")

	for i := 0; i < 30; i++ {
		resp, err := client.Get(serverURL + "/health")
		if err == nil {
			if resp.StatusCode == http.StatusOK {
				resp.Body.Close()
				fmt.Println("Connected.")
				return
			}
			resp.Body.Close()
		}

		if i == 0 {
			fmt.Println("\nServer is asleep or starting up. Waking it up (this may take 1-2 minutes)...")
		}

		fmt.Printf("Waiting for server to wake up... (Attempt %d/30)\r", i+1)
		time.Sleep(5 * time.Second)
	}

	fmt.Println("\nWarning: Could not verify server is awake. Attempting operation anyway...")
}

type progressReader struct {
	r       io.Reader
	total   int64
	current int64
}

func (pr *progressReader) Read(p []byte) (n int, err error) {
	n, err = pr.r.Read(p)
	if n > 0 {
		pr.current += int64(n)
		percent := float64(pr.current) / float64(pr.total) * 100
		mbCurrent := float64(pr.current) / (1024 * 1024)
		mbTotal := float64(pr.total) / (1024 * 1024)
		fmt.Printf("Uploading: %.1f%% (%.1f/%.1f MB)\r", percent, mbCurrent, mbTotal)
	}
	return
}

type writeCounter struct {
	total   int64
	current int64
	name    string
}

func (wc *writeCounter) Write(p []byte) (int, error) {
	n := len(p)
	wc.current += int64(n)
	percent := float64(wc.current) / float64(wc.total) * 100
	mbCurrent := float64(wc.current) / (1024 * 1024)
	mbTotal := float64(wc.total) / (1024 * 1024)
	fmt.Printf("Downloading %s: %.1f%% (%.1f/%.1f MB)\r", wc.name, percent, mbCurrent, mbTotal)
	return n, nil
}

var quoteEscaper = strings.NewReplacer("\\", "\\\\", "\"", "\\\"")

func escapeQuotes(s string) string {
	return quoteEscaper.Replace(s)
}

func calculateUploadSize(prefix string, files []uploadFile, boundary string) (int64, error) {
	var totalSize int64

	// 1. Prefix field size
	var prefixBuf bytes.Buffer
	w := multipart.NewWriter(&prefixBuf)
	w.SetBoundary(boundary)
	_ = w.WriteField("prefix", prefix)
	totalSize += int64(prefixBuf.Len())

	// 2. Files size
	for _, file := range files {
		var fileHeaderBuf bytes.Buffer
		fw := multipart.NewWriter(&fileHeaderBuf)
		fw.SetBoundary(boundary)
		_, err := fw.CreateFormFile("files", file.relPath)
		if err != nil {
			return 0, err
		}
		totalSize += int64(fileHeaderBuf.Len())

		info, err := os.Stat(file.absPath)
		if err != nil {
			return 0, err
		}
		totalSize += info.Size()
		totalSize += 2 // trailing \r\n after file part data
	}

	// 3. Closing boundary size (\r\n--boundary--\r\n)
	totalSize += int64(len("\r\n--" + boundary + "--\r\n"))

	return totalSize, nil
}

// bytes helper import for size calculation
type bytesBufferWriter struct {
	bytes.Buffer
}

func runSave(targets []string) {
	cfg := loadConfig()
	ensureServerAwake(cfg.ServerURL)

	for _, target := range targets {
		saveSingleTarget(cfg, target)
	}
}

func saveSingleTarget(cfg Config, target string) {
	cleanDir := filepath.Clean(target)
	info, err := os.Stat(cleanDir)
	if err != nil {
		fmt.Printf("Error: Target '%s' not found, skipping.\n", target)
		return
	}

	var filesToUpload []uploadFile

	if !info.IsDir() {
		filesToUpload = append(filesToUpload, uploadFile{
			relPath: filepath.Base(cleanDir),
			absPath: cleanDir,
		})
	} else {
		err = filepath.Walk(cleanDir, func(path string, fileInfo os.FileInfo, walkErr error) error {
			if walkErr != nil {
				return walkErr
			}
			if fileInfo.IsDir() {
				return nil
			}

			parts := strings.Split(path, string(filepath.Separator))
			for _, part := range parts {
				if part == ".git" || part == "node_modules" || part == "__pycache__" || part == ".venv" || part == "venv" {
					return nil
				}
			}

			rel, err := filepath.Rel(filepath.Dir(cleanDir), path)
			if err == nil {
				filesToUpload = append(filesToUpload, uploadFile{
					relPath: filepath.ToSlash(rel),
					absPath: path,
				})
			}
			return nil
		})
		if err != nil {
			fmt.Printf("Error scanning target %s: %v\n", target, err)
			return
		}
	}

	if len(filesToUpload) == 0 {
		fmt.Printf("No files found to save in target %s.\n", target)
		return
	}

	dateStr := time.Now().Format("2006-01-02")
	timeStr := time.Now().Format("15-04-05")
	folderSuffix := filepath.Base(cleanDir)
	prefix := fmt.Sprintf("%s/%s_%s", dateStr, timeStr, folderSuffix)

	fmt.Printf("\n--- Saving target: %s (%d files) ---\n", target, len(filesToUpload))
	fmt.Printf("Uploading snapshot: %s\n", prefix)

	// Create a dummy multipart writer to generate a consistent boundary string
	dummyWriter := multipart.NewWriter(nil)
	boundary := dummyWriter.Boundary()

	totalSize, err := calculateUploadSize(prefix, filesToUpload, boundary)
	if err != nil {
		fmt.Printf("Failed to calculate upload size: %v\n", err)
		return
	}

	pipeReader, pipeWriter := io.Pipe()
	writer := multipart.NewWriter(pipeWriter)
	writer.SetBoundary(boundary)

	go func() {
		defer pipeWriter.Close()
		defer writer.Close()

		err := writer.WriteField("prefix", prefix)
		if err != nil {
			return
		}

		for _, file := range filesToUpload {
			fileWriter, err := writer.CreateFormFile("files", file.relPath)
			if err != nil {
				return
			}

			f, err := os.Open(file.absPath)
			if err != nil {
				return
			}

			_, err = io.Copy(fileWriter, f)
			f.Close()
			if err != nil {
				return
			}
		}
	}()

	reqURL := fmt.Sprintf("%s/upload", cfg.ServerURL)
	pr := &progressReader{
		r:     pipeReader,
		total: totalSize,
	}

	req, err := http.NewRequest(http.MethodPost, reqURL, pr)
	if err != nil {
		fmt.Printf("Error creating upload request: %v\n", err)
		return
	}

	req.Header.Set("Content-Type", writer.FormDataContentType())
	req.ContentLength = totalSize
	if cfg.APIKey != "" {
		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", cfg.APIKey))
	}

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		fmt.Printf("\nUpload failed: %v\n", err)
		return
	}
	defer resp.Body.Close()
	fmt.Println() // print newline after progress percentage

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		fmt.Printf("Upload failed with status %s: %s\n", resp.Status, string(body))
		return
	}

	var result map[string]interface{}
	if err := json.NewDecoder(resp.Body).Decode(&result); err == nil {
		fmt.Printf("Successfully saved %s! (%v files)\n", target, result["file_count"])
	} else {
		fmt.Printf("Successfully saved %s!\n", target)
	}
}

func runList(date string) {
	cfg := loadConfig()
	ensureServerAwake(cfg.ServerURL)

	reqURL := fmt.Sprintf("%s/list", cfg.ServerURL)
	if date != "" {
		reqURL += "?date=" + date
	}

	resp, err := http.Get(reqURL)
	if err != nil {
		fmt.Printf("Request failed: %v\n", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		fmt.Printf("Request failed: %s\n", resp.Status)
		os.Exit(1)
	}

	var results []string
	if err := json.NewDecoder(resp.Body).Decode(&results); err != nil {
		fmt.Printf("Failed to decode response: %v\n", err)
		os.Exit(1)
	}

	if len(results) == 0 {
		fmt.Println("No backups found.")
		return
	}

	for _, item := range results {
		fmt.Println(item)
	}
}

func runMount(args []string) {
	cfg := loadConfig()
	ensureServerAwake(cfg.ServerURL)

	var snapshot string

	if len(args) == 1 {
		folderName := args[0]
		dates, err := fetchList(cfg.ServerURL, "")
		if err != nil {
			fmt.Printf("Error: Failed to fetch date indexes: %v\n", err)
			os.Exit(1)
		}

		found := false
		for i := len(dates) - 1; i >= 0; i-- {
			runs, err := fetchList(cfg.ServerURL, dates[i])
			if err != nil {
				continue
			}

			for j := len(runs) - 1; j >= 0; j-- {
				if strings.HasSuffix(runs[j], "_"+folderName) {
					snapshot = dates[i] + "/" + runs[j]
					found = true
					break
				}
			}
			if found {
				break
			}
		}

		if !found {
			fmt.Printf("Error: No backups found matching folder name: %s\n", folderName)
			os.Exit(1)
		}

	} else if len(args) == 2 {
		snapshot = args[0] + "/" + args[1]
	} else {
		fmt.Println("Error: Invalid arguments for mount command.")
		os.Exit(1)
	}

	fmt.Printf("Mounting snapshot: %s\n", snapshot)

	manifestURL := fmt.Sprintf("%s/download?snapshot=%s", cfg.ServerURL, snapshot)
	resp, err := http.Get(manifestURL)
	if err != nil {
		fmt.Printf("Failed to download snapshot manifest: %v\n", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		fmt.Printf("Failed to fetch manifest: %s - %s\n", resp.Status, string(body))
		os.Exit(1)
	}

	var files []string
	if err := json.NewDecoder(resp.Body).Decode(&files); err != nil {
		fmt.Printf("Failed to decode snapshot list: %v\n", err)
		os.Exit(1)
	}

	if len(files) == 0 {
		fmt.Println("Snapshot is empty.")
		return
	}

	fmt.Printf("Restoring %d files...\n", len(files))

	for _, file := range files {
		targetLocalPath := filepath.FromSlash(file)
		localDir := filepath.Dir(targetLocalPath)
		if localDir != "." {
			if err := os.MkdirAll(localDir, 0755); err != nil {
				fmt.Printf("Failed to create folder %s: %v\n", localDir, err)
				os.Exit(1)
			}
		}

		fileURL := fmt.Sprintf("%s/download?file=%s/%s", cfg.ServerURL, snapshot, file)
		err = downloadFile(fileURL, targetLocalPath, filepath.Base(targetLocalPath))
		if err != nil {
			fmt.Printf("Error downloading file %s: %v\n", file, err)
			os.Exit(1)
		}
	}

	fmt.Println("Restore completed successfully!")
}

func fetchList(serverURL, date string) ([]string, error) {
	reqURL := fmt.Sprintf("%s/list", serverURL)
	if date != "" {
		reqURL += "?date=" + date
	}

	resp, err := http.Get(reqURL)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("unexpected status %s", resp.Status)
	}

	var res []string
	err = json.NewDecoder(resp.Body).Decode(&res)
	return res, err
}

func downloadFile(url, dest string, fileName string) error {
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("server returned %s - %s", resp.Status, string(body))
	}

	out, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer out.Close()

	if resp.ContentLength > 5*1024*1024 {
		counter := &writeCounter{total: resp.ContentLength, name: fileName}
		_, err = io.Copy(out, io.TeeReader(resp.Body, counter))
		fmt.Println()
	} else {
		fmt.Printf("Downloading %s...\n", fileName)
		_, err = io.Copy(out, resp.Body)
	}
	return err
}