package main

import (
	"flag"
	"fmt"
	"os"
	"runtime"
	"strings"
)

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

	var err error
	switch os.Args[1] {
	case "sync":
		err = cmdSync(os.Args[2:])
	case "clone-disk":
		err = cmdCloneDisk(os.Args[2:])
	case "agent":
		err = cmdAgent(os.Args[2:])
	case "version", "--version":
		fmt.Printf("clonetool %s/%s\n", runtime.GOOS, runtime.GOARCH)
		return
	case "-h", "--help", "help":
		usage()
		return
	default:
		usage()
		os.Exit(2)
	}
	if err != nil {
		fmt.Fprintln(os.Stderr, "clonetool: "+err.Error())
		os.Exit(1)
	}
}

func usage() {
	fmt.Fprint(os.Stderr, `clonetool - block-level file/device sync

Usage:
  clonetool sync --source LOC --dest LOC [options]
  clonetool clone-disk --source LOC --dest LOC [options]
  clonetool version
  clonetool agent --role {control|sink|source-stream|fs-send|fs-recv} ...   (internal, spawned automatically)

LOC is either a local path, or [user@]host:path for a path reached over SSH.
Source, destination, and the machine running "sync" (the manager) may all be
different machines: the manager only orchestrates, it never reads or writes
a single block itself.

clonetool keeps no state between runs: every sync re-reads and re-hashes both
the source and the destination and transfers only the blocks that differ.

Options for clone-disk (whole-disk: boot record + partition table + per-filesystem data):
  --parts LIST            only clone these partition numbers (comma-separated; default all)
  --raw LIST              force a raw block clone for these partitions
  --file-level LIST       mkfs + rsync these partitions instead of a filesystem image
  --file-level-auto       use file-level for filesystem types with no image cloner
  --allow-shrink          permit shrinking trailing partitions to fit a smaller target
                          (resizes the SOURCE filesystem in place before imaging)
  --no-shrink             never shrink; fail instead if the target is too small
  --new-ids               randomize the GPT disk GUID / MBR signature on the target
  --reinstall-bootloader  after copy, run grub-install / grub-mkconfig on the target (Linux)
  --vss                   Windows source: take a Volume Shadow Copy per NTFS volume (default on)
  --image-size SIZE       file target: image size (default = enough for the layout)
  (also accepts --block-size --job --yes --sudo --deploy --ssh --ssh-opt --remote-bin
   --connect-timeout --manager-host, same meaning as sync)

Options for sync:
  --block-size SIZE       block size, e.g. 4M (default 4M)
  --job NAME             optional label shown in progress/log output
  --yes                   don't prompt before shrinking an existing destination file
  --sudo MODE             device-access privilege escalation: auto (escalate on a
                          permission error, default), always, or never
  --deploy               copy this binary to remote hosts that lack it
                          (default true; --deploy=false to disable)
  --connect-timeout SEC   ssh connect timeout for the push/pull direction probe (default 8)
  --ssh PATH              ssh binary to use (default "ssh")
  --ssh-opt OPT           extra "-o OPT" passed to ssh (repeatable)
  --remote-bin PATH       path to clonetool on remote hosts (default "clonetool")
  --manager-host HOST     address peers should use to reach this machine, when
                          source or dest has no host part (defaults to the local hostname)
`)
}

type stringSlice []string

func (s *stringSlice) String() string { return fmt.Sprint([]string(*s)) }
func (s *stringSlice) Set(v string) error {
	*s = append(*s, v)
	return nil
}

func cmdSync(args []string) error {
	fs := flag.NewFlagSet("sync", flag.ContinueOnError)
	job := fs.String("job", "", "optional label for progress/log output")
	source := fs.String("source", "", "source location (required)")
	dest := fs.String("dest", "", "destination location (required)")
	blockSizeStr := fs.String("block-size", "4M", "block size, e.g. 4M")
	yes := fs.Bool("yes", false, "don't prompt before shrinking destination")
	sudoMode := fs.String("sudo", "auto", "device-access privilege escalation: auto|always|never")
	deploy := fs.Bool("deploy", true, "copy this binary to remote hosts that lack it")
	connectTimeout := fs.Int("connect-timeout", defaultConnectTimeoutSec, "ssh connect timeout (seconds)")
	sshBin := fs.String("ssh", "ssh", "ssh binary")
	remoteBin := fs.String("remote-bin", "clonetool", "clonetool path on remote hosts")
	managerHost := fs.String("manager-host", "", "address peers use to reach this machine")
	var sshOpts stringSlice
	fs.Var(&sshOpts, "ssh-opt", `extra "-o OPT" passed to ssh (repeatable)`)
	if err := fs.Parse(args); err != nil {
		return err
	}

	if *source == "" || *dest == "" {
		fs.Usage()
		return fmt.Errorf("--source and --dest are required")
	}
	blockSize, err := parseSize(*blockSizeStr)
	if err != nil {
		return fmt.Errorf("--block-size: %w", err)
	}
	switch *sudoMode {
	case "auto", "always", "never":
	default:
		return fmt.Errorf("--sudo: want auto|always|never, got %q", *sudoMode)
	}

	return runSync(SyncConfig{
		Job: *job, Source: *source, Dest: *dest, BlockSize: blockSize,
		Yes: *yes, Sudo: *sudoMode, Deploy: *deploy, ConnectTimeoutSec: *connectTimeout,
		SSHBin: *sshBin, SSHOpts: sshOpts, RemoteBin: *remoteBin, ManagerHost: *managerHost,
	})
}

func parseIntList(s string) ([]int, error) {
	if strings.TrimSpace(s) == "" {
		return nil, nil
	}
	var out []int
	for _, f := range strings.Split(s, ",") {
		f = strings.TrimSpace(f)
		if f == "" {
			continue
		}
		var n int
		if _, err := fmt.Sscanf(f, "%d", &n); err != nil || n <= 0 {
			return nil, fmt.Errorf("invalid partition number %q", f)
		}
		out = append(out, n)
	}
	return out, nil
}

func cmdCloneDisk(args []string) error {
	fs := flag.NewFlagSet("clone-disk", flag.ContinueOnError)
	job := fs.String("job", "", "optional label for progress/log output")
	source := fs.String("source", "", "source whole-disk location (required)")
	dest := fs.String("dest", "", "destination whole-disk or image location (required)")
	blockSizeStr := fs.String("block-size", "4M", "block size for raw/boot-region transfers, e.g. 4M")
	yes := fs.Bool("yes", false, "don't prompt before shrinking an existing destination image")
	sudoMode := fs.String("sudo", "auto", "device-access privilege escalation: auto|always|never")
	deploy := fs.Bool("deploy", true, "copy this binary to remote hosts that lack it")
	connectTimeout := fs.Int("connect-timeout", defaultConnectTimeoutSec, "ssh connect timeout (seconds)")
	sshBin := fs.String("ssh", "ssh", "ssh binary")
	remoteBin := fs.String("remote-bin", "clonetool", "clonetool path on remote hosts")
	managerHost := fs.String("manager-host", "", "address peers use to reach this machine")
	partsStr := fs.String("parts", "", "only clone these partition numbers (comma-separated)")
	rawStr := fs.String("raw", "", "force raw block clone for these partition numbers")
	fileLevelStr := fs.String("file-level", "", "mkfs + rsync these partition numbers")
	fileAuto := fs.Bool("file-level-auto", false, "file-level for fs types with no image cloner")
	allowShrink := fs.Bool("allow-shrink", false, "permit shrinking trailing partitions to fit a smaller target")
	noShrink := fs.Bool("no-shrink", false, "never shrink; fail if the target is too small")
	newIDs := fs.Bool("new-ids", false, "randomize the GPT disk GUID / MBR signature on the target")
	reinstallBoot := fs.Bool("reinstall-bootloader", false, "run grub-install / grub-mkconfig on the target after copy")
	vss := fs.Bool("vss", true, "Windows source: take a Volume Shadow Copy per NTFS volume")
	imageSizeStr := fs.String("image-size", "", "file target: image size (default: enough for the layout)")
	var sshOpts stringSlice
	fs.Var(&sshOpts, "ssh-opt", `extra "-o OPT" passed to ssh (repeatable)`)
	if err := fs.Parse(args); err != nil {
		return err
	}
	if *source == "" || *dest == "" {
		fs.Usage()
		return fmt.Errorf("--source and --dest are required")
	}
	blockSize, err := parseSize(*blockSizeStr)
	if err != nil {
		return fmt.Errorf("--block-size: %w", err)
	}
	switch *sudoMode {
	case "auto", "always", "never":
	default:
		return fmt.Errorf("--sudo: want auto|always|never, got %q", *sudoMode)
	}
	parts, err := parseIntList(*partsStr)
	if err != nil {
		return fmt.Errorf("--parts: %w", err)
	}
	raw, err := parseIntList(*rawStr)
	if err != nil {
		return fmt.Errorf("--raw: %w", err)
	}
	fileLevel, err := parseIntList(*fileLevelStr)
	if err != nil {
		return fmt.Errorf("--file-level: %w", err)
	}
	var imageSize int64
	if *imageSizeStr != "" {
		if imageSize, err = parseSize(*imageSizeStr); err != nil {
			return fmt.Errorf("--image-size: %w", err)
		}
	}

	return runCloneDisk(CloneDiskConfig{
		SyncConfig: SyncConfig{
			Job: *job, Source: *source, Dest: *dest, BlockSize: blockSize,
			Yes: *yes, Sudo: *sudoMode, Deploy: *deploy, ConnectTimeoutSec: *connectTimeout,
			SSHBin: *sshBin, SSHOpts: sshOpts, RemoteBin: *remoteBin, ManagerHost: *managerHost,
		},
		Parts: parts, Raw: raw, FileLevel: fileLevel, FileAuto: *fileAuto,
		AllowShrink: *allowShrink, NoShrink: *noShrink, NewIDs: *newIDs,
		ReinstallBoot: *reinstallBoot, VSS: *vss, ImageSize: imageSize,
	})
}

func parseSize(s string) (int64, error) {
	if s == "" {
		return 0, fmt.Errorf("empty size")
	}
	mult := int64(1)
	numPart := s
	switch s[len(s)-1] {
	case 'k', 'K':
		mult = 1024
		numPart = s[:len(s)-1]
	case 'm', 'M':
		mult = 1024 * 1024
		numPart = s[:len(s)-1]
	case 'g', 'G':
		mult = 1024 * 1024 * 1024
		numPart = s[:len(s)-1]
	}
	var n int64
	if _, err := fmt.Sscanf(numPart, "%d", &n); err != nil {
		return 0, fmt.Errorf("invalid size %q", s)
	}
	if n <= 0 {
		return 0, fmt.Errorf("size must be positive")
	}
	return n * mult, nil
}