// proxmox-backup-watch
//
// Waits for the vzdump task of the given VMID(s) to finish, then restarts the
// guest (QEMU VM or LXC container). By default it restarts regardless of the
// backup result; use -only-failed to restart only when the backup did not
// finish OK.
//
// Usage: proxmox-backup-watch [options] VMID [VMID...]
package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

// Set at build time with -ldflags (see Makefile).
var (
	version = "dev"
	commit  = "none"
	date    = "unknown"
)

type task struct {
	ID        string `json:"id"`
	StartTime int64  `json:"starttime"`
	EndTime   *int64 `json:"endtime"`
	Status    string `json:"status"`
	UPID      string `json:"upid"`
}

type config struct {
	node       string
	dryRun     bool
	timeout    time.Duration
	interval   time.Duration
	onlyFailed bool
	warnOK     bool
	startedAt  int64
}

func main() {
	os.Exit(run())
}

func run() int {
	cfg := config{startedAt: time.Now().Unix()}

	fs := flag.NewFlagSet("proxmox-backup-watch", flag.ContinueOnError)
	fs.BoolVar(&cfg.dryRun, "dry-run", false, "show what would be done, don't restart anything")
	fs.BoolVar(&cfg.dryRun, "n", false, "shorthand for -dry-run")
	fs.DurationVar(&cfg.timeout, "timeout", 6*time.Hour, "max time to wait per VMID (0 = forever)")
	fs.DurationVar(&cfg.timeout, "t", 6*time.Hour, "shorthand for -timeout")
	fs.DurationVar(&cfg.interval, "interval", 30*time.Second, "poll interval")
	fs.DurationVar(&cfg.interval, "i", 30*time.Second, "shorthand for -interval")
	fs.BoolVar(&cfg.onlyFailed, "only-failed", false, "restart only if the backup did NOT finish OK")
	fs.BoolVar(&cfg.onlyFailed, "f", false, "shorthand for -only-failed")
	fs.BoolVar(&cfg.warnOK, "warn-ok", false, `with -only-failed, treat "WARNINGS: n" as success`)
	fs.BoolVar(&cfg.warnOK, "w", false, "shorthand for -warn-ok")
	fs.StringVar(&cfg.node, "node", "", "Proxmox node name (default: short hostname)")
	var showVersion bool
	fs.BoolVar(&showVersion, "version", false, "print version and exit")
	fs.BoolVar(&showVersion, "v", false, "shorthand for -version")
	fs.Usage = func() {
		fmt.Fprintln(os.Stderr, "Usage: proxmox-backup-watch [options] VMID [VMID...]")
		fs.PrintDefaults()
	}

	// Allow flags before, between and after VMIDs.
	var vmids []string
	args := os.Args[1:]
	for {
		if err := fs.Parse(args); err != nil {
			if errors.Is(err, flag.ErrHelp) {
				return 0
			}
			return 2
		}
		rest := fs.Args()
		if len(rest) == 0 {
			break
		}
		vmids = append(vmids, rest[0])
		args = rest[1:]
	}

	if showVersion {
		fmt.Printf("proxmox-backup-watch %s (commit %s, built %s)\n", version, commit, date)
		return 0
	}

	if len(vmids) == 0 {
		fmt.Fprintln(os.Stderr, "Error: at least one VMID is required")
		fs.Usage()
		return 2
	}
	for _, id := range vmids {
		if !isNumeric(id) {
			fmt.Fprintf(os.Stderr, "Error: invalid VMID %q\n", id)
			return 2
		}
	}

	if cfg.node == "" {
		h, err := os.Hostname()
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: cannot determine hostname: %v\n", err)
			return 1
		}
		cfg.node, _, _ = strings.Cut(h, ".")
	}

	fmt.Printf("proxmox-backup-watch %s\n", version)

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	failed := false
	for _, vmid := range vmids {
		if err := handle(ctx, &cfg, vmid); err != nil {
			fmt.Fprintf(os.Stderr, "VM %s: %v\n", vmid, err)
			failed = true
			if ctx.Err() != nil {
				break
			}
		}
	}
	if failed {
		return 1
	}
	return 0
}

func handle(ctx context.Context, cfg *config, vmid string) error {
	fmt.Printf("VM %s: waiting for backup task to finish...\n", vmid)

	t, err := waitForTask(ctx, cfg, vmid)
	if err != nil {
		return err
	}

	if cfg.onlyFailed {
		ok := t.Status == "OK" || (cfg.warnOK && strings.HasPrefix(t.Status, "WARNINGS"))
		if ok {
			fmt.Printf("VM %s: backup finished (%s), no restart needed\n", vmid, t.Status)
			return nil
		}
	}

	fmt.Printf("VM %s: backup finished (%s) -> restarting\n", vmid, t.Status)
	return restartGuest(ctx, cfg, vmid)
}

// waitForTask polls until the newest vzdump task for vmid (running now, or
// finished after this program started) has an end time.
func waitForTask(ctx context.Context, cfg *config, vmid string) (*task, error) {
	var deadline time.Time
	if cfg.timeout > 0 {
		deadline = time.Now().Add(cfg.timeout)
	}

	for {
		t, err := latestTask(ctx, cfg, vmid)
		if err != nil {
			// Transient pvesh/JSON errors shouldn't abort a long wait.
			fmt.Fprintf(os.Stderr, "VM %s: warning: %v\n", vmid, err)
		} else if t != nil && t.EndTime != nil {
			return t, nil
		}

		if !deadline.IsZero() && time.Now().After(deadline) {
			return nil, fmt.Errorf("timed out after %s, no restart", cfg.timeout)
		}

		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-time.After(cfg.interval):
		}
	}
}

func latestTask(ctx context.Context, cfg *config, vmid string) (*task, error) {
	out, err := output(ctx, "pvesh", "get", "/nodes/"+cfg.node+"/tasks",
		"--typefilter", "vzdump", "--limit", "500", "--output-format", "json")
	if err != nil {
		return nil, err
	}

	var tasks []task
	if err := json.Unmarshal(out, &tasks); err != nil {
		return nil, fmt.Errorf("parsing pvesh output: %w", err)
	}

	var latest *task
	for i := range tasks {
		t := &tasks[i]
		if t.ID != vmid {
			continue
		}
		// Running (no endtime) or finished after we started.
		if t.EndTime != nil && *t.EndTime < cfg.startedAt {
			continue
		}
		if latest == nil || t.StartTime > latest.StartTime {
			latest = t
		}
	}
	return latest, nil
}

func restartGuest(ctx context.Context, cfg *config, vmid string) error {
	var tool string
	switch {
	case exec.CommandContext(ctx, "qm", "status", vmid).Run() == nil:
		tool = "qm"
	case exec.CommandContext(ctx, "pct", "status", vmid).Run() == nil:
		tool = "pct"
	default:
		return errors.New("not found on this node")
	}

	// Clear a leftover 'backup' lock; failure here is not fatal.
	_ = runCmd(ctx, cfg, tool, "unlock", vmid)

	if err := runCmd(ctx, cfg, tool, "shutdown", vmid, "--timeout", "120", "--forceStop", "1"); err != nil {
		return fmt.Errorf("shutdown failed: %w", err)
	}
	if err := runCmd(ctx, cfg, tool, "start", vmid); err != nil {
		return fmt.Errorf("start failed: %w", err)
	}
	return nil
}

func runCmd(ctx context.Context, cfg *config, name string, args ...string) error {
	if cfg.dryRun {
		fmt.Printf("[dry-run] %s %s\n", name, strings.Join(args, " "))
		return nil
	}
	cmd := exec.CommandContext(ctx, name, args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func output(ctx context.Context, name string, args ...string) ([]byte, error) {
	var stdout, stderr bytes.Buffer
	cmd := exec.CommandContext(ctx, name, args...)
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf("%s: %w: %s", name, err, strings.TrimSpace(stderr.String()))
	}
	return stdout.Bytes(), nil
}

func isNumeric(s string) bool {
	if s == "" {
		return false
	}
	for _, r := range s {
		if r < '0' || r > '9' {
			return false
		}
	}
	return true
}